使用Psycopg预获取查询结果大小遇ProgrammingError问题排查

1. 你遗漏了什么？ #

核心问题是Psycopg与PGAdmin对PL/pgSQL函数结果的处理逻辑存在差异：

• PGAdmin会自动展示函数内部最后一次查询的输出，但Psycopg要求函数明确返回结果集或标量值，且你必须通过fetchone()/fetchall()等方法主动获取返回值。
• 如果你的PL/pgSQL函数仅执行了创建临时表、查询大小的操作，但没有用RETURN或RETURN QUERY语句返回最终的大小值，Psycopg就会判定“无结果产生”并抛出ProgrammingError。
• 另外，若你在Psycopg中调用函数的方式错误（比如对返回标量的函数用SELECT * FROM 函数名()），也会导致无法正确获取结果。

举个反例：如果函数内部仅执行SELECT pg_total_relation_size('temp_table');却不把结果返回，PGAdmin会显示这个查询的结果，但Psycopg无法识别，必须改成RETURN 结果值;才能被Psycopg捕获。

2. 替代实现方案 #

方案一：直接用事务+临时表（无需PL/pgSQL函数） #

通过Psycopg直接执行临时表创建、大小查询，用事务包裹确保临时表自动清理：

import psycopg2

# 建立连接
conn = psycopg2.connect("dbname=你的数据库名 user=你的用户名")
cur = conn.cursor()
conn.autocommit = False  # 开启事务

try:
    # 创建临时表，ON COMMIT DROP确保事务结束后自动删除
    cur.execute("CREATE TEMP TABLE temp_result ON COMMIT DROP AS SELECT * FROM 你的表 WHERE 你的条件;")
    # 查询临时表总大小（含索引、TOAST表）
    cur.execute("SELECT pg_total_relation_size('temp_result');")
    result_size = cur.fetchone()[0]
    print(f"查询结果总大小：{result_size} 字节")
finally:
    # 回滚事务，自动清理临时表
    conn.rollback()
    cur.close()
    conn.close()

方案二：用EXPLAIN估算结果大小（无需创建临时表） #

如果不需要精确值，可通过EXPLAIN或EXPLAIN ANALYZE从执行计划中提取估算的行数和行宽，计算大致大小：

import psycopg2

conn = psycopg2.connect("dbname=你的数据库名 user=你的用户名")
cur = conn.cursor()

# 执行EXPLAIN（不带ANALYZE则不实际执行查询，估算精度稍低）
cur.execute("EXPLAIN ANALYZE SELECT * FROM 你的表 WHERE 你的条件;")
plan_lines = [line[0] for line in cur.fetchall()]

# 解析执行计划中的行数和行宽
for line in plan_lines:
    if "rows" in line and "width" in line:
        row_info = line.split()
        rows = int(row_info[0].split('=')[1])
        row_width = int(row_info[1].split('=')[1])
        estimated_size = rows * row_width
        print(f"估算结果大小：{estimated_size} 字节")
        break

cur.close()
conn.close()

方案三：改进PL/pgSQL函数并正确调用 #

修改函数使其明确返回结果，再用Psycopg正确获取：
首先创建函数：

CREATE OR REPLACE FUNCTION calculate_query_size(query_text text) RETURNS bigint AS $$
DECLARE
    temp_table text := 'temp_' || md5(random()::text);  # 生成唯一临时表名避免冲突
    total_size bigint;
BEGIN
    -- 动态执行传入的查询创建临时表
    EXECUTE format('CREATE TEMP TABLE %I AS %s', temp_table, query_text);
    -- 查询临时表大小
    EXECUTE format('SELECT pg_total_relation_size(%I)', temp_table) INTO total_size;
    -- 删除临时表
    EXECUTE format('DROP TABLE %I', temp_table);
    -- 返回结果
    RETURN total_size;
END;
$$ LANGUAGE plpgsql;

然后在Python中调用：

import psycopg2

conn = psycopg2.connect("dbname=你的数据库名 user=你的用户名")
cur = conn.cursor()

target_query = "SELECT * FROM 你的表 WHERE 你的条件;"
cur.execute("SELECT calculate_query_size(%s);", (target_query,))
result_size = cur.fetchone()[0]
print(f"查询结果总大小：{result_size} 字节")

cur.close()
conn.close()

内容的提问来源于stack exchange，提问作者Pietro D'Antuono