一、架构背景与痛点

在自动化测试用例执行、测试数据准备以及后端数据核对过程中，PostgreSQL 是最常用的关系型数据库之一。

在实际编写数据脚本时，很多初学者容易写出低效甚至危险的代码：

1. 每次查询新建物理连接：高频测试场景下快速耗尽数据库 max_connections；
2. 字符串拼接 SQL：产生严重的 SQL 注入安全隐患，且无法享受数据库预编译执行计划缓存；
3. 循环单条插入（逐条 execute）：插入 1 万条数据需要发起 1 万次网络往返（RTT），耗时数十秒；
4. 一次性 fetchall()：面对数十万条大表查询时，进程直接把几百兆数据全部载入内存，引发客户端内存溢出（OOM）。

---

二、使用 SimpleConnectionPool 管理连接池

避免频繁握手创建物理连接，使用连接池统一分配与回收：

import psycopg2
from psycopg2 import pool
from contextlib import contextmanager

class PgDatabaseManager:
    _pool = None

    @classmethod
    def init_pool(cls, host="127.0.0.1", port=5432, dbname="testdb", user="postgres", password="secure_password"):
        if cls._pool is None:
            cls._pool = pool.SimpleConnectionPool(
                minconn=2,
                maxconn=10,
                host=host,
                port=port,
                dbname=dbname,
                user=user,
                password=password
            )

    @classmethod
    @contextmanager
    def get_connection(cls):
        """利用上下文管理器自动获取并归还连接"""
        conn = cls._pool.getconn()
        try:
            yield conn
        finally:
            cls._pool.putconn(conn)

---

三、参数化查询与动态构造

严禁用 Python 字符串 f"SELECT ... WHERE id = '{user_id}'" 拼接，必须使用元组传递参数，由驱动底层完成预编译与参数化转义：

def query_user_metrics(user_id: int, status_code: str) -> list[dict]:
    sql = """
        SELECT id, username, email, created_at
        FROM sys_user
        WHERE id = %s AND status = %s;
    """
    with PgDatabaseManager.get_connection() as conn:
        with conn.cursor() as cur:
            cur.execute(sql, (user_id, status_code))
            rows = cur.fetchall()
            
            # 将查询结果与字段名拼装为字典列表
            columns = [desc[0] for desc in cur.description]
            return [dict(zip(columns, row)) for row in rows]

---

四、高性能批量数据写入：execute_batch vs copy_from

针对批量构造 1 万 ~ 100 万条测试数据的场景，对比三种写入性能：

1. 推荐方案：psycopg2.extras.execute_batch

相比于循环单条 execute，execute_batch 将多条记录打包在单个网络请求中分批下发，效率提升 20 倍以上：

from psycopg2.extras import execute_batch

def batch_insert_device_metrics(records: list[tuple]):
    """
    records 示例：[(101, 'TEMP', 45.2), (102, 'SPEED', 1200.0), ...]
    """
    insert_sql = """
        INSERT INTO device_metric (device_id, metric_code, value)
        VALUES (%s, %s, %s);
    """
    with PgDatabaseManager.get_connection() as conn:
        with conn.cursor() as cur:
            # page_size: 每批次打包 1000 条发送给服务端
            execute_batch(cur, insert_sql, records, page_size=1000)
        conn.commit()

2. 极致性能：copy_expert 拷贝流

若需要单次导入数十万条以上的离线数据集，使用基于 PostgreSQL 原生 COPY 协议的流式导入是性能极限方案：

import io

def stream_copy_large_dataset(csv_buffer: io.StringIO):
    with PgDatabaseManager.get_connection() as conn:
        with conn.cursor() as cur:
            cur.copy_expert("COPY device_metric (device_id, metric_code, value) FROM STDIN WITH CSV", csv_buffer)
        conn.commit()

---

五、大表流式读取防 OOM：服务端游标（Server-side Cursor）

当需要对几百万条历史指标数据进行核对分析时，默认客户端游标会尝试把所有结果全部下载到本地内存。通过给游标命名，可强制开启 PostgreSQL 服务端游标，按需分批拉取：

def iterate_large_metrics_safely(batch_size=5000):
    sql = "SELECT id, device_id, raw_payload FROM large_history_logs;"
    
    with PgDatabaseManager.get_connection() as conn:
        # 传入 name 参数，此时游标转为服务端游标（Server-side Cursor）
        with conn.cursor(name="stream_cursor_01") as cur:
            cur.itersize = batch_size
            cur.execute(sql)
            
            while True:
                batch = cur.fetchmany(batch_size)
                if not batch:
                    break
                for row in batch:
                    yield row