kaxil commented on code in PR #74586: URL: https://github.com/apache/airflow/pull/74586#discussion_r4250620784
########## task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py: ########## @@ -0,0 +1,404 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# 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. + +""" +Build a validated :class:`~.models.DagDocument` into a :class:`~airflow.sdk.DAG`. + +The parser produces a validated document; this turns it into real DAG + operator +objects. Ready-made operators are instantiated from their ``uses`` import path; +tasks with ``run`` are built into stub operators and handled by AIP-108 task +handlers at execution time. + +Templating is explicit, realized importer-side (a la AIP-80). The format is +literal by default; only values the author marked render as templates. Value +grammar markers are only accepted on fields declared in ``template_fields``. + +Tasks are built in topological order so a task referenced by an upstream's XCom +already exists when the downstream operator is constructed. +""" + +from __future__ import annotations + +import datetime +import functools +import graphlib +import inspect +import keyword +import linecache +import re +from typing import TYPE_CHECKING, Any, NamedTuple, cast + +import attrs + +from airflow.sdk import DAG, BaseOperator, XComArg, literal, task as task_decorators +from airflow.sdk._shared.module_loading import import_string +from airflow.sdk.importers.yaml_importer.models import ( + ConstRef, + FileRef, + Literal, + TemplateRef, + TimetableSchedule, + XComRef, +) + +if TYPE_CHECKING: + from collections.abc import Callable + + from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 + from airflow.sdk.importers.base import DagDefinition + from airflow.sdk.importers.yaml_importer.models import DagDocument, Task, Value + + +class YamlDagBuildError(ValueError): + """A validated document could not be built into a DAG.""" + + +_BASE_TASK_ARGS = frozenset(BaseOperator.__init__._BaseOperatorMeta__param_names) # type: ignore[attr-defined] +_UNSUPPORTED_TASK_ARGS = frozenset({"task_id", "dag", "task_group", "default_args"}) + + [email protected](kw_only=True) +class DagBuilder: + """ + Build a :class:`DagDocument` into a :class:`DAG`. + + This holds the state a single build shares. The state belongs to one build + and is discarded with the instance. + """ + + _definition: DagDefinition + _bundle: BaseDagBundle + _operators: dict[str, BaseOperator] = attrs.field(factory=dict, init=False) + _stub_callables: dict[tuple[str, ...], Callable[..., None]] = attrs.field(factory=dict, init=False) + + def build(self, doc: DagDocument) -> DAG: + tasks = {t.id_: t for t in doc.tasks} + extras = _coerce_temporal(doc.model_extra or {}, DAG) + with DAG(dag_id=doc.dag_id, schedule=_build_schedule(doc.schedule), **extras) as dag: + for task_id in _topological_order(doc.tasks): + self._operators[task_id] = self._build_operator(tasks[task_id]) + + for task in doc.tasks: + for upstream in sorted(task.get_dependency_ids()): + self._operators[upstream] >> self._operators[task.id_] + dag.bundle_name = self._bundle.name + dag.fileloc = repr(self._definition) + dag.relative_fileloc = self._definition.get_relative_loc(self._bundle.path) + return dag + + def _build_stub_task(self, task: Task, args: dict[str, Value]) -> BaseOperator: + """ + Build a stub from a "run:" task. + + The task body is a registered SDK task handler resolved at run time; the + run mapping binds as that handler's arguments, where XCom references + resolve natively. The stub is fed a callable whose parameters are the + run argument names, so serialization can build and later validate + (AIP-108) the handler's argument spec. + """ + run_kwargs: dict[str, Any] = {} + for name, value in args.items(): + if not name.isidentifier() or keyword.iskeyword(name): + raise YamlDagBuildError(f"task {task.id_!r}: run argument {name!r} is not a valid identifier") + _reject_unbindable_run_marker(task.id_, value) + run_kwargs[name] = self._lower(value, protect=False) + decorator = task_decorators.stub( + task_id=task.id_, **_coerce_temporal(task.model_extra or {}, BaseOperator) + ) + return decorator(self._stub_callable(tuple(run_kwargs)))(**run_kwargs).operator + + def _build_operator(self, task: Task) -> BaseOperator: + """Instantiate *task* as an operator.""" + if reserved := sorted( + k for k in (*task.with_, *(task.model_extra or {})) if k in _UNSUPPORTED_TASK_ARGS + ): + raise YamlDagBuildError( + f"task {task.id_!r}: {reserved} cannot be set on a task; the task id comes from 'id', and " + "dag/task_group/default_args are managed by the importer" + ) + if misplaced := sorted(k for k in task.with_ if k in _BASE_TASK_ARGS): + raise YamlDagBuildError( + f"task {task.id_!r}: {misplaced} are task-level arguments; set them at task level, not under 'with'" + ) + if unknown := sorted(k for k in (task.model_extra or {}) if k not in _BASE_TASK_ARGS): + raise YamlDagBuildError( + f"task {task.id_!r}: {unknown} are not task-level (BaseOperator) arguments; " + "operator-specific arguments belong under 'with'" + ) + + if (args := task.run) is not None: Review Comment: A `run:` task never reads `task.with_`, so `{id: t, run: {n: 1}, with: {region: eu}}` builds without error and `region` is gone. A `$x` under that `with:` still adds the upstream edge but passes nothing. The likelier way in is `extends` from a template that carries `with:`, since `_merge_task` merges it into the task. Could this raise `YamlDagBuildError` when a `run:` task has a non-empty `with:`? ########## task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py: ########## @@ -0,0 +1,404 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# 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. + +""" +Build a validated :class:`~.models.DagDocument` into a :class:`~airflow.sdk.DAG`. + +The parser produces a validated document; this turns it into real DAG + operator +objects. Ready-made operators are instantiated from their ``uses`` import path; +tasks with ``run`` are built into stub operators and handled by AIP-108 task +handlers at execution time. + +Templating is explicit, realized importer-side (a la AIP-80). The format is +literal by default; only values the author marked render as templates. Value +grammar markers are only accepted on fields declared in ``template_fields``. + +Tasks are built in topological order so a task referenced by an upstream's XCom +already exists when the downstream operator is constructed. +""" + +from __future__ import annotations + +import datetime +import functools +import graphlib +import inspect +import keyword +import linecache +import re +from typing import TYPE_CHECKING, Any, NamedTuple, cast + +import attrs + +from airflow.sdk import DAG, BaseOperator, XComArg, literal, task as task_decorators +from airflow.sdk._shared.module_loading import import_string +from airflow.sdk.importers.yaml_importer.models import ( + ConstRef, + FileRef, + Literal, + TemplateRef, + TimetableSchedule, + XComRef, +) + +if TYPE_CHECKING: + from collections.abc import Callable + + from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 + from airflow.sdk.importers.base import DagDefinition + from airflow.sdk.importers.yaml_importer.models import DagDocument, Task, Value + + +class YamlDagBuildError(ValueError): + """A validated document could not be built into a DAG.""" + + +_BASE_TASK_ARGS = frozenset(BaseOperator.__init__._BaseOperatorMeta__param_names) # type: ignore[attr-defined] +_UNSUPPORTED_TASK_ARGS = frozenset({"task_id", "dag", "task_group", "default_args"}) + + [email protected](kw_only=True) +class DagBuilder: + """ + Build a :class:`DagDocument` into a :class:`DAG`. + + This holds the state a single build shares. The state belongs to one build + and is discarded with the instance. + """ + + _definition: DagDefinition + _bundle: BaseDagBundle + _operators: dict[str, BaseOperator] = attrs.field(factory=dict, init=False) + _stub_callables: dict[tuple[str, ...], Callable[..., None]] = attrs.field(factory=dict, init=False) + + def build(self, doc: DagDocument) -> DAG: + tasks = {t.id_: t for t in doc.tasks} + extras = _coerce_temporal(doc.model_extra or {}, DAG) Review Comment: Dag-level extras never go through the value grammar, so `params: {p: {$f: README.md}}` reaches `DAG()` as a Param whose default is the literal dict `{'$f': 'README.md'}`, with no error. Rejecting `$`-prefixed keys in the Dag- and task-level extras, the way `Literal._reject_reserved_keys` does under `with:`, would catch the typo-shaped case. ########## task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py: ########## @@ -0,0 +1,552 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# 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. + +from __future__ import annotations + +import datetime +import inspect +import textwrap +import zipfile +from types import SimpleNamespace + +import pytest + +from airflow.sdk import BaseOperator, XComArg +from airflow.sdk.definitions._internal.templater import LiteralValue +from airflow.sdk.definitions.timetables.simple import OnceTimetable +from airflow.sdk.importers import ZipMemberDagDefinition +from airflow.sdk.importers.base import FilesystemDagDefinition +from airflow.sdk.importers.yaml_importer.importer import YamlDagImporter + +HEAD_SCHEMA = "https://airflow.apache.org/schemas/dag/2026-10-30.json"; + + +class _RecordingOperator(BaseOperator): + """A trivial operator with one templated field (``payload``) and one plain field (``note``).""" + + template_fields = ("payload",) + template_ext = (".sql",) # non-empty so the importer emptying it to () is observable, not vacuous + + def __init__(self, *, payload=None, note=None, **kwargs): + super().__init__(**kwargs) + self.payload = payload + self.note = note + + def execute(self, context): + return self.payload + + +OP = f"{__name__}.{_RecordingOperator.__name__}" + + +class _DefaultTemplateOperator(BaseOperator): + """Templates ``preset`` by default (set in __init__, not by the author) plus author-set ``payload``.""" + + template_fields = ("payload", "preset") + + def __init__(self, *, payload=None, preset="{{ ds }}", **kwargs): + super().__init__(**kwargs) + self.payload = payload + self.preset = preset + + def execute(self, context): + return self.payload + + +OP_DEFAULT_TEMPLATE = f"{__name__}.{_DefaultTemplateOperator.__name__}" + + +class _TemporalOperator(BaseOperator): + """An operator with a datetime-typed argument (``cutoff``).""" + + def __init__(self, *, cutoff: datetime.datetime | None = None, **kwargs): + super().__init__(**kwargs) + self.cutoff = cutoff + + +OP_TEMPORAL = f"{__name__}.{_TemporalOperator.__name__}" + + [email protected] +def mock_bundle(tmp_path): + bundle_dir = tmp_path / "bundle" + bundle_dir.mkdir(parents=True, exist_ok=True) + return SimpleNamespace(name="test_bundle", path=bundle_dir) + + +def _def(bundle, name: str, body: str) -> FilesystemDagDefinition: + path = bundle.path / name + path.write_text(textwrap.dedent(body)) + return FilesystemDagDefinition(path=path) + + +def _doc(tasks: str, *, dag_id: str = "d", extra: str = "") -> str: + return f"$schema: {HEAD_SCHEMA}\ndag_id: {dag_id}\n{extra}tasks:\n{textwrap.indent(textwrap.dedent(tasks), ' ')}" + + +def test_can_handle(): + importer = YamlDagImporter() + assert importer.can_handle("dag.yaml") + assert importer.can_handle("dag.yml") + assert not importer.can_handle("dag.py") + assert not importer.can_handle("dag.json") + + +def test_get_source_code(mock_bundle): + definition = _def(mock_bundle, "d.yaml", f"$schema: {HEAD_SCHEMA}\ndag_id: d\ntasks: []\n") + src = YamlDagImporter().get_source_code(definition) + assert src.language == "yaml" + assert "$schema" in src.source_code + + +def test_might_contain_dag(mock_bundle): + importer = YamlDagImporter() + has = _def(mock_bundle, "has.yaml", f"$schema: {HEAD_SCHEMA}\ndag_id: d\ntasks: []\n") + nope = _def(mock_bundle, "nope.yaml", "just: config\n") + # A $schema that is not the canonical Dag URL (e.g. a JSON-Schema file) must not be picked up. + other = _def(mock_bundle, "schema.yaml", "$schema: https://json-schema.org/draft/2020-12/schema\n";) + # The canonical URL merely mentioned, without a $schema declaration, must not be picked up either. + mention = _def(mock_bundle, "mention.yaml", f"note: see {HEAD_SCHEMA}\n") + assert importer.might_contain_dag(has, safe_mode=True) + assert not importer.might_contain_dag(nope, safe_mode=True) + assert not importer.might_contain_dag(other, safe_mode=True) + assert not importer.might_contain_dag(mention, safe_mode=True) + assert importer.might_contain_dag(nope, safe_mode=False) # sniff skipped + + +def test_list_dag_definitions_skips_non_dag_yaml(mock_bundle): + _def(mock_bundle, "dag.yaml", f"$schema: {HEAD_SCHEMA}\ndag_id: d\ntasks: []\n") + _def(mock_bundle, "config.yaml", "just: config\n") # no $schema -> skipped + _def(mock_bundle, "notes.txt", "hello") # wrong extension + names = {d.path.name for d in YamlDagImporter().list_dag_definitions(mock_bundle)} + assert names == {"dag.yaml"} + + +def test_import_builds_operator_dag(mock_bundle): + definition = _def( + mock_bundle, + "dag.yaml", + _doc( + f""" + - {{id: a, uses: {OP}, with: {{payload: hi}}}} + - {{id: b, uses: {OP}, with: {{payload: {{$x: a}}}}}} + """, + dag_id="pipeline", + extra="catchup: false\n", + ), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + (dag,) = result.dags + assert dag.dag_id == "pipeline" + assert set(dag.task_ids) == {"a", "b"} + assert dag.catchup is False # Dag-level attribute passed through + assert dag.bundle_name == "test_bundle" + a, b = dag.get_task("a"), dag.get_task("b") + assert a.payload == "hi" # literal + assert isinstance(b.payload, XComArg) # $x -> XComArg + assert b.upstream_task_ids == {"a"} # XCom wired the dependency + assert "payload" in b.template_fields # $x opted the field into rendering + assert "payload" not in a.template_fields # a's unmarked literal was dropped from rendering + + +def test_import_wires_needs(mock_bundle): + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f""" + - {{id: a, uses: {OP}, with: {{payload: x}}}} + - {{id: b, uses: {OP}, with: {{payload: y}}, needs: [a]}} + """), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + assert result.dags[0].get_task("b").upstream_task_ids == {"a"} + + +def test_import_bodyless_task_becomes_empty_operator(mock_bundle): + definition = _def(mock_bundle, "dag.yaml", _doc("- {id: a}")) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + assert type(result.dags[0].get_task("a")).__name__ == "EmptyOperator" + + +def test_import_bodyless_task_with_args_errors(mock_bundle): + definition = _def(mock_bundle, "dag.yaml", _doc("- {id: a, with: {x: 1}}")) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert any( + "Invalid arguments were passed to EmptyOperator" in e.message and "'x'" in e.message + for e in result.errors + ) + + +def test_import_schedule_cron(mock_bundle): + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f"- {{id: a, uses: {OP}, with: {{payload: x}}}}", extra="schedule: '0 3 * * *'\n"), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + assert result.dags[0].schedule == "0 3 * * *" + + +def test_import_schedule_timetable(mock_bundle): + definition = _def( + mock_bundle, + "dag.yaml", + _doc( + f"- {{id: a, uses: {OP}, with: {{payload: x}}}}", + extra="schedule: {uses: airflow.sdk.definitions.timetables.simple.OnceTimetable}\n", + ), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + assert isinstance(result.dags[0].timetable, OnceTimetable) + + +def test_import_keyed_xcom_on_uses_task(mock_bundle): + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f""" + - {{id: a, uses: {OP}, with: {{payload: x}}}} + - {{id: b, uses: {OP}, with: {{payload: {{$x: {{task: a, key: out}}}}}}}} + """), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + b = result.dags[0].get_task("b") + assert isinstance(b.payload, XComArg) + assert b.payload.key == "out" + assert b.upstream_task_ids == {"a"} + + [email protected]( + ("value", "hour"), + [("2024-01-01", 0), ("'2024-01-01'", 0), ("2024-01-01T03:30:00", 3)], # YAML date, ISO string, datetime +) +def test_import_coerces_temporal_start_date(mock_bundle, value, hour): + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f"- {{id: a, uses: {OP}, with: {{payload: x}}}}", extra=f"start_date: {value}\n"), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + start = result.dags[0].start_date + assert isinstance(start, datetime.datetime) # promoted from date/string + assert start.tzinfo is not None # localized by DAG() + assert (start.year, start.month, start.day, start.hour) == (2024, 1, 1, hour) + + +def test_import_does_not_coerce_nontemporal_arg(mock_bundle): + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f"- {{id: a, uses: {OP}, owner: '2024-01-01', with: {{payload: x}}}}"), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + assert result.dags[0].get_task("a").owner == "2024-01-01" + + +def test_import_coerces_temporal_with_arg(mock_bundle): + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f"- {{id: a, uses: {OP_TEMPORAL}, with: {{cutoff: 2024-01-01}}}}"), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + cutoff = result.dags[0].get_task("a").cutoff + assert isinstance(cutoff, datetime.datetime) + assert (cutoff.year, cutoff.month, cutoff.day) == (2024, 1, 1) + + +def test_import_code_task_builds_stub_operator(mock_bundle): + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f""" + - {{id: a, uses: {OP}, with: {{payload: hi}}}} + - {{id: b, run: {{value: {{$x: a}}, label: done}}, queue: handlers}} + """), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + b = result.dags[0].get_task("b") + assert type(b).__name__ == "_StubOperator" + assert getattr(b, "is_stub", False) is True # the real stub operator Airflow serializes + assert b.queue == "handlers" # task-level attribute passed through + assert isinstance(b.op_kwargs["value"], XComArg) # $x -> XComArg, bound as a handler argument + assert b.op_kwargs["label"] == "done" # literal bound verbatim (run args are never rendered) + assert b.upstream_task_ids == {"a"} # $x wired the dependency + # The stub is fed a callable whose ordered parameters are the run argument names, so core + # serialization can build the handler arg-binding spec (AIP-108) from a real signature. + assert list(inspect.signature(b.python_callable).parameters) == ["value", "label"] + + +def test_import_mixed_container_keeps_literal_siblings_literal(mock_bundle): Review Comment: Only a plain string sits next to the `$x` here, so nothing checks that a `$const` or `$f` sibling stays wrapped. Those are the `ConstRef` and `FileRef` branches of `_lower` that call `_protect`; adding both to this payload would cover them. ########## task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py: ########## @@ -0,0 +1,404 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# 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. + +""" +Build a validated :class:`~.models.DagDocument` into a :class:`~airflow.sdk.DAG`. + +The parser produces a validated document; this turns it into real DAG + operator +objects. Ready-made operators are instantiated from their ``uses`` import path; +tasks with ``run`` are built into stub operators and handled by AIP-108 task +handlers at execution time. + +Templating is explicit, realized importer-side (a la AIP-80). The format is +literal by default; only values the author marked render as templates. Value +grammar markers are only accepted on fields declared in ``template_fields``. + +Tasks are built in topological order so a task referenced by an upstream's XCom +already exists when the downstream operator is constructed. +""" + +from __future__ import annotations + +import datetime +import functools +import graphlib +import inspect +import keyword +import linecache +import re +from typing import TYPE_CHECKING, Any, NamedTuple, cast + +import attrs + +from airflow.sdk import DAG, BaseOperator, XComArg, literal, task as task_decorators +from airflow.sdk._shared.module_loading import import_string +from airflow.sdk.importers.yaml_importer.models import ( + ConstRef, + FileRef, + Literal, + TemplateRef, + TimetableSchedule, + XComRef, +) + +if TYPE_CHECKING: + from collections.abc import Callable + + from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 + from airflow.sdk.importers.base import DagDefinition + from airflow.sdk.importers.yaml_importer.models import DagDocument, Task, Value + + +class YamlDagBuildError(ValueError): + """A validated document could not be built into a DAG.""" + + +_BASE_TASK_ARGS = frozenset(BaseOperator.__init__._BaseOperatorMeta__param_names) # type: ignore[attr-defined] +_UNSUPPORTED_TASK_ARGS = frozenset({"task_id", "dag", "task_group", "default_args"}) + + [email protected](kw_only=True) +class DagBuilder: + """ + Build a :class:`DagDocument` into a :class:`DAG`. + + This holds the state a single build shares. The state belongs to one build + and is discarded with the instance. + """ + + _definition: DagDefinition + _bundle: BaseDagBundle + _operators: dict[str, BaseOperator] = attrs.field(factory=dict, init=False) + _stub_callables: dict[tuple[str, ...], Callable[..., None]] = attrs.field(factory=dict, init=False) + + def build(self, doc: DagDocument) -> DAG: + tasks = {t.id_: t for t in doc.tasks} + extras = _coerce_temporal(doc.model_extra or {}, DAG) + with DAG(dag_id=doc.dag_id, schedule=_build_schedule(doc.schedule), **extras) as dag: + for task_id in _topological_order(doc.tasks): + self._operators[task_id] = self._build_operator(tasks[task_id]) + + for task in doc.tasks: + for upstream in sorted(task.get_dependency_ids()): + self._operators[upstream] >> self._operators[task.id_] + dag.bundle_name = self._bundle.name + dag.fileloc = repr(self._definition) + dag.relative_fileloc = self._definition.get_relative_loc(self._bundle.path) + return dag + + def _build_stub_task(self, task: Task, args: dict[str, Value]) -> BaseOperator: + """ + Build a stub from a "run:" task. + + The task body is a registered SDK task handler resolved at run time; the + run mapping binds as that handler's arguments, where XCom references + resolve natively. The stub is fed a callable whose parameters are the + run argument names, so serialization can build and later validate + (AIP-108) the handler's argument spec. + """ + run_kwargs: dict[str, Any] = {} + for name, value in args.items(): + if not name.isidentifier() or keyword.iskeyword(name): + raise YamlDagBuildError(f"task {task.id_!r}: run argument {name!r} is not a valid identifier") + _reject_unbindable_run_marker(task.id_, value) + run_kwargs[name] = self._lower(value, protect=False) + decorator = task_decorators.stub( + task_id=task.id_, **_coerce_temporal(task.model_extra or {}, BaseOperator) + ) + return decorator(self._stub_callable(tuple(run_kwargs)))(**run_kwargs).operator + + def _build_operator(self, task: Task) -> BaseOperator: + """Instantiate *task* as an operator.""" + if reserved := sorted( + k for k in (*task.with_, *(task.model_extra or {})) if k in _UNSUPPORTED_TASK_ARGS + ): + raise YamlDagBuildError( + f"task {task.id_!r}: {reserved} cannot be set on a task; the task id comes from 'id', and " + "dag/task_group/default_args are managed by the importer" + ) + if misplaced := sorted(k for k in task.with_ if k in _BASE_TASK_ARGS): + raise YamlDagBuildError( + f"task {task.id_!r}: {misplaced} are task-level arguments; set them at task level, not under 'with'" + ) + if unknown := sorted(k for k in (task.model_extra or {}) if k not in _BASE_TASK_ARGS): + raise YamlDagBuildError( + f"task {task.id_!r}: {unknown} are not task-level (BaseOperator) arguments; " + "operator-specific arguments belong under 'with'" + ) + + if (args := task.run) is not None: + return self._build_stub_task(task, args) + + op_cls: Any + if task.uses is None: # uses: null (or omitted) becomes EmptyOperator. + from airflow.providers.standard.operators.empty import EmptyOperator as op_cls # noqa: SDK002 + else: + op_cls = import_string(task.uses) + + template_fields = set(getattr(op_cls, "template_fields", ())) + temporal = _temporal_fields(op_cls) + kwargs: dict[str, Any] = {} + render_fields: list[str] = [] + off_template: list[str] = [] + for name, value in task.with_.items(): + needs_render, has_marker = _field_markers(value) + lowered = self._lower(value, protect=needs_render) + if name in temporal and not has_marker: # a literal date/ISO string for a datetime arg + lowered = _to_datetime(lowered) + kwargs[name] = lowered + if has_marker and name not in template_fields: + off_template.append(name) + if needs_render: + render_fields.append(name) + if off_template: + raise YamlDagBuildError( + f"task {task.id_!r}: value grammar markers are only allowed on templated fields, not {off_template}" + ) + if task.model_extra: # Merge in task-level arguments (raw, never templated). + kwargs.update(_coerce_temporal(task.model_extra, BaseOperator)) + + return _restrict_template_fields(op_cls(task_id=task.id_, **kwargs), task.with_, render_fields) + + def _lower(self, value: Any, *, protect: bool) -> Any: + """ + Lower a resolved value-grammar node into an operator argument. + + :param protect: This should be set when the enclosing field renders to + avoid double-rendering. :class:`LiteralValue` is used to wrap + rendered values so they stay literal. + """ + if isinstance(value, XComRef): + return self._xcom_arg(value) # resolves at render time; never wrapped + if isinstance(value, TemplateRef): + # $t over a string renders that string; $t over $f is a template file (read, then render). + if isinstance(value.source, FileRef): + # A constructed FileRef has a str path (its validator rejects a templated filename). + return self._read_file(cast("str", value.source.path)) + return value.source + if isinstance(value, FileRef): + return self._protect(self._read_file(cast("str", value.path)), protect) + if isinstance(value, ConstRef): + return self._protect(value.value, protect) + if isinstance(value, Literal): + return self._lower(value.root, protect=protect) + if isinstance(value, dict): + return {k: self._lower(v, protect=protect) for k, v in value.items()} + if isinstance(value, list): + return [self._lower(v, protect=protect) for v in value] + return self._protect(value, protect) + + @staticmethod + def _protect(value: Any, protect: bool) -> Any: + return literal(value) if protect else value + + def _xcom_arg(self, value: XComRef) -> XComArg: + """Resolve an :class:`XComRef` against the already-built operators.""" + target = value.target + if isinstance(target, str): + return XComArg(self._operators[target]) + if target.key is None: + return XComArg(self._operators[target.task]) + return XComArg(self._operators[target.task], key=target.key) + + def _read_file(self, rel_path: str) -> str: + """Read a ``$f`` at build time, relative to the document.""" + try: + return self._definition.sibling(rel_path).read_bytes().decode("utf-8") + except (OSError, ValueError) as exc: + raise YamlDagBuildError(f"cannot read $f {rel_path!r}: {exc}") from exc + + def _stub_callable(self, param_names: tuple[str, ...]) -> Callable[..., None]: + """ + Build a placeholder callable to feed to ``@task.stub``. + + The stub needs a real callable with a meaningful signature; core + serialization reads it to build the TaskFlow arg-binding spec a lang-SDK + handler is validated against (AIP-108), and ``@task.stub`` reads the + callable's source to check the body is empty. + + The callable declares exactly the run arguments as a fixed, ordered + parameter list with an empty body, and its source is registered with + :mod:`linecache` so :func:`inspect.getsource` can read it back. + """ + if (cached := self._stub_callables.get(param_names)) is not None: + return cached + source = f"def _stub({', '.join(param_names)}):\n ...\n" + filename = f"<yaml-dag stub ({', '.join(param_names)})>" + linecache.cache[filename] = (len(source), None, source.splitlines(keepends=True), filename) + namespace: dict[str, Any] = {} + exec(compile(source, filename, "exec"), namespace) + self._stub_callables[param_names] = stub = namespace["_stub"] + return stub + + +def _restrict_template_fields( + op: BaseOperator, with_: dict[str, Any], render_fields: list[str] +) -> BaseOperator: + """ + Make *op* render literal-by-default, per instance. + + Only fields the author set in ``with:`` without a marker are dropped from + ``template_fields``; fields the operator templates by default or marked as + templates (e.g. uses ``$t``) are kept, so an operator's own defaults (e.g. + ``ExternalTaskMarker.logical_date``) still render. + + ``template_ext`` is emptied: a file is read only when asked for explicitly + (``$f``), never by a rendered value's suffix. + """ + op.template_fields = tuple(f for f in op.template_fields if f not in with_ or f in render_fields) + op.template_ext = () + return op + + +def _reject_unbindable_run_marker(task_id: str, value: Any) -> None: + """ + Reject value-grammar a task handler's argument binding cannot carry. + + The binding carries only literals and an upstream task's return value, so a + template (``$t``) or a keyed XCom (``$x`` with a ``key``) anywhere in a run + argument is an error. A plain ``$x`` (return value), ``$f``, ``$const`` and + literals are fine. Review Comment: Three more `run:` inputs build fine here and then fail at serialization with a `@task.stub` message from `stub_arg_bindings.py`: a nested plain `$x` like `items: [{$x: a}]` (`_ensure_json_literal`), an unquoted YAML date like `day: 2024-01-01` (not JSON-serializable), and an argument named after a context key such as `ds` or `params` (`_validate_stub_signature`). Checking them here would give a YAML-worded error at import time, and a plain `$x` is only fine at the top level of an argument, not nested as this docstring says. ########## task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py: ########## @@ -0,0 +1,552 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# 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. + +from __future__ import annotations + +import datetime +import inspect +import textwrap +import zipfile +from types import SimpleNamespace + +import pytest + +from airflow.sdk import BaseOperator, XComArg +from airflow.sdk.definitions._internal.templater import LiteralValue +from airflow.sdk.definitions.timetables.simple import OnceTimetable +from airflow.sdk.importers import ZipMemberDagDefinition +from airflow.sdk.importers.base import FilesystemDagDefinition +from airflow.sdk.importers.yaml_importer.importer import YamlDagImporter + +HEAD_SCHEMA = "https://airflow.apache.org/schemas/dag/2026-10-30.json"; + + +class _RecordingOperator(BaseOperator): + """A trivial operator with one templated field (``payload``) and one plain field (``note``).""" + + template_fields = ("payload",) + template_ext = (".sql",) # non-empty so the importer emptying it to () is observable, not vacuous + + def __init__(self, *, payload=None, note=None, **kwargs): + super().__init__(**kwargs) + self.payload = payload + self.note = note + + def execute(self, context): + return self.payload + + +OP = f"{__name__}.{_RecordingOperator.__name__}" + + +class _DefaultTemplateOperator(BaseOperator): + """Templates ``preset`` by default (set in __init__, not by the author) plus author-set ``payload``.""" + + template_fields = ("payload", "preset") + + def __init__(self, *, payload=None, preset="{{ ds }}", **kwargs): + super().__init__(**kwargs) + self.payload = payload + self.preset = preset + + def execute(self, context): + return self.payload + + +OP_DEFAULT_TEMPLATE = f"{__name__}.{_DefaultTemplateOperator.__name__}" + + +class _TemporalOperator(BaseOperator): + """An operator with a datetime-typed argument (``cutoff``).""" + + def __init__(self, *, cutoff: datetime.datetime | None = None, **kwargs): + super().__init__(**kwargs) + self.cutoff = cutoff + + +OP_TEMPORAL = f"{__name__}.{_TemporalOperator.__name__}" + + [email protected] +def mock_bundle(tmp_path): + bundle_dir = tmp_path / "bundle" + bundle_dir.mkdir(parents=True, exist_ok=True) + return SimpleNamespace(name="test_bundle", path=bundle_dir) + + +def _def(bundle, name: str, body: str) -> FilesystemDagDefinition: + path = bundle.path / name + path.write_text(textwrap.dedent(body)) + return FilesystemDagDefinition(path=path) + + +def _doc(tasks: str, *, dag_id: str = "d", extra: str = "") -> str: + return f"$schema: {HEAD_SCHEMA}\ndag_id: {dag_id}\n{extra}tasks:\n{textwrap.indent(textwrap.dedent(tasks), ' ')}" + + +def test_can_handle(): + importer = YamlDagImporter() + assert importer.can_handle("dag.yaml") + assert importer.can_handle("dag.yml") + assert not importer.can_handle("dag.py") + assert not importer.can_handle("dag.json") + + +def test_get_source_code(mock_bundle): + definition = _def(mock_bundle, "d.yaml", f"$schema: {HEAD_SCHEMA}\ndag_id: d\ntasks: []\n") + src = YamlDagImporter().get_source_code(definition) + assert src.language == "yaml" + assert "$schema" in src.source_code + + +def test_might_contain_dag(mock_bundle): + importer = YamlDagImporter() + has = _def(mock_bundle, "has.yaml", f"$schema: {HEAD_SCHEMA}\ndag_id: d\ntasks: []\n") + nope = _def(mock_bundle, "nope.yaml", "just: config\n") + # A $schema that is not the canonical Dag URL (e.g. a JSON-Schema file) must not be picked up. + other = _def(mock_bundle, "schema.yaml", "$schema: https://json-schema.org/draft/2020-12/schema\n";) + # The canonical URL merely mentioned, without a $schema declaration, must not be picked up either. + mention = _def(mock_bundle, "mention.yaml", f"note: see {HEAD_SCHEMA}\n") + assert importer.might_contain_dag(has, safe_mode=True) + assert not importer.might_contain_dag(nope, safe_mode=True) + assert not importer.might_contain_dag(other, safe_mode=True) + assert not importer.might_contain_dag(mention, safe_mode=True) + assert importer.might_contain_dag(nope, safe_mode=False) # sniff skipped + + +def test_list_dag_definitions_skips_non_dag_yaml(mock_bundle): + _def(mock_bundle, "dag.yaml", f"$schema: {HEAD_SCHEMA}\ndag_id: d\ntasks: []\n") + _def(mock_bundle, "config.yaml", "just: config\n") # no $schema -> skipped + _def(mock_bundle, "notes.txt", "hello") # wrong extension + names = {d.path.name for d in YamlDagImporter().list_dag_definitions(mock_bundle)} + assert names == {"dag.yaml"} + + +def test_import_builds_operator_dag(mock_bundle): + definition = _def( + mock_bundle, + "dag.yaml", + _doc( + f""" + - {{id: a, uses: {OP}, with: {{payload: hi}}}} + - {{id: b, uses: {OP}, with: {{payload: {{$x: a}}}}}} + """, + dag_id="pipeline", + extra="catchup: false\n", + ), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + (dag,) = result.dags + assert dag.dag_id == "pipeline" + assert set(dag.task_ids) == {"a", "b"} + assert dag.catchup is False # Dag-level attribute passed through + assert dag.bundle_name == "test_bundle" + a, b = dag.get_task("a"), dag.get_task("b") + assert a.payload == "hi" # literal + assert isinstance(b.payload, XComArg) # $x -> XComArg + assert b.upstream_task_ids == {"a"} # XCom wired the dependency + assert "payload" in b.template_fields # $x opted the field into rendering + assert "payload" not in a.template_fields # a's unmarked literal was dropped from rendering + + +def test_import_wires_needs(mock_bundle): + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f""" + - {{id: a, uses: {OP}, with: {{payload: x}}}} + - {{id: b, uses: {OP}, with: {{payload: y}}, needs: [a]}} + """), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + assert result.dags[0].get_task("b").upstream_task_ids == {"a"} + + +def test_import_bodyless_task_becomes_empty_operator(mock_bundle): + definition = _def(mock_bundle, "dag.yaml", _doc("- {id: a}")) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + assert type(result.dags[0].get_task("a")).__name__ == "EmptyOperator" + + +def test_import_bodyless_task_with_args_errors(mock_bundle): + definition = _def(mock_bundle, "dag.yaml", _doc("- {id: a, with: {x: 1}}")) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert any( + "Invalid arguments were passed to EmptyOperator" in e.message and "'x'" in e.message + for e in result.errors + ) + + +def test_import_schedule_cron(mock_bundle): + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f"- {{id: a, uses: {OP}, with: {{payload: x}}}}", extra="schedule: '0 3 * * *'\n"), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + assert result.dags[0].schedule == "0 3 * * *" + + +def test_import_schedule_timetable(mock_bundle): + definition = _def( + mock_bundle, + "dag.yaml", + _doc( + f"- {{id: a, uses: {OP}, with: {{payload: x}}}}", + extra="schedule: {uses: airflow.sdk.definitions.timetables.simple.OnceTimetable}\n", + ), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + assert isinstance(result.dags[0].timetable, OnceTimetable) + + +def test_import_keyed_xcom_on_uses_task(mock_bundle): + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f""" + - {{id: a, uses: {OP}, with: {{payload: x}}}} + - {{id: b, uses: {OP}, with: {{payload: {{$x: {{task: a, key: out}}}}}}}} + """), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + b = result.dags[0].get_task("b") + assert isinstance(b.payload, XComArg) + assert b.payload.key == "out" + assert b.upstream_task_ids == {"a"} + + [email protected]( + ("value", "hour"), + [("2024-01-01", 0), ("'2024-01-01'", 0), ("2024-01-01T03:30:00", 3)], # YAML date, ISO string, datetime +) +def test_import_coerces_temporal_start_date(mock_bundle, value, hour): + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f"- {{id: a, uses: {OP}, with: {{payload: x}}}}", extra=f"start_date: {value}\n"), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + start = result.dags[0].start_date + assert isinstance(start, datetime.datetime) # promoted from date/string + assert start.tzinfo is not None # localized by DAG() + assert (start.year, start.month, start.day, start.hour) == (2024, 1, 1, hour) + + +def test_import_does_not_coerce_nontemporal_arg(mock_bundle): + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f"- {{id: a, uses: {OP}, owner: '2024-01-01', with: {{payload: x}}}}"), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + assert result.dags[0].get_task("a").owner == "2024-01-01" + + +def test_import_coerces_temporal_with_arg(mock_bundle): + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f"- {{id: a, uses: {OP_TEMPORAL}, with: {{cutoff: 2024-01-01}}}}"), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + cutoff = result.dags[0].get_task("a").cutoff + assert isinstance(cutoff, datetime.datetime) + assert (cutoff.year, cutoff.month, cutoff.day) == (2024, 1, 1) + + +def test_import_code_task_builds_stub_operator(mock_bundle): + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f""" + - {{id: a, uses: {OP}, with: {{payload: hi}}}} + - {{id: b, run: {{value: {{$x: a}}, label: done}}, queue: handlers}} + """), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + b = result.dags[0].get_task("b") + assert type(b).__name__ == "_StubOperator" + assert getattr(b, "is_stub", False) is True # the real stub operator Airflow serializes + assert b.queue == "handlers" # task-level attribute passed through + assert isinstance(b.op_kwargs["value"], XComArg) # $x -> XComArg, bound as a handler argument + assert b.op_kwargs["label"] == "done" # literal bound verbatim (run args are never rendered) + assert b.upstream_task_ids == {"a"} # $x wired the dependency + # The stub is fed a callable whose ordered parameters are the run argument names, so core + # serialization can build the handler arg-binding spec (AIP-108) from a real signature. + assert list(inspect.signature(b.python_callable).parameters) == ["value", "label"] + + +def test_import_mixed_container_keeps_literal_siblings_literal(mock_bundle): + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f""" + - {{id: up, uses: {OP}, with: {{payload: x}}}} + - {{id: t, uses: {OP}, with: {{payload: {{a: {{$x: up}}, b: '{{{{ not a template }}}}'}}}}}} + """), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + payload = result.dags[0].get_task("t").payload + assert isinstance(payload["a"], XComArg) # $x still resolves + assert isinstance(payload["b"], LiteralValue) # literal sibling protected from rendering + assert payload["b"].value == "{{ not a template }}" + + +def test_import_code_task_rejects_non_identifier_arg(mock_bundle): + definition = _def(mock_bundle, "dag.yaml", _doc("- {id: a, run: {'bad-name': 1}}")) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert any("not a valid identifier" in e.message for e in result.errors) + + +def test_import_run_rejects_template(mock_bundle): + definition = _def(mock_bundle, "dag.yaml", _doc("- {id: a, run: {day: {$t: '{{ ds }}'}}}")) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert any("template ($t)" in e.message for e in result.errors) + + +def test_import_run_rejects_keyed_xcom(mock_bundle): + definition = _def( + mock_bundle, + "dag.yaml", + _doc(""" + - {id: a, run: {}} + - {id: b, run: {v: {$x: {task: a, key: k}}}} + """), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert any("keyed XCom" in e.message for e in result.errors) + + [email protected]("marker", ["{$x: a}", "{$t: 'v'}", "{$f: q.sql}"]) +def test_import_marker_on_non_template_field_errors(mock_bundle, marker): + (mock_bundle.path / "q.sql").write_text("SELECT 1") # exists, so this is not a missing-file error + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f""" + - {{id: a, uses: {OP}, with: {{payload: hi}}}} + - {{id: b, uses: {OP}, with: {{payload: ok, note: {marker}}}}} + """), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert any("only allowed on template_fields" in e.message for e in result.errors) Review Comment: This still expects `only allowed on template_fields`, but `builder.py` now says `only allowed on templated fields`, so all three cases fail (3 failed when I ran the file in breeze). Matching the new wording keeps the off-template guard covered. ########## task-sdk/src/airflow/sdk/importers/zip_importer.py: ########## @@ -101,6 +102,18 @@ def read_bytes(self) -> bytes: self._content = z.read(self.file_path) return self._content + def sibling(self, rel_path: str) -> zipfile.Path: Review Comment: The base and filesystem versions name this parameter `relpath`, so a keyword call `definition.sibling(relpath=...)` fails for zip members only. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
