modify scripts

src/db_utils/db_common.py

@@ -0,0 +1,121 @@
+import sqlite3
+import json
+import logging
+from datetime import datetime
+import src.config.config as config
+
+# 连接 SQLite 数据库
+DB_PATH = f"{config.global_share_data_dir}/sqlite/shared.db"  # 替换为你的数据库文件
+
+# 检查 SQLite 版本
+lower_sqlite_version = False
+sqlite_version = sqlite3.sqlite_version_info
+if sqlite_version < (3, 24, 0):
+    lower_sqlite_version = True
+
+# 获取表的列名和默认值
+def get_table_columns_and_defaults(cursor, tbl_name):
+    try:
+        cursor.execute(f"PRAGMA table_info({tbl_name})")
+        columns = cursor.fetchall()
+        column_info = {}
+        for col in columns:
+            col_name = col[1]
+            default_value = col[4]
+            column_info[col_name] = default_value
+        return column_info
+    except sqlite3.Error as e:
+        logging.error(f"Error getting table columns: {e}")
+        return None
+
+# 检查并处理数据
+def check_and_process_data(cursor, data, tbl_name):
+    column_info = get_table_columns_and_defaults(cursor=cursor, tbl_name=tbl_name)
+    if column_info is None:
+        return None
+    processed_data = {}
+    for col, default in column_info.items():
+        if col == 'id' or col == 'created_at':      # 自增主键，不需要用户提供; 创建日期，使用建表默认值
+            continue
+        if col == 'updated_at':                     # 日期函数，用户自己指定即可
+            processed_data[col] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        if col in data:
+            processed_data[col] = data[col]
+
+    return processed_data
+
+        
+# 插入或更新数据
+def insert_or_update_common(cursor, conn, data, tbl_name, uniq_key='url'):
+    if lower_sqlite_version:
+        return insert_or_update_common_lower(cursor, conn, data, tbl_name, uniq_key)
+    
+    try:
+        processed_data = check_and_process_data(cursor, data, tbl_name)
+        if processed_data is None:
+            return None
+        
+        columns = ', '.join(processed_data.keys())
+        values = list(processed_data.values())
+        placeholders = ', '.join(['?' for _ in values])
+        update_clause = ', '.join([f"{col}=EXCLUDED.{col}" for col in processed_data.keys() if col != uniq_key]) 
+
+        sql = f'''
+            INSERT INTO {tbl_name} ({columns})
+            VALUES ({placeholders})
+            ON CONFLICT ({uniq_key}) DO UPDATE SET {update_clause}
+        '''
+        cursor.execute(sql, values)
+        conn.commit()
+
+        # 获取插入或更新后的 report_id
+        cursor.execute(f"SELECT id FROM {tbl_name} WHERE {uniq_key} = ?", (data[uniq_key],))
+        report_id = cursor.fetchone()[0]
+        return report_id
+    except sqlite3.Error as e:
+        logging.error(f"Error inserting or updating data: {e}")
+        return None
+
+# 插入或更新数据
+def insert_or_update_common_lower(cursor, conn, data, tbl_name, uniq_key='url'):
+    try:
+        processed_data = check_and_process_data(cursor, data, tbl_name)
+        if processed_data is None:
+            return None
+        
+        columns = ', '.join(processed_data.keys())
+        values = list(processed_data.values())
+        placeholders = ', '.join(['?' for _ in values])
+
+        # 先尝试插入数据
+        try:
+            sql = f'''
+                INSERT INTO {tbl_name} ({columns})
+                VALUES ({placeholders})
+            '''
+            cursor.execute(sql, values)
+            conn.commit()
+        except sqlite3.IntegrityError:  # 唯一键冲突，执行更新操作
+            update_clause = ', '.join([f"{col}=?" for col in processed_data.keys() if col != uniq_key]) 
+            update_values = [processed_data[col] for col in processed_data.keys() if col != uniq_key]
+            update_values.append(data[uniq_key])
+            sql = f"UPDATE {tbl_name} SET {update_clause} WHERE {uniq_key} = ?"
+            cursor.execute(sql, update_values)
+            conn.commit()
+
+        # 获取插入或更新后的 report_id
+        cursor.execute(f"SELECT id FROM {tbl_name} WHERE {uniq_key} = ?", (data[uniq_key],))
+        report_id = cursor.fetchone()[0]
+        return report_id
+    except sqlite3.Error as e:
+        logging.error(f"Error inserting or updating data: {e}")
+        return None
+
+
+# 测试代码
+if __name__ == "__main__":
+    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
+    cursor = conn.cursor()
+
+    tbl_name_actors = 'javhd_models'
+    print(get_table_columns_and_defaults(cursor, tbl_name_actors))