fix(airflow-apis): avoid DagContext autoregister race in build_dag (#28500) (#28744)

Pipeline deploy intermittently failed with
`KeyError: 'unusual_prefix_<hash>_<other dag>'` on Airflow 3.x.

`common.build_dag` used `with DAG(...) as dag:`, which relies on Airflow
3.x's process-global `DagContext.current_autoregister_module_name`. On
`DAG.__exit__` -> `DagContext.pop()` it dereferences
`sys.modules[current_autoregister_module_name]`. When multiple DAG files
parse concurrently in the same api-server process (a deploy racing the
DAG-processor scan / status polls / other deploys), that global is
clobbered and may point to a module already evicted from `sys.modules`,
raising the KeyError on a *different* DAG than the one being deployed.

Build the DAG and attach the task with an explicit `dag=` reference
instead of the context manager. Airflow's BaseOperator only consults
DagContext when `dag is None`, so this bypasses autoregister entirely.
The DAG is already registered into module globals by
`WorkflowFactory.register_dag`, so autoregister was never needed. The
change is version-agnostic (Airflow 2.x and 3.x).

Adds tests/unit/.../test_build_dag.py: a deterministic regression that
poisons DagContext.current_autoregister_module_name and asserts build_dag
still succeeds (fails with the KeyError before this fix).

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit 079761f)

Author: 2 people

openmetadata-airflow-apis/openmetadata_managed_apis/workflows/ingestion/common.py

@@ -425,26 +425,34 @@ def build_dag(
     :return: DAG
     """
 
-    with DAG(**build_dag_configs(ingestion_pipeline)) as dag:
-        # Initialize with random UUID4. Will be used by the callback instead of
-        # generating it inside the Workflow itself.
-        workflow_config.pipelineRunId = Uuid(uuid.uuid4())
-
-        CustomPythonOperator(
-            task_id=task_name,
-            python_callable=workflow_fn,
-            op_kwargs={
-                "workflow_config": workflow_config,
-            },
-            # There's no need to retry if we have had an error. Wait until the next schedule or manual rerun.
-            retries=ingestion_pipeline.airflowConfig.retries or 0,
-            # each DAG will call its own OpenMetadataWorkflowConfig
-            on_failure_callback=partial(send_failed_status_callback, workflow_config),
-            # Add tag and ownership to easily identify DAGs generated by OM
-            owner=ingestion_pipeline.owners.root[0].name
-            if (ingestion_pipeline.owners and ingestion_pipeline.owners.root)
-            else "openmetadata",
-            params=params,
-        )
+    # Build the DAG and attach the task with an explicit `dag=` reference instead of
+    # the `with DAG(...) as dag:` context manager. The context manager relies on
+    # Airflow 3.x's process-global DagContext autoregister, which races when multiple
+    # DAG files are parsed concurrently in the same process and raises a KeyError on
+    # __exit__ (see issue #28500). The DAG is registered into the module globals
+    # explicitly by WorkflowFactory.register_dag, so autoregister is not needed here.
+    dag = DAG(**build_dag_configs(ingestion_pipeline))
+
+    # Initialize with random UUID4. Will be used by the callback instead of
+    # generating it inside the Workflow itself.
+    workflow_config.pipelineRunId = Uuid(uuid.uuid4())
+
+    CustomPythonOperator(
+        task_id=task_name,
+        python_callable=workflow_fn,
+        op_kwargs={
+            "workflow_config": workflow_config,
+        },
+        # There's no need to retry if we have had an error. Wait until the next schedule or manual rerun.
+        retries=ingestion_pipeline.airflowConfig.retries or 0,
+        # each DAG will call its own OpenMetadataWorkflowConfig
+        on_failure_callback=partial(send_failed_status_callback, workflow_config),
+        # Add tag and ownership to easily identify DAGs generated by OM
+        owner=ingestion_pipeline.owners.root[0].name
+        if (ingestion_pipeline.owners and ingestion_pipeline.owners.root)
+        else "openmetadata",
+        params=params,
+        dag=dag,
+    )
 
-        return dag
+    return dag

openmetadata-airflow-apis/tests/unit/ingestion_pipeline/test_build_dag.py

@@ -0,0 +1,140 @@
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+Regression tests for ``build_dag`` (issue #28500).
+
+``build_dag`` must NOT rely on Airflow 3.x's ``DagContext`` autoregister context
+manager. That context manager keeps a process-global
+``current_autoregister_module_name`` and, on ``DAG.__exit__`` -> ``DagContext.pop()``,
+dereferences ``sys.modules[current_autoregister_module_name]``. When multiple DAG
+files are parsed concurrently in the same process that global gets clobbered and the
+referenced module may already be evicted from ``sys.modules``, raising a ``KeyError``.
+
+Building the DAG with an explicit ``dag=`` reference (instead of ``with DAG(...)``)
+bypasses ``DagContext`` entirely, so these tests assert both that the task attaches
+correctly and that a poisoned ``DagContext`` no longer breaks the build.
+"""
+
+import uuid
+
+import pytest
+
+from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
+    OpenMetadataConnection,
+)
+from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
+    AirflowConfig,
+    IngestionPipeline,
+    PipelineType,
+)
+from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline import (
+    DatabaseServiceMetadataPipeline,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    OpenMetadataWorkflowConfig,
+    Sink,
+    Source,
+    SourceConfig,
+    WorkflowConfig,
+)
+from metadata.generated.schema.security.client.openMetadataJWTClientConfig import (
+    OpenMetadataJWTClientConfig,
+)
+from metadata.generated.schema.type.entityReference import EntityReference
+from openmetadata_managed_apis.workflows.ingestion.common import build_dag
+
+TASK_NAME = "ingestion_task"
+
+
+def _noop_workflow_fn(*_, **__):
+    return None
+
+
+def _server_config() -> OpenMetadataConnection:
+    return OpenMetadataConnection(
+        hostPort="http://localhost:8585/api",
+        authProvider="openmetadata",
+        securityConfig=OpenMetadataJWTClientConfig(jwtToken="x.y.z"),
+    )
+
+
+def _ingestion_pipeline(name: str) -> IngestionPipeline:
+    return IngestionPipeline(
+        name=name,
+        pipelineType=PipelineType.metadata,
+        fullyQualifiedName=f"svc.{name}",
+        sourceConfig=SourceConfig(config=DatabaseServiceMetadataPipeline()),
+        openMetadataServerConnection=_server_config(),
+        airflowConfig=AirflowConfig(),
+        service=EntityReference(id=str(uuid.uuid4()), type="databaseService", name="svc"),
+    )
+
+
+def _workflow_config(name: str) -> OpenMetadataWorkflowConfig:
+    return OpenMetadataWorkflowConfig(
+        source=Source(
+            type="mysql",
+            serviceName="svc",
+            sourceConfig=SourceConfig(config=DatabaseServiceMetadataPipeline()),
+        ),
+        sink=Sink(type="metadata-rest", config={}),
+        workflowConfig=WorkflowConfig(openMetadataServerConfig=_server_config()),
+        ingestionPipelineFQN=f"svc.{name}",
+    )
+
+
+def _build(name: str):
+    return build_dag(
+        task_name=TASK_NAME,
+        ingestion_pipeline=_ingestion_pipeline(name),
+        workflow_config=_workflow_config(name),
+        workflow_fn=_noop_workflow_fn,
+    )
+
+
+def test_build_dag_attaches_task():
+    dag = _build("attach_dag")
+
+    assert dag.dag_id == "attach_dag"
+    assert len(dag.tasks) == 1
+    assert TASK_NAME in dag.task_dict
+    assert dag.task_dict[TASK_NAME].dag is dag
+
+
+def test_build_dag_survives_poisoned_dag_context():
+    """
+    Simulate the concurrent-parse race deterministically: point the process-global
+    autoregister module name at a module that is not in ``sys.modules``. Before the
+    fix this raised ``KeyError`` on ``with DAG(...) as dag:`` __exit__; the explicit
+    ``dag=`` build must be immune to it.
+    """
+    dag_context = pytest.importorskip("airflow.sdk.definitions._internal.contextmanager").DagContext
+
+    if not hasattr(dag_context, "current_autoregister_module_name"):
+        pytest.skip("DagContext autoregister not present on this Airflow version")
+
+    original = dag_context.current_autoregister_module_name
+    dag_context.current_autoregister_module_name = f"missing_module_{uuid.uuid4().hex}"
+    try:
+        dag = _build("poisoned_dag")
+    finally:
+        dag_context.current_autoregister_module_name = original
+
+    assert dag.dag_id == "poisoned_dag"
+    assert TASK_NAME in dag.task_dict
+
+
+def test_build_dag_does_not_use_dag_context_stack():
+    dag_context = pytest.importorskip("airflow.sdk.definitions._internal.contextmanager").DagContext
+
+    _build("stack_dag")
+
+    assert dag_context.get_current_dag() is None