移除过时的db handler

src/KSP_tools/EngineDataProcess_DB_upload.py

@@ -1,5 +1,5 @@
 import pandas as pd
-from src.PostgresSQL.databaseHandler import DatabaseHandler
+from src.PostgresSQL.db_handler import PgHandler
 
 
 def process_excel():
@@ -31,9 +31,12 @@ def upload_to_database(database_config):
             table_structure[key] = 'FLOAT4'
         else:
             table_structure[key] = 'TEXT'
-    handler = DatabaseHandler(database_config)
+    handler = PgHandler(host=database_config.get('host'), port=int(database_config.get('port', 5432)),
+                        user=database_config.get('user'), password=database_config.get('password'),
+                        database=database_config.get('database'))
 
-    handler.create_table(table_structure, 'Engine_Data', force_create=True)
+    # 使用 PgHandler.create_table_from_dict 来从 dict 创建表
+    handler.create_table_from_dict(table_structure, 'Engine_Data', force_create=True, primary_key_list=None)
     handler.insert_into_table('Engine_Data', data_content_list)
 
     handler.close()

src/PostgresSQL/databaseHandler.py

@@ -1,123 +0,0 @@
-import psycopg2
-import logging
-from src.My_Logger.project_logger import LoggerClass
-
-
-class DatabaseHandler:
-    def __init__(self, config, logger=None):
-        if not logger or not isinstance(logger, logging.RootLogger):
-            log = LoggerClass()
-            self.logger = log.get_log(logging.DEBUG)
-        else:
-            self.logger = logger
-        self.config = config
-        # {database="postgres", user="postgres", password="123456", host="localhost", port="5432"}
-        self.connection = psycopg2.connect(database=config['database'], user=config['user'],
-                                           password=config['password'], host=config['host'], port=config['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 create_table(self, table_structure, table_name, force_create=False, primary_key_list=None):
-        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):
-        table_name = table_name.lower()
-        if not self.exists_table(table_name):
-            if info_flag:
-                self.logger.error(f'table: {table_name} not exists in database.')
-            return
-        self.cursor.execute(f'drop table {table_name}')
-        self.connection.commit()
-        if info_flag:
-            self.logger.info(f'table: {table_name} dropped')
-        return
-
-    def insert_into_table(self, table_name, data_dict_list, batch=1000):
-        """
-
-        :param batch:
-        :param table_name:
-        :param data_dict_list:
-        :return:
-        """
-        table_name = table_name.lower()
-        if not self.exists_table(table_name):
-            self.logger.error(f'table: {table_name} not exists in database.')
-            return
-
-        table_structure = data_dict_list[0]
-        content_list = []
-        for key, values in table_structure.items():
-            content_list.append(f'{key.replace(" ", "_").replace("-", "_")}'.lower())
-
-        sql = f"INSERT INTO {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()
-            self.logger.info(f'Inserted into table {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:
-        """
-        table_name = table_name.lower()
-        if not self.exists_table(table_name):
-            self.logger.error(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')
-        sql = (f"INSERT INTO {table_name}({','.join(table_structure)}) VALUES({','.join(['%s'] * len(table_structure))})"
-               f" ON CONFLICT ({unique_columns}) DO NOTHING")
-        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()
-            self.logger.info(f'Inserted into table {table_name}, rows: {min(i+batch, len(target))}/{len(target)}')
-
-    def close(self):
-        self.connection.close()

src/PostgresSQL/db_handler.py

@@ -1,17 +1,6 @@
 """Postgres helper for reading tables into pandas DataFrame.
 
-Provides PgHandler class with methods:
+Provides PgHandler class with methods for reading and basic DDL/DML.
-- read_table(table, schema='public', columns=None)
-- read_table_by_conditions(table, conditions, schema='public', columns=None)
-
-Conditions format (single dict or list of dicts):
-{
-    'column': 'aaa',
-    'condition': 'larger than',  # 'equal', 'larger than', 'lower than'
-    'target': ['relative', 'bbb'] or ['direct', 123]
-}
-
-All returned results are pandas.DataFrame.
 """
 from typing import Optional, List, Any
 import pandas as pd
@@ -24,9 +13,6 @@ class PgHandler:
 
     Parameters (pass to constructor):
       - host, port, user, password, database (db)
-
-    Behavior:
-      - Connects directly to the given host:port with provided credentials.
     """
 
     def __init__(self, *, host: str = "localhost", port: int = 5432, user: str = "postgres",
@@ -45,19 +31,14 @@ class PgHandler:
         if self._conn is None:
             if self.database is None:
                 raise ValueError("database must be provided")
-
+            self._conn = psycopg2.connect(host=self.host, port=self.port, user=self.user,
-            connect_host = self.host
-            connect_port = self.port
-
-            self._conn = psycopg2.connect(host=connect_host, port=connect_port, user=self.user,
                                           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
+        This returns a new connection instance (caller must close it).
-        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)
@@ -72,8 +53,6 @@ class PgHandler:
 
     def read_table(self, table: str, schema: str = "public", columns: Optional[List[str]] = None) -> pd.DataFrame:
         """Read a table (or specific columns) and return a DataFrame.
-
-        If columns is None or empty -> select *.
         """
         conn = self._get_conn()
         cols = "*" if not columns else ", ".join([f'"{c}"' for c in columns])
@@ -81,15 +60,7 @@ class PgHandler:
         return pd.read_sql_query(q, con=conn)
 
     def read_table_by_conditions(self, table: str, conditions: Any, schema: str = "public", columns: Optional[List[str]] = None) -> pd.DataFrame:
-        """Read rows filtered by condition dict or list of such dicts.
+        """Read rows filtered by condition dict or list of such dicts."""
-
-        conditions can be a single dict or a list of dicts. Each dict has keys:
-          - column: column name to compare
-          - condition: 'equal', 'larger than', 'lower than'
-          - target: ['relative', other_column] or ['direct', value]
-
-        Multiple conditions are AND-ed together.
-        """
         if isinstance(conditions, dict):
             condition_list = [conditions]
         elif isinstance(conditions, list):
@@ -97,7 +68,6 @@ class PgHandler:
         else:
             raise ValueError("conditions must be a dict or list of dicts")
 
-        # build where clauses
         where_clauses = []
         params_list: List[Any] = []
 
@@ -147,7 +117,6 @@ class PgHandler:
             q = q + " WHERE " + where_sql
 
         conn = self._get_conn()
-        # pandas.read_sql_query accepts DB-API connection and positional params
         return pd.read_sql_query(q, con=conn, params=tuple(params_list) if params_list else None)
 
     def query(self, query: str) -> list[dict]:
@@ -160,36 +129,29 @@ class PgHandler:
         return result_list
 
     def get_all_tables(self, schema: str = "public") -> list[str]:
-        """Get a list of all table names in the specified schema."""
         conn = self._get_conn()
         cursor = conn.cursor()
-        cursor.execute(f"select tablename from pg_tables where schemaname='{schema}'", (schema,))
+        cursor.execute("select tablename from pg_tables where schemaname=%s", (schema,))
         tables = [row[0] for row in cursor.fetchall()]
         return tables
 
-    # -------------------- 新增的数据库/模式/表 操作 --------------------
+    # -------------------- DDL / DB operations --------------------
     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 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:
@@ -200,21 +162,17 @@ class PgHandler:
             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()
@@ -226,16 +184,13 @@ class PgHandler:
         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
@@ -256,11 +211,12 @@ class PgHandler:
         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:
+    # alias to match some older call sites
-        """Create a table with provided columns definition SQL (e.g. 'id serial PRIMARY KEY, name text').
+    def exists_table(self, table_name: str, schema: str = 'public') -> bool:
+        return self.table_exists(table_name, schema=schema)
 
-        Returns True if created, False if already existed and if_force_create is False.
+    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')."""
         conn = self._get_conn()
         cur = conn.cursor()
         if self.table_exists(table, schema):
@@ -268,72 +224,99 @@ class PgHandler:
                 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 create_table_from_dict(self, table_structure: dict, table_name: str, force_create: bool = False, primary_key_list: Optional[list] = None, schema: str = 'public') -> None:
+        """Create table from dict mapping column name -> sql type. Preserves primary keys if provided."""
+        table_name_l = table_name.lower()
+        # build column definitions
+        cols = []
+        for key, val in table_structure.items():
+            colname = key.replace(' ', '_').replace('-', '_')
+            coldef = f'"{colname}" {val}'
+            if primary_key_list and key in primary_key_list:
+                coldef += ' PRIMARY KEY'
+            cols.append(coldef)
+        columns_def = ', '.join(cols)
+        # delegate to create_table
+        self.create_table(table_name_l, columns_def, schema=schema, if_force_create=force_create)
+
     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:
+    def insert_into_table(self, table_name: str, data_dict_list: list, batch: int = 1000, schema: str = 'public') -> None:
-        """
+        """Insert list of dicts into table in batches."""
-        批量将json（字典）列表插入到指定schema下的table中。
+        if not data_dict_list:
-        :param table: 表名
+            return
-        :param json_list: 字典列表，每个字典为一行数据
+        table_name_l = table_name.lower()
-        :param schema: schema名，默认为public
+        if not self.table_exists(table_name_l, schema=schema):
-        """
+            raise ValueError(f'table: {schema}.{table_name_l} not exists in database.')
+        conn = self._get_conn()
+        cur = conn.cursor()
+        columns = [k.replace(' ', '_').replace('-', '_') for k in data_dict_list[0].keys()]
+        col_names = ','.join([f'"{c}"' for c in columns])
+        placeholders = ','.join(['%s'] * len(columns))
+        sql_stmt = f'INSERT INTO "{schema}"."{table_name_l}" ({col_names}) VALUES ({placeholders})'
+        target = [list(item.values()) for item in data_dict_list]
+        for i in range(0, len(target), batch):
+            cur.executemany(sql_stmt, target[i:i+batch])
+            conn.commit()
+
+    def insert_json_into_table(self, table: str, json_list: list, schema: str = "public", batch: int = 1000, unique_columns: Optional[str] = None) -> None:
+        """批量将 json（字典）列表插入到指定 schema 下的 table 中，支持 batch 和 unique_columns（ON CONFLICT ...）。"""
         if not json_list:
             return
         conn = self._get_conn()
-        columns = list(json_list[0].keys())
+        cur = conn.cursor()
-        values = [[row.get(col) for col in columns] for row in json_list]
+        table_name_l = table.lower()
+        if not self.table_exists(table_name_l, schema=schema):
+            raise ValueError(f'table: {schema}.{table_name_l} not exists in database.')
+        columns = [k.replace(' ', '_').replace('-', '_') for k in json_list[0].keys()]
+        col_names = ','.join([f'"{c}"' for c in columns])
         placeholders = ','.join(['%s'] * len(columns))
-        col_names = ', '.join([f'"{col}"' for col in columns])
+        conflict_sql = f' ON CONFLICT ({unique_columns}) DO NOTHING' if unique_columns else ''
-        sql = f'INSERT INTO "{schema}"."{table}" ({col_names}) VALUES ({placeholders})'
+        base_sql = f'INSERT INTO "{schema}"."{table_name_l}" ({col_names}) VALUES ({placeholders}){conflict_sql}'
-        with conn.cursor() as cur:
+        target = [list(item.values()) for item in json_list]
-            cur.executemany(sql, values)
+        for i in range(0, len(target), batch):
+            try:
+                cur.executemany(base_sql, target[i:i+batch])
+            except Exception:
+                # fall back to row-by-row to avoid losing all rows on partial error
+                for row in target[i:i+batch]:
+                    try:
+                        cur.execute(base_sql, row)
+                    except Exception:
+                        pass
             conn.commit()
 
     def __enter__(self):
-        # ensure connection is created and return self for use in `with` blocks
         self._get_conn()
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        # always close connection on exit
         self.close()
-        # do not suppress exceptions
         return False
 
 
 if __name__ == '__main__':
     from src.config.load_config import project_config
 
-
+    host = project_config.get('db_host') or project_config.get('host')
-
+    port = project_config.get('db_port') or project_config.get('port')
-    host = project_config['db_host']
+    user = project_config.get('db_user') or project_config.get('user')
-    port = project_config['db_port']
+    password = project_config.get('db_password') or project_config.get('password')
-    user = project_config['db_user']
+    database = project_config.get('db_name') or project_config.get('database')
-    password = project_config['db_password']
-    database = project_config['db_name']
 
     pg_handler = PgHandler(
-        host=host, port=port, user=user, password=password or None, database=database,
+        host=host, port=int(port or 5432), user=user, password=password or None, database=database,
     )
 
     with pg_handler as pg:
         xxx = pg.get_all_tables()
         print(xxx)
-
-
-    #
-    # with PgHandler(host=host, port=port, user=user, password=password or None, database=database) as pg:
-    #     sql = f'select "CEC manager" from public.cec_manager where "Vendor Name" = \'Sutherland\' and "Segment" = \'Standard\''
-    #     df = pg.query(sql)
-    #     print('111')

src/TG-tools/history_message_upload.py

@@ -10,7 +10,7 @@ from pathlib import Path
 from tqdm import tqdm
 
 from src.My_Logger.project_logger import LoggerClass
-from src.PostgresSQL.databaseHandler import DatabaseHandler
+from src.PostgresSQL.db_handler import PgHandler
 from src.utils.utils import size_of_file
 from src.config.load_config import project_config
 
@@ -46,7 +46,10 @@ class MessageUploader:
         self.config = project_config
         self.load_path = local_load_path
         self.file_param_list = ['photo', 'file']
-        self.db_handler = DatabaseHandler(self.config, self.logger)
+        # 使用 PgHandler 替代旧的 DatabaseHandler，按 config 构造连接
+        self.db_handler = PgHandler(host=self.config['host'], port=int(self.config.get('port', 5432)),
+                                    user=self.config.get('user'), password=self.config.get('password'),
+                                    database=self.config.get('database'))
         self.file_db_path = local_file_db_path
 
         self.size_limit = size_limit * 1024 * 1024 * 1024
@@ -82,7 +85,8 @@ class MessageUploader:
     def prepare_database(self, table_name):
         # self.db_handler.create_database(table_name, force_create=False)
         if not self.db_handler.exists_table(table_name):
-            self.db_handler.create_table(self.table_format, table_name, False, ['message_id'])
+            # 使用 PgHandler 的 create_table_from_dict 接口创建表
+            self.db_handler.create_table_from_dict(self.table_format, table_name, force_create=False, primary_key_list=['message_id'])
             time.sleep(0.5)
             self.logger.info(f'Table {table_name} created.')
 
@@ -94,8 +98,8 @@ class MessageUploader:
 
     def upload_to_database(self, target_list, table_name):
         db_handler = self.db_handler
-        db_handler.insert_json_into_table(table_name=table_name, data_list=target_list,
+        # 调用 PgHandler.insert_json_into_table(table, json_list, schema, batch, unique_columns)
-                                          batch=10000, unique_columns='message_id')
+        db_handler.insert_json_into_table(table=table_name, json_list=target_list, schema='public', batch=10000, unique_columns='message_id')
 
     def process_loaded_message(self, loaded_messages, table_name):
         self.logger.info('Start processing messages')
@@ -209,4 +213,3 @@ if __name__ == "__main__":
     finally:
         tg.db_handler.close()
         print('done')
-