import os import time import datetime import json import shutil import pandas as pd import logging from pathlib import Path from tqdm import tqdm from src.My_Logger.project_logger import LoggerClass from src.PostgresSQL.db_handler import PgHandler 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') config_path = os.path.join(project_path, 'config') 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, local_logger, local_load_path, local_file_db_path, size_limit=0.5): self.logger = local_logger self.config = project_config self.load_path = local_load_path self.file_param_list = ['photo', 'file'] # 使用 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 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 '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): # 使用 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.') 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 # 调用 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') start_time = time.time() conversation_id = loaded_messages['id'] messages = [item for item in loaded_messages['messages']] target_list = [] for message in tqdm(messages, desc=f'Processing 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) def logger_test(self): self.logger.info('this is a test') 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: 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') if __name__ == "__main__": 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() finally: tg.db_handler.close() print('done')