當 InterruptionKey 時，使用 psycopg2 殺死未完成的 SQL 提交查詢

Question

我使用 Python 腳本，該腳本使用 psycopg2 庫在 AWS Redshift 數據庫上運行各種 SQL 自動提交查詢。 該腳本是從我的本地工作站手動執行的。 流程如下：

• 使用 psycopg2.connect() 創建數據庫連接
• 使用 execute() 對數據庫執行自動提交的查詢
• 關閉連接。

由於各種原因，數據庫可能不可用（網絡問題，許多查詢已經在運行......），最好停止 Python 腳本。 在這一點上，我然后通過檢索與這些查詢關聯的 pid 通過 SQL 客戶端（SQL 工作台）終止已經提交（和未完成）的查詢。 當用戶停止它（ctrl + c）時，我想直接在 Python 腳本中自動執行最后一步。 流程是：

• 使用 psycopg2.connect() 創建數據庫連接
• 使用 execute() 對數據庫執行自動提交的查詢
• 使用 info.backend_pid Connection 屬性存儲與查詢關聯的當前 PID
• 如果收到 InterruptionKey 異常，則使用之前存儲的 PID 終止正在運行的查詢
• 關閉連接。

我在筆記本上做了一些測試，以檢查是否可以檢索 back_pid 信息：

log = logging.getLogger(__name__)
session = psycopg2.connect(
    connection_factory=LoggingConnection,
        host=host,
        port=port,
        dbname=database,
        user=user,
        password=password,
        sslmode="require",
)
session.initialize(log)
session.set_session(autocommit=True)

query = """
CREATE OR REPLACE FUNCTION janky_sleep (x float) RETURNS bool IMMUTABLE as $$
    from time import sleep
    sleep(x)
    return True
$$ LANGUAGE plpythonu;
"""

cur = session.cursor()
cur.execute(query)
cur.execute("select janky_sleep(60.0)")

我使用 sleep 函數來復制需要 60 秒才能完成的查詢的行為。 當獲取 backend_pid 如下：

session.info.backend_pid

問題是會話對象已被 execute() 方法（運行查詢）使用，並且 backend_pid 信息僅在會話空閑時產生，即當查詢完成時。

我想旋轉一個並發 Python 進程來監視父進程。 一旦父進程停止，子進程將通過第二個數據庫連接獲取 backend_pid，然后運行終止查詢。 然而，這種方法似乎有點矯枉過正。

處理這種情況的正確方法是什么？

謝謝

Answer 1

我終於使用了從文檔中找到的資源： http : //initd.org/psycopg/docs/faq.html#faq-interrupt-query 。 它使 Psycopg2 能夠獲取 SIGINT 信號並終止后續查詢。

>>> psycopg2.extensions.set_wait_callback(psycopg2.extras.wait_select)
>>> cnn = psycopg2.connect('')
>>> cur = cnn.cursor()
>>> cur.execute("select pg_sleep(10)")
^C
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  QueryCanceledError: canceling statement due to user request

>>> cnn.rollback()
>>> # You can use the connection and cursor again from here