DG47 opened a new pull request, #75975:
URL: https://github.com/apache/airflow/pull/75975
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
**Depends on #74279.** This branch is stacked on #74279's branch, so until
that PR is merged the diff here also shows its commits. The only new commit is
00ea612612, which changes `PostgresHook`, its tests and the provider's operator
guide. I'll rebase onto main once #74279 is merged.
This is the follow-up promised in #74279, which added
`DbApiHook.cancel_query()` and `SQLExecuteQueryOperator.on_kill` /
`cancel_on_kill`. As shown there, it has no effect when `PostgresHook` runs on
psycopg2. That's the driver it uses whenever psycopg 3 or SQLAlchemy 2 is
missing, which includes every Airflow 2.11 install. psycopg2 waits for a
statement inside libpq, so Python can't run the `SIGTERM` handler (`on_kill`)
or the `SIGALRM` handler of `execution_timeout` until the statement finishes.
On mark failed, the supervisor SIGKILLs the task after 5 s and the statement
keeps running on the server. On `execution_timeout`, the timeout isn't even
raised until the statement returns.
**Change:** while a statement runs, `PostgresHook` installs a psycopg2 wait
callback that waits in `select()` (`psycopg2.extras.wait_select`), where Python
runs signal handlers:
* **Mark failed:** `on_kill` runs mid-statement, and the existing
`cancel_query()` sends `connection.cancel()`.
* **`execution_timeout`:** `AirflowTaskTimeout` is raised inside the
callback. psycopg2 closes the connection when a wait callback raises, so the
callback cancels the statement first. `cancel_query()` then reports the
statement as already cancelled, instead of trying a closed connection.
The wait callback is a process-wide psycopg2 setting, and `copy_expert`
refuses to run while one is set, so the hook keeps it narrow. It installs the
callback only:
* for the duration of the statement (`_track_running_cursor`),
* in the main thread, which is the only one where Python runs signal
handlers,
* when no other wait callback is set (for example one for gevent),
* when `cancel_on_interrupt` / `cancel_on_kill` is enabled.
psycopg 3 is unchanged.
**Verified on a real Airflow 3.3.2** (LocalExecutor, Postgres 17, psycopg 3
uninstalled so the hook falls back to psycopg2, the same harness as in #74279):
| kill | before (#74279 only) | with this PR |
|---|---|---|
| mark failed, `cancel_on_kill=True` | `on_kill` never runs; task SIGKILLed
after 5 s; statement keeps running on the server | cancelled through
`connection.cancel()` 0.4 s after SIGTERM (`QueryCanceled`) |
| `execution_timeout`, `cancel_on_kill=True` | timeout only raised when the
statement returns | cancelled 0.1 s after the timeout; task fails with
`AirflowTaskTimeout` |
| either, `cancel_on_kill=False` | statement keeps running | statement keeps
running, as designed |
**Tests:**
* Unit tests (mocked psycopg2) check that the callback:
* is installed only while a statement runs, and removed after a failure,
* is not installed with `cancel_on_interrupt=False`, from a non-main
thread, or over an existing callback,
* cancels on interruption,
* makes `cancel_query()` handle a connection closed by the callback.
* Two `backend("postgres")` tests run `pg_sleep(20)` through psycopg2 and
interrupt it from a signal handler after 0.5 s: once raising
`AirflowTaskTimeout`, once calling `cancel_query()`. They check that the
statement is gone from `pg_stat_activity`. Without this change both take the
full 20 s and fail.
* `providers/postgres/tests/unit/postgres`: 177 passed with psycopg 3 and
177 passed with psycopg2 only. The hook tests with `--backend postgres`: 23
passed with psycopg 3 and 30 passed with psycopg2 only.
The postgres operator guide gets a short section on this. The provider's
`common-sql` lower bound doesn't need to change for it to install: the
overrides only run when the base hook calls them, which needs the `common-sql`
release that ships #74279.
related: #74279, #27314
--
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]
