From 4fd3fe49bc642c77d9aa8209c55b68d168697fdb Mon Sep 17 00:00:00 2001 From: Armor00 <2654988228@qq.com> Date: Wed, 10 Jul 2024 10:19:50 +0800 Subject: [PATCH] TG tools update --- src/PostgresSQL/databaseHandler.py | 47 ++++++++++++++++++++++++-- src/TG-tools/history_message_upload.py | 31 ++++++++++------- src/utils/utils.py | 5 +++ 3 files changed, 68 insertions(+), 15 deletions(-) diff --git a/src/PostgresSQL/databaseHandler.py b/src/PostgresSQL/databaseHandler.py index 900da26..0d9475d 100644 --- a/src/PostgresSQL/databaseHandler.py +++ b/src/PostgresSQL/databaseHandler.py @@ -16,18 +16,22 @@ class DatabaseHandler: 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): + 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(): - content_list.append(f'{key.replace(" ", "_").replace("-", "_")} {values}') + 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() @@ -74,5 +78,44 @@ class DatabaseHandler: 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/TG-tools/history_message_upload.py b/src/TG-tools/history_message_upload.py index 2653638..5ea16dd 100644 --- a/src/TG-tools/history_message_upload.py +++ b/src/TG-tools/history_message_upload.py @@ -6,10 +6,16 @@ import shutil import pandas as pd import logging +from pathlib import Path + from src.My_Logger.project_logger import LoggerClass from src.PostgresSQL.databaseHandler import DatabaseHandler from src.utils.utils import load_config, size_of_file +project_path = Path(__file__).parent.parent.parent.absolute() +data_path = os.path.join(project_path, 'data') +config_path = os.path.join(project_path, 'config') + def process_message_text(text): if isinstance(text, str): @@ -157,20 +163,8 @@ class MessageUploader: return file_path.replace(file_name, new_file_name) -if __name__ == "__main__": - config = load_config('../config') - config['database'] = "TestDB" - tg = MessageUploader(config_path='../config', - local_logger=LoggerClass().get_log(logging.DEBUG), - local_load_path='E:\\SNSFiles\\ChatExport_2023-12-27', - local_file_db_path=config['file_db_path'], - size_limit=0.25) - - - # with open('../data/Rynco_Group.txt', 'r', encoding='utf-8') as f: - # file_content = f.readlines() +def upload_didi_saved_message_to_db(): msg_list = [] - csv_content = pd.read_csv('../data/Rynco_Group.csv', encoding='utf-8') for row, content in csv_content.iterrows(): try: @@ -194,6 +188,17 @@ if __name__ == "__main__": } msg_list.append(target) tg.upload_to_database(msg_list, 'Rynco_Group_private_supergroup') + + +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_2023-12-27', + local_file_db_path=config['file_db_path'], + size_limit=0.25) + tg.do_upload() tg.db_handler.close() print('done') diff --git a/src/utils/utils.py b/src/utils/utils.py index 530c850..3f45ee5 100644 --- a/src/utils/utils.py +++ b/src/utils/utils.py @@ -3,6 +3,11 @@ 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: