diff --git a/src/KSP_tools/EngineDataProcess_DB_upload.py b/src/KSP_tools/EngineDataProcess_DB_upload.py index a2f6922..e1573ce 100644 --- a/src/KSP_tools/EngineDataProcess_DB_upload.py +++ b/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() diff --git a/src/PostgresSQL/databaseHandler.py b/src/PostgresSQL/databaseHandler.py deleted file mode 100644 index 86fe6dd..0000000 --- a/src/PostgresSQL/databaseHandler.py +++ /dev/null @@ -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() diff --git a/src/PostgresSQL/db_handler.py b/src/PostgresSQL/db_handler.py index 55afa29..311a816 100644 --- a/src/PostgresSQL/db_handler.py +++ b/src/PostgresSQL/db_handler.py @@ -1,17 +1,6 @@ """Postgres helper for reading tables into pandas DataFrame. -Provides PgHandler class with methods: -- 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. +Provides PgHandler class with methods for reading and basic DDL/DML. """ 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") - - connect_host = self.host - connect_port = self.port - - self._conn = psycopg2.connect(host=connect_host, port=connect_port, user=self.user, + self._conn = psycopg2.connect(host=self.host, port=self.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 - CREATE/DROP DATABASE, autocommit should be set to True by the caller. + This returns a new connection instance (caller must close it). """ 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. - - 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. - """ + """Read rows filtered by condition dict or list of such dicts.""" 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 exists + return cur.fetchone() is not None 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: - """Create a table with provided columns definition SQL (e.g. 'id serial PRIMARY KEY, name text'). + # alias to match some older call sites + 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: - """ - 批量将json（字典）列表插入到指定schema下的table中。 - :param table: 表名 - :param json_list: 字典列表，每个字典为一行数据 - :param schema: schema名，默认为public - """ + 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.""" + if not data_dict_list: + return + table_name_l = table_name.lower() + 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()) - 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) + cur = conn.cursor() + 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)) + conflict_sql = f' ON CONFLICT ({unique_columns}) DO NOTHING' if unique_columns else '' + base_sql = f'INSERT INTO "{schema}"."{table_name_l}" ({col_names}) VALUES ({placeholders}){conflict_sql}' + target = [list(item.values()) for item in json_list] + 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['db_host'] - port = project_config['db_port'] - user = project_config['db_user'] - password = project_config['db_password'] - database = project_config['db_name'] + host = project_config.get('db_host') or project_config.get('host') + port = project_config.get('db_port') or project_config.get('port') + user = project_config.get('db_user') or project_config.get('user') + password = project_config.get('db_password') or project_config.get('password') + database = project_config.get('db_name') or project_config.get('database') 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') diff --git a/src/TG-tools/history_message_upload.py b/src/TG-tools/history_message_upload.py index 9454dc1..8767a65 100644 --- a/src/TG-tools/history_message_upload.py +++ b/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, - batch=10000, unique_columns='message_id') + # 调用 PgHandler.insert_json_into_table(table, json_list, schema, batch, unique_columns) + 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') -