pandas确实蛮强的,记录使用的代码

from sqlalchemy import create_engine
import pandas as pd
import os

# 从环境变量中读取数据库连接信息
user = os.getenv('DB_USER', 'xxx')              # 数据库用户名
password = os.getenv('DB_PASSWORD', 'xxx')      # 数据库密码
host = os.getenv('DB_HOST', '192.168.0.1')      # 数据库主机地址
port = os.getenv('DB_PORT', '5432')             # 数据库端口
dbname = os.getenv('DB_NAME', 'xxx')            # 数据库名称

class DatabaseManager():
    def __init__(self, 
                 user, 
                 password,
                 host, 
                 port, 
                 dbname):
        # 从环境变量或默认值中读取数据库连接信息
        self.user = user
        self.password = password
        self.host = host
        self.port = port
        self.dbname = dbname
        
        # 创建数据库引擎
        self.engine = self.create_engine_()

    def create_engine_(self):
        return create_engine(
            f'postgresql://{self.user}:{self.password}@{self.host}:{self.port}/{self.dbname}'
        )
    
    def insert_data(self, df, sql_name, schema, operation='append'):
        """
        插入数据到数据库
        :param df: pandas DataFrame,包含要插入的数据
        :operation: 'append' 或 'replace',指定插入操作的行为
        """
        df.to_sql(name=sql_name, schema=schema, con=self.engine, if_exists=operation, index=False)
    
    def query_data(self, qureysql):
        """
        查询数据库中的数据
        :return: 查询结果,作为一个 pandas DataFrame 返回
        """
        return pd.read_sql_query(text(qureysql), self.engine)
    
    def update_data_fail(self, project_id, schema, sql_name):
        """
        更新数据库中的数据
        """
        with self.engine.connect() as conn:
            with conn.begin():  # 显式开始事务
                sql_text = f"UPDATE {schema}.{sql_name} SET run_status=-1 WHERE id = '{project_id}';"
                conn.execute(text(sql_text))

    def update_data_success(self, project_id, schema, sql_name):
        """
        更新数据库中的数据
        """
        with self.engine.connect() as conn:
            with conn.begin():  # 显式开始事务
                sql_text = f"UPDATE {schema}.{sql_name} SET run_status=2 WHERE id = '{project_id}';"
                conn.execute(text(sql_text))

    def dispose(self):
        """
        关闭数据库连接
        """
        self.engine.dispose()

# 主程序
if __name__ == '__main__':
    # 创建一个实例
    a = DatabaseManager(user=sql_config.user,
                        password=sql_config.password,
                        host=sql_config.host,
                        port=sql_config.port,
                        dbname=sql_config.dbname)
    
    qureysql = r"select * from aa_algo.employees where department like '%e%'"
    # qureysql = args.qureysql

    df = a.query_data(text(qureysql))
    print(df)

更多推荐