216 lines
8.8 KiB
Python
216 lines
8.8 KiB
Python
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')
|