Refactor save and remove methods of BaseObject subclasses.

[plomtask] / plomtask / db.py

diff --git a/plomtask/db.py b/plomtask/db.py
index 982ddfe3b96915d4ffc8bd8d4bce264671c09069..e4d5f6e9a42fe4726c5dff35e9488e13b61eea13 100644 (file)
--- 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]
 
@@ -245,11 +247,11 @@ class BaseModel(Generic[BaseModelId]):
                 items[item.id_] = item
         return list(items.values())
 
-    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
@@ -265,10 +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 and cache."""
+        """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_)