Browse Source

Update database operation: support insert ignore

StrayWarrior 4 months ago
parent
commit
69edfc2d97
2 changed files with 18 additions and 5 deletions
  1. 12 1
      run_category_model_v1.py
  2. 6 4
      src/common/database.py

+ 12 - 1
run_category_model_v1.py

@@ -163,7 +163,18 @@ def main():
         compare_version(db_manager, dt_version, records_to_save, account_id_map)
         return
 
-    db_manager.batch_insert('account_category', records_to_save)
+    rows = db_manager.batch_insert('account_category', records_to_save, ignore=True)
+    if rows != len(records_to_save):
+        for record in records_to_save:
+            sql = f"""
+                UPDATE account_category
+                SET category_map = '{record['category_map']}',
+                    update_timestamp = {record['update_timestamp']}
+                WHERE dt = {record['dt']} AND gh_id = '{record['gh_id']}'
+                    AND category_map != '{record['category_map']}'
+            """
+            update_rows = db_manager.execute(sql)
+            print(f"updated rows: {update_rows}, {record['gh_id']}")
     clear_old_version(db_manager, dt_version)
 
     # 过滤空账号

+ 6 - 4
src/common/database.py

@@ -37,14 +37,14 @@ class MySQLManager:
         finally:
             conn.close()
 
-    def batch_insert(self, table, data, columns=None):
+    def batch_insert(self, table, data, columns=None, ignore=False):
         """
         table: table name, string
         data: data, list[tuple] or list[dict]
         columns: column names, list, required if data is list[tuple]
         """
         if data is None or len(data) == 0:
-            return
+            return None
         conn = pymysql.connect(**self.config)
         try:
             if isinstance(data[0], dict):
@@ -56,10 +56,12 @@ class MySQLManager:
                     raise Exception("data length != column length")
                 columns_str = ','.join(columns)
                 placeholders_str = ','.join(['%s'] * len(data[0]))
+            ignore_keyword = 'IGNORE' if ignore else ''
             with conn.cursor() as cursor:
-                sql_str = f"INSERT INTO {table} ({columns_str}) VALUES ({placeholders_str})"
-                cursor.executemany(sql_str, data)
+                sql_str = f"INSERT {ignore_keyword} INTO {table} ({columns_str}) VALUES ({placeholders_str})"
+                rows = cursor.executemany(sql_str, data)
                 conn.commit()
+                return rows
         except pymysql.MySQLError as e:
             print(f"Error in batch_insert: {e}")
             conn.rollback()