From a40e5d4d1da010fb4b7ca9e861a95ddab5816440 Mon Sep 17 00:00:00 2001 From: Iurii Kuchits Date: Sat, 25 Jul 2026 00:04:26 +0200 Subject: [PATCH 1/2] Document get_parsing_context for custom Dag bundle initialize Custom bundle authors need to know when Airflow sets dag_id and task_id during initialize(), and when they must prepare the full bundle instead. --- .../dag-bundles.rst | 23 ++++++ .../execution_time/test_task_runner.py | 75 ++++++++++++++++++- 2 files changed, 96 insertions(+), 2 deletions(-) diff --git a/airflow-core/docs/administration-and-deployment/dag-bundles.rst b/airflow-core/docs/administration-and-deployment/dag-bundles.rst index 354e6ecda1626..2045bd030060a 100644 --- a/airflow-core/docs/administration-and-deployment/dag-bundles.rst +++ b/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,27 @@ 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. diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 0031974b86646..430bcc5d09ee0 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/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") From 50091dc26d60e8bbf8407242adc4c574b3cfb730 Mon Sep 17 00:00:00 2001 From: Iurii Kuchits Date: Sat, 25 Jul 2026 09:18:33 +0200 Subject: [PATCH 2/2] Document get_parsing_context for custom Dag bundle initialize --- airflow-core/docs/administration-and-deployment/dag-bundles.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/airflow-core/docs/administration-and-deployment/dag-bundles.rst b/airflow-core/docs/administration-and-deployment/dag-bundles.rst index 2045bd030060a..b93c6bfac8644 100644 --- a/airflow-core/docs/administration-and-deployment/dag-bundles.rst +++ b/airflow-core/docs/administration-and-deployment/dag-bundles.rst @@ -404,6 +404,7 @@ Other Considerations from airflow.sdk import get_parsing_context + class MyBundle(BaseDagBundle): def initialize(self): dag_id = get_parsing_context().dag_id