X-Git-Url: https://plomlompom.com/repos/%7B%7Bprefix%7D%7D/balance?a=blobdiff_plain;f=plomtask%2Fdb.py;h=e4d5f6e9a42fe4726c5dff35e9488e13b61eea13;hb=5e3c633f1994329297999899790e69d28516934b;hp=8e5529062ae18c60a2ed222deb27305f9c18e54c;hpb=3a74c67ef298d16e7a7e5dd030e1fd813f42f0b4;p=plomtask

diff --git a/plomtask/db.py b/plomtask/db.py
index 8e55290..e4d5f6e 100644
--- a/plomtask/db.py
+++ b/plomtask/db.py
@@ -75,7 +75,7 @@ class DatabaseConnection:
         """Close DB connection."""
         self.conn.close()
 
-    def rewrite_relations(self, table_name: str, key: str, target: int,
+    def rewrite_relations(self, table_name: str, key: str, target: int | str,
                           rows: list[list[Any]]) -> None:
         """Rewrite relations in table_name to target, with rows values."""
         self.delete_where(table_name, key, target)
@@ -121,6 +121,8 @@ class BaseModel(Generic[BaseModelId]):
     """Template for most of the models we use/derive from the DB."""
     table_name = ''
     to_save: list[str] = []
+    to_save_versioned: list[str] = []
+    to_save_relations: list[tuple[str, str, str]] = []
     id_: None | BaseModelId
     cache_: dict[BaseModelId, Self]
 
@@ -130,6 +132,23 @@ class BaseModel(Generic[BaseModelId]):
             raise HandledException(msg)
         self.id_ = id_
 
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, self.__class__):
+            return False
+        to_hash_me = tuple([self.id_] +
+                           [getattr(self, name) for name in self.to_save])
+        to_hash_other = tuple([other.id_] +
+                              [getattr(other, name) for name in other.to_save])
+        return hash(to_hash_me) == hash(to_hash_other)
+
+    def __lt__(self, other: Any) -> bool:
+        if not isinstance(other, self.__class__):
+            msg = 'cannot compare to object of different class'
+            raise HandledException(msg)
+        assert isinstance(self.id_, int)
+        assert isinstance(other.id_, int)
+        return self.id_ < other.id_
+
     @classmethod
     def get_cached(cls: type[BaseModelInstance],
                    id_: BaseModelId) -> BaseModelInstance | None:
@@ -209,7 +228,13 @@ class BaseModel(Generic[BaseModelId]):
     @classmethod
     def all(cls: type[BaseModelInstance],
             db_conn: DatabaseConnection) -> list[BaseModelInstance]:
-        """Collect all objects of class."""
+        """Collect all objects of class into list.
+
+        Note that this primarily returns the contents of the cache, and only
+        _expands_ that by additional findings in the DB. This assumes the
+        cache is always instantly cleaned of any items that would be removed
+        from the DB.
+        """
         items: dict[BaseModelId, BaseModelInstance] = {}
         for k, v in cls.get_cache().items():
             assert isinstance(v, cls)
@@ -222,21 +247,11 @@ class BaseModel(Generic[BaseModelId]):
                 items[item.id_] = item
         return list(items.values())
 
-    def __eq__(self, other: object) -> bool:
-        if not isinstance(other, self.__class__):
-            msg = 'cannot compare to object of different class'
-            raise HandledException(msg)
-        to_hash_me = tuple([self.id_] +
-                           [getattr(self, name) for name in self.to_save])
-        to_hash_other = tuple([other.id_] +
-                              [getattr(other, name) for name in other.to_save])
-        return hash(to_hash_me) == hash(to_hash_other)
-
-    def save_core(self, db_conn: DatabaseConnection) -> None:
-        """Write bare-bones self (sans connected items), ensuring self.id_.
+    def save(self, db_conn: DatabaseConnection) -> None:
+        """Write self to DB and cache and ensure .id_.
 
         Write both to DB, and to cache. To DB, write .id_ and attributes
-        listed in cls.to_save.
+        listed in cls.to_save[_versioned|_relations].
 
         Ensure self.id_ by setting it to what the DB command returns as the
         last saved row's ID (cursor.lastrowid), EXCEPT if self.id_ already
@@ -252,9 +267,21 @@ class BaseModel(Generic[BaseModelId]):
         if not isinstance(self.id_, str):
             self.id_ = cursor.lastrowid  # type: ignore[assignment]
         self.cache()
+        for attr_name in self.to_save_versioned:
+            getattr(self, attr_name).save(db_conn)
+        for table, column, attr_name in self.to_save_relations:
+            assert isinstance(self.id_, (int, str))
+            db_conn.rewrite_relations(table, column, self.id_,
+                                      [[i.id_] for i
+                                       in getattr(self, attr_name)])
 
     def remove(self, db_conn: DatabaseConnection) -> None:
-        """Remove from DB."""
-        assert isinstance(self.id_, int | str)
+        """Remove from DB and cache, including dependencies."""
+        if self.id_ is None or self.__class__.get_cached(self.id_) is None:
+            raise HandledException('cannot remove unsaved item')
+        for attr_name in self.to_save_versioned:
+            getattr(self, attr_name).remove(db_conn)
+        for table, column, attr_name in self.to_save_relations:
+            db_conn.delete_where(table, column, self.id_)
         self.uncache()
         db_conn.delete_where(self.table_name, 'id', self.id_)