项目重构
This commit is contained in:
@@ -0,0 +1,203 @@
|
||||
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')
|
||||
Reference in New Issue
Block a user