如何使用psycopg终止运行超时的PostgreSQL查询？

一、利用psycopg2的原生查询超时参数 #

psycopg2从2.8版本开始，支持在cursor.execute()中直接传入timeout参数（单位：秒），超时后会自动抛出psycopg2.errors.QueryCanceled异常，捕获后即可跳过当前查询继续执行下一个。

示例代码：

import psycopg2
from psycopg2.errors import QueryCanceled

# 建立数据库连接
conn = psycopg2.connect(
    dbname="your_db",
    user="your_user",
    password="your_pwd",
    host="your_host"
)
cursor = conn.cursor()

text_list = ["text1", "text2", "text3"]

for text in text_list:
    try:
        # 设置10秒查询超时
        cursor.execute("SELECT * FROM your_table WHERE content = %s;", (text,), timeout=10)
        result = cursor.fetchall()
        # 处理查询结果
        print(f"查询「{text}」成功：{result}")
    except QueryCanceled:
        print(f"查询「{text}」超时，已终止")
        conn.rollback()  # 超时后需回滚事务，避免影响后续查询
    except Exception as e:
        print(f"查询「{text}」出错：{str(e)}")
        conn.rollback()

# 关闭资源
cursor.close()
conn.close()

这个参数本质是临时设置PostgreSQL的statement_timeout参数，查询结束后会自动恢复原有配置，无需修改数据库全局设置。

二、手动设置会话级查询超时（适配旧版psycopg2） #

如果使用的是2.8以下版本的psycopg2，可以手动在会话中设置statement_timeout（单位：毫秒），对当前连接的所有查询生效：

示例代码：

import psycopg2
from psycopg2.errors import QueryCanceled

conn = psycopg2.connect(...)
cursor = conn.cursor()

# 设置会话超时为10秒（10000毫秒）
cursor.execute("SET statement_timeout = 10000;")
conn.commit()

text_list = [...]

for text in text_list:
    try:
        cursor.execute("SELECT * FROM your_table WHERE content = %s;", (text,))
        result = cursor.fetchall()
        print(f"查询「{text}」成功：{result}")
    except QueryCanceled:
        print(f"查询「{text}」超时")
        conn.rollback()
    except Exception as e:
        print(f"查询「{text}」出错：{str(e)}")
        conn.rollback()

# 若后续需要取消超时，可执行：
# cursor.execute("SET statement_timeout = 0;")
# conn.commit()

三、Python线程/进程配合超时控制（兜底方案） #

如果上述数据库端的方法无法生效，可以用Python的线程配合超时等待实现，但注意线程无法强制终止，需结合数据库超时设置避免阻塞：

示例代码：

import psycopg2
import threading

def run_query(text, cursor, result_container):
    try:
        cursor.execute("SELECT * FROM your_table WHERE content = %s;", (text,))
        result_container["data"] = cursor.fetchall()
    except Exception as e:
        result_container["data"] = f"错误：{str(e)}"

conn = psycopg2.connect(...)
cursor = conn.cursor()

# 先设置数据库端超时，避免线程永久阻塞
cursor.execute("SET statement_timeout = 10000;")
conn.commit()

text_list = [...]

for text in text_list:
    result = {}
    query_thread = threading.Thread(target=run_query, args=(text, cursor, result))
    query_thread.start()
    query_thread.join(timeout=10)  # 等待10秒
    
    if query_thread.is_alive():
        print(f"查询「{text}」超时，已终止")
        # 超时后需重建连接，避免当前连接处于异常状态
        conn.close()
        conn = psycopg2.connect(...)
        cursor = conn.cursor()
        cursor.execute("SET statement_timeout = 10000;")
        conn.commit()
    else:
        print(f"查询「{text}」完成：{result.get('data')}")

cursor.close()
conn.close()

四、从根源优化查询（减少超时发生） #

超时本质是查询效率低下，建议优先优化：

• 给查询字段（比如content）建立B树索引：CREATE INDEX idx_your_table_content ON your_table(content);
• 若为模糊查询（如LIKE '%xxx%'），改用PostgreSQL全文索引（GIN/GIST类型）
• 用EXPLAIN ANALYZE分析查询计划，定位慢查询瓶颈

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