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