import psycopg2 import json class DatabaseHandler: def __init__(self, config): self.config = config self.connection = psycopg2.connect(database=config['db_database'], user=config['db_user'], password=config['db_password'], host=config['db_host'], port=config['db_port']) self.cursor = self.connection.cursor() self.cursor.execute(f"select tablename from pg_tables where schemaname='public'") self.table_list = [item[0] for item in self.cursor.fetchall()] def ensure_connection(self): try: self.cursor.execute(f"select tablename from pg_tables where schemaname='public'") except Exception as e: print(e) config = self.config self.connection = psycopg2.connect(database=config['db_database'], user=config['db_user'], password=config['db_password'], host=config['db_host'], port=config['db_port']) self.cursor = self.connection.cursor() def create_database(self, db_name, force_create=False): self.ensure_connection() if force_create: self.drop_database(db_name) try: self.connection.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT) self.cursor.execute(f'CREATE DATABASE {db_name}') print(f'Database {db_name} successfully created.') except psycopg2.errors.DuplicateDatabase as e: print(f'Database {db_name} already exists: {e}') def drop_database(self, db_name): self.ensure_connection() self.connection.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT) self.cursor.execute(f'drop database if exists {db_name}') print(f'Database {db_name} successfully dropped.') self.table_list.remove(db_name) def create_table(self, table_structure, table_name, force_create=False, primary_key_list=None): self.ensure_connection() 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(): 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() if table_name in self.table_list: return True else: return False def drop_table(self, table_name, info_flag=False): self.ensure_connection() table_name = table_name.lower() if not self.exists_table(table_name): if info_flag: print(f'table: {table_name} not exists in database.') return self.cursor.execute(f'drop table {table_name}') self.connection.commit() if info_flag: print(f'table: {table_name} dropped') return def insert_into_table(self, table_name, data_dict_list, batch=1000, schema='public'): """ :param schema: :param batch: :param table_name: :param data_dict_list: :return: """ self.ensure_connection() full_table_name = f"{schema}.{table_name.lower()}" # 检查表是否存在时也要加上schema self.cursor.execute( "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_schema=%s AND table_name=%s)", (schema, table_name.lower()) ) if not self.cursor.fetchone()[0]: print(f'table: {full_table_name} not exists in database.') return table_structure = data_dict_list[0] content_list = [f'{key.replace(" ", "_").replace("-", "_")}'.lower() for key in table_structure.keys()] sql = f"INSERT INTO {full_table_name}({','.join(content_list)}) VALUES({','.join(['%s'] * len(content_list))})" target = [list(item.values()) for item in data_dict_list] for i in range(0, len(target), batch): self.cursor.executemany(sql, target[i: i + batch]) self.connection.commit() print(f'Inserted into table {full_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: """ self.ensure_connection() table_name = table_name.lower() if not self.exists_table(table_name): print(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') if unique_columns: sql = (f"INSERT INTO {table_name}({','.join(table_structure)}) " f"VALUES({','.join(['%s'] * len(table_structure))})" f" ON CONFLICT ({unique_columns}) DO NOTHING") else: sql = (f"INSERT INTO {table_name}({','.join(table_structure)}) " f"VALUES({','.join(['%s'] * len(table_structure))})") 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() print(f'Inserted into table {table_name}, rows: {min(i+batch, len(target))}/{len(target)}') def get_data(self, table_name, field_name_list=None): self.ensure_connection() if not field_name_list: sql = f"SELECT * FROM {table_name}" elif isinstance(field_name_list, str): sql = f"SELECT {field_name_list} FROM {table_name}" else: sql = f"select {','.join(field_name_list)} from {table_name}" self.cursor.execute(sql) res = self.cursor.fetchall() return res def update_table(self, table_name, update_data: dict, key_column): self.ensure_connection() table_name = table_name.lower() if not self.exists_table(table_name): print(f'table: {table_name} not exists in database.') return if key_column not in update_data: raise ValueError(f'{key_column} not in data') target_list = [] for key, value in update_data.items(): if key == key_column: continue if isinstance(value, dict): target_list.append(f"""{key} = '{json.dumps(value).replace("'", "''")}'""") else: target_list.append(f"""{key} = '{str(value).replace("'", "''")}'""") target = ','.join(target_list) sql = f"update {table_name} set {target} where {key_column}='{update_data[key_column]}'" self.cursor.execute(sql) self.connection.commit() def close(self): self.connection.close() if __name__ == "__main__": config = { "db_user": "app", "db_password": "Letmein1", "db_host": "pgm-8vb3061igsxn50is9o.pgsql.zhangbei.rds.aliyuncs.com", "db_port": "5432", "db_database": "flex_desk", } # db_handler = DatabaseHandler(config) # print(db_handler.get_data("wiki.feishu_doc_public")) # db_handler.insert_into_table('feishu_doc_public', [table_schema], batch=1, schema='wiki')