You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "DG47 (via GitHub)" <gi...@apache.org> on 2026/10/05 17:03:30 UTC
[PR] Cancel the running statement in SQLExecuteQueryOperator.on_kill [airflow]
DG47 opened a new pull request, #74279:
URL: https://github.com/apache/airflow/pull/74279
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
`SQLExecuteQueryOperator` had no `on_kill` implementation, so killing a task (marking it as failed in the UI, a `SIGTERM` from the executor, or an elapsed `execution_timeout`) left the SQL statement running on the database until it completed on its own.
This PR adds a generic, hook-level cancellation mechanism to `common.sql` and wires it into the operator:
* `DbApiHook._run_command` keeps a reference to the cursor while `cursor.execute()` is in flight (cleared in a `finally`), so the hook knows which statement is currently running.
* New public method `DbApiHook.cancel_query() -> bool`. The default implementation uses the `cancel()` extension of the DB-API cursor, falling back to `cursor.connection.cancel()`. That covers `psycopg2`, `psycopg`, `oracledb`, `pyodbc`, `trino`, `databricks-sql-connector` and others out of the box. When no statement is running, or the driver has no cancel extension (e.g. `pymysql`, `mysqlclient`), it logs a message and returns `False` instead of raising. Hooks for databases that need a server-side command (`KILL QUERY`, `SYSTEM$CANCEL_QUERY`, ...) can override it.
* `SQLExecuteQueryOperator.on_kill` delegates to `self.get_db_hook().cancel_query()` and never raises, so it is safe on the kill path (Airflow calls `on_kill` from the task process `SIGTERM` handler while `execute()` may still be blocked inside the driver, and after `AirflowTaskTimeout`).
Design notes, relative to the earlier attempt in #27514:
* That PR tracked query ids *after* `execute()` returned, which (as pointed out in review) means the "running" queries had already finished. Tracking the live cursor during `execute()` is what makes cancellation from the signal handler actually reach a running statement.
* The state lives on the hook rather than the operator because the hook owns the connection/cursor lifecycle; the operator only has the hook object. Hooks with a custom `run()` that still call `self._run_command` (e.g. `DatabricksSqlHook`) are covered automatically; hooks that bypass `_run_command` entirely (e.g. `SnowflakeHook`) get the safe no-op until they override `cancel_query`.
* No locks: `on_kill` runs in the main thread's signal handler, so attribute assignment is sufficient and a lock could deadlock against the interrupted frame.
Tests mock the hook (operator) and the DB-API cursor/connection (hook), including the case where `cancel_query` is invoked from inside `cursor.execute` to simulate a kill mid-statement. The `hooks/sql.pyi` public API stub gains `cancel_query` and the operator guide documents the behaviour. No changelog entry is added since the provider changelog is maintained by the release manager and provider newsfragments are not used.
closes: #27314
---
* Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment: `{pr_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments). You can add this file in a follow-up commit after the PR is created so you know the PR number.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Cancel the running statement in SQLExecuteQueryOperator.on_kill [airflow]
Posted by "DG47 (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Cancel the running statement in SQLExecuteQueryOperator.on_kill [airflow]
Posted by "DG47 (via GitHub)" <gi...@apache.org>.
DG47 commented on PR #74279:
URL: https://github.com/apache/airflow/pull/74279#issuecomment-6448570613
@kaxil would you be able to review this? It closes #27314: `SQLExecuteQueryOperator.on_kill` now cancels the running statement through a new `DbApiHook.cancel_query()`, so a killed or timed-out task stops the query on the database. I merged main last week after a test-file conflict, and the common-sql hook and operator tests pass.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Cancel the running statement in SQLExecuteQueryOperator.on_kill [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74279:
URL: https://github.com/apache/airflow/pull/74279#discussion_r4393906833
##########
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:
On the `execution_timeout` path this `finally` clears the cursor before `on_kill` gets a chance. `AirflowTaskTimeout` is raised from SIGALRM inside `cur.execute`, unwinds through here and through the `closing(cur)` / `closing(conn)` in `run()`, and only then does the task runner call `task.on_kill()` from its `except AirflowTaskTimeout:`. I ran it on Postgres with psycopg 3, `timeout(2)` around `hook.run("select pg_sleep(10)")`: `cancel_query()` logged "nothing to cancel" and returned `False`, and the statement was still `active` in `pg_stat_activity` 7s after the client closed.
Could `_run_command` send the cancel itself when a non-`Exception` `BaseException` interrupts `execute` (that's what `AirflowTaskTimeout` is), while the cursor is still open, then re-raise? A test where `execute` raises `AirflowTaskTimeout` and asserts the cancel went out would cover it. `test_running_cursor_is_cleared_when_statement_fails` currently pins the opposite for that case. If that's out of scope, the timeout wording in the `on_kill` docstring and in `operators.rst` should go.
##########
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:
I don't think psycopg2 belongs on this list. Its synchronous `execute` blocks in libpq and doesn't return to Python for a signal, so the SIGTERM handler only runs after the statement is done. I sent a real SIGTERM 2s into a `pg_sleep(8)` going through `_run_command`: with psycopg2 the handler ran once the 8s were up, and `cancel_query()` still logged "Cancelling the running SQL statement" and returned `True` for a statement that had already finished. psycopg 3 got `QueryCanceled` at 2.1s. `PostgresHook` only uses psycopg 3 when it's installed and SQLAlchemy is 2.x, so psycopg2 setups get nothing from this.
##########
providers/common/sql/tests/unit/common/sql/hooks/test_sql.py:
##########
@@ -503,3 +503,95 @@ def get_sqlalchemy_engine_with_listeners(engine_kwargs=None):
assert table_names == ["users"]
assert opened
assert closed == opened
+
+
+class _CancelQueryHook(DbApiHook):
Review Comment:
`mock_db_hook` from `tests_common.test_utils.common_sql` already gives a `DbApiHook` with a mocked `get_conn`, and this module imports it. Could these tests use it instead of a new helper? The cursors in `test_nothing_to_cancel_when_no_statement_is_running` and `test_running_cursor_is_cleared_when_statement_fails` are bare `MagicMock()`s too, which have `.cancel` and `.connection.cancel` by default; a spec list like the one in `test_cancel_through_cursor` would keep them honest.
##########
providers/common/sql/tests/unit/common/sql/operators/test_sql.py:
##########
@@ -209,6 +210,25 @@ def test_output_processor(self, mock_get_db_hook):
assert descriptions == ("id", "name")
assert result == [(1, "Alice"), (2, "Bob")]
+ @mock.patch.object(SQLExecuteQueryOperator, "get_db_hook")
+ def test_on_kill_cancels_running_query(self, mock_get_db_hook):
Review Comment:
`get_db_hook` returns a bare `MagicMock` here, so this passes even if `cancel_query` were renamed. `mock.create_autospec(DbApiHook, instance=True)` as the return value would pin it. It would also help to have one test that goes through `execute()` with a real hook and calls `on_kill()` from the cursor's `execute` side effect, since the fix depends on `on_kill` reaching the same hook object that's still holding the cursor.
##########
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:
This is also what gets logged while a statement is running on hooks that don't go through the base `_run_command`. `SnowflakeHook._run_command` overrides it without calling super, and `ExasolHook.run` executes through `conn.execute`, so killing a `SQLExecuteQueryOperator` on either logs at INFO that nothing is running. `operators.rst` says a warning is logged when cancelling isn't supported. Maybe "no running statement is tracked by %s" at WARNING?
##########
providers/common/sql/src/airflow/providers/common/sql/operators/sql.py:
##########
@@ -580,6 +580,22 @@ def execute(self, context):
self.log.info("result: %s", result)
return result
+ def on_kill(self) -> None:
+ """
+ Cancel the SQL statement that is still running when the task instance is killed.
+
+ Airflow calls this method from the ``SIGTERM`` handler of the task process, i.e. while ``execute``
+ may still be blocked in the database driver waiting for the statement to finish, and after the
+ ``execution_timeout`` of the task elapsed. Cancellation is delegated to
+ :meth:`~airflow.providers.common.sql.hooks.sql.DbApiHook.cancel_query`, so every database hook can
+ use the mechanism its driver offers. It is best effort and never raises.
+ """
+ self.log.info("Task instance is being killed, cancelling the running SQL statement if there is one")
+ try:
+ self.get_db_hook().cancel_query()
Review Comment:
`get_db_hook()` can build a new hook here, with a connection lookup, inside the SIGTERM handler. That happens if the kill lands before `_hook` is cached, and on every kill for subclasses that don't cache it (`DuckDBExecuteQueryOperator.get_db_hook` builds a fresh one per call). A fresh hook has no running cursor, so the lookup is wasted. `hook = self.__dict__.get("_hook")` with an early return on `None` keeps this to the hook `execute()` actually ran on.
##########
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:
The window opened here closes before `handler(cur)` runs in `run()`, so a kill while results are being fetched isn't covered. For Trino that's most of a SELECT: reading the client, `TrinoQuery.execute` returns once the first page of rows arrives, and `fetchall()` pulls the rest while the query is still running on the server. `trino.dbapi.Cursor.cancel()` would work during that phase. Should the cursor stay tracked until the handler has consumed it?
##########
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)
Review Comment:
Is the `getattr` needed? `__init__` sets `_running_cursor`, and a hook that skipped `__init__` would already fail in `run()` on `self.log_sql` and `self.descriptions`, which the same `__init__` sets. Plain `self._running_cursor` and dropping `test_nothing_to_cancel_when_init_was_not_called` would be simpler.
##########
providers/common/sql/src/airflow/providers/common/sql/operators/sql.py:
##########
@@ -580,6 +580,22 @@ def execute(self, context):
self.log.info("result: %s", result)
return result
+ def on_kill(self) -> None:
Review Comment:
The other operators that cancel remote work on kill take `cancel_on_kill: bool = True` (`BigQueryInsertJobOperator`, `RedshiftDataOperator`, the EMR operators). Should this one too? It makes cancellation unconditional for every `SQLExecuteQueryOperator` subclass, and an autocommit load someone would rather let finish after the task is killed seems like a real case.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Cancel the running statement in SQLExecuteQueryOperator.on_kill [airflow]
Posted by "DG47 (via GitHub)" <gi...@apache.org>.
DG47 commented on PR #74279:
URL: https://github.com/apache/airflow/pull/74279#issuecomment-6269937878
Following up in case this slipped through. It closes #27314: `SQLExecuteQueryOperator.on_kill` now cancels the statement that's still running, through a new `DbApiHook.cancel_query()`, so a killed or timed-out task no longer leaves the query running on the database. It was rebased on 10-13. Happy to split the hook and operator parts into separate PRs if that's easier to review.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Cancel the running statement in SQLExecuteQueryOperator.on_kill [airflow]
Posted by "DG47 (via GitHub)" <gi...@apache.org>.
DG47 commented on code in PR #74279:
URL: https://github.com/apache/airflow/pull/74279#discussion_r4394035578
##########
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)
Review Comment:
Not needed, you're right. It's plain `self._running_cursor` in cd780893b4, and `test_nothing_to_cancel_when_init_was_not_called` is removed.
##########
providers/common/sql/src/airflow/providers/common/sql/operators/sql.py:
##########
@@ -580,6 +580,22 @@ def execute(self, context):
self.log.info("result: %s", result)
return result
+ def on_kill(self) -> None:
+ """
+ Cancel the SQL statement that is still running when the task instance is killed.
+
+ Airflow calls this method from the ``SIGTERM`` handler of the task process, i.e. while ``execute``
+ may still be blocked in the database driver waiting for the statement to finish, and after the
+ ``execution_timeout`` of the task elapsed. Cancellation is delegated to
+ :meth:`~airflow.providers.common.sql.hooks.sql.DbApiHook.cancel_query`, so every database hook can
+ use the mechanism its driver offers. It is best effort and never raises.
+ """
+ self.log.info("Task instance is being killed, cancelling the running SQL statement if there is one")
+ try:
+ self.get_db_hook().cancel_query()
Review Comment:
Done in cd780893b4. `on_kill` reads `self.__dict__.get("_hook")` and returns early (with an INFO log) when `execute` hasn't created a hook yet. Nothing is built or looked up in the signal handler. `test_on_kill_does_not_create_a_hook` asserts that `get_hook` isn't called. Subclasses that build a fresh hook in `get_db_hook` (like DuckDB) get no cancellation from this, which matches what their hook could do anyway.
##########
providers/common/sql/src/airflow/providers/common/sql/operators/sql.py:
##########
@@ -580,6 +580,22 @@ def execute(self, context):
self.log.info("result: %s", result)
return result
+ def on_kill(self) -> None:
Review Comment:
Good point, added in cd780893b4: `cancel_on_kill: bool = True`, documented in the docstring and in `operators.rst`. Since the timeout case is now cancelled inside the hook, the operator also passes the flag to the hook (`hook._cancel_on_interrupt = self.cancel_on_kill` in `execute`), so `cancel_on_kill=False` lets the statement finish both on SIGTERM and on `execution_timeout`. Both paths are tested with the flag on and off.
##########
providers/common/sql/tests/unit/common/sql/operators/test_sql.py:
##########
@@ -209,6 +210,25 @@ def test_output_processor(self, mock_get_db_hook):
assert descriptions == ("id", "name")
assert result == [(1, "Alice"), (2, "Bob")]
+ @mock.patch.object(SQLExecuteQueryOperator, "get_db_hook")
+ def test_on_kill_cancels_running_query(self, mock_get_db_hook):
Review Comment:
Changed in cd780893b4. The operator tests now use `mock.create_autospec(DbApiHook, instance=True)` through a patched `get_hook`, so a renamed `cancel_query` would fail them. I also added `test_on_kill_during_execute_reaches_the_running_statement`: it runs `execute()` against a real `DbApiHook` from `mock_db_hook`, and the cursor's `execute` side effect calls `operator.on_kill()`. It asserts that `cursor.cancel()` was called (or not, with `cancel_on_kill=False`), so it checks that `on_kill` reaches the same hook that holds the cursor.
##########
providers/common/sql/tests/unit/common/sql/hooks/test_sql.py:
##########
@@ -503,3 +503,95 @@ def get_sqlalchemy_engine_with_listeners(engine_kwargs=None):
assert table_names == ["users"]
assert opened
assert closed == opened
+
+
+class _CancelQueryHook(DbApiHook):
Review Comment:
Done in cd780893b4. `_CancelQueryHook` is gone and the tests use `mock_db_hook(DbApiHook)` with the cursor swapped in through `hook.get_conn().cursor`. Every cursor now has an explicit spec, and `cancel` / `connection` are only added where the test needs them.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Cancel the running statement in SQLExecuteQueryOperator.on_kill [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on PR #74279:
URL: https://github.com/apache/airflow/pull/74279#issuecomment-6453508900
Please post few screenshot/demo video with actual Airflow of running and cancellation with Postgres, sqlite and some DBs
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org