This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch v3-4-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-4-test by this push:
new de4c2047883 [v3-4-test] Fix `Trigger.clean_unused` rejecting a
caller-provided session (#75096) (#75122)
de4c2047883 is described below
commit de4c20478834347fb82fea7a362a92baad6b06be
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Tue Oct 20 09:41:44 2026 +0530
[v3-4-test] Fix `Trigger.clean_unused` rejecting a caller-provided session
(#75096) (#75122)
Restore the keyword-only session argument that #70961 removed, so callers
that run the cleanup inside their own transaction keep working. The model
method now deletes one batch in the caller's transaction without retrying:
a retry would have to roll back that transaction and silently discard the
caller's other writes. The triggerer supervisor owns the retry instead and
runs each attempt in a new transaction.
(cherry picked from commit e3110a3092de1c68596520e60b66e6bf84936479)
Co-authored-by: Kaxil Naik <[email protected]>
---
.../src/airflow/jobs/triggerer_job_runner.py | 7 ++-
airflow-core/src/airflow/models/trigger.py | 31 +++++-----
airflow-core/tests/unit/jobs/test_triggerer_job.py | 69 +++++++++++++++++++++-
airflow-core/tests/unit/models/test_trigger.py | 63 ++++++--------------
4 files changed, 107 insertions(+), 63 deletions(-)
diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py
b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
index 0bcc790224b..4c6cc56ceea 100644
--- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py
+++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
@@ -139,6 +139,7 @@ from airflow.triggers.shared_stream import
SharedStreamManager
from airflow.utils.helpers import prune_dict
from airflow.utils.log.logging_mixin import LoggingMixin
from airflow.utils.log.task_log_address import prepare_task_log_contexts,
render_task_log_filename
+from airflow.utils.retries import run_with_db_retries
from airflow.utils.session import create_session, provide_session
if TYPE_CHECKING:
@@ -818,8 +819,10 @@ class TriggerRunnerSupervisor(WatchedSubprocess):
Trigger.submit_event(trigger_id=trigger_id, event=event)
def clean_unused(self) -> None:
- """Remove triggers that are no longer needed."""
- Trigger.clean_unused()
+ """Remove triggers that are no longer needed, retrying database errors
in a new transaction."""
+ for attempt in run_with_db_retries():
+ with attempt, create_session() as session:
+ Trigger.clean_unused(session=session)
def handle_failed_triggers(self):
"""
diff --git a/airflow-core/src/airflow/models/trigger.py
b/airflow-core/src/airflow/models/trigger.py
index 1713f9aebb7..a7193074f6f 100644
--- a/airflow-core/src/airflow/models/trigger.py
+++ b/airflow-core/src/airflow/models/trigger.py
@@ -37,8 +37,7 @@ from airflow.models.base import Base
from airflow.models.taskinstance import TaskInstance
from airflow.serialization.enums import stringify_encoding_keys
from airflow.triggers.base import BaseTaskEndEvent
-from airflow.utils.retries import run_with_db_retries
-from airflow.utils.session import NEW_SESSION, create_session, provide_session
+from airflow.utils.session import NEW_SESSION, provide_session
from airflow.utils.sqlalchemy import UtcDateTime, with_row_locks
from airflow.utils.state import TaskInstanceState
@@ -249,30 +248,30 @@ class Trigger(Base):
return set(session.scalars(query))
@classmethod
- def clean_unused(cls) -> None:
+ @provide_session
+ def clean_unused(cls, *, session: Session = NEW_SESSION) -> None:
"""
Delete up to one batch of triggers that have no dependent tasks,
assets, or callbacks.
Triggers have a one-to-many relationship to task instances, so we need
to clean those up first.
Afterward we can drop unreferenced triggers; later triggerer loops
drain any remaining batches.
+
+ Database errors are not retried; to retry, run the call again in a new
transaction.
"""
batch_size = conf.getint("triggerer",
"unreferenced_triggers_cleanup_batch_size", fallback=500)
if batch_size <= 0:
raise ValueError("[triggerer]
unreferenced_triggers_cleanup_batch_size must be at least 1")
- for attempt in run_with_db_retries():
- with attempt:
- with create_session(scoped=False) as session:
- # Update all task instances with trigger IDs that are not
DEFERRED to remove them
- session.execute(
- update(TaskInstance)
- .where(
- TaskInstance.state != TaskInstanceState.DEFERRED,
- TaskInstance.trigger_id.is_not(None),
- )
- .values(trigger_id=None)
- )
- cls._delete_unused_batch(batch_size, session=session)
+ # Update all task instances with trigger IDs that are not DEFERRED to
remove them
+ session.execute(
+ update(TaskInstance)
+ .where(
+ TaskInstance.state != TaskInstanceState.DEFERRED,
+ TaskInstance.trigger_id.is_not(None),
+ )
+ .values(trigger_id=None)
+ )
+ cls._delete_unused_batch(batch_size, session=session)
@classmethod
def _delete_unused_batch(cls, batch_size: int, *, session: Session) ->
None:
diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py
b/airflow-core/tests/unit/jobs/test_triggerer_job.py
index 345a5859743..bbf01eba0f1 100644
--- a/airflow-core/tests/unit/jobs/test_triggerer_job.py
+++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py
@@ -47,7 +47,9 @@ from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import
InMemorySpanExporter
from opentelemetry.trace.propagation.tracecontext import
TraceContextTextMapPropagator
from pydantic import TypeAdapter
-from sqlalchemy import delete, update
+from sqlalchemy import create_engine, delete, event, func, select, text, update
+from sqlalchemy.exc import OperationalError
+from sqlalchemy.pool import NullPool
from structlog.typing import FilteringBoundLogger
from airflow._shared.timezones import timezone
@@ -685,6 +687,71 @@ def
test_heartbeat_watchdog_disabled_when_threshold_is_zero(supervisor_builder,
perform_heartbeat_mock.assert_called_once()
+def _is_trigger_delete(statement: str) -> bool:
+ return "delete from trigger" in " ".join(statement.lower().replace('"',
"").replace("`", "").split())
+
+
+def
test_clean_unused_retries_failed_batches_in_new_transactions(jobless_supervisor,
session):
+ session.add(Trigger(classpath="airflow.triggers.testing.SuccessTrigger",
kwargs={}))
+ session.commit()
+
+ delete_attempt_count = 0
+ begin_count = 0
+
+ def fail_first_two_deletes(connection, cursor, statement, parameters,
context, executemany):
+ nonlocal delete_attempt_count
+ if _is_trigger_delete(statement):
+ delete_attempt_count += 1
+ if delete_attempt_count <= 2:
+ raise OperationalError(statement, parameters,
RuntimeError("injected cleanup failure"))
+
+ def count_begin(cleanup_session, transaction, connection):
+ nonlocal begin_count
+ begin_count += 1
+
+ event.listen(session.bind, "before_cursor_execute", fail_first_two_deletes)
+ event.listen(session.__class__, "after_begin", count_begin)
+ try:
+ jobless_supervisor.clean_unused()
+ finally:
+ event.remove(session.bind, "before_cursor_execute",
fail_first_two_deletes)
+ event.remove(session.__class__, "after_begin", count_begin)
+
+ assert (delete_attempt_count, begin_count) == (3, 3)
+ session.expire_all()
+ assert session.scalar(select(func.count()).select_from(Trigger)) == 0
+
+
[email protected]("postgres")
+def
test_clean_unused_recovers_after_the_database_connection_drops(jobless_supervisor,
session):
+ session.add(Trigger(classpath="airflow.triggers.testing.SuccessTrigger",
kwargs={}))
+ session.commit()
+
+ killer_engine = create_engine(session.bind.url, poolclass=NullPool)
+ terminated = []
+
+ def terminate_connection_before_delete(connection, cursor, statement,
parameters, context, executemany):
+ if not terminated and _is_trigger_delete(statement):
+ backend_pid =
connection.connection.dbapi_connection.info.backend_pid
+ with killer_engine.connect() as killer:
+ terminated.append(
+ killer.execute(
+ text("SELECT pg_terminate_backend(:pid, 5000)"),
{"pid": backend_pid}
+ ).scalar()
+ )
+
+ event.listen(session.bind, "before_cursor_execute",
terminate_connection_before_delete)
+ try:
+ jobless_supervisor.clean_unused()
+ finally:
+ event.remove(session.bind, "before_cursor_execute",
terminate_connection_before_delete)
+ killer_engine.dispose()
+
+ assert terminated == [True]
+ session.expire_all()
+ assert session.scalar(select(func.count()).select_from(Trigger)) == 0
+
+
def test_metric_tags_default_uses_job_hostname(supervisor_builder):
"""metric_tags() defaults to {"hostname": job.hostname} when a job is
present."""
supervisor = supervisor_builder()
diff --git a/airflow-core/tests/unit/models/test_trigger.py
b/airflow-core/tests/unit/models/test_trigger.py
index 3e673126d03..54a89a8ce69 100644
--- a/airflow-core/tests/unit/models/test_trigger.py
+++ b/airflow-core/tests/unit/models/test_trigger.py
@@ -343,24 +343,26 @@ def
test_clean_unused_rejects_non_positive_batch_size_before_transaction(session
assert commit_count == 0
-@conf_vars({("triggerer", "unreferenced_triggers_cleanup_batch_size"): "2"})
-def test_clean_unused_retries_failed_batch_in_fresh_transactions(session):
- session.add_all(
- [
-
Trigger(classpath=f"airflow.triggers.testing.SuccessTrigger{index}", kwargs={})
- for index in range(5)
- ]
- )
+def test_clean_unused_runs_in_the_caller_transaction(session):
+ session.add(Trigger(classpath="airflow.triggers.testing.SuccessTrigger",
kwargs={}))
session.commit()
- Trigger.clean_unused()
- session.expire_all()
- assert session.scalar(select(func.count()).select_from(Trigger)) == 3
+ Trigger.clean_unused(session=session)
+ assert session.scalar(select(func.count()).select_from(Trigger)) == 0
+ session.rollback()
+
+ assert session.scalar(select(func.count()).select_from(Trigger)) == 1
+
+
+def
test_clean_unused_raises_database_errors_to_the_caller_without_retrying(session):
+ session.add(Trigger(classpath="airflow.triggers.testing.SuccessTrigger",
kwargs={}))
+ session.commit()
+ # The caller's own pending write must survive the cleanup failing.
+ job = Job()
+ session.add(job)
+ session.flush()
delete_attempt_count = 0
- cleanup_sessions = []
- commit_count = 0
- rollback_count = 0
def fail_delete(connection, cursor, statement, parameters, context,
executemany):
nonlocal delete_attempt_count
@@ -369,42 +371,15 @@ def
test_clean_unused_retries_failed_batch_in_fresh_transactions(session):
delete_attempt_count += 1
raise OperationalError(statement, parameters,
RuntimeError("injected cleanup failure"))
- def capture_session(cleanup_session, transaction, connection):
- cleanup_sessions.append(cleanup_session)
-
- def count_commit(committed_session):
- nonlocal commit_count
- commit_count += 1
-
- def count_rollback(rolled_back_session):
- nonlocal rollback_count
- rollback_count += 1
-
event.listen(session.bind, "before_cursor_execute", fail_delete)
- event.listen(session.__class__, "after_begin", capture_session)
- event.listen(session.__class__, "after_commit", count_commit)
- event.listen(session.__class__, "after_rollback", count_rollback)
try:
with pytest.raises(OperationalError, match="injected cleanup failure"):
- Trigger.clean_unused()
+ Trigger.clean_unused(session=session)
finally:
event.remove(session.bind, "before_cursor_execute", fail_delete)
- event.remove(session.__class__, "after_begin", capture_session)
- event.remove(session.__class__, "after_commit", count_commit)
- event.remove(session.__class__, "after_rollback", count_rollback)
- assert len({id(cleanup_session) for cleanup_session in cleanup_sessions})
== 3
- assert (delete_attempt_count, commit_count, rollback_count) == (3, 0, 3)
- session.expire_all()
- assert session.scalar(select(func.count()).select_from(Trigger)) == 3
-
- Trigger.clean_unused()
- session.expire_all()
- assert session.scalar(select(func.count()).select_from(Trigger)) == 1
-
- Trigger.clean_unused()
- session.expire_all()
- assert session.scalar(select(func.count()).select_from(Trigger)) == 0
+ assert delete_attempt_count == 1
+ assert job in session
@patch.object(TriggererCallback, "handle_event")
