This commit is contained in:
2024-07-09 17:16:03 +08:00
parent 2bbfadb420
commit 20fc3073d3
8 changed files with 475 additions and 6 deletions
View File
+199
View File
@@ -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')
+234
View File
@@ -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')
+13
View File
@@ -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')