DG47 commented on PR #74279: URL: https://github.com/apache/airflow/pull/74279#issuecomment-6472404678
Thanks for asking for this. Running it on a real Airflow showed the driver matters a lot more than the PR description let on, so here are the full results. **Setup:** Airflow 3.3.2 from PyPI (constraints-3.3.2, Python 3.12), api-server + scheduler + dag-processor with LocalExecutor, Postgres 17 and MySQL 26.10 running locally, and `common.sql` installed from this branch. The baseline is `common-sql` 2.4.0 from PyPI, which has the same version and no `on_kill`. Each DAG runs four `SQLExecuteQueryOperator` tasks with a 5-minute statement: `pg_sleep(300)`, `SLEEP(300)` or a recursive CTE on MySQL, and a recursive CTE on SQLite. Two tasks are marked failed while running, and two have `execution_timeout=15s`. In each pair, one task uses `cancel_on_kill=True` and the other `cancel_on_kill=False`. I checked on the DB side (`pg_stat_activity`, `SHOW PROCESSLIST`, or the task process for SQLite) whether the statement actually stopped. ### Postgres with psycopg 3: cancelled psycopg 3 is what `PostgresHook` uses by default on Airflow 3.3 with SQLAlchemy 2. - Marking the task failed in the UI: the statement is gone from `pg_stat_activity` about 0.2 s after the supervisor sends SIGTERM, and the task fails with `QueryCanceled`. - `execution_timeout`: cancelled within 0.6 s, and the task fails with `AirflowTaskTimeout`. - On main, and with `cancel_on_kill=False`, the statement keeps running on the server after the task is failed. Video of the whole flow (66 s): [pr_psycopg3_mark_failed_demo.mp4](https://github.com/DG47/airflow/blob/8d044d3ec852788550363440211818938e4e1d92/pr_psycopg3_mark_failed_demo.mp4) <details><summary>Screenshots</summary> Before: all four statements running  Marking `mark_failed` as failed  Task log after the kill  `execution_timeout` (taken before the follow-up commit, so it still shows the "nothing to cancel" warning discussed below)  After: only the two `cancel_on_kill=False` statements are still running  Same run on main: the killed statements are all still running  </details> ### psycopg2, SQLite and MySQL (mysqlclient and mysql-connector-python): not cancelled by this PR These drivers stay in C for the whole statement, so Python never runs the signal handlers. - Mark failed: `on_kill` never runs, and the supervisor SIGKILLs the task 5 s after SIGTERM. - Postgres keeps running the statement after the client is gone. - SQLite stops only because it dies with the process. - MySQL `SLEEP()` happens to notice the dropped connection, but a CPU-bound statement keeps running. - `execution_timeout`: the timeout isn't raised until the statement returns. On SQLite the 15 s timeout fired after 5 min 23 s, on main and on this branch alike. <details><summary>Screenshots</summary> psycopg2, mark failed: killed 5 s after SIGTERM, `on_kill` never ran  SQLite, `execution_timeout` of 15 s fired after 5 min 23 s  MySQL CPU-bound CTE, mark failed  </details> So the PR as it stands helps drivers that hand control back to Python while a statement runs. psycopg 3 is the one I verified. For the others it's a no-op and doesn't regress anything. I've corrected the description, which said psycopg2 handles the kill "afterwards" (in practice it never does) and named the wrong MySQL drivers. The `cancel_query` docstring and the operator guide now say the same. I also tried hook-level overrides for the two blocking cases, on a separate local branch: - **psycopg2:** a wait callback (`wait_select`, which also cancels on interrupt), installed only while the statement runs, because `copy_expert` refuses to run with one. - **SQLite:** a progress handler, plus `connection.interrupt()` in `cancel_query`. With those, psycopg2 and SQLite cancel too: about 1.6 s from the click, and 0.4 s after the timeout. <details><summary>Screenshots</summary>   </details> They're provider-specific, so I'd rather send them as follow-up PRs to the postgres and sqlite providers than grow this one. I can fold them in here if you'd prefer. MySQL would need a `KILL QUERY` on a second connection, which I haven't attempted. One small fix is in 06293042af: after the hook has already cancelled a statement interrupted by the timeout, `on_kill` used to log a warning that there was nothing to cancel, which read like an error. That message is now logged at info level. -- 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]
