update
This commit is contained in:
@@ -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'],
|
||||
|
||||
@@ -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')
|
||||
|
||||
|
||||
@@ -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')
|
||||
|
||||
@@ -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')
|
||||
Reference in New Issue
Block a user