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')