Fix dag_tag unique-constraint crash when multiple dag processors parse concurrently

airflow-core/newsfragments/70002.bugfix.rst

@@ -0,0 +1 @@
+Fix dag processor crash when multiple dag processors run in HA and parse the same file concurrently. Both processors could INSERT the same ``dag_tag`` row, raising a unique-constraint violation (surfacing as ``PendingRollbackError``) that crashed the processor. Tags are now inserted with a dialect-aware conflict-ignoring statement (``ON CONFLICT DO NOTHING`` on PostgreSQL/SQLite, a no-op ``ON DUPLICATE KEY UPDATE`` on MySQL).

airflow-core/src/airflow/dag_processing/collection.py

@@ -215,34 +215,53 @@ def calculate(cls, dag: LazyDeserializedDAG, *, session: Session) -> Self:
         return cls(latest_run, active_run_counts.get(dag.dag_id, 0))
 
 
-def _update_dag_tags(tag_names: set[str], dm: DagModel, *, session: Session) -> None:
-    orm_tags = {t.name: t for t in dm.tags}
-    tags_to_delete = []
-    for name, orm_tag in orm_tags.items():
-        if name not in tag_names:
-            session.delete(orm_tag)
-            tags_to_delete.append(orm_tag)
-
-    tags_to_add = tag_names.difference(orm_tags)
-    if tags_to_delete:
-        # Remove deleted tags from the collection to keep it in sync
-        for tag in tags_to_delete:
-            dm.tags.remove(tag)
-
-        # Check if there's a potential case-only rename on MySQL (e.g., 'tag' -> 'TAG').
-        # MySQL uses case-insensitive collation for the (name, dag_id) primary key by default,
-        # which can cause duplicate key errors when renaming tags with only case changes.
-        if get_dialect_name(session) == "mysql":
-            orm_tags_lower = {name.lower(): name for name in orm_tags}
-            has_case_only_change = any(tag.lower() in orm_tags_lower for tag in tags_to_add)
-
-            if has_case_only_change:
-                # Force DELETE operations to execute before INSERT operations.
-                session.flush()
-                # Refresh the tags relationship from the database to reflect the deletions.
-                session.expire(dm, ["tags"])
-
-    dm.tags.extend(DagTag(name=name, dag_id=dm.dag_id) for name in tags_to_add)
+def _update_dag_tags(updates: Iterable[tuple[DagModel, set[str]]], *, session: Session) -> None:
+    rows_to_add: list[dict[str, str]] = []
+    dms_to_expire: list[DagModel] = []
+    for dm, tag_names in updates:
+        orm_tags = {t.name: t for t in dm.tags}
+        for name, orm_tag in orm_tags.items():
+            if name not in tag_names:
+                session.delete(orm_tag)
+                # Remove the deleted tag from the collection to keep it in sync
+                dm.tags.remove(orm_tag)
+        if tags_to_add := tag_names.difference(orm_tags):
+            rows_to_add.extend({"name": name, "dag_id": dm.dag_id} for name in tags_to_add)
+            dms_to_expire.append(dm)
+
+    if not rows_to_add:
+        return
+
+    # Flush so the DELETE operations above hit the database before the INSERT.
+    # This is required for case-only renames (e.g. 'tag' -> 'TAG') on MySQL,
+    # where the (name, dag_id) primary key is case-insensitive by default, and
+    # it also guarantees newly-added DagModel rows exist to satisfy the
+    # dag_tag foreign key.
+    session.flush()
+
+    # Insert with a conflict-ignoring statement instead of via the ORM
+    # relationship: multiple dag processors running in HA can parse the same
+    # file concurrently, and a plain INSERT of a tag that another processor
+    # has just committed raises a unique-constraint violation that poisons the
+    # session and crashes the processor.
+    if (dialect_name := get_dialect_name(session)) == "postgresql":
+        from sqlalchemy.dialects.postgresql import insert as postgresql_insert
+
+        stmt: Any = postgresql_insert(DagTag).on_conflict_do_nothing()
+    elif dialect_name == "mysql":
+        from sqlalchemy.dialects.mysql import insert as mysql_insert
+
+        # MySQL does not support "do nothing"; this updates the row in
+        # conflict with its own value to achieve the same idea.
+        stmt = mysql_insert(DagTag).on_duplicate_key_update(name=DagTag.name)
+    else:
+        from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+        stmt = sqlite_insert(DagTag).on_conflict_do_nothing()
+    session.execute(stmt, rows_to_add)
+    # The new rows bypassed the ORM relationships; reload them on next access.
+    for dm in dms_to_expire:
+        session.expire(dm, ["tags"])
 
 
 def _update_dag_owner_links(dag_owner_links: dict[str, str], dm: DagModel, *, session: Session) -> None:
@@ -629,6 +648,7 @@ def update_dags(
         session: Session,
     ) -> None:
         # we exclude backfill from active run counts since their concurrency is separate
+        tag_updates: list[tuple[DagModel, set[str]]] = []
         for dag_id, dm in sorted(orm_dags.items()):
             run_info = _RunInfo.calculate(dag=self.dags[dag_id], session=session)
             dag = self.dags[dag_id]
@@ -704,14 +724,18 @@ def update_dags(
             # FIXME: STORE NEW REFERENCES.
 
             if dag.tags:
-                _update_dag_tags(set(dag.tags), dm, session=session)
+                # Collected here and applied in one batch below so the whole
+                # bulk write costs a constant number of extra queries.
+                tag_updates.append((dm, set(dag.tags)))
             else:  # Optimization: no references at all, just clear everything.
                 dm.tags = []
             if dag.owner_links:
                 _update_dag_owner_links(dag.owner_links, dm, session=session)
             else:  # Optimization: no references at all, just clear everything.
                 dm.dag_owner_links = []
 
+        _update_dag_tags(tag_updates, session=session)
+
     def update_dag_asset_expression(
         self,
         *,

airflow-core/tests/unit/dag_processing/test_collection.py

@@ -27,7 +27,7 @@
 from unittest.mock import patch
 
 import pytest
-from sqlalchemy import delete, func, inspect as sa_inspect, select
+from sqlalchemy import delete, func, insert, inspect as sa_inspect, select
 from sqlalchemy.exc import OperationalError, SAWarning
 
 import airflow.dag_processing.collection
@@ -1442,7 +1442,7 @@ class TestUpdateDagTags:
     @pytest.fixture(autouse=True)
     def setup_teardown(self, session):
         yield
-        session.execute(delete(DagModel).where(DagModel.dag_id == "test_dag"))
+        session.execute(delete(DagModel).where(DagModel.dag_id.in_(("test_dag", "test_dag_2"))))
         session.commit()
 
     @pytest.mark.parametrize(
@@ -1461,11 +1461,64 @@ def test_update_dag_tags(self, testing_dag_bundle, session, initial_tags, new_ta
         session.add(dag_model)
         session.commit()
 
-        _update_dag_tags(new_tags, dag_model, session=session)
+        _update_dag_tags([(dag_model, new_tags)], session=session)
         session.commit()
 
         assert {t.name for t in dag_model.tags} == expected_tags
 
+    def test_update_dag_tags_ignores_concurrently_inserted_tag(self, testing_dag_bundle, session):
+        """
+        A tag inserted by another dag processor must not crash the update.
+
+        With multiple dag processors running in HA, two processors can parse the
+        same file concurrently: both read the (empty) tags relationship, then both
+        insert the same tag row. The plain ORM INSERT raised a unique-constraint
+        violation (e.g. ``UniqueViolation`` on ``dag_tag_pkey``), poisoning the
+        session and crashing the processor. The insert must tolerate the conflict
+        and keep the remaining tags.
+        """
+        dag_model = DagModel(dag_id="test_dag", bundle_name="testing")
+        session.add(dag_model)
+        session.commit()
+
+        # Load the (empty) tags relationship, as update_dags does...
+        assert dag_model.tags == []
+        # ...then simulate another dag processor committing the same tag in the
+        # meantime, bypassing the already-loaded relationship.
+        session.execute(insert(DagTag).values(name="shared", dag_id="test_dag"))
+
+        _update_dag_tags([(dag_model, {"shared", "extra"})], session=session)
+        session.commit()
+
+        assert {t.name for t in dag_model.tags} == {"shared", "extra"}
+
+    def test_update_dag_tags_batches_multiple_dags(self, testing_dag_bundle, session):
+        """Tag changes for several dags are applied in one batched statement."""
+        dag_model_1 = DagModel(dag_id="test_dag", bundle_name="testing")
+        dag_model_1.tags = [DagTag(name="stale", dag_id="test_dag")]
+        dag_model_2 = DagModel(dag_id="test_dag_2", bundle_name="testing")
+        session.add_all([dag_model_1, dag_model_2])
+        session.commit()
+
+        _update_dag_tags(
+            [(dag_model_1, {"fresh"}), (dag_model_2, {"shared", "other"})],
+            session=session,
+        )
+        session.commit()
+
+        assert {t.name for t in dag_model_1.tags} == {"fresh"}
+        assert {t.name for t in dag_model_2.tags} == {"shared", "other"}
+
+    def test_update_dag_tags_flushes_new_dag_before_insert(self, testing_dag_bundle, session):
+        """A not-yet-flushed DagModel is flushed first so the dag_tag FK is satisfied."""
+        dag_model = DagModel(dag_id="test_dag", bundle_name="testing")
+        session.add(dag_model)
+
+        _update_dag_tags([(dag_model, {"new"})], session=session)
+        session.commit()
+
+        assert {t.name for t in dag_model.tags} == {"new"}
+
 
 @pytest.mark.db_test
 class TestPartitionMapperInfoSync: