PostgreSQL

类

import psycopg2
from psycopg2.extras import execute_values
from typing import List

def save_data_to_pg(data_list: List[dict], table_name: str, dbname: str, user: str, password: str, host: str, port: int) -> None:
    """
    将字典列表数据存储到 PostgreSQL 数据库中

    Args:
        data_list (List[dict]): 包含字典数据的列表
        table_name (str): 要插入数据的表名
        dbname (str): 数据库名称
        user (str): 数据库用户名
        password (str): 数据库密码
        host (str): 数据库主机
        port (int): 数据库端口

    Returns:
        None
    """
    # 连接数据库
    conn = psycopg2.connect(
        dbname=dbname,
        user=user,
        password=password,
        host=host,
        port=port
    )
    # 创建游标
    cur = conn.cursor()
    try:
        # 获取字段名列表
        fields = list(data_list[0].keys())
        # 构造 SQL 语句
        sql = "INSERT INTO {} ({}) VALUES %s".format(
            table_name,
            ", ".join(fields)
        )
        # 获取数据值列表
        values = [[data[field] for field in fields] for data in data_list]
        # 执行 SQL 语句批量插入数据
        execute_values(cur, sql, values)
        # 提交更改
        conn.commit()
        print("成功插入{}条数据".format(len(data_list)))
    except Exception as e:
        # 发生异常时回滚
        conn.rollback()
        print("插入数据时发生异常：", e)
    finally:
        # 关闭游标和连接
        cur.close()
        conn.close()

from typing import List, Dict, Union, Tuple
import psycopg2
from psycopg2.extras import execute_values


class PostgreSQL:
    def __init__(self, dbname: str, user: str, password: str, host: str, port: int):
        self.dbname = dbname
        self.user = user
        self.password = password
        self.host = host
        self.port = port

    def connect(self):
        # 建立数据库连接
        self.conn = psycopg2.connect(
            dbname=self.dbname,
            user=self.user,
            password=self.password,
            host=self.host,
            port=self.port
        )
        self.cur = self.conn.cursor()

    def disconnect(self):
        # 断开数据库连接
        self.cur.close()
        self.conn.close()

    def execute_query(self, sql: str, params: Union[List, Tuple] = None, fetch_one: bool = False) -> Union[None, List, Tuple]:
        """
        执行查询语句

        Args:
            sql (str): 查询语句
            params (Union[List, Tuple], optional): 参数列表，默认为 None
            fetch_one (bool, optional): 是否只获取一条数据，默认为 False

        Returns:
            Union[None, List, Tuple]: 查询结果
        """
        self.connect()
        if params is None:
            self.cur.execute(sql)
        else:
            self.cur.execute(sql, params)
        if fetch_one:
            result = self.cur.fetchone()
        else:
            result = self.cur.fetchall()
        self.disconnect()
        return result

    def execute_update(self, sql: str, params: Union[List, Tuple] = None) -> int:
        """
        执行更新语句

        Args:
            sql (str): 更新语句
            params (Union[List, Tuple], optional): 参数列表，默认为 None

        Returns:
            int: 更新的行数
        """
        self.connect()
        if params is None:
            self.cur.execute(sql)
        else:
            self.cur.execute(sql, params)
        rowcount = self.cur.rowcount
        self.conn.commit()
        self.disconnect()
        return rowcount

    def execute_insert(self, table_name: str, data: Dict) -> None:
        """
        插入一条数据

        Args:
            table_name (str): 表名
            data (Dict): 数据字典

        Returns:
            None
        """
        self.connect()
        columns = list(data.keys())
        values = tuple(data.values())
        placeholders = ", ".join(["%s" for _ in range(len(values))])
        sql = f"INSERT INTO {table_name} ({', '.join(columns)}) VALUES ({placeholders})"
        self.cur.execute(sql, values)
        self.conn.commit()
        self.disconnect()

    def execute_batch_insert(self, table_name: str, data_list: List[Dict]) -> None:
        """
        批量插入数据

        Args:
            table_name (str): 表名
            data_list (List[Dict]): 数据字典列表

        Returns:
            None
        """
        self.connect()
        columns = list(data_list[0].keys())
        values = [[data[field] for field in columns] for data in data_list]
        placeholders = ", ".join(["%s" for _ in range(len(columns))])
        sql = f"INSERT INTO {table_name} ({', '.join(columns)}) VALUES ({placeholders})"
        execute_values(self.cur, sql, values)
        self.conn.commit()
        self.disconnect()

    def execute_native_query(self, sql: str, params: Union[List, Tuple] = None) -> Union[None, List, Tuple]:
        """
        执行原生 SQL 查询语句

        Args:
            sql (str): 原生 SQL 查询语句
            params (Union[List, Tuple], optional): 参数列表，默认为 None

        Returns:
            Union[None, List, Tuple]: 查询结果
        """
        self.connect()
        if params is None:
            self.cur.execute(sql)
        else:
            self.cur.execute(sql, params)
        result = self.cur.fetchall()
        self.disconnect()
        return result

    def execute_native_update(self, sql: str, params: Union[List, Tuple] = None) -> int:
        """
        执行原生 SQL 更新语句

        Args:
            sql (str): 原生 SQL 更新语句
            params (Union[List, Tuple], optional): 参数列表，默认为 None

        Returns:
            int: 更新的行数
        """
        self.connect()
        if params is None:
            self.cur.execute(sql)
        else:
            self.cur.execute(sql, params)
        rowcount = self.cur.rowcount
        self.conn.commit()
        self.disconnect()
        return rowcount