204 lines
8.4 KiB
Python
204 lines
8.4 KiB
Python
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')
|