TG tools update

src/PostgresSQL/databaseHandler.py

@@ -16,18 +16,22 @@ class DatabaseHandler:
         self.cursor.execute(f"select tablename from pg_tables where schemaname='public'")
         self.table_list = [item[0] for item in self.cursor.fetchall()]
 
-    def create_table(self, table_structure, table_name, force_create=False):
+    def create_table(self, table_structure, table_name, force_create=False, primary_key_list=None):
         table_name = table_name.lower()
         if force_create and self.exists_table(table_name):
             self.drop_table(table_name)
         sql = f'Create Table {table_name}('
         content_list = []
         for key, values in table_structure.items():
-            content_list.append(f'{key.replace(" ", "_").replace("-", "_")}           {values}')
+            key_string = f'{key.replace(" ", "_").replace("-", "_")}           {values}'
+            if primary_key_list and key in primary_key_list:
+                key_string += ' PRIMARY KEY'
+            content_list.append(key_string)
         sql += ',\n'.join(content_list)
         sql += ');'
         self.connection.cursor().execute(sql)
         self.connection.commit()
+        self.table_list.append(table_name)
 
     def exists_table(self, table_name):
         table_name = table_name.lower()
@@ -74,5 +78,44 @@ class DatabaseHandler:
             self.connection.commit()
             self.logger.info(f'Inserted into table {table_name}, rows: {min(i+batch, len(target))}/{len(target)}')
 
+    def insert_json_into_table(self, table_name, data_list, batch=1000, unique_columns=None):
+        """
+        insert a list of json into table. The json content should contain all the required columns as its keys.
+        :param batch: batch size
+        :param table_name:  which table to insert into
+        :param data_list:   target json
+        :param unique_columns:  the unique column of the table. If set, that column will be
+        checked to ensure no duplicate entries
+        :return:
+        """
+        table_name = table_name.lower()
+        if not self.exists_table(table_name):
+            self.logger.error(f'table: {table_name} not exists in database.')
+            return
+
+        table_structure = list(data_list[0].keys())
+        if unique_columns and unique_columns not in table_structure:
+            raise ValueError(f'{unique_columns} not in data')
+        sql = (f"INSERT INTO {table_name}({','.join(table_structure)}) VALUES({','.join(['%s'] * len(table_structure))})"
+               f" ON CONFLICT ({unique_columns}) DO NOTHING")
+        target = [list(item.values()) for item in data_list]
+
+        for i in range(0, len(target), batch):
+            try:
+                self.cursor.executemany(sql, target[i: i+batch])
+            except psycopg2.errors.InternalError_ as e:
+                print(e)
+                size = 0
+                for line in target[i:i+batch]:
+                    if line[-1]:
+                        size += line[-1]
+                print(size / 1024 / 1024)
+                print('error')
+            except Exception as e:
+                print(sql)
+                print(e)
+            self.connection.commit()
+            self.logger.info(f'Inserted into table {table_name}, rows: {min(i+batch, len(target))}/{len(target)}')
+
     def close(self):
         self.connection.close()

src/TG-tools/history_message_upload.py

@@ -6,10 +6,16 @@ import shutil
 import pandas as pd
 import logging
 
+from pathlib import Path
+
 from src.My_Logger.project_logger import LoggerClass
 from src.PostgresSQL.databaseHandler import DatabaseHandler
 from src.utils.utils import load_config, size_of_file
 
+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):
@@ -157,20 +163,8 @@ class MessageUploader:
         return file_path.replace(file_name, new_file_name)
 
 
-if __name__ == "__main__":
+def upload_didi_saved_message_to_db():
-    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:
@@ -194,6 +188,17 @@ if __name__ == "__main__":
         }
         msg_list.append(target)
     tg.upload_to_database(msg_list, 'Rynco_Group_private_supergroup')
+
+
+if __name__ == "__main__":
+    config = load_config(config_path)
+    config['database'] = "TestDB"
+    tg = MessageUploader(config_path=config_path,
+                         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)
+    tg.do_upload()
     tg.db_handler.close()
     print('done')
 

src/utils/utils.py

@@ -3,6 +3,11 @@ import json
 
 
 def load_config(config_path):
+    """
+
+    :param config_path: config folder path
+    :return:
+    """
     if os.path.exists(os.path.join(config_path, 'config_local.txt')):
         config_path = os.path.join(config_path, 'config_local.txt')
     else: