项目重构

src/PostgresSQL/db_handler.py

@@ -16,10 +16,11 @@ All returned results are pandas.DataFrame.
 from typing import Optional, List, Any
 import pandas as pd
 import psycopg2
+from psycopg2 import sql
 
 
 class PgHandler:
-    """Simple Postgres handler for reading tables into pandas DataFrames.
+    """Simple Postgres handler for reading tables.
 
     Parameters (pass to constructor):
       - host, port, user, password, database (db)
@@ -52,6 +53,15 @@ class PgHandler:
                                           password=self.password, dbname=self.database)
         return self._conn
 
+    def _get_admin_conn(self, dbname: str = 'postgres'):
+        """Create a fresh connection to a given database name (useful for creating/dropping DBs).
+
+        This returns a new connection instance (caller must close it). For operations like
+        CREATE/DROP DATABASE, autocommit should be set to True by the caller.
+        """
+        return psycopg2.connect(host=self.host, port=self.port, user=self.user,
+                                password=self.password, dbname=dbname)
+
     def close(self):
         """Close the underlying DB connection."""
         if self._conn is not None:
@@ -157,6 +167,138 @@ class PgHandler:
         tables = [row[0] for row in cursor.fetchall()]
         return tables
 
+    # -------------------- 新增的数据库/模式/表 操作 --------------------
+    def database_exists(self, db_name: str) -> bool:
+        conn = self._get_admin_conn('postgres')
+        try:
+            cur = conn.cursor()
+            cur.execute("SELECT 1 FROM pg_database WHERE datname = %s", (db_name,))
+            exists = cur.fetchone() is not None
+            return exists
+        finally:
+            conn.close()
+
+    def create_database(self, db_name: str, if_force_create: bool = False) -> bool:
+        """Create a database. Returns True if created, False if not created because it existed and if_force_create is False."""
+        if self.database == db_name:
+            # if current handler is bound to this database, operate via admin connection instead
+            pass
+
+        if self.database_exists(db_name):
+            if not if_force_create:
+                return False
+            # terminate connections and drop db first
+            self.drop_database(db_name)
+
+        conn = self._get_admin_conn('postgres')
+        conn.autocommit = True
+        try:
+            cur = conn.cursor()
+            cur.execute(sql.SQL('CREATE DATABASE {}').format(sql.Identifier(db_name)))
+            return True
+        finally:
+            conn.close()
+
+    def drop_database(self, db_name: str) -> None:
+        """Drop a database. Terminates active connections first."""
+        conn = self._get_admin_conn('postgres')
+        conn.autocommit = True
+        try:
+            cur = conn.cursor()
+            # terminate other backends connected to the target database
+            cur.execute("SELECT pid FROM pg_stat_activity WHERE datname = %s AND pid <> pg_backend_pid()", (db_name,))
+            rows = cur.fetchall()
+            for (pid,) in rows:
+                try:
+                    cur.execute("SELECT pg_terminate_backend(%s)", (pid,))
+                except Exception:
+                    # best-effort
+                    pass
+
+            cur.execute(sql.SQL('DROP DATABASE IF EXISTS {}').format(sql.Identifier(db_name)))
+        finally:
+            conn.close()
+
+    def schema_exists(self, schema: str) -> bool:
+        conn = self._get_conn()
+        cur = conn.cursor()
+        cur.execute("SELECT 1 FROM pg_namespace WHERE nspname = %s", (schema,))
+        return cur.fetchone() is not None
+
+    def create_schema(self, schema: str, if_force_create: bool = False) -> bool:
+        """Create schema in the current database. Returns True if created."""
+        conn = self._get_conn()
+        cur = conn.cursor()
+        if self.schema_exists(schema):
+            if not if_force_create:
+                return False
+            else:
+                # drop then create
+                self.drop_schema(schema)
+
+        cur.execute(sql.SQL('CREATE SCHEMA {}').format(sql.Identifier(schema)))
+        conn.commit()
+        return True
+
+    def drop_schema(self, schema: str) -> None:
+        conn = self._get_conn()
+        cur = conn.cursor()
+        cur.execute(sql.SQL('DROP SCHEMA IF EXISTS {} CASCADE').format(sql.Identifier(schema)))
+        conn.commit()
+
+    def table_exists(self, table: str, schema: str = 'public') -> bool:
+        conn = self._get_conn()
+        cur = conn.cursor()
+        cur.execute(
+            "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_schema = %s AND table_name = %s)",
+            (schema, table)
+        )
+        row = cur.fetchone()
+        return bool(row and row[0])
+
+    def create_table(self, table: str, columns_def: str, schema: str = 'public', if_force_create: bool = False) -> bool:
+        """Create a table with provided columns definition SQL (e.g. 'id serial PRIMARY KEY, name text').
+
+        Returns True if created, False if already existed and if_force_create is False.
+        """
+        conn = self._get_conn()
+        cur = conn.cursor()
+        if self.table_exists(table, schema):
+            if not if_force_create:
+                return False
+            else:
+                self.drop_table(table, schema=schema)
+
+        q = sql.SQL('CREATE TABLE {}.{} ({})').format(sql.Identifier(schema), sql.Identifier(table), sql.SQL(columns_def))
+        cur.execute(q)
+        conn.commit()
+        return True
+
+    def drop_table(self, table: str, schema: str = 'public') -> None:
+        conn = self._get_conn()
+        cur = conn.cursor()
+        cur.execute(sql.SQL('DROP TABLE IF EXISTS {}.{} CASCADE').format(sql.Identifier(schema), sql.Identifier(table)))
+        conn.commit()
+    # -------------------- 新增结束 --------------------
+
+    def insert_json_into_table(self, table: str, json_list: list, schema: str = "public") -> None:
+        """
+        批量将json（字典）列表插入到指定schema下的table中。
+        :param table: 表名
+        :param json_list: 字典列表，每个字典为一行数据
+        :param schema: schema名，默认为public
+        """
+        if not json_list:
+            return
+        conn = self._get_conn()
+        columns = list(json_list[0].keys())
+        values = [[row.get(col) for col in columns] for row in json_list]
+        placeholders = ', '.join(['%s'] * len(columns))
+        col_names = ', '.join([f'"{col}"' for col in columns])
+        sql = f'INSERT INTO "{schema}"."{table}" ({col_names}) VALUES ({placeholders})'
+        with conn.cursor() as cur:
+            cur.executemany(sql, values)
+            conn.commit()
 
     def __enter__(self):
         # ensure connection is created and return self for use in `with` blocks

