From 20fc3073d33572ec0151e114aab4001abb707993 Mon Sep 17 00:00:00 2001 From: Armor <2654988228@qq.com> Date: Tue, 9 Jul 2024 17:16:03 +0800 Subject: [PATCH] update --- config/config.txt | 1 + src/My_Logger/project_logger.py | 12 +- src/TG-tools/__init__.py | 0 src/TG-tools/history_message_upload.py | 199 +++++++++++++++++++++ src/TG-tools/message_upload.py | 234 +++++++++++++++++++++++++ src/TG-tools/run_TG_message_manager.py | 13 ++ src/utils/__init__.py | 0 src/utils/utils.py | 22 +++ 8 files changed, 475 insertions(+), 6 deletions(-) create mode 100644 config/config.txt create mode 100644 src/TG-tools/__init__.py create mode 100644 src/TG-tools/history_message_upload.py create mode 100644 src/TG-tools/message_upload.py create mode 100644 src/TG-tools/run_TG_message_manager.py create mode 100644 src/utils/__init__.py create mode 100644 src/utils/utils.py diff --git a/config/config.txt b/config/config.txt new file mode 100644 index 0000000..6ff7160 --- /dev/null +++ b/config/config.txt @@ -0,0 +1 @@ +{"user": "", "password": "", "host": "", "port": "", "file_db_path": ""} \ No newline at end of file diff --git a/src/My_Logger/project_logger.py b/src/My_Logger/project_logger.py index a13fcfc..34bd7db 100644 --- a/src/My_Logger/project_logger.py +++ b/src/My_Logger/project_logger.py @@ -11,10 +11,8 @@ log_path = os.path.join(project_path, 'logs') class LoggerClass: - def __init__(self, level="DEBUG"): + def __init__(self): # 创建日志器对象 - self.logger = logging.getLogger(__name__) - self.logger.setLevel(level) self.format = logging.Formatter(f'[%(filename)s: %(funcName)s:%(lineno)4d][%(asctime)s][%(levelname)-.5s]: ' f'%(message)s', datefmt='%Y-%m-%d %H:%M:%S') @@ -46,13 +44,15 @@ class LoggerClass: return file_handler def get_log(self, level, log_file_path=log_path): + logger = logging.getLogger() + logger.setLevel(level) # 日志器中添加控制台处理器 - self.logger.addHandler(self.console_handler(level)) + logger.addHandler(self.console_handler(level)) # 日志器中添加文件处理器 - self.logger.addHandler(self.file_handler(level, log_file_path)) + logger.addHandler(self.file_handler(level, log_file_path)) # 返回日志实例对象 - return self.logger + return logger if __name__ == '__main__': diff --git a/src/TG-tools/__init__.py b/src/TG-tools/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/TG-tools/history_message_upload.py b/src/TG-tools/history_message_upload.py new file mode 100644 index 0000000..2653638 --- /dev/null +++ b/src/TG-tools/history_message_upload.py @@ -0,0 +1,199 @@ +import os +import time +import datetime +import json +import shutil +import pandas as pd +import logging + +from src.My_Logger.project_logger import LoggerClass +from src.PostgresSQL.databaseHandler import DatabaseHandler +from src.utils.utils import load_config, size_of_file + + +def process_message_text(text): + if isinstance(text, str): + if '[Sticker ' in text and ']' in text: + return 'Sticker' + return text + elif isinstance(text, list): + res = '' + for item in text: + if isinstance(item, str): + res += item + '\n' + elif isinstance(item, dict): + res += f"{item.get('text')}\n" + elif not item: + continue + else: + raise TypeError + return res + else: + raise TypeError + + +class MessageUploader: + def __init__(self, config_path, local_logger, local_load_path, local_file_db_path, size_limit=0.5): + self.logger = local_logger + self.config = load_config(config_path) + self.load_path = local_load_path + self.file_param_list = ['photo', 'file'] + self.db_handler = DatabaseHandler(self.config, self.logger) + self.file_db_path = local_file_db_path + + self.size_limit = size_limit * 1024 * 1024 * 1024 + self.single_file_size_limit = 0.24 * 1024 * 1024 * 1024 # 240MB + + self.table_format = { + 'message_id': 'text', + 'message_type': 'text', + 'message_datetime': 'text', + 'message_unix': 'text', + 'message_from': 'text', + 'message_from_id': 'text', + 'message_text': 'text', + 'forward_from': 'text', + 'reply_message_id': 'text', + 'file_location': 'text', + 'file_type': 'text', + 'file_size': 'text', + 'file_size_raw': 'numeric', + } + + def do_upload(self): + start_time = time.time() + loaded_messages = self.load_message() + table_name = f'{loaded_messages["name"]}_{loaded_messages["type"]}' + # table_name = 'Rynco_Group_private_supergroup' + self.prepare_database(table_name) + target_list = self.process_loaded_message(loaded_messages, table_name) + self.upload_to_database(target_list, table_name) + self.logger.info(f'Message upload completed in {round(time.time() - start_time, 3)}s') + + 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']) + time.sleep(0.5) + self.logger.info(f'Table {table_name} created.') + + def load_message(self): + with open(os.path.join(self.load_path, 'result.json'), 'r', encoding='utf-8') as file: + file_content = file.read() + self.logger.info('Message loaded') + return json.loads(file_content) + + 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') + + def process_loaded_message(self, loaded_messages, table_name): + self.logger.info('Start processing messages') + start_time = time.time() + conversation_id = loaded_messages['id'] + messages = [item for item in loaded_messages['messages']] + target_list = [] + for message in messages: + if message['type'] == 'service': + continue + target = { + 'message_id': f"{conversation_id}_{message['id']}", + 'message_type': message['type'], + 'message_datetime': message['date'], + 'message_unix': message['date_unixtime'], + 'message_from': message['from'], + 'message_from_id': message['from_id'], + 'message_text': process_message_text(message['text']), + 'forward_from': message.get('forwarded_from', None), + 'reply_message_id': message.get('reply_to_message_id', None), + 'file_location': None, + 'file_type': None, + 'file_size': None, + 'file_size_raw': None + } + for file_param in self.file_param_list: + if file_param in message: + file_location = self.transfer_file_to_db(message[file_param]) + target['file_location'] = file_location + target['file_size_raw'] = os.path.getsize(os.path.join(self.load_path, message[file_param])) + target['file_size'] = size_of_file(target['file_size_raw']) + if file_param == 'photo': + target['file_type'] = file_param + elif 'sticker_emoji' in message: + target['file_type'] = 'sticker' + elif 'media_type' in message: + target['file_type'] = message['media_type'] + else: + target['file_type'] = 'file' + + target_list.append(target) + self.logger.info(f'Data process complete, time: {time.time() - start_time}') + return target_list + + def transfer_file_to_db(self, file_path): + source_path = os.path.join(self.load_path, file_path) + destination_path = os.path.join(self.file_db_path, file_path) + destination_folder, file_name = os.path.split(destination_path) + new_file_name = file_name + # 路径不存在则创建路径 + if not os.path.exists(destination_folder): + os.mkdir(destination_folder) + + # 已存在相同文件名的情况 + if os.path.exists(destination_path): + # 如果文件大小不相等，或者大小相同但不是sticker文件，直接进行复制操作。 + if (os.path.getsize(source_path) != os.path.getsize(destination_path) or + ('sticker' not in source_path) and + os.path.getsize(source_path) == os.path.getsize(destination_path)): + i = 0 + while os.path.exists(destination_path): + new_file_name = f"{i}_{file_name}" + destination_path = os.path.join(destination_folder, new_file_name) + i += 1 + # 如果是sticker文件，则直接进行映射（sticker文件出错问题应该不大） + + shutil.copy2(source_path, destination_path) + 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() + msg_list = [] + + csv_content = pd.read_csv('../data/Rynco_Group.csv', encoding='utf-8') + for row, content in csv_content.iterrows(): + try: + text = process_message_text(content['text']) + except TypeError as e: + text = '' + target = { + 'message_id': content['id'], + 'message_type': '' + 'message', + 'message_datetime': content['date'], + 'message_unix': int(datetime.datetime.strptime(content['date'], '%Y-%m-%dT%H:%M:%S').timestamp()), + 'message_from': content['from_name'], + 'message_from_id': content['from_id'], + 'message_text': text, + 'forward_from': content.get('forwarded_from'), + 'reply_message_id': content.get('reply_to_message_id', None), + 'file_location': None, + 'file_type': 'sticker' if 'Sticker' == text else None, + 'file_size': None, + 'file_size_raw': None + } + msg_list.append(target) + tg.upload_to_database(msg_list, 'Rynco_Group_private_supergroup') + tg.db_handler.close() + print('done') + diff --git a/src/TG-tools/message_upload.py b/src/TG-tools/message_upload.py new file mode 100644 index 0000000..3f40c7b --- /dev/null +++ b/src/TG-tools/message_upload.py @@ -0,0 +1,234 @@ +import os +import time +import json +import shutil +import logging + +from src.My_Logger.project_logger import LoggerClass +from src.PostgresSQL.databaseHandler import DatabaseHandler +from src.utils.utils import load_config, size_of_file + + +def process_message_text(text): + if isinstance(text, str): + return text + elif isinstance(text, list): + res = '' + for item in text: + if isinstance(item, str): + res += item + '\n' + elif isinstance(item, dict): + res += f"{item.get('text')}\n" + elif not item: + continue + else: + raise TypeError + return res + else: + raise TypeError + + +class MessageUploader: + def __init__(self, config_path, local_logger, local_load_path, local_file_db_path, size_limit=0.5): + self.logger = local_logger + self.config = load_config(config_path) + self.load_path = local_load_path + self.file_param_list = ['photo', 'file'] + self.db_handler = DatabaseHandler(self.config, self.logger) + self.file_db_path = local_file_db_path + + self.size_limit = size_limit * 1024 * 1024 * 1024 + self.single_file_size_limit = 0.24 * 1024 * 1024 * 1024 # 240MB + + self.table_format = { + 'message_id': 'text', + 'message_type': 'text', + 'message_datetime': 'text', + 'message_unix': 'text', + 'message_from': 'text', + 'message_from_id': 'text', + 'message_text': 'text', + 'forward_from': 'text', + 'reply_message_id': 'text', + 'file_location': 'text', + 'file_type': 'text', + 'file_size': 'text', + 'file_size_raw': 'numeric', + } + + def do_upload(self): + start_time = time.time() + loaded_messages = self.load_message() + table_name = f'{loaded_messages["name"]}_{loaded_messages["type"]}' + if 'wss://rynco.me' in table_name: + table_name = 'Rynco_Group_private_supergroup' + self.prepare_database(table_name) + target_list = self.process_loaded_message(loaded_messages, table_name) + self.upload_to_database(target_list, table_name) + self.logger.info(f'Message upload completed in {round(time.time() - start_time, 3)}s') + + 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']) + time.sleep(0.5) + self.logger.info(f'Table {table_name} created.') + + def load_message(self): + with open(os.path.join(self.load_path, 'result.json'), 'r', encoding='utf-8') as file: + file_content = file.read() + self.logger.info('Message loaded') + return json.loads(file_content) + + 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') + + def process_loaded_message(self, loaded_messages, table_name): + self.logger.info('Start processing messages') + start_time = time.time() + conversation_id = loaded_messages['id'] + messages = [item for item in loaded_messages['messages']] + target_list = [] + for message in messages: + if message['type'] == 'service': + continue + target = { + 'message_id': message['id'], + 'message_type': message['type'], + 'message_datetime': message['date'], + 'message_unix': message['date_unixtime'], + 'message_from': message['from'], + 'message_from_id': message['from_id'], + 'message_text': process_message_text(message['text']), + 'forward_from': message.get('forwarded_from', None), + 'reply_message_id': message.get('reply_to_message_id', None), + 'file_location': None, + 'file_type': None, + 'file_size': None, + 'file_size_raw': None + } + for file_param in self.file_param_list: + if file_param in message: + file_location = self.transfer_file_to_db(message[file_param]) + target['file_location'] = file_location + target['file_size_raw'] = os.path.getsize(os.path.join(self.load_path, message[file_param])) + target['file_size'] = size_of_file(target['file_size_raw']) + if file_param == 'photo': + target['file_type'] = file_param + elif 'sticker_emoji' in message: + target['file_type'] = 'sticker' + elif 'media_type' in message: + target['file_type'] = message['media_type'] + else: + target['file_type'] = 'file' + + target_list.append(target) + self.logger.info(f'Data process complete, time: {time.time() - start_time}') + return target_list + + def transfer_file_to_db(self, file_path): + source_path = os.path.join(self.load_path, file_path) + destination_path = os.path.join(self.file_db_path, file_path) + destination_folder, file_name = os.path.split(destination_path) + new_file_name = file_name + # 路径不存在则创建路径 + if not os.path.exists(destination_folder): + os.mkdir(destination_folder) + + # 已存在相同文件名的情况 + if os.path.exists(destination_path): + # 如果文件大小不相等，或者大小相同但不是sticker文件，直接进行复制操作。 + if (os.path.getsize(source_path) != os.path.getsize(destination_path) or + ('sticker' not in source_path) and + os.path.getsize(source_path) == os.path.getsize(destination_path)): + i = 0 + while os.path.exists(destination_path): + new_file_name = f"{i}_{file_name}" + destination_path = os.path.join(destination_folder, new_file_name) + i += 1 + # 如果是sticker文件，则直接进行映射（sticker文件出错问题应该不大） + + shutil.copy2(source_path, destination_path) + 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) + + try: + tg.do_upload() + + finally: + tg.db_handler.close() + + + # load_path = '../data/ChatExport_2023-12-27' + # with open(os.path.join(load_path, 'result.json'), 'r', encoding='utf-8') as file: + # content = file.read() + # + # target_content = json.loads(content) + # conversation_id = target_content['id'] + # conversation_name = target_content['name'] + # conversation_type = target_content['type'] + # messages = [item for item in target_content['messages']] + # + # first_msg, last_msg = messages[0], messages[-1] + # first_datetime = datetime.datetime.strptime(first_msg['date'], '%Y-%m-%dT%H:%M:%S') + # last_datetime = datetime.datetime.strptime(last_msg['date'], '%Y-%m-%dT%H:%M:%S') + # + # time_interval = 180 + # time_list = [] + # start_time = first_datetime + # while start_time <= last_datetime: + # time_list.append(start_time) + # start_time += datetime.timedelta(days=time_interval) + # + # file_param_list = ['photo', 'file'] + # target_list = [] + # for message in messages: + # target = { + # 'message_id': f"{conversation_id}_{message['id']}", + # 'message_type': message['type'], + # 'message_datetime': message['date'], + # 'message_unix': message['date_unixtime'], + # 'message_from': message['from'], + # 'message_from_id': message['from_id'], + # 'message_text': process_message_text(message['text']), + # 'forward_from': message.get('forwarded_from', None), + # 'reply_message_id': message.get('reply_to_message_id', None), + # 'file_data': None, + # 'file_type': None, + # 'file_size': None, + # 'file_size_raw': None + # } + # for file_param in file_param_list: + # if file_param in message: + # target_file_path = message[file_param] + # with open(os.path.join(load_path, target_file_path), 'rb') as f: + # file_raw_data = f.read() + # file_data = Binary(file_raw_data) + # target['file_data'] = file_data + # target['file_size_raw'] = len(file_raw_data) + # target['file_size'] = size_of_file(len(file_raw_data)) + # if file_param == 'photo': + # target['file_type'] = file_param + # elif 'sticker_emoji' in message: + # target['file_type'] = 'sticker' + # target_list.append(target) + # + # db_logger = logger.Logger('../logs').get_log() + # db_handler = DatabaseHandler(config, db_logger) + # db_handler.insert_json_into_table(table_name='sample_table', data_list=target_list, + # batch=1000, unique_columns='message_id') + # + # db_handler.close() + print('done') + diff --git a/src/TG-tools/run_TG_message_manager.py b/src/TG-tools/run_TG_message_manager.py new file mode 100644 index 0000000..c5178de --- /dev/null +++ b/src/TG-tools/run_TG_message_manager.py @@ -0,0 +1,13 @@ +import os +import time +import datetime +import json + + +from src.My_Logger.project_logger import LoggerClass +from src.PostgresSQL.databaseHandler import DatabaseHandler + +if __name__ == "__main__": + + print('done') + diff --git a/src/utils/__init__.py b/src/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/utils/utils.py b/src/utils/utils.py new file mode 100644 index 0000000..530c850 --- /dev/null +++ b/src/utils/utils.py @@ -0,0 +1,22 @@ +import os +import json + + +def load_config(config_path): + 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 + while file_size > 1024.0 and i < len(units) - 1: + file_size /= 1024.0 + i += 1 + return f'{round(file_size, 2)} {units[i]}'