energy-data-trans/utils/db/trans_pg.py

# -*- coding: utf-8 -*-
# @Time    : 2024/5/15
# @Author  : 魏志亮

import pandas as pd
# 建立数据库连接
from sqlalchemy import create_engine, Engine, text

from utils.log.trans_log import trans_print

trans_user = 'postgres'
trans_password = 'admin123456'
trans_host = '192.168.50.235'
trans_port = 5432
trans_database = 'postgres'

trans_engine = create_engine(
    f'postgresql://{trans_user}:{trans_password}@{trans_host}:{trans_port}/{trans_database}', echo=True)


def __query(engine: Engine, sql):
    trans_print('开始执行SQL:', sql)
    with engine.connect() as conn:
        df = pd.read_sql(text(sql), conn)
    return df


def __ddl_sql(engine: Engine, sql):
    trans_print('开始执行SQL:', sql)
    with engine.connect() as conn:
        conn.execute(text(sql))
        conn.commit()


def creat_table_and_add_partition(table_name, partition):
    lower_partition = partition.lower()
    query_table = f"""
    select  c.relname
    from pg_class c
             join pg_inherits pi on pi.inhrelid = c. oid
             join pg_class c2 on c2.oid = pi.inhparent
    where
            c2.relname = '{table_name}' and c.relname = '{table_name}_{lower_partition}'
    """
    df = __query(trans_engine, query_table)

    if df.empty:
        add_partition_sql = f"""
            create table {table_name}_{partition} PARTITION OF {table_name} FOR VALUES  IN ('{partition}');
        """
        __ddl_sql(trans_engine, add_partition_sql)


if __name__ == '__main__':
    # creat_table_and_add_partition("test_11", "123_002")
    df = __query(trans_engine, """
    select  c.relname
from pg_class c
         join pg_inherits pi on pi.inhrelid = c. oid
         join pg_class c2 on c2.oid = pi.inhparent
where
        c2.relname = 'energy_data_second' and c.relname = 'energy_data_second_202405201453_a006';
    """)
    print(df)
    print(df.empty)