Roll back session between retries in Trigger.clean_unused

airflow-core/src/airflow/models/trigger.py

@@ -25,8 +25,10 @@
 from typing import TYPE_CHECKING, Any
 
 from sqlalchemy import ForeignKey, Index, Integer, String, Text, delete, func, or_, select, update
+from sqlalchemy.exc import DBAPIError
 from sqlalchemy.ext.associationproxy import association_proxy
 from sqlalchemy.orm import Mapped, Session, mapped_column, relationship, selectinload
+from sqlalchemy.orm.exc import StaleDataError
 from sqlalchemy.sql.functions import coalesce
 
 from airflow._shared.timezones import timezone
@@ -239,16 +241,27 @@ def clean_unused(cls, *, session: Session = NEW_SESSION) -> None:
         Triggers have a one-to-many relationship to task instances, so we need to clean those up first.
         Afterward we can drop the triggers not referenced by anyone.
         """
-        # Update all task instances with trigger IDs that are not DEFERRED to remove them
+        # Update all task instances with trigger IDs that are not DEFERRED to remove them.
+        # A transient DB disconnect raises DBAPIError and leaves the session in a
+        # "pending rollback" state; without an explicit rollback here, the next retry
+        # attempt reuses that invalidated transaction and raises PendingRollbackError,
+        # which run_with_db_retries() does not match. Mirrors retry_db_transaction in
+        # airflow/utils/retries.py, which cannot be used here because clean_unused is
+        # already decorated with @provide_session.
         for attempt in run_with_db_retries():
             with attempt:
-                session.execute(
-                    update(TaskInstance)
-                    .where(
-                        TaskInstance.state != TaskInstanceState.DEFERRED, TaskInstance.trigger_id.is_not(None)
+                try:
+                    session.execute(
+                        update(TaskInstance)
+                        .where(
+                            TaskInstance.state != TaskInstanceState.DEFERRED,
+                            TaskInstance.trigger_id.is_not(None),
+                        )
+                        .values(trigger_id=None)
                     )
-                    .values(trigger_id=None)
-                )
+                except (DBAPIError, StaleDataError):
+                    session.rollback()
+                    raise
 
         # Get all triggers that have no task instances, assets, or callbacks depending on them and delete them
         ids = select(cls.id).where(

airflow-core/tests/unit/models/test_trigger.py

@@ -20,7 +20,7 @@
 import json
 from collections.abc import AsyncIterator
 from typing import TYPE_CHECKING, Any
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pendulum
 import pytest
@@ -197,6 +197,42 @@ def test_clean_unused(session, dag_maker):
     assert {result.id for result in results} == {trigger1.id, trigger4.id, trigger5.id, trigger6.id}
 
 
+def test_clean_unused_rolls_back_session_on_transient_db_error():
+    """
+    Regression for https://github.com/apache/airflow/issues/74614.
+
+    A transient DB disconnect raises DBAPIError on the UPDATE inside
+    run_with_db_retries(). Without an explicit session rollback between retry
+    attempts, SQLAlchemy's pending-rollback state causes the next attempt to
+    raise PendingRollbackError, which is not matched by the retry policy, so
+    the triggerer process exits. The fix calls session.rollback() before
+    re-raising inside the retry attempt so Tenacity drives the next retry on
+    a clean session.
+    """
+    import tenacity
+    from sqlalchemy.exc import DBAPIError
+
+    mock_session = MagicMock()
+    mock_session.bind.dialect.name = "postgresql"
+    mock_session.execute.side_effect = [
+        DBAPIError("UPDATE task_instance", {}, Exception("connection reset")),
+        MagicMock(),
+        MagicMock(),
+    ]
+
+    fast_retry = tenacity.Retrying(
+        retry=tenacity.retry_if_exception_type((DBAPIError,)),
+        wait=tenacity.wait_none(),
+        stop=tenacity.stop_after_attempt(3),
+        reraise=True,
+    )
+    with patch("airflow.models.trigger.run_with_db_retries", return_value=fast_retry):
+        Trigger.clean_unused(session=mock_session)
+
+    mock_session.rollback.assert_called_once()
+    assert mock_session.execute.call_count >= 2
+
+
 @patch.object(TriggererCallback, "handle_event")
 def test_submit_event(mock_callback_handle_event, session, create_task_instance):
     """