X-Git-Url: https://plomlompom.com/repos/feed.xml?a=blobdiff_plain;f=plomtask%2Fdb.py;h=f45d2b6adac0e109938d9f8018653cf1caa73fcf;hb=7adcf651f0053f0dc5d719457a016a0d5b12253b;hp=abd8f6181db946b357902a0ec851f51919f37344;hpb=5a5d713ce0b223ab2f6ef34c15bb82b614bdda98;p=plomtask

diff --git a/plomtask/db.py b/plomtask/db.py
index abd8f61..f45d2b6 100644
--- a/plomtask/db.py
+++ b/plomtask/db.py
@@ -67,6 +67,25 @@ class DatabaseConnection:
         """Close DB connection."""
         self.conn.close()
 
+    def rewrite_relations(self, table_name: str, key: str, target: int,
+                          rows: list[list[Any]]) -> None:
+        """Rewrite relations in table_name to target, with rows values."""
+        self.exec(f'DELETE FROM {table_name} WHERE {key} = ?', (target,))
+        for row in rows:
+            values = tuple([target] + row)
+            q_marks = self.__class__.q_marks_from_values(values)
+            self.exec(f'INSERT INTO {table_name} VALUES {q_marks}', values)
+
+    def all_where(self, table_name: str, key: str, target: int) -> list[Row]:
+        """Return list of Rows at table where key == target."""
+        return list(self.exec(f'SELECT * FROM {table_name} WHERE {key} = ?',
+                              (target,)))
+
+    @staticmethod
+    def q_marks_from_values(values: tuple[Any]) -> str:
+        """Return placeholder to insert values into SQL code."""
+        return '(' + ','.join(['?'] * len(values)) + ')'
+
 
 class BaseModel:
     """Template for most of the models we use/derive from the DB."""
@@ -114,11 +133,11 @@ class BaseModel:
     def save_core(self, db_conn: DatabaseConnection,
                   update_with_lastrowid: bool = True) -> None:
         """Write bare-bones self (sans connected items), ensuring self.id_."""
-        q_marks = ','.join(['?'] * (len(self.to_save) + 1))
         values = tuple([self.id_] + [getattr(self, key)
                                      for key in self.to_save])
+        q_marks = DatabaseConnection.q_marks_from_values(values)
         table_name = self.table_name
-        cursor = db_conn.exec(f'REPLACE INTO {table_name} VALUES ({q_marks})',
+        cursor = db_conn.exec(f'REPLACE INTO {table_name} VALUES {q_marks}',
                               values)
         if update_with_lastrowid:
             self.id_ = cursor.lastrowid