Pass dag_id to Dag bundle construction on task execution

airflow-core/docs/administration-and-deployment/dag-bundles.rst

@@ -383,6 +383,8 @@ In addition to the abstract methods, you may choose to override the following me
 **initialize**
     This method is called before the bundle is first used in the Dag processor or worker. It allows you to perform expensive operations only when the bundle's content is accessed.
 
+    When a worker starts a task, Airflow sets the Dag parsing context before calling ``initialize()``. See **Parsing context** under Other Considerations.
+
 **view_url**
     This method should return a URL as a string to view the bundle on an external system (e.g., a Git repository's web interface).
 
@@ -391,6 +393,28 @@ Other Considerations
 
 - **Versioning**: If your bundle supports versioning, ensure that ``initialize``, ``get_current_version`` and ``refresh`` are implemented to handle version-specific logic.
 
+- **Parsing context**: If expensive bundle setup can be scoped to a single Dag (for example, a sparse checkout of only the paths that Dag needs), worker task startup can do less work in ``initialize()``.
+  Call :py:func:`~airflow.sdk.get_parsing_context` inside ``initialize()`` and branch on ``dag_id``.
+
+  Airflow sets ``dag_id`` and ``task_id`` only on the worker task startup path (``startup()`` wraps ``parse()``, which calls ``initialize()``).
+  The Dag processor, CLI (for example ``airflow dags reserialize`` or ``dag report``), callbacks, and other call sites invoke ``initialize()`` without setting this parsing context, so custom bundles should prepare the full bundle on those paths.
+  ``refresh()`` is not invoked with Dag-specific parsing context either; use it to update the entire bundle, not for per-Dag preparation.
+
+  .. code-block:: python
+
+      from airflow.sdk import get_parsing_context
+
+
+      class MyBundle(BaseDagBundle):
+          def initialize(self):
+              dag_id = get_parsing_context().dag_id
+              if dag_id:
+                  # Worker task startup only: prepare just this Dag (e.g. sparse checkout)
+                  ...
+              else:
+                  # Dag processor, CLI, callbacks, and other paths: prepare the full bundle
+                  ...
+
 - **Concurrency**: Workers may create many bundles simultaneously, and does nothing to serialize calls to the bundle objects. Thus, the bundle class must handle locking if
   that is problematic for the underlying technology. For example, if you are cloning a git repo, the bundle class is responsible for locking to ensure only 1 bundle
   object is cloning at a time. There is a ``lock`` method in the base class that can be used for this purpose, if necessary.

task-sdk/tests/task_sdk/execution_time/test_task_runner.py

@@ -829,6 +829,79 @@ def test_parse_module_in_bundle_root(tmp_path: Path, make_ti_context):
     assert ti.task.dag.dag_id == "dag_name"
 
 
+@pytest.mark.parametrize(
+    ("use_startup", "expected_dag_id", "expected_task_id"),
+    [
+        pytest.param(True, "expected_dag", "my_task", id="startup"),
+        pytest.param(False, None, None, id="parse"),
+    ],
+)
+@mock.patch("airflow.sdk.execution_time.task_runner.get_listener_manager")
+@mock.patch("airflow.dag_processing.dagbag.BundleDagBag")
+@mock.patch("airflow.sdk.execution_time.task_runner.DagBundlesManager")
+def test_bundle_initialize_parsing_context(
+    mock_manager_cls,
+    mock_bag_cls,
+    mock_listener_manager,
+    use_startup,
+    expected_dag_id,
+    expected_task_id,
+    make_ti_context,
+    tmp_path: Path,
+    monkeypatch,
+):
+    from airflow.sdk import get_parsing_context
+    from airflow.sdk.definitions.context import (
+        _AIRFLOW_PARSING_CONTEXT_DAG_ID,
+        _AIRFLOW_PARSING_CONTEXT_TASK_ID,
+    )
+
+    captured: dict[str, str | None] = {}
+    mock_bundle = mock.Mock()
+    mock_bundle.path = tmp_path
+
+    def capture_parsing_context():
+        ctx = get_parsing_context()
+        captured["dag_id"] = ctx.dag_id
+        captured["task_id"] = ctx.task_id
+
+    mock_bundle.initialize.side_effect = capture_parsing_context
+    mock_manager = mock_manager_cls.return_value
+    mock_manager.get_bundle.return_value = mock_bundle
+
+    with DAG("expected_dag") as dag:
+        BaseOperator(task_id="my_task")
+    mock_bag_cls.return_value.dags = {"expected_dag": dag}
+
+    what = StartupDetails(
+        ti=TaskInstance(
+            id=uuid7(),
+            task_id="my_task",
+            dag_id="expected_dag",
+            run_id="test_run",
+            try_number=1,
+            dag_version_id=uuid7(),
+            queue="default",
+        ),
+        dag_rel_path="dag.py",
+        bundle_info=BundleInfo(name="my-bundle", version=None),
+        ti_context=make_ti_context(),
+        start_date=timezone.utcnow(),
+        sentry_integration="",
+    )
+
+    monkeypatch.delenv(_AIRFLOW_PARSING_CONTEXT_DAG_ID, raising=False)
+    monkeypatch.delenv(_AIRFLOW_PARSING_CONTEXT_TASK_ID, raising=False)
+
+    if use_startup:
+        startup(what)
+    else:
+        parse(what, mock.Mock())
+
+    assert captured == {"dag_id": expected_dag_id, "task_id": expected_task_id}
+    mock_bundle.initialize.assert_called_once()
+
+
 def test_verify_bundle_access_raises_when_not_accessible(tmp_path: Path, make_ti_context):
     """Test that _verify_bundle_access raises AirflowException when bundle path is not accessible."""
     from airflow.sdk.execution_time.task_runner import _verify_bundle_access
@@ -6421,7 +6494,6 @@ class TestRegisterDeserializationAllowedClasses:
     """
 
     def test_registers_real_and_mapped_operators(self):
-
         with DAG("walker_dag") as dag:
             # Non-mapped producer: output_type is a plain attribute.
             _WalkerOperator(task_id="real", output_type=_WalkerModelA)
@@ -6437,7 +6509,6 @@ def test_registers_real_and_mapped_operators(self):
         assert _WalkerModelB in registered, "mapped operator output_type not registered"
 
     def test_default_operator_registers_nothing(self):
-
         with DAG("walker_dag_plain") as dag:
             BaseOperator(task_id="plain")