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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