src/PostgresSQL/pg_db_handler.py

@@ -0,0 +1,203 @@
+import psycopg2
+import json
+
+
+class DatabaseHandler:
+    def __init__(self, config):
+        self.config = config
+        self.connection = psycopg2.connect(database=config['db_database'], user=config['db_user'],
+                                           password=config['db_password'], host=config['db_host'], port=config['db_port'])
+        self.cursor = self.connection.cursor()
+        self.cursor.execute(f"select tablename from pg_tables where schemaname='public'")
+        self.table_list = [item[0] for item in self.cursor.fetchall()]
+
+    def ensure_connection(self):
+        try:
+            self.cursor.execute(f"select tablename from pg_tables where schemaname='public'")
+        except Exception as e:
+            print(e)
+            config = self.config
+            self.connection = psycopg2.connect(database=config['db_database'], user=config['db_user'],
+                                               password=config['db_password'], host=config['db_host'],
+                                               port=config['db_port'])
+            self.cursor = self.connection.cursor()
+
+    def create_database(self, db_name, force_create=False):
+        self.ensure_connection()
+        if force_create:
+            self.drop_database(db_name)
+        try:
+            self.connection.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+            self.cursor.execute(f'CREATE DATABASE {db_name}')
+            print(f'Database {db_name} successfully created.')
+        except psycopg2.errors.DuplicateDatabase as e:
+            print(f'Database {db_name} already exists: {e}')
+
+    def drop_database(self, db_name):
+        self.ensure_connection()
+        self.connection.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+        self.cursor.execute(f'drop database if exists {db_name}')
+        print(f'Database {db_name} successfully dropped.')
+        self.table_list.remove(db_name)
+
+    def create_table(self, table_structure, table_name, force_create=False, primary_key_list=None):
+        self.ensure_connection()
+        table_name = table_name.lower()
+        if force_create and self.exists_table(table_name):
+            self.drop_table(table_name)
+        sql = f'Create Table {table_name}('
+        content_list = []
+        for key, values in table_structure.items():
+            key_string = f'{key.replace(" ", "_").replace("-", "_")}           {values}'
+            if primary_key_list and key in primary_key_list:
+                key_string += ' PRIMARY KEY'
+            content_list.append(key_string)
+        sql += ',\n'.join(content_list)
+        sql += ');'
+        self.connection.cursor().execute(sql)
+        self.connection.commit()
+        self.table_list.append(table_name)
+
+    def exists_table(self, table_name):
+        table_name = table_name.lower()
+        if table_name in self.table_list:
+            return True
+        else:
+            return False
+
+    def drop_table(self, table_name, info_flag=False):
+        self.ensure_connection()
+        table_name = table_name.lower()
+        if not self.exists_table(table_name):
+            if info_flag:
+                print(f'table: {table_name} not exists in database.')
+            return
+        self.cursor.execute(f'drop table {table_name}')
+        self.connection.commit()
+        if info_flag:
+            print(f'table: {table_name} dropped')
+        return
+
+    def insert_into_table(self, table_name, data_dict_list, batch=1000, schema='public'):
+        """
+
+        :param schema:
+        :param batch:
+        :param table_name:
+        :param data_dict_list:
+        :return:
+        """
+        self.ensure_connection()
+        full_table_name = f"{schema}.{table_name.lower()}"
+        # 检查表是否存在时也要加上schema
+        self.cursor.execute(
+            "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_schema=%s AND table_name=%s)",
+            (schema, table_name.lower())
+        )
+        if not self.cursor.fetchone()[0]:
+            print(f'table: {full_table_name} not exists in database.')
+            return
+
+        table_structure = data_dict_list[0]
+        content_list = [f'{key.replace(" ", "_").replace("-", "_")}'.lower() for key in table_structure.keys()]
+        sql = f"INSERT INTO {full_table_name}({','.join(content_list)}) VALUES({','.join(['%s'] * len(content_list))})"
+        target = [list(item.values()) for item in data_dict_list]
+
+        for i in range(0, len(target), batch):
+            self.cursor.executemany(sql, target[i: i + batch])
+            self.connection.commit()
+            print(f'Inserted into table {full_table_name}, rows: {min(i + batch, len(target))}/{len(target)}')
+
+    def insert_json_into_table(self, table_name, data_list, batch=1000, unique_columns=None):
+        """
+        insert a list of json into table. The json content should contain all the required columns as its keys.
+        :param batch: batch size
+        :param table_name:  which table to insert into
+        :param data_list:   target json
+        :param unique_columns:  the unique column of the table. If set, that column will be
+        checked to ensure no duplicate entries
+        :return:
+        """
+        self.ensure_connection()
+        table_name = table_name.lower()
+        if not self.exists_table(table_name):
+            print(f'table: {table_name} not exists in database.')
+            return
+
+        table_structure = list(data_list[0].keys())
+        if unique_columns and unique_columns not in table_structure:
+            raise ValueError(f'{unique_columns} not in data')
+        if unique_columns:
+            sql = (f"INSERT INTO {table_name}({','.join(table_structure)}) "
+                   f"VALUES({','.join(['%s'] * len(table_structure))})"
+                   f" ON CONFLICT ({unique_columns}) DO NOTHING")
+        else:
+            sql = (f"INSERT INTO {table_name}({','.join(table_structure)}) "
+                   f"VALUES({','.join(['%s'] * len(table_structure))})")
+        target = [list(item.values()) for item in data_list]
+
+        for i in range(0, len(target), batch):
+            try:
+                self.cursor.executemany(sql, target[i: i+batch])
+            except psycopg2.errors.InternalError_ as e:
+                print(e)
+                size = 0
+                for line in target[i:i+batch]:
+                    if line[-1]:
+                        size += line[-1]
+                print(size / 1024 / 1024)
+                print('error')
+            except Exception as e:
+                print(sql)
+                print(e)
+            self.connection.commit()
+            print(f'Inserted into table {table_name}, rows: {min(i+batch, len(target))}/{len(target)}')
+
+    def get_data(self, table_name, field_name_list=None):
+        self.ensure_connection()
+        if not field_name_list:
+            sql = f"SELECT * FROM {table_name}"
+        elif isinstance(field_name_list, str):
+            sql = f"SELECT {field_name_list} FROM {table_name}"
+        else:
+            sql = f"select {','.join(field_name_list)} from {table_name}"
+        self.cursor.execute(sql)
+        res = self.cursor.fetchall()
+        return res
+
+    def update_table(self, table_name, update_data: dict, key_column):
+        self.ensure_connection()
+        table_name = table_name.lower()
+        if not self.exists_table(table_name):
+            print(f'table: {table_name} not exists in database.')
+            return
+        if key_column not in update_data:
+            raise ValueError(f'{key_column} not in data')
+        target_list = []
+        for key, value in update_data.items():
+            if key == key_column:
+                continue
+            if isinstance(value, dict):
+                target_list.append(f"""{key} = '{json.dumps(value).replace("'", "''")}'""")
+            else:
+                target_list.append(f"""{key} = '{str(value).replace("'", "''")}'""")
+        target = ','.join(target_list)
+        sql = f"update {table_name} set {target} where {key_column}='{update_data[key_column]}'"
+        self.cursor.execute(sql)
+        self.connection.commit()
+
+    def close(self):
+        self.connection.close()
+
+
+if __name__ == "__main__":
+    config = {
+"db_user": "app",
+"db_password": "Letmein1",
+"db_host": "pgm-8vb3061igsxn50is9o.pgsql.zhangbei.rds.aliyuncs.com",
+"db_port": "5432",
+"db_database": "flex_desk",
+}
+    # db_handler = DatabaseHandler(config)
+    # print(db_handler.get_data("wiki.feishu_doc_public"))
+    # db_handler.insert_into_table('feishu_doc_public', [table_schema], batch=1, schema='wiki')

