Armor/general_tools
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

移除过时的db handler

Browse Source
This commit is contained in:
Armor
2025-12-28 18:43:30 +08:00
parent 4aa014aff0
commit 483122a497
4 changed files with 88 additions and 222 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/KSP_tools/EngineDataProcess_DB_upload.py
+6 -3
View File
@@ -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
-123
View File
@@ -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
+72 -89
View File
@@ -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]
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))
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)
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')
src/TG-tools/history_message_upload.py
+9 -6
View File
@@ -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')