项目重构

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": ""
+
+}

config/config.txt

@@ -1 +0,0 @@
-{"user": "", "password": "", "host": "", "port": "", "file_db_path": ""}

data/note.txt

@@ -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.

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,
+    tg = MessageUploader(
         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)
+        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