diff --git a/src/PostgresSQL/databaseHandler.py b/src/PostgresSQL/databaseHandler.py index 0d9475d..7732a95 100644 --- a/src/PostgresSQL/databaseHandler.py +++ b/src/PostgresSQL/databaseHandler.py @@ -5,9 +5,11 @@ from src.My_Logger.project_logger import LoggerClass class DatabaseHandler: def __init__(self, config, logger=None): - if not logger: + if not logger or not isinstance(logger, LoggerClass): log = LoggerClass() self.logger = log.get_log(logging.DEBUG) + else: + self.logger = logger self.config = config # {database="postgres", user="postgres", password="123456", host="localhost", port="5432"} self.connection = psycopg2.connect(database=config['database'], user=config['user'], diff --git a/src/TG-tools/history_message_upload.py b/src/TG-tools/history_message_upload.py index 5ea16dd..e907421 100644 --- a/src/TG-tools/history_message_upload.py +++ b/src/TG-tools/history_message_upload.py @@ -7,6 +7,7 @@ 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.databaseHandler import DatabaseHandler @@ -70,7 +71,8 @@ class MessageUploader: start_time = time.time() loaded_messages = self.load_message() table_name = f'{loaded_messages["name"]}_{loaded_messages["type"]}' - # table_name = 'Rynco_Group_private_supergroup' + 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) @@ -100,7 +102,7 @@ class MessageUploader: conversation_id = loaded_messages['id'] messages = [item for item in loaded_messages['messages']] target_list = [] - for message in messages: + for message in tqdm(messages, desc=f'Processing messages'): if message['type'] == 'service': continue target = { @@ -191,6 +193,7 @@ def upload_didi_saved_message_to_db(): if __name__ == "__main__": + config = load_config(config_path) config['database'] = "TestDB" tg = MessageUploader(config_path=config_path, @@ -198,7 +201,9 @@ if __name__ == "__main__": 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') + try: + tg.do_upload() + finally: + tg.db_handler.close() + print('done') diff --git a/src/TG-tools/message_upload.py b/src/TG-tools/message_upload.py deleted file mode 100644 index 3f40c7b..0000000 --- a/src/TG-tools/message_upload.py +++ /dev/null @@ -1,234 +0,0 @@ -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/tools/filename_length_scan.py b/src/tools/filename_length_scan.py new file mode 100644 index 0000000..922a6a4 --- /dev/null +++ b/src/tools/filename_length_scan.py @@ -0,0 +1,24 @@ +import os + + +def list_all_files(path): + file_list = [] + for item in os.walk(path): + for file in item[2]: + file_list.append(item[0] + file) + return file_list + + +def scan_path_length(path, threshold=180): + file_path_list = list_all_files(path) + result_list = [] + for file in file_path_list: + if len(file) >= threshold: + result_list.append(file) + return result_list + + +if __name__ == '__main__': + path = r'H:\Mass Storage Drive\classified documents' + target = scan_path_length(path) + print('done')