DG47 commented on code in PR #74279:
URL: https://github.com/apache/airflow/pull/74279#discussion_r4394034564
##########
providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py:
##########
@@ -900,13 +902,23 @@ def _run_command(self, cur, sql_statement, parameters):
if self.log_sql:
self.log.info("Running statement: %s, parameters: %s",
sql_statement, parameters)
- if parameters:
- # If we're using psycopg3, we might need to handle parameters
differently
- if hasattr(cur, "__module__") and "psycopg" in cur.__module__ and
isinstance(parameters, list):
- parameters = tuple(parameters)
- cur.execute(sql_statement, parameters)
- else:
- cur.execute(sql_statement)
+ # Keep a reference to the cursor while the statement is in flight, so
that ``cancel_query`` can
+ # reach it when the task is killed in the middle of a long-running
statement.
+ self._running_cursor = cur
+ try:
+ if parameters:
+ # If we're using psycopg3, we might need to handle parameters
differently
+ if (
+ hasattr(cur, "__module__")
+ and "psycopg" in cur.__module__
+ and isinstance(parameters, list)
+ ):
+ parameters = tuple(parameters)
+ cur.execute(sql_statement, parameters)
+ else:
+ cur.execute(sql_statement)
+ finally:
+ self._running_cursor = None
Review Comment:
Thanks for running it, that's a real gap. Fixed in cd780893b4.
`_run_command` and `run` now go through a `_track_running_cursor` context
manager. When a non-`Exception` `BaseException` (`AirflowTaskTimeout`,
`KeyboardInterrupt`) leaves the block, it calls `cancel_query()` while the
cursor is still open and then re-raises. A failing cancel is logged and doesn't
replace the timeout. `test_running_cursor_is_cleared_when_statement_fails` now
only covers a plain `RuntimeError`, which doesn't cancel. The new
`test_statement_is_cancelled_when_interrupted_by_timeout` makes `execute` (and
`fetchall`) raise `AirflowTaskTimeout` and checks that `cursor.cancel()` went
out. On the operator side,
`test_execution_timeout_cancels_the_running_statement` covers the same thing
through `execute()`. The `on_kill` docstring now says the timeout case is
handled by the hook before `on_kill` runs.
##########
providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py:
##########
@@ -918,6 +930,42 @@ def _run_command(self, cur, sql_statement, parameters):
if (row_count := handlers.get_row_count(cur)) is not None:
self.log.info("Rows affected: %s", row_count)
+ def cancel_query(self) -> bool:
+ """
+ Cancel the SQL statement that this hook is currently executing, if
there is one.
+
+
:class:`~airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator`
calls this method from
+ ``on_kill`` when the task instance is killed while a statement is
still running, so that the
+ database stops working on a query whose result nobody is waiting for
anymore.
+
+ The default implementation relies on the cancellation extensions
offered by many DB-API drivers:
+ ``cursor.cancel()`` is used when the cursor provides it, otherwise
``connection.cancel()`` is used
+ when the cursor exposes its connection and the connection provides it
(e.g. ``psycopg2``,
Review Comment:
Agreed, thanks for checking with a real SIGTERM. In cd780893b4 psycopg2 is
gone from the list (it's now `psycopg`, `oracledb`, `trino`). The docstring and
`operators.rst` now say this only helps with drivers that return to Python to
handle signals during a statement, and name psycopg2 as one that doesn't. I'll
update the PR description too.
##########
providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py:
##########
@@ -900,13 +902,23 @@ def _run_command(self, cur, sql_statement, parameters):
if self.log_sql:
self.log.info("Running statement: %s, parameters: %s",
sql_statement, parameters)
- if parameters:
- # If we're using psycopg3, we might need to handle parameters
differently
- if hasattr(cur, "__module__") and "psycopg" in cur.__module__ and
isinstance(parameters, list):
- parameters = tuple(parameters)
- cur.execute(sql_statement, parameters)
- else:
- cur.execute(sql_statement)
+ # Keep a reference to the cursor while the statement is in flight, so
that ``cancel_query`` can
+ # reach it when the task is killed in the middle of a long-running
statement.
+ self._running_cursor = cur
Review Comment:
Yes. In cd780893b4 `run()` keeps the cursor tracked from `_run_command`
through `handler(cur)`, so a kill during `fetchall()` reaches
`cursor.cancel()`. The tracking is re-entrant for the same cursor:
`_run_command` still tracks on its own for hooks that call it directly (e.g.
`DatabricksSqlHook.run`), and when `run()` already tracks the cursor it leaves
the clearing to `run()`. Covered by `test_cancel_while_handler_fetches_results`
and by the `fetchall` case of the timeout test.
##########
providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py:
##########
@@ -918,6 +930,42 @@ def _run_command(self, cur, sql_statement, parameters):
if (row_count := handlers.get_row_count(cur)) is not None:
self.log.info("Rows affected: %s", row_count)
+ def cancel_query(self) -> bool:
+ """
+ Cancel the SQL statement that this hook is currently executing, if
there is one.
+
+
:class:`~airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator`
calls this method from
+ ``on_kill`` when the task instance is killed while a statement is
still running, so that the
+ database stops working on a query whose result nobody is waiting for
anymore.
+
+ The default implementation relies on the cancellation extensions
offered by many DB-API drivers:
+ ``cursor.cancel()`` is used when the cursor provides it, otherwise
``connection.cancel()`` is used
+ when the cursor exposes its connection and the connection provides it
(e.g. ``psycopg2``,
+ ``psycopg`` or ``oracledb``). Hooks for databases whose driver has no
such extension can override
+ this method and cancel the statement in a database specific way, for
example by issuing a
+ ``KILL QUERY`` command on a separate connection.
+
+ :return: ``True`` if a cancellation request was sent to the database,
``False`` if no statement
+ is running or the driver does not support cancelling a running
statement.
+ """
+ # getattr: a subclass that does not call ``DbApiHook.__init__`` should
still get the no-op path
+ cur = getattr(self, "_running_cursor", None)
+ if cur is None:
+ self.log.info("No SQL statement is currently running, there is
nothing to cancel.")
Review Comment:
Done in cd780893b4: it's now `"No running SQL statement is tracked by %s,
there is nothing to cancel."` at WARNING, with the hook class name. The
`cancel_query` docstring also says that only statements run through `run` /
`_run_command` are tracked, and `operators.rst` mentions hooks that run
statements another way.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