src/TG-tools/history_message_upload.py

@@ -11,7 +11,8 @@ from tqdm import tqdm
 
 from src.My_Logger.project_logger import LoggerClass
 from src.PostgresSQL.databaseHandler import DatabaseHandler
-from src.utils.utils import load_config, size_of_file
+from src.utils.utils import size_of_file
+from src.config.load_config import project_config
 
 project_path = Path(__file__).parent.parent.parent.absolute()
 data_path = os.path.join(project_path, 'data')
@@ -40,9 +41,9 @@ def process_message_text(text):
 
 
 class MessageUploader:
-    def __init__(self, config_path, local_logger, local_load_path, local_file_db_path, size_limit=0.5):
+    def __init__(self, local_logger, local_load_path, local_file_db_path, size_limit=0.5):
         self.logger = local_logger
-        self.config = load_config(config_path)
+        self.config = project_config
         self.load_path = local_load_path
         self.file_param_list = ['photo', 'file']
         self.db_handler = DatabaseHandler(self.config, self.logger)
@@ -196,14 +197,12 @@ def upload_didi_saved_message_to_db():
 
 
 if __name__ == "__main__":
-
-    config = load_config(config_path)
-    config['database'] = "TestDB"
-    tg = MessageUploader(config_path=config_path,
-                         local_logger=LoggerClass().get_log(logging.DEBUG),
-                         local_load_path='E:\\SNSFiles\\ChatExport_2024-07-10',
-                         local_file_db_path=config['file_db_path'],
-                         size_limit=0.25)
+    tg = MessageUploader(
+        local_logger=LoggerClass().get_log(logging.DEBUG),
+        local_load_path='E:\\SNSFiles\\ChatExport_2024-07-10',
+        local_file_db_path=project_config['file_db_path'],
+        size_limit=0.25
+    )
     try:
         tg.do_upload()
         # tg.logger_test()

src/config/load_config.py

@@ -19,7 +19,3 @@ else:
 
 with open(CONFIG_PATH, 'r', encoding='utf-8') as config_file:
     project_config = json.load(config_file)
-
-
-LLM = ChatOpenAI(model=project_config['llm_model_name'], api_key= SecretStr(project_config['llm_secret_key']),
-                 base_url=project_config['llm_base_url'], temperature=0.01)

src/utils/utils.py

@@ -2,22 +2,6 @@ import os
 import json
 
 
-def load_config(config_path):
-    """
-
-    :param config_path: config folder path
-    :return:
-    """
-    if os.path.exists(os.path.join(config_path, 'config_local.txt')):
-        config_path = os.path.join(config_path, 'config_local.txt')
-    else:
-        config_path = os.path.join(config_path, 'config.txt')
-    with open(config_path, 'r', encoding='utf-8') as config_file:
-        config_json = json.loads(config_file.read())
-    config_json['database'] = 'Telegram_Message'
-    return config_json
-
-
 def size_of_file(file_size):
     units = ['B', 'KB', 'MB', 'GB', 'TB', 'PB']
     i = 0