如何从Python调用PostgreSQL的count_estimate函数？

问题背景 #

在PostgreSQL的ais_history schema下创建了一个用于获取近似行数的自定义函数，函数定义如下：

CREATE FUNCTION ais_history.count_estimate(query text)
        RETURNS integer
        LANGUAGE plpgsql AS
        $func$
                DECLARE
                        rec   record;
                        rows  integer;
                BEGIN
                FOR rec IN EXECUTE 'EXPLAIN ' || query LOOP
                        rows := substring(rec."QUERY PLAN" FROM ' rows=([[:digit:]]+)');
                        EXIT WHEN rows IS NOT NULL;
                END LOOP;

                RETURN rows;
                END
        $func$;

在psql命令行中调用该函数完全正常：

db=> SELECT count_estimate('SELECT * FROM schema.table_name');
 count_estimate
----------------
              4
(1 row)

但使用Python通过psycopg2调用时，却抛出UndefinedFunction错误，代码示例：

conn.autocommit = True
cursor = conn.cursor()
cursor.execute("SELECT count_estimate('SELECT * FROM schema.table_name');")
conn.commit()
conn.close()

错误信息：

UndefinedFunction                         Traceback (most recent call last)
Input In [91], in <cell line: 8>()
      5 conn.autocommit = True
      6 cursor = conn.cursor()
----> 8 cursor.execute("SELECT count_estimate('SELECT * FROM schema.table_name');")
     10 conn.commit()
     11 conn.close()

UndefinedFunction: function count_estimate(unknown) does not exist
LINE 1: SELECT count_estimate('SELECT * FROM schema.table_name...
               ^
HINT:  No function matches the given name and argument types. You might need to add explicit type casts.

解决方法 #

1. 显式指定函数所属schema并强制参数类型转换 #

报错核心原因：一是Python连接未将ais_history加入搜索路径，导致找不到函数；二是psycopg2未自动将字符串字面量转为text类型，和函数定义的参数类型不匹配。

修改调用语句，同时指定schema和强制类型转换：

conn.autocommit = True
cursor = conn.cursor()
# 显式指定schema+将参数强制转为text类型
cursor.execute("SELECT ais_history.count_estimate('SELECT * FROM schema.table_name'::text);")
# 获取查询结果
result = cursor.fetchone()
print(f"近似行数：{result[0]}")
conn.commit()
conn.close()

2. 配置连接的搜索路径（search_path） #

如果不想每次调用都写schema，可以在建立连接后先设置搜索路径，包含ais_history：

conn.autocommit = True
cursor = conn.cursor()
# 将ais_history加入搜索路径，顺序按需调整
cursor.execute("SET search_path TO public, ais_history;")
# 强制参数转为text类型调用函数
cursor.execute("SELECT count_estimate('SELECT * FROM schema.table_name'::text);")
result = cursor.fetchone()
print(f"近似行数：{result[0]}")
conn.commit()
conn.close()

额外优化：连接时指定搜索路径 #

也可以在建立数据库连接时直接指定search_path，避免每次执行额外语句：

import psycopg2
# 在connect参数中加入options配置search_path
conn = psycopg2.connect(
    dbname="your_db",
    user="your_user",
    password="your_pwd",
    host="your_host",
    options="-c search_path=public,ais_history"
)
conn.autocommit = True
cursor = conn.cursor()
cursor.execute("SELECT count_estimate('SELECT * FROM schema.table_name'::text);")
result = cursor.fetchone()
print(f"近似行数：{result[0]}")
conn.close()

内容的提问来源于stack exchange，提问作者user3367601