From 4aa014aff0edb98dea47098a928aaa2bbf6ebe94 Mon Sep 17 00:00:00 2001 From: Armor <2654988228@qq.com> Date: Sun, 28 Dec 2025 16:59:29 +0800 Subject: [PATCH] =?UTF-8?q?=E9=A1=B9=E7=9B=AE=E9=87=8D=E6=9E=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- config/config.json | 22 +++ config/config.txt | 1 - data/note.txt | 66 -------- src/PostgresSQL/db_handler.py | 144 +++++++++++++++++- src/PostgresSQL/pg_db_handler.py | 203 +++++++++++++++++++++++++ src/TG-tools/history_message_upload.py | 21 ++- src/config/load_config.py | 4 - src/utils/utils.py | 16 -- 8 files changed, 378 insertions(+), 99 deletions(-) delete mode 100644 config/config.txt delete mode 100644 data/note.txt create mode 100644 src/PostgresSQL/pg_db_handler.py diff --git a/config/config.json b/config/config.json index e69de29..48b9291 100644 --- a/config/config.json +++ b/config/config.json @@ -0,0 +1,22 @@ +{ + "user": "", + "password": "", + "host": "", + "port": "", + "file_db_path": "", + + "db_host": "", + "db_port": "", + "db_user": "", + "db_password": "", + "db_name": "", + + "bailian_API_key": "", + "bailian_MODEL_NAME": "", + "bailian_API_URL": "", + + "xiaomi_API_KEY": "", + "xiaomi_MODEL_NAME": "", + "xiaomi_API_URL": "" + +} diff --git a/config/config.txt b/config/config.txt deleted file mode 100644 index 6ff7160..0000000 --- a/config/config.txt +++ /dev/null @@ -1 +0,0 @@ -{"user": "", "password": "", "host": "", "port": "", "file_db_path": ""} \ No newline at end of file diff --git a/data/note.txt b/data/note.txt deleted file mode 100644 index 64e0c53..0000000 --- a/data/note.txt +++ /dev/null @@ -1,66 +0,0 @@ -Soooo, lets think about it. -在银河系，一个离银心不近也不远的地方，诞生了希尔文明。 -他们蹒跚起步，征服了陆地，迈向了海洋。他们很幸运， - -一个路过的黑洞扯碎了希尔文明的母星系，但是希尔文明没有心灰气馁，他们抢在灾难降临之前，倾尽全部的智慧，把文明的种子播撒了出去。 -现在，距离希尔舰队启航，已经过去了15个千年。新生的希尔人已经习惯了在群星中的生活。他们出生在星舰，在星河间穿行，以探索为荣耀。 -舰队偶尔会在某个星系落脚，补充星际旅行所必需的燃料和物质。但他们不会长久的停留。 -为了保证舰队的出现不会干涉到孕育中的文明种子，在舰队泊入一个恒星系之前，需要派出探险队确认目标星系没有生命诞生。 -这是一个重要而有趣的工作，在年轻人中很受欢迎，但只有最优秀的探星者才能参与这项充满荣耀的工作，驾驶着【星火】级星系探索飞船，在舰载AI的帮助下对恒星系展开细致而全面的探索。 -以及为后续舰队补充资源建立前哨基地。 -这是严垣从小就梦寐以求的工作。能够离开【此处需要插入某个形容狭小的形容词】的星舰，前往未知的世界探索，这样的生活实在是令人心潮澎湃。 -在经过了10余年的勤学苦读，12个月的艰苦训练后，严垣如愿以偿，成为了一名光荣的探星者。 - - - - - - - - - - - - -We need to ensure a 1600 arrival in the destination, leaving us with around 30 minutes of redundancy. -The current estimation of travel time is around 90 minutes. -Which means a 1400 departure is optimal. - -We need to check the traffic status 10 minutes before departure. Anything longer than 120 minutes will forbid the driving plan. -In which case we will transition into public traffic. - -We will need to pack the consumables, including water and necessary food. Which I intend to be prepared before lunchtime. -We also need to pack 2 set of personal terminals, power bank, Etc - -1200 Consumables preparation complete -1350 Final travel plan decision point. - In the event of the selection of public transportation, the plan is bike-9-1-5-bus. With a 1400 departure we should catch the 1540 bus. -1355 Donning the Doctor's gear. -1400 Departure. - - - - - - - - - - - - -I am listening to the OST of Lone-trail while I type this down. -Kristen. -I'm not going to deny, that I like her. -I like the idea of exploration. I like the idea of pushing the boundary. And she managed to do just that. -Comparing to her colleague, she is just a mortal. She is not as indestructible as Seria. She does not have the noble blood like Muelsyse. -But her none-stopping spirit of advance is the purest of them all. And under such spirit, Rhine Lab was found and united. -Her sins cannot be neglected, that I admit. But I can not resist my desire to salute her. -I salute to a pioneer, a scientist, a visionary, who gave all, risk all to challenge the unknown. -And now, even sky is not the limit. -Mortals are weak, but in that weakness, greatness is born. I find this phrase very suitable for her. -Ad astra, to the stars. Kristen, good night. - -And about Lone-trail, there's more I enjoy. -Kristen is a great character that I like so much, but the storyline did not just focus on how great her is. -It did not just shape a great character and call her sins necessary evil. \ No newline at end of file diff --git a/src/PostgresSQL/db_handler.py b/src/PostgresSQL/db_handler.py index 80eb9b8..55afa29 100644 --- a/src/PostgresSQL/db_handler.py +++ b/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 diff --git a/src/PostgresSQL/pg_db_handler.py b/src/PostgresSQL/pg_db_handler.py new file mode 100644 index 0000000..1861234 --- /dev/null +++ b/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') diff --git a/src/TG-tools/history_message_upload.py b/src/TG-tools/history_message_upload.py index 90c3b78..9454dc1 100644 --- a/src/TG-tools/history_message_upload.py +++ b/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() diff --git a/src/config/load_config.py b/src/config/load_config.py index 9eee28d..bdb1406 100644 --- a/src/config/load_config.py +++ b/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) diff --git a/src/utils/utils.py b/src/utils/utils.py index 3f45ee5..0e6576b 100644 --- a/src/utils/utils.py +++ b/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