From 7cebda58c4a3c78634cb7816fbac66fc0893ff6f Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 12:15:43 +0100 Subject: [PATCH 01/25] Add missing PyYAML dependency to Task SDK --- task-sdk/pyproject.toml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/task-sdk/pyproject.toml b/task-sdk/pyproject.toml index 0ce16d4b713e9..9a4c9d44ea97b 100644 --- a/task-sdk/pyproject.toml +++ b/task-sdk/pyproject.toml @@ -54,14 +54,15 @@ dependencies = [ "babel>=2.17.0", "cadwyn>=6.1.1", "fsspec>=2023.10.0", + "greenback>=1.2.1", "httpx>=0.27.0", "jinja2>=3.1.5", "methodtools>=0.4.7", "msgspec>=0.19.0", - "python-dateutil>=2.7.0", "psutil>=6.1.0", + "python-dateutil>=2.7.0", + "pyyaml>=6.0.3", "structlog>=26.1.0", - "greenback>=1.2.1", "tenacity>=8.3.0", # Start of shared timezones dependencies "pendulum>=3.1.0", From 07503aa14802a3afa440b8dcde655ce890c08c55 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 12:34:08 +0100 Subject: [PATCH 02/25] Implement file-loading syntax $f --- .../sdk/importers/yaml_importer/models.py | 69 ++++++++++++++++--- .../sdk/importers/yaml_importer/schema.json | 62 +++++++++++++++-- .../importers/yaml_importer/test_parser.py | 10 +++ 3 files changed, 128 insertions(+), 13 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/models.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/models.py index 94fdd50f852f4..4fc628e1f5626 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/models.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/models.py @@ -50,6 +50,7 @@ XCOM_KEYS = ("$x", "$xcom") TEMPLATE_KEYS = ("$t", "$template") +FILE_KEYS = ("$f", "$file") CONST_KEY = "$const" # Keys deferred to a later edition... @@ -106,13 +107,26 @@ class _MarkerRef(BaseModel): _marker_keys: ClassVar[tuple[str, ...]] + @classmethod + def _process_value_schema(cls, field_schema: JsonSchemaValue) -> JsonSchemaValue: + """ + Post-process published schema for each marker key. + + The default publishes the field's schema as a disjoint union. + :class:`FileRef` overrides it to reject a templated filename in the + schema too. + """ + return _as_exclusive_union(field_schema) + @classmethod def __get_pydantic_json_schema__( cls, core_schema: CoreSchema, handler: GetJsonSchemaHandler ) -> JsonSchemaValue: schema = handler(core_schema) - value_schema = _as_exclusive_union(next(iter(schema.get("properties", {}).values()))) - schema["properties"] = {key: value_schema for key in cls._marker_keys} + schema["properties"] = { + key: cls._process_value_schema(next(iter(schema.get("properties", {}).values()))) + for key in cls._marker_keys + } schema["additionalProperties"] = False if len(cls._marker_keys) > 1: # Pydantic sets the correct 'required' for the single-variant case. schema.pop("required", None) @@ -132,11 +146,39 @@ class XComRef(_MarkerRef): class TemplateRef(_MarkerRef): - """A Jinja template.""" + """ + A Jinja template, rendered server-side. + + ``source`` is a Jinja string, or a :class:`FileRef` to load a template file. + """ _marker_keys = TEMPLATE_KEYS - source: str = Field(serialization_alias="$t", validation_alias=AliasChoices(*TEMPLATE_KEYS)) + source: str | FileRef = Field(serialization_alias="$t", validation_alias=AliasChoices(*TEMPLATE_KEYS)) + + +class FileRef(_MarkerRef): + """ + A file input (AIP-80's FileTemplate). + + ``path`` is a static path. The inverse nesting (templated filename) is + currently reserved and rejected since it does not work with current Airflow + internals. + """ + + _marker_keys = FILE_KEYS + + path: str | TemplateRef = Field(serialization_alias="$f", validation_alias=AliasChoices(*FILE_KEYS)) + + @classmethod + def _process_value_schema(cls, field_schema: JsonSchemaValue) -> JsonSchemaValue: + return {"type": "string"} + + @model_validator(mode="after") + def _reject_templated_filename(self) -> FileRef: + if isinstance(self.path, TemplateRef): + raise ValueError("templated filename ($f wrapping $t) is not supported yet") + return self class ConstRef(_MarkerRef): @@ -163,12 +205,16 @@ def _value_discriminator(v: Any) -> str: return "xcom" if key in TEMPLATE_KEYS: return "template" + if key in FILE_KEYS: + return "file" if key == CONST_KEY: return "const" if isinstance(v, XComRef): return "xcom" if isinstance(v, TemplateRef): return "template" + if isinstance(v, FileRef): + return "file" if isinstance(v, ConstRef): return "const" return "literal" @@ -222,11 +268,14 @@ def __get_pydantic_json_schema__( Value = Annotated[ Annotated[XComRef, Tag("xcom")] | Annotated[TemplateRef, Tag("template")] + | Annotated[FileRef, Tag("file")] | Annotated[ConstRef, Tag("const")] | Annotated[Literal, Tag("literal")], Discriminator(_value_discriminator), ] +TemplateRef.model_rebuild() +FileRef.model_rebuild() Literal.model_rebuild() @@ -265,6 +314,13 @@ def _check_task(cls, data: Any) -> Any: raise ValueError(f"task {data.get('id')!r} may not set both 'uses' and 'run'") return data + def get_dependency_ids(self) -> set[str]: + """Ids of the tasks this task depends on: explicit ``needs`` plus any XCom references.""" + ids = set(self.needs) + for value in itertools.chain(self.with_.values(), (self.run or {}).values()): + ids.update(_iter_xcom_task_ids(value)) + return ids + @classmethod def __get_pydantic_json_schema__( cls, core_schema: CoreSchema, handler: GetJsonSchemaHandler @@ -449,10 +505,7 @@ def _check_refs(self): raise ValueError(f"duplicate task id(s): {duplicates}") validated_task_ids = set(task_id_counts) for t in self.tasks: - refs = set(t.needs) - for v in itertools.chain(t.with_.values(), (t.run or {}).values()): - refs.update(_iter_xcom_task_ids(v)) - if t.id_ in refs: + if t.id_ in (refs := t.get_dependency_ids()): raise ValueError(f"task {t.id_!r} depends on itself") if missing := (refs - validated_task_ids): raise ValueError(f"task {t.id_!r} references unknown task(s): {sorted(missing)}") diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/schema.json b/task-sdk/src/airflow/sdk/importers/yaml_importer/schema.json index 953c61b6ccca4..10f5a0f559c1f 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/schema.json +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/schema.json @@ -14,6 +14,32 @@ "title": "ConstRef", "type": "object" }, + "FileRef": { + "additionalProperties": false, + "description": "A file input (AIP-80's FileTemplate).\n\n``path`` is a static path. The inverse nesting (templated filename) is\ncurrently reserved and rejected since it does not work with current Airflow\ninternals.", + "oneOf": [ + { + "required": [ + "$f" + ] + }, + { + "required": [ + "$file" + ] + } + ], + "properties": { + "$f": { + "type": "string" + }, + "$file": { + "type": "string" + } + }, + "title": "FileRef", + "type": "object" + }, "Literal": { "description": "A literal value.", "not": { @@ -34,6 +60,9 @@ { "$ref": "#/$defs/TemplateRef" }, + { + "$ref": "#/$defs/FileRef" + }, { "$ref": "#/$defs/ConstRef" }, @@ -53,6 +82,9 @@ { "$ref": "#/$defs/TemplateRef" }, + { + "$ref": "#/$defs/FileRef" + }, { "$ref": "#/$defs/ConstRef" }, @@ -140,6 +172,9 @@ { "$ref": "#/$defs/TemplateRef" }, + { + "$ref": "#/$defs/FileRef" + }, { "$ref": "#/$defs/ConstRef" }, @@ -167,6 +202,9 @@ { "$ref": "#/$defs/TemplateRef" }, + { + "$ref": "#/$defs/FileRef" + }, { "$ref": "#/$defs/ConstRef" }, @@ -196,7 +234,7 @@ }, "TemplateRef": { "additionalProperties": false, - "description": "A Jinja template.", + "description": "A Jinja template, rendered server-side.\n\n``source`` is a Jinja string, or a :class:`FileRef` to load a template file.", "oneOf": [ { "required": [ @@ -211,12 +249,26 @@ ], "properties": { "$t": { - "title": "$T", - "type": "string" + "oneOf": [ + { + "type": "string" + }, + { + "$ref": "#/$defs/FileRef" + } + ], + "title": "$T" }, "$template": { - "title": "$T", - "type": "string" + "oneOf": [ + { + "type": "string" + }, + { + "$ref": "#/$defs/FileRef" + } + ], + "title": "$T" } }, "title": "TemplateRef", diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_parser.py b/task-sdk/tests/task_sdk/importers/yaml_importer/test_parser.py index 22a69df2c753f..91b0dd33f17cd 100644 --- a/task-sdk/tests/task_sdk/importers/yaml_importer/test_parser.py +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_parser.py @@ -531,6 +531,14 @@ def test_published_schema_accepts_null_and_empty_bodies(task): "tasks:\n - {id: t, uses: X, with: {k: {$t: 'v-{{ ds }}'}}}", id="template", ), + pytest.param( + "tasks:\n - {id: t, uses: X, with: {k: {$f: 'reports/daily.sql'}}}", + id="file", + ), + pytest.param( + "tasks:\n - {id: t, uses: X, with: {k: {$t: {$f: 'reports/daily.sql'}}}}", + id="template-file", + ), pytest.param( "templates: {base: {run: {r: 1}}}\ntasks:\n - {id: t, extends: base}", id="extends-bare", @@ -564,6 +572,8 @@ def test_valid_shapes_round_trip(body): pytest.param("{cfg: {$weird: 1}}", id="dollar-key-nested"), pytest.param("{$x: up, $xcom: up}", id="two-marker-keys"), pytest.param("{$x: {task: up, nope: 1}}", id="xcom-target-extra-key"), + pytest.param("{$f: {$t: 'q-{{ ds }}.sql'}}", id="templated-filename-reserved"), + pytest.param("{$t: {$f: {$t: x}}}", id="templated-filename-nested-reserved"), ], ) def test_invalid_shapes_rejected_by_both(value): From 45c27baeee1d1622a8d427744f71a806529fda65 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 13:09:03 +0100 Subject: [PATCH 03/25] Implement the YAML importer shell --- .../sdk/importers/yaml_importer/builder.py | 254 ++++++++++++++ .../sdk/importers/yaml_importer/importer.py | 144 ++++++++ .../importers/yaml_importer/test_importer.py | 310 ++++++++++++++++++ 3 files changed, 708 insertions(+) create mode 100644 task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py create mode 100644 task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py create mode 100644 task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py new file mode 100644 index 0000000000000..37870367835af --- /dev/null +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py @@ -0,0 +1,254 @@ +# 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 keyword +import linecache +from typing import TYPE_CHECKING, Any, cast + +from airflow.sdk import DAG, XComArg +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 pathlib import Path + + from airflow.sdk import BaseOperator + from airflow.sdk.importers.yaml_importer.models import DagDocument, Task, Value + + +class YamlDagBuildError(ValueError): + """A validated document could not be built into a DAG.""" + + +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. + """ + + def __init__(self, source_dir: Path | None) -> None: + self._source_dir = source_dir + self._operators: dict[str, BaseOperator] = {} + self._stub_callables: dict[tuple[str, ...], Callable[..., None]] = {} + + def build(self, doc: DagDocument) -> DAG: + tasks = {t.id_: t for t in doc.tasks} + with DAG(dag_id=doc.dag_id, schedule=_build_schedule(doc.schedule), **(doc.model_extra or {})) as dag: + # A task depended by another task must be built first. + for task_id in _topological_order(doc.tasks): + self._operators[task_id] = self._build_operator(tasks[task_id]) + + # Wire dependency edges explicitly. + for task in doc.tasks: + for upstream in sorted(task.get_dependency_ids()): + self._operators[upstream] >> self._operators[task.id_] + 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 and template + 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. + """ + from airflow.sdk import task as task_decorators # Lazy to avoid provider registration. + + 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") + run_kwargs[name] = self._lower(value)[0] + decorator = task_decorators.stub(task_id=task.id_, **(task.model_extra or {})) + return decorator(self._stub_callable(tuple(run_kwargs)))(**run_kwargs).operator + + def _build_operator(self, task: Task) -> BaseOperator: + """Instantiate *task* as an operator.""" + 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) + + # Validate template fields. + template_fields = set(getattr(op_cls, "template_fields", ())) + kwargs: dict[str, Any] = {} + render_fields: list[str] = [] + off_template: list[str] = [] + for name, value in task.with_.items(): + kwargs[name], needs_render, has_marker = self._lower(value) + 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 only allowed on template_fields") + if extras := task.model_extra: # Merge in task-level arguments (raw, never templated). + kwargs.update(extras) + + return _apply_template_lowering(op_cls(task_id=task.id_, **kwargs), tuple(render_fields)) + + def _lower(self, value: Any) -> tuple[Any, bool, bool]: + """ + Lower a resolved value-grammar node. + + The return value is ``(operator argument, needs_render, has_marker)``. + + * ``needs_render`` is True if the node carries ``$t`` or ``$x``, and we + must validate it's in the operator's ``template_fields``. + reference anywhere, i.e. + * ``has_marker`` is True when it carries any value-grammar markers, and + we should only be allowed if the field can be templated. + """ + if isinstance(value, XComRef): + return self._xcom_arg(value), True, True + 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)), True, True + return value.source, True, True + if isinstance(value, FileRef): + return self._read_file(cast("str", value.path)), False, True # verbatim; not a template + if isinstance(value, ConstRef): + return value.value, False, False + if isinstance(value, Literal): + return self._lower(value.root) + if isinstance(value, dict): + items = {k: self._lower(v) for k, v in value.items()} + return ( + {k: lv for k, (lv, _, _) in items.items()}, + any(render for _, render, _ in items.values()), + any(marker for _, _, marker in items.values()), + ) + if isinstance(value, list): + elements = [self._lower(v) for v in value] + return ( + [lv for lv, _, _ in elements], + any(render for _, render, _ in elements), + any(marker for _, _, marker in elements), + ) + return value, False, False + + 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, resolved against ``source_dir``.""" + if self._source_dir is None: + raise YamlDagBuildError(f"cannot read $f {rel_path!r}: the document has no source directory") + try: + return (self._source_dir / rel_path).read_text(encoding="utf-8") + except OSError 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"" + 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 _apply_template_lowering(op: BaseOperator, templated_fields: tuple[str, ...]) -> BaseOperator: + """ + Make *op* render exactly *templated_fields*. + + This allows us to implement AIP-80 file-reading semantics without changing + core; fortunately Airflow reads them per-instance so a simple patch works. + """ + op.template_fields = templated_fields + op.template_ext = () + return op + + +def _build_schedule(schedule: Any) -> Any: + """Build *schedule* into a schedule construct for :class:`DAG`.""" + if isinstance(schedule, TimetableSchedule): + return import_string(schedule.uses)(**schedule.with_) + return schedule + + +def _topological_order(tasks: list[Task]) -> list[str]: + """Return task ids ordered so every task follows the tasks it depends on.""" + upstreams = {task.id_: task.get_dependency_ids() for task in tasks} + order: list[str] = [] + done: set[str] = set() + remaining = [task.id_ for task in tasks] + while remaining: + ready = [task_id for task_id in remaining if upstreams[task_id] <= done] + if not ready: + raise YamlDagBuildError(f"dependency cycle among tasks: {sorted(remaining)}") + order.extend(ready) + done.update(ready) + remaining = [task_id for task_id in remaining if task_id not in done] + return order diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py new file mode 100644 index 0000000000000..84dfa04abaff2 --- /dev/null +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py @@ -0,0 +1,144 @@ +# 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. + +""" +AIP-85 YAML/JSON Dag importer. + +This discovers ``.yaml`` and ``.yml`` Dag files in a bundle, and builds the Dag +objects they describe. It is not turned on by default; enable with:: + + [dag_processor] + dag_importer_configs = [{"classpath": "airflow.sdk.importers.yaml_importer.importer.YamlDagImporter"}] +""" + +from __future__ import annotations + +import logging +import traceback +from typing import TYPE_CHECKING + +from airflow.sdk.importers.base import ( + AbstractDagImporter, + DagImportError, + DagImportResult, + DagSourceCode, + FileDagDefinition, + _normalize_extensions, + find_file_dag_definitions, + get_file_suffix, +) +from airflow.sdk.importers.yaml_importer.builder import DagBuilder +from airflow.sdk.importers.yaml_importer.parser import YamlDagParseError, parse_documents + +if TYPE_CHECKING: + from collections.abc import Iterator + from pathlib import Path + + from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 + from airflow.sdk.importers.base import DagDefinition + +log = logging.getLogger(__name__) + + +class YamlDagImporter(AbstractDagImporter[FileDagDefinition]): + """Importer for YAML/JSON DAG sources. Opt-in via ``dag_importer_configs`` (not a default).""" + + supported_extensions = [".yaml", ".yml"] + + def __init__(self, extensions: list[str] | None = None) -> None: + if extensions is not None: + self.supported_extensions = _normalize_extensions(extensions) + + def can_handle(self, definition: DagDefinition | str | Path) -> bool: + suffix = get_file_suffix(definition) + return suffix in self.supported_extensions if suffix else False + + def list_dag_definitions( + self, + bundle: BaseDagBundle, + *, + safe_mode: bool = True, + ) -> Iterator[FileDagDefinition | DagImportError]: + for definition in find_file_dag_definitions(bundle.path, self.supported_extensions): + if self.might_contain_dag(definition, safe_mode): + yield definition + else: + log.debug("Skipping %r: no '$schema' marker found", definition) + + def import_definition( + self, + definition: FileDagDefinition, + bundle: BaseDagBundle, + ) -> DagImportResult: + result = DagImportResult(definition=definition) + try: + source = definition.read_text(encoding="utf-8") + except Exception as exc: + result.errors.append( + DagImportError( + source_reference=repr(definition), + message=str(exc), + error_type="import", + stacktrace=traceback.format_exc(), + ) + ) + return result + + # A parse failure fails the whole file (like a Python module that won't import); a build + # failure is per-document so the file's other DAGs still load. + try: + documents = list(parse_documents(source, source=repr(definition))) + except YamlDagParseError as exc: + result.errors.append( + DagImportError(source_reference=repr(definition), message=str(exc), error_type="parse") + ) + return result + + # ``$f`` file references resolve against the document's own directory, when it has one. + source_dir = path.parent if (path := getattr(definition, "path", None)) is not None else None + for document in documents: + try: + dag = DagBuilder(source_dir).build(document) + except Exception as exc: + result.errors.append( + DagImportError( + source_reference=f"{definition!r} (dag_id={document.dag_id!r})", + message=str(exc), + error_type="import", + stacktrace=traceback.format_exc(), + ) + ) + continue + dag.bundle_name = bundle.name + dag.fileloc = repr(definition) + dag.relative_fileloc = definition.get_relative_loc(bundle.path) + result.dags.append(dag) + log.debug("Built DAG %s from %r", dag.dag_id, definition) + + return result + + def get_source_code(self, definition: DagDefinition, dag_id: str | None = None) -> DagSourceCode: + return DagSourceCode(source_code=definition.read_text(encoding="utf-8"), language="yaml") + + def might_contain_dag(self, definition: DagDefinition, safe_mode: bool) -> bool: + """Cheap sniff: a DAG document declares a ``$schema``. Keep the file if the read fails.""" + if not safe_mode: + return True + try: + return b"$schema" in definition.read_bytes() + except OSError: + return True diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py new file mode 100644 index 0000000000000..374dabe0d4e1a --- /dev/null +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py @@ -0,0 +1,310 @@ +# 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 inspect +import textwrap +from types import SimpleNamespace + +import pytest + +from airflow.sdk import BaseOperator, XComArg +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",) + + 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__}" + + +@pytest.fixture +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") + assert importer.might_contain_dag(has, safe_mode=True) + assert not importer.might_contain_dag(nope, 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 + + +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): + # A body-less task becomes an EmptyOperator, which rejects the extra 'with' arguments + # (BaseOperator raises on unknown kwargs) -- so it fails to build rather than silently dropping. + 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 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_code_task_builds_stub_operator(mock_bundle): + # A code task (run:) builds the stub operator Airflow understands (@task.stub). Its body is a + # registered SDK task handler; the run mapping binds as that handler's arguments (op_kwargs), + # and $x inside it wires the dependency natively. + 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 preserved + 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_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) + + +@pytest.mark.parametrize("marker", ["{$x: a}", "{$t: 'v'}", "{$f: q.sql}"]) +def test_import_marker_on_non_template_field_errors(mock_bundle, marker): + # A marker only works where the operator class already templates. ``note`` is not in + # _RecordingOperator.template_fields, so any $f/$t/$x on it is an error -- the importer never + # adds a field to template_fields to make a marker work. + (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) + + +def test_import_literal_by_default(mock_bundle): + # A bare Jinja-looking string on a field the operator class *does* template is NOT rendered: + # literal by default, and the instance's template_fields is emptied of it. + definition = _def( + mock_bundle, "dag.yaml", _doc(f"- {{id: a, uses: {OP}, with: {{payload: '{{{{ ds }}}}'}}}}") + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + a = result.dags[0].get_task("a") + assert a.payload == "{{ ds }}" # kept literal + assert tuple(a.template_fields) == () # not rendered + assert tuple(a.template_ext) == () # and the suffix-based file-load foot-gun is closed + + +def test_import_template_marker_sets_template_fields(mock_bundle): + definition = _def( + mock_bundle, "dag.yaml", _doc(f"- {{id: a, uses: {OP}, with: {{payload: {{$t: '{{{{ ds }}}}'}}}}}}") + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + a = result.dags[0].get_task("a") + assert a.payload == "{{ ds }}" + assert set(a.template_fields) == {"payload"} # $t opts the field into rendering + + +def test_import_file_ref_reads_verbatim(mock_bundle): + (mock_bundle.path / "query.sql").write_text("SELECT 1") + definition = _def( + mock_bundle, "dag.yaml", _doc(f"- {{id: a, uses: {OP}, with: {{payload: {{$f: query.sql}}}}}}") + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + a = result.dags[0].get_task("a") + assert a.payload == "SELECT 1" # file contents, read at build time + assert tuple(a.template_fields) == () # $f is verbatim, not a template + + +def test_import_template_file_reads_then_templates(mock_bundle): + (mock_bundle.path / "query.sql").write_text("SELECT {{ ds }}") + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f"- {{id: a, uses: {OP}, with: {{payload: {{$t: {{$f: query.sql}}}}}}}}"), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + a = result.dags[0].get_task("a") + assert a.payload == "SELECT {{ ds }}" # file contents, to be rendered at run time + assert set(a.template_fields) == {"payload"} # $t over $f is a template file + + +def test_import_templated_filename_is_reserved(mock_bundle): + # $f wrapping $t (a templated filename) needs run-time file resolution; the parser reserves it, + # so the whole file fails to parse. + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f"- {{id: a, uses: {OP}, with: {{payload: {{$f: {{$t: '{{{{ ds }}}}.sql'}}}}}}}}"), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert result.errors + assert result.errors[0].error_type == "parse" + + +def test_import_parse_error(mock_bundle): + definition = _def(mock_bundle, "dag.yaml", "dag_id: d\ntasks: []\n") # missing $schema + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert result.errors + assert result.errors[0].error_type == "parse" + + +def test_import_unknown_operator_errors(mock_bundle): + definition = _def(mock_bundle, "dag.yaml", _doc("- {id: a, uses: does.not.Exist, with: {}}")) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert result.errors + + +def test_import_dependency_cycle_errors(mock_bundle): + # The parser allows a cycle (full cycle detection is the builder's job); the build rejects it. + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f""" + - {{id: a, uses: {OP}, with: {{payload: x}}, needs: [b]}} + - {{id: b, uses: {OP}, with: {{payload: y}}, needs: [a]}} + """), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert any("cycle" in e.message for e in result.errors) From 83663b053c5e2f44b8f57ad24e97b2fcbc5aa050 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 14:16:32 +0100 Subject: [PATCH 04/25] Fix double-rendering --- .../sdk/importers/yaml_importer/builder.py | 79 ++++++++++++------- .../importers/yaml_importer/test_importer.py | 24 +++++- 2 files changed, 72 insertions(+), 31 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py index 37870367835af..c79f9c70e04cb 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py @@ -39,6 +39,7 @@ from airflow.sdk import DAG, XComArg from airflow.sdk._shared.module_loading import import_string +from airflow.sdk.definitions._internal.templater import LiteralValue from airflow.sdk.importers.yaml_importer.models import ( ConstRef, FileRef, @@ -102,7 +103,7 @@ def _build_stub_task(self, task: Task, args: dict[str, Value]) -> BaseOperator: 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") - run_kwargs[name] = self._lower(value)[0] + run_kwargs[name] = self._lower(value, protect=True) decorator = task_decorators.stub(task_id=task.id_, **(task.model_extra or {})) return decorator(self._stub_callable(tuple(run_kwargs)))(**run_kwargs).operator @@ -123,7 +124,8 @@ def _build_operator(self, task: Task) -> BaseOperator: render_fields: list[str] = [] off_template: list[str] = [] for name, value in task.with_.items(): - kwargs[name], needs_render, has_marker = self._lower(value) + needs_render, has_marker = _field_markers(value) + kwargs[name] = self._lower(value, protect=needs_render) if has_marker and name not in template_fields: off_template.append(name) if needs_render: @@ -135,47 +137,37 @@ def _build_operator(self, task: Task) -> BaseOperator: return _apply_template_lowering(op_cls(task_id=task.id_, **kwargs), tuple(render_fields)) - def _lower(self, value: Any) -> tuple[Any, bool, bool]: + def _lower(self, value: Any, *, protect: bool) -> Any: """ - Lower a resolved value-grammar node. + Lower a resolved value-grammar node into an operator argument. - The return value is ``(operator argument, needs_render, has_marker)``. - - * ``needs_render`` is True if the node carries ``$t`` or ``$x``, and we - must validate it's in the operator's ``template_fields``. - reference anywhere, i.e. - * ``has_marker`` is True when it carries any value-grammar markers, and - we should only be allowed if the field can be templated. + :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), True, True + 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)), True, True - return value.source, True, True + return self._read_file(cast("str", value.source.path)) + return value.source if isinstance(value, FileRef): - return self._read_file(cast("str", value.path)), False, True # verbatim; not a template + return self._protect(self._read_file(cast("str", value.path)), protect) if isinstance(value, ConstRef): - return value.value, False, False + return self._protect(value.value, protect) if isinstance(value, Literal): - return self._lower(value.root) + return self._lower(value.root, protect=protect) if isinstance(value, dict): - items = {k: self._lower(v) for k, v in value.items()} - return ( - {k: lv for k, (lv, _, _) in items.items()}, - any(render for _, render, _ in items.values()), - any(marker for _, _, marker in items.values()), - ) + return {k: self._lower(v, protect=protect) for k, v in value.items()} if isinstance(value, list): - elements = [self._lower(v) for v in value] - return ( - [lv for lv, _, _ in elements], - any(render for _, render, _ in elements), - any(marker for _, _, marker in elements), - ) - return value, False, False + return [self._lower(v, protect=protect) for v in value] + return self._protect(value, protect) + + @staticmethod + def _protect(value: Any, protect: bool) -> Any: + return LiteralValue(value) if protect else value def _xcom_arg(self, value: XComRef) -> XComArg: """Resolve an :class:`XComRef` against the already-built operators.""" @@ -238,6 +230,33 @@ def _build_schedule(schedule: Any) -> Any: return schedule +def _field_markers(value: Any) -> tuple[bool, bool]: + """ + Scan a value-grammar node for ``(needs_render, has_marker)``. + + ``needs_render`` is True when it carries a ``$t`` or ``$x`` anywhere (the field must render); + ``has_marker`` is True when it carries any of ``$x``/``$t``/``$f`` (only allowed on a templated + field). ``$const`` is a force-literal escape, not a marker. + """ + if isinstance(value, XComRef): + return True, True + if isinstance(value, TemplateRef): + return True, True + if isinstance(value, FileRef): + return False, True + if isinstance(value, ConstRef): + return False, False + if isinstance(value, Literal): + return _field_markers(value.root) + if isinstance(value, dict): + flags = [_field_markers(v) for v in value.values()] + return any(r for r, _ in flags), any(m for _, m in flags) + if isinstance(value, list): + flags = [_field_markers(v) for v in value] + return any(r for r, _ in flags), any(m for _, m in flags) + return False, False + + def _topological_order(tasks: list[Task]) -> list[str]: """Return task ids ordered so every task follows the tasks it depends on.""" upstreams = {task.id_: task.get_dependency_ids() for task in tasks} diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py index 374dabe0d4e1a..1506570cd6576 100644 --- a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py @@ -24,6 +24,7 @@ import pytest from airflow.sdk import BaseOperator, XComArg +from airflow.sdk.definitions._internal.templater import LiteralValue from airflow.sdk.importers.base import FilesystemDagDefinition from airflow.sdk.importers.yaml_importer.importer import YamlDagImporter @@ -182,13 +183,34 @@ def test_import_code_task_builds_stub_operator(mock_bundle): 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 preserved + # op_kwargs is a template field, so the literal is protected to stay literal when rendered. + assert isinstance(b.op_kwargs["label"], LiteralValue) + assert b.op_kwargs["label"].value == "done" 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): + # When a container field renders (it holds a $x), its literal siblings must stay literal rather + # than being rendered as Jinja. The importer wraps them in LiteralValue; only $x/$t render. + 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) From ebc8b3530c5976da8aa165c6cb0216d2eca04ce7 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 14:40:18 +0100 Subject: [PATCH 05/25] Protect file read using the same logic as Python templates --- task-sdk/src/airflow/sdk/importers/base.py | 28 ++++++++++++++++++- .../sdk/importers/yaml_importer/builder.py | 14 ++++------ .../sdk/importers/yaml_importer/importer.py | 4 +-- .../src/airflow/sdk/importers/zip_importer.py | 13 +++++++++ .../importers/yaml_importer/test_importer.py | 12 ++++++++ 5 files changed, 59 insertions(+), 12 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/base.py b/task-sdk/src/airflow/sdk/importers/base.py index e594e15ad2402..b1c04842c187d 100644 --- a/task-sdk/src/airflow/sdk/importers/base.py +++ b/task-sdk/src/airflow/sdk/importers/base.py @@ -34,11 +34,17 @@ if TYPE_CHECKING: from collections.abc import Generator, Iterable, Iterator - from typing import Self + from typing import Protocol, Self from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 from airflow.sdk import DAG + class _SiblingResource(Protocol): + """A readable handle to a resource, e.g. a :class:`pathlib.Path` or :class:`zipfile.Path`.""" + + def read_bytes(self) -> bytes: ... + + log = logging.getLogger(__name__) DefT = TypeVar("DefT", bound="DagDefinition") @@ -64,6 +70,19 @@ def read_text(self, encoding: str = "utf-8") -> str: """Read and return the content of the resource as a string.""" return self.read_bytes().decode(encoding) + def sibling(self, relpath: str) -> _SiblingResource: + """ + Return a readable handle to a resource stored next to this definition. + + The handle supports ``read_bytes()``. Resolution matches a Python Dag's + default Jinja search path (relative to the directory of the definition's + own file) and raises if *relpath* escapes it. + + The base cannot locate siblings; a definition type that can + hold them overrides this. + """ + raise ValueError(f"{type(self).__name__} cannot resolve sibling resources") + @abstractmethod def as_file(self) -> contextlib.AbstractContextManager[Path]: """ @@ -133,6 +152,13 @@ def suffix(self) -> str: def read_bytes(self) -> bytes: return self.path.read_bytes() + def sibling(self, relpath: str) -> Path: + base = self.path.parent.resolve() + target = base.joinpath(relpath).resolve() + if not target.is_relative_to(base): + raise ValueError(f"{relpath!r} resolves outside the Dag file's directory {str(base)!r}") + return target + @contextlib.contextmanager def as_file(self) -> Generator[Path, None, None]: yield self.path diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py index c79f9c70e04cb..0d49b362f8803 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py @@ -51,9 +51,9 @@ if TYPE_CHECKING: from collections.abc import Callable - from pathlib import Path from airflow.sdk import BaseOperator + from airflow.sdk.importers.base import _SiblingResource from airflow.sdk.importers.yaml_importer.models import DagDocument, Task, Value @@ -69,8 +69,8 @@ class DagBuilder: and is discarded with the instance. """ - def __init__(self, source_dir: Path | None) -> None: - self._source_dir = source_dir + def __init__(self, sibling: Callable[[str], _SiblingResource]) -> None: + self._sibling = sibling self._operators: dict[str, BaseOperator] = {} self._stub_callables: dict[tuple[str, ...], Callable[..., None]] = {} @@ -179,12 +179,10 @@ def _xcom_arg(self, value: XComRef) -> XComArg: return XComArg(self._operators[target.task], key=target.key) def _read_file(self, rel_path: str) -> str: - """Read a ``$f`` at build time, resolved against ``source_dir``.""" - if self._source_dir is None: - raise YamlDagBuildError(f"cannot read $f {rel_path!r}: the document has no source directory") + """Read a ``$f`` at build time, relative to the document.""" try: - return (self._source_dir / rel_path).read_text(encoding="utf-8") - except OSError as exc: + return self._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]: diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py index 84dfa04abaff2..2a7e70328cbd3 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py @@ -108,11 +108,9 @@ def import_definition( ) return result - # ``$f`` file references resolve against the document's own directory, when it has one. - source_dir = path.parent if (path := getattr(definition, "path", None)) is not None else None for document in documents: try: - dag = DagBuilder(source_dir).build(document) + dag = DagBuilder(definition.sibling).build(document) except Exception as exc: result.errors.append( DagImportError( diff --git a/task-sdk/src/airflow/sdk/importers/zip_importer.py b/task-sdk/src/airflow/sdk/importers/zip_importer.py index 1b7c53ad104e6..ee51f24ff3b8c 100644 --- a/task-sdk/src/airflow/sdk/importers/zip_importer.py +++ b/task-sdk/src/airflow/sdk/importers/zip_importer.py @@ -21,6 +21,7 @@ import contextlib import logging import os +import posixpath import sys import tempfile import threading @@ -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: + base_dir = posixpath.dirname(self.file_path) + member = posixpath.normpath(posixpath.join(base_dir, rel_path)) + within = ( + (member == base_dir or member.startswith(f"{base_dir}/")) + if base_dir + else not (member == ".." or member.startswith("../")) + ) + if posixpath.isabs(member) or not within: + raise ValueError(f"{rel_path!r} resolves outside the Dag file's directory") + return zipfile.Path(self.zip_path, member) + @contextlib.contextmanager def as_file(self) -> Generator[Path, None, None]: # Keep the member's file name: importers may derive identity from it, such as a Dag id. diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py index 1506570cd6576..4e4659fcb4c25 100644 --- a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py @@ -274,6 +274,18 @@ def test_import_file_ref_reads_verbatim(mock_bundle): assert tuple(a.template_fields) == () # $f is verbatim, not a template +@pytest.mark.parametrize("bad_path", ["../../etc/passwd", "/etc/passwd"]) +def test_import_file_ref_outside_dag_dir_is_rejected(mock_bundle, bad_path): + # $f resolves relative to the Dag file's directory, like a Python Dag's template search path; + # a path that escapes it (traversal or absolute) is rejected rather than read. + definition = _def( + mock_bundle, "dag.yaml", _doc(f"- {{id: a, uses: {OP}, with: {{payload: {{$f: '{bad_path}'}}}}}}") + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert any("resolves outside" in e.message for e in result.errors) + + def test_import_template_file_reads_then_templates(mock_bundle): (mock_bundle.path / "query.sql").write_text("SELECT {{ ds }}") definition = _def( From 14a12647478d0bf1be6a502ab2857e835b102a41 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 14:47:21 +0100 Subject: [PATCH 06/25] Re-export to reduce import levels --- .../src/airflow/sdk/importers/yaml_importer/__init__.py | 9 +++++---- .../src/airflow/sdk/importers/yaml_importer/importer.py | 2 +- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/__init__.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/__init__.py index 2f229838802e7..a08ed337b706d 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/__init__.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/__init__.py @@ -19,14 +19,15 @@ from __future__ import annotations +from airflow.sdk.importers.yaml_importer.builder import YamlDagBuildError +from airflow.sdk.importers.yaml_importer.importer import YamlDagImporter from airflow.sdk.importers.yaml_importer.models import DagDocument -from airflow.sdk.importers.yaml_importer.parser import ( - YamlDagParseError, - parse_documents, -) +from airflow.sdk.importers.yaml_importer.parser import YamlDagParseError, parse_documents __all__ = [ "DagDocument", + "YamlDagBuildError", + "YamlDagImporter", "YamlDagParseError", "parse_documents", ] diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py index 2a7e70328cbd3..f1294d95e6534 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py @@ -22,7 +22,7 @@ objects they describe. It is not turned on by default; enable with:: [dag_processor] - dag_importer_configs = [{"classpath": "airflow.sdk.importers.yaml_importer.importer.YamlDagImporter"}] + dag_importer_configs = [{"classpath": "airflow.sdk.importers.yaml_importer.YamlDagImporter"}] """ from __future__ import annotations From 91224e1b660eb999d4b27f591246b39caeddcac5 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 14:48:16 +0100 Subject: [PATCH 07/25] Correctly import in an import context --- .../sdk/importers/yaml_importer/importer.py | 35 ++++++++--------- .../importers/yaml_importer/test_importer.py | 38 +++++++++++++++++++ 2 files changed, 56 insertions(+), 17 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py index f1294d95e6534..ba1e7a90f0e0e 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py @@ -108,24 +108,25 @@ def import_definition( ) return result - for document in documents: - try: - dag = DagBuilder(definition.sibling).build(document) - except Exception as exc: - result.errors.append( - DagImportError( - source_reference=f"{definition!r} (dag_id={document.dag_id!r})", - message=str(exc), - error_type="import", - stacktrace=traceback.format_exc(), + with definition.import_context(): + for document in documents: + try: + dag = DagBuilder(definition.sibling).build(document) + except Exception as exc: + result.errors.append( + DagImportError( + source_reference=f"{definition!r} (dag_id={document.dag_id!r})", + message=str(exc), + error_type="import", + stacktrace=traceback.format_exc(), + ) ) - ) - continue - dag.bundle_name = bundle.name - dag.fileloc = repr(definition) - dag.relative_fileloc = definition.get_relative_loc(bundle.path) - result.dags.append(dag) - log.debug("Built DAG %s from %r", dag.dag_id, definition) + continue + dag.bundle_name = bundle.name + dag.fileloc = repr(definition) + dag.relative_fileloc = definition.get_relative_loc(bundle.path) + result.dags.append(dag) + log.debug("Built DAG %s from %r", dag.dag_id, definition) return result diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py index 4e4659fcb4c25..c749bb453acd0 100644 --- a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py @@ -19,12 +19,14 @@ 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.importers import ZipMemberDagDefinition from airflow.sdk.importers.base import FilesystemDagDefinition from airflow.sdk.importers.yaml_importer.importer import YamlDagImporter @@ -286,6 +288,42 @@ def test_import_file_ref_outside_dag_dir_is_rejected(mock_bundle, bad_path): assert any("resolves outside" in e.message for e in result.errors) +_ZIP_OP_SOURCE = """ +from airflow.sdk import BaseOperator + + +class ZipOp(BaseOperator): + template_fields = ("payload",) + + def __init__(self, *, payload=None, **kwargs): + super().__init__(**kwargs) + self.payload = payload + + def execute(self, context): + return self.payload +""" + + +def test_import_yaml_inside_zip_bundle(tmp_path): + # A YAML Dag inside a zip: `uses:` imports a module next to it (import_context puts the archive on + # sys.path), and `$f` reads a sibling member from the same archive. + zip_path = tmp_path / "bundle.zip" + with zipfile.ZipFile(zip_path, "w") as z: + z.writestr("_ziptest_ops.py", _ZIP_OP_SOURCE) + z.writestr("q.sql", "SELECT 42") + z.writestr( + "dag.yaml", + _doc("- {id: a, uses: _ziptest_ops.ZipOp, with: {payload: {$f: q.sql}}}", dag_id="z"), + ) + definition = ZipMemberDagDefinition(zip_path=zip_path, file_path="dag.yaml") + bundle = SimpleNamespace(name="zbundle", path=tmp_path) + result = YamlDagImporter().import_definition(definition, bundle=bundle) + assert result.success, result.errors + a = result.dags[0].get_task("a") + assert type(a).__name__ == "ZipOp" # imported from a module inside the zip (import_context) + assert a.payload == "SELECT 42" # $f read a sibling member in the same archive + + def test_import_template_file_reads_then_templates(mock_bundle): (mock_bundle.path / "query.sql").write_text("SELECT {{ ds }}") definition = _def( From 19978d95b50d8464ff8ee275227a4d7691355066 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 14:58:29 +0100 Subject: [PATCH 08/25] Refactor DagBuilder interface --- .../sdk/importers/yaml_importer/builder.py | 24 ++++++++++++------- .../sdk/importers/yaml_importer/importer.py | 5 +--- 2 files changed, 16 insertions(+), 13 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py index 0d49b362f8803..4defc3ede0144 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py @@ -37,9 +37,10 @@ import linecache from typing import TYPE_CHECKING, Any, cast -from airflow.sdk import DAG, XComArg +import attrs + +from airflow.sdk import DAG, XComArg, literal from airflow.sdk._shared.module_loading import import_string -from airflow.sdk.definitions._internal.templater import LiteralValue from airflow.sdk.importers.yaml_importer.models import ( ConstRef, FileRef, @@ -52,8 +53,9 @@ if TYPE_CHECKING: from collections.abc import Callable + from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 from airflow.sdk import BaseOperator - from airflow.sdk.importers.base import _SiblingResource + from airflow.sdk.importers.base import DagDefinition from airflow.sdk.importers.yaml_importer.models import DagDocument, Task, Value @@ -61,6 +63,7 @@ class YamlDagBuildError(ValueError): """A validated document could not be built into a DAG.""" +@attrs.define(kw_only=True) class DagBuilder: """ Build a :class:`DagDocument` into a :class:`DAG`. @@ -69,10 +72,10 @@ class DagBuilder: and is discarded with the instance. """ - def __init__(self, sibling: Callable[[str], _SiblingResource]) -> None: - self._sibling = sibling - self._operators: dict[str, BaseOperator] = {} - self._stub_callables: dict[tuple[str, ...], Callable[..., None]] = {} + _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} @@ -85,6 +88,9 @@ def build(self, doc: DagDocument) -> DAG: 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: @@ -167,7 +173,7 @@ def _lower(self, value: Any, *, protect: bool) -> Any: @staticmethod def _protect(value: Any, protect: bool) -> Any: - return LiteralValue(value) if protect else value + return literal(value) if protect else value def _xcom_arg(self, value: XComRef) -> XComArg: """Resolve an :class:`XComRef` against the already-built operators.""" @@ -181,7 +187,7 @@ def _xcom_arg(self, value: XComRef) -> XComArg: def _read_file(self, rel_path: str) -> str: """Read a ``$f`` at build time, relative to the document.""" try: - return self._sibling(rel_path).read_bytes().decode("utf-8") + 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 diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py index ba1e7a90f0e0e..0c9f8cec9377c 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py @@ -111,7 +111,7 @@ def import_definition( with definition.import_context(): for document in documents: try: - dag = DagBuilder(definition.sibling).build(document) + dag = DagBuilder(definition=definition, bundle=bundle).build(document) except Exception as exc: result.errors.append( DagImportError( @@ -122,9 +122,6 @@ def import_definition( ) ) continue - dag.bundle_name = bundle.name - dag.fileloc = repr(definition) - dag.relative_fileloc = definition.get_relative_loc(bundle.path) result.dags.append(dag) log.debug("Built DAG %s from %r", dag.dag_id, definition) From 8be499eb1dd7af338a80ebb3c0f4e056843d005d Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 15:17:31 +0100 Subject: [PATCH 09/25] Stricter top-level and with argument checks --- .../sdk/importers/yaml_importer/builder.py | 16 +++++++++++-- .../importers/yaml_importer/test_importer.py | 23 +++++++++++++++++++ 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py index 4defc3ede0144..4e7ad5256c5a6 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py @@ -39,7 +39,7 @@ import attrs -from airflow.sdk import DAG, XComArg, literal +from airflow.sdk import DAG, BaseOperator, XComArg, literal from airflow.sdk._shared.module_loading import import_string from airflow.sdk.importers.yaml_importer.models import ( ConstRef, @@ -54,7 +54,6 @@ from collections.abc import Callable from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 - from airflow.sdk import BaseOperator from airflow.sdk.importers.base import DagDefinition from airflow.sdk.importers.yaml_importer.models import DagDocument, Task, Value @@ -63,6 +62,9 @@ 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] + + @attrs.define(kw_only=True) class DagBuilder: """ @@ -115,6 +117,16 @@ def _build_stub_task(self, task: Task, args: dict[str, Value]) -> BaseOperator: def _build_operator(self, task: Task) -> BaseOperator: """Instantiate *task* as an operator.""" + 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) diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py index c749bb453acd0..146faecaba114 100644 --- a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py @@ -239,6 +239,29 @@ def test_import_marker_on_non_template_field_errors(mock_bundle, marker): assert any("only allowed on template_fields" in e.message for e in result.errors) +def test_import_base_operator_arg_in_with_errors(mock_bundle): + # `with:` is the operator-specific namespace; a BaseOperator (task-level) arg like `retries` + # there is rejected rather than silently accepted. + definition = _def( + mock_bundle, "dag.yaml", _doc(f"- {{id: a, uses: {OP}, with: {{payload: x, retries: 2}}}}") + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert any("'retries'" in e.message and "not under 'with'" in e.message for e in result.errors) + + +@pytest.mark.parametrize( + "body", ["- {{id: a, uses: {op}, with: {{payload: x}}, note: 1}}", "- {{id: a, run: {{}}, note: 1}}"] +) +def test_import_operator_arg_at_task_level_errors(mock_bundle, body): + # Task-level keys must be BaseOperator args; an operator-specific key (`note`) there is rejected + # (for both `uses:` and `run:` tasks) instead of being passed through raw. + definition = _def(mock_bundle, "dag.yaml", _doc(body.format(op=OP))) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert any("not task-level" in e.message and "'note'" in e.message for e in result.errors) + + def test_import_literal_by_default(mock_bundle): # A bare Jinja-looking string on a field the operator class *does* template is NOT rendered: # literal by default, and the instance's template_fields is emptied of it. From 7857b248b8e163a47402090597f2d98a931c9460 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 15:36:14 +0100 Subject: [PATCH 10/25] Tighten template field patching --- .../sdk/importers/yaml_importer/builder.py | 65 +++++++++++++++---- .../importers/yaml_importer/test_importer.py | 30 +++++++++ 2 files changed, 83 insertions(+), 12 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py index 4e7ad5256c5a6..7f6946804c825 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py @@ -99,11 +99,11 @@ 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 and template - 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. + The task body is a registered SDK task handler resolved at run time; each run argument binds + as a handler argument -- a literal value, or an upstream task's return value (``$x``). The + binding carries no rendering, so templates (``$t``) and keyed XCom (``$x`` with a ``key``) + are rejected. 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. """ from airflow.sdk import task as task_decorators # Lazy to avoid provider registration. @@ -111,7 +111,10 @@ def _build_stub_task(self, task: Task, args: dict[str, Value]) -> BaseOperator: 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") - run_kwargs[name] = self._lower(value, protect=True) + _reject_unbindable_run_marker(task.id_, value) + # Not protected: run arguments bind as literal/XCom, they are never rendered, and a + # LiteralValue wrapper is not JSON-serializable for the handler arg binding. + run_kwargs[name] = self._lower(value, protect=False) decorator = task_decorators.stub(task_id=task.id_, **(task.model_extra or {})) return decorator(self._stub_callable(tuple(run_kwargs)))(**run_kwargs).operator @@ -153,7 +156,7 @@ def _build_operator(self, task: Task) -> BaseOperator: if extras := task.model_extra: # Merge in task-level arguments (raw, never templated). kwargs.update(extras) - return _apply_template_lowering(op_cls(task_id=task.id_, **kwargs), tuple(render_fields)) + return _restrict_template_fields(op_cls(task_id=task.id_, **kwargs), task.with_, render_fields) def _lower(self, value: Any, *, protect: bool) -> Any: """ @@ -227,18 +230,56 @@ def _stub_callable(self, param_names: tuple[str, ...]) -> Callable[..., None]: return stub -def _apply_template_lowering(op: BaseOperator, templated_fields: tuple[str, ...]) -> BaseOperator: +def _restrict_template_fields( + op: BaseOperator, with_: dict[str, Any], render_fields: list[str] +) -> BaseOperator: """ - Make *op* render exactly *templated_fields*. + Make *op* render literal-by-default, per instance. - This allows us to implement AIP-80 file-reading semantics without changing - core; fortunately Airflow reads them per-instance so a simple patch works. + 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 = templated_fields + 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. + """ + if isinstance(value, TemplateRef): + raise YamlDagBuildError( + f"task {task_id!r}: a template ($t) reference is not supported in a run: task; the task " + "handler binding carries literal values and XCom return values, not rendered templates" + ) + if isinstance(value, XComRef): + target = value.target + if not isinstance(target, str) and target.key is not None: + raise YamlDagBuildError( + f"task {task_id!r}: a keyed XCom ($x with 'key') is not supported in a run: task; " + "only a task's return value crosses the handler boundary" + ) + return + if isinstance(value, Literal): + _reject_unbindable_run_marker(task_id, value.root) + elif isinstance(value, dict): + for item in value.values(): + _reject_unbindable_run_marker(task_id, item) + elif isinstance(value, list): + for item in value: + _reject_unbindable_run_marker(task_id, item) + + def _build_schedule(schedule: Any) -> Any: """Build *schedule* into a schedule construct for :class:`DAG`.""" if isinstance(schedule, TimetableSchedule): diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py index 146faecaba114..c73f4da69a025 100644 --- a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py @@ -50,6 +50,23 @@ def execute(self, context): 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__}" + + @pytest.fixture def mock_bundle(tmp_path): bundle_dir = tmp_path / "bundle" @@ -287,6 +304,19 @@ def test_import_template_marker_sets_template_fields(mock_bundle): assert set(a.template_fields) == {"payload"} # $t opts the field into rendering +def test_import_keeps_operator_default_template_fields(mock_bundle): + # The operator templates ``preset`` by default; the author only sets ``payload`` (unmarked). + # Literal-by-default drops the author's unmarked field but keeps the operator's own default. + definition = _def( + mock_bundle, "dag.yaml", _doc(f"- {{id: a, uses: {OP_DEFAULT_TEMPLATE}, with: {{payload: hi}}}}") + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert result.success, result.errors + a = result.dags[0].get_task("a") + assert set(a.template_fields) == {"preset"} # operator default kept; author's unmarked payload dropped + assert a.payload == "hi" # author-set unmarked field no longer renders + + def test_import_file_ref_reads_verbatim(mock_bundle): (mock_bundle.path / "query.sql").write_text("SELECT 1") definition = _def( From b4742e789b5c85d619e298288ad30437f59c6718 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 15:37:49 +0100 Subject: [PATCH 11/25] Tighten literal handling in run --- .../sdk/importers/yaml_importer/builder.py | 19 +++++++------ .../importers/yaml_importer/test_importer.py | 27 ++++++++++++++++--- 2 files changed, 33 insertions(+), 13 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py index 7f6946804c825..18c65100f5ce7 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py @@ -99,11 +99,11 @@ 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; each run argument binds - as a handler argument -- a literal value, or an upstream task's return value (``$x``). The - binding carries no rendering, so templates (``$t``) and keyed XCom (``$x`` with a ``key``) - are rejected. 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. + 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. """ from airflow.sdk import task as task_decorators # Lazy to avoid provider registration. @@ -112,8 +112,6 @@ def _build_stub_task(self, task: Task, args: dict[str, Value]) -> BaseOperator: 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) - # Not protected: run arguments bind as literal/XCom, they are never rendered, and a - # LiteralValue wrapper is not JSON-serializable for the handler arg binding. run_kwargs[name] = self._lower(value, protect=False) decorator = task_decorators.stub(task_id=task.id_, **(task.model_extra or {})) return decorator(self._stub_callable(tuple(run_kwargs)))(**run_kwargs).operator @@ -253,9 +251,10 @@ 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. + 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. """ if isinstance(value, TemplateRef): raise YamlDagBuildError( diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py index c73f4da69a025..db4043a41c2e8 100644 --- a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py @@ -202,9 +202,7 @@ def test_import_code_task_builds_stub_operator(mock_bundle): 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 - # op_kwargs is a template field, so the literal is protected to stay literal when rendered. - assert isinstance(b.op_kwargs["label"], LiteralValue) - assert b.op_kwargs["label"].value == "done" + 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. @@ -237,6 +235,29 @@ def test_import_code_task_rejects_non_identifier_arg(mock_bundle): assert any("not a valid identifier" in e.message for e in result.errors) +def test_import_run_rejects_template(mock_bundle): + # A $t in run: cannot be carried by the handler binding (it would arrive unrendered); reject it. + 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): + # Only a task's return value crosses the handler boundary; a keyed $x is rejected at build time. + 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) + + @pytest.mark.parametrize("marker", ["{$x: a}", "{$t: 'v'}", "{$f: q.sql}"]) def test_import_marker_on_non_template_field_errors(mock_bundle, marker): # A marker only works where the operator class already templates. ``note`` is not in From b39b0e7682dc343969a2a700762cfd4b448b79b4 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 15:55:10 +0100 Subject: [PATCH 12/25] Consolidate import-time errors from a YAML file --- .../sdk/importers/yaml_importer/importer.py | 43 +++++++++++++------ .../importers/yaml_importer/test_importer.py | 13 +++++- 2 files changed, 41 insertions(+), 15 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py index 0c9f8cec9377c..8813eac46149b 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py @@ -49,6 +49,7 @@ from pathlib import Path from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 + from airflow.sdk import DAG from airflow.sdk.importers.base import DagDefinition log = logging.getLogger(__name__) @@ -98,8 +99,8 @@ def import_definition( ) return result - # A parse failure fails the whole file (like a Python module that won't import); a build - # failure is per-document so the file's other DAGs still load. + # A parse failure fails the whole file (like a Python module that won't + # import). All documents' failures are reported together. try: documents = list(parse_documents(source, source=repr(definition))) except YamlDagParseError as exc: @@ -108,23 +109,37 @@ def import_definition( ) return result + dags: list[DAG] = [] + build_errors: list[Exception] = [] with definition.import_context(): for document in documents: try: - dag = DagBuilder(definition=definition, bundle=bundle).build(document) + dags.append(DagBuilder(definition=definition, bundle=bundle).build(document)) except Exception as exc: - result.errors.append( - DagImportError( - source_reference=f"{definition!r} (dag_id={document.dag_id!r})", - message=str(exc), - error_type="import", - stacktrace=traceback.format_exc(), - ) - ) - continue - result.dags.append(dag) - log.debug("Built DAG %s from %r", dag.dag_id, definition) + exc.add_note(f"dag_id={document.dag_id!r}") + build_errors.append(exc) + + if build_errors: + summary = f"{len(build_errors)} of {len(documents)} document(s) failed to build" + detail = "\n".join(f"{err.__notes__[-1]}: {err}" for err in build_errors) + stacktrace = "\n".join( + f"{err.__notes__[-1]}:\n" + + "".join(traceback.format_exception(type(err), err, err.__traceback__)) + for err in build_errors + ) + result.errors.append( + DagImportError( + source_reference=repr(definition), + message=f"{summary}:\n{detail}", + error_type="import", + stacktrace=stacktrace, + ) + ) + return result + result.dags.extend(dags) + for dag in dags: + log.debug("Built Dag %s from %r", dag.dag_id, definition) return result def get_source_code(self, definition: DagDefinition, dag_id: str | None = None) -> DagSourceCode: diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py index db4043a41c2e8..25fa9f0938ca8 100644 --- a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py @@ -441,8 +441,19 @@ def test_import_unknown_operator_errors(mock_bundle): assert result.errors +def test_import_one_bad_document_fails_whole_file(mock_bundle): + good = _doc(f"- {{id: a, uses: {OP}, with: {{payload: x}}}}", dag_id="good") + bad = _doc("- {id: a, uses: does.not.Exist, with: {}}", dag_id="bad") + definition = _def(mock_bundle, "dag.yaml", f"{good}\n---\n{bad}") + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert len(result.errors) == 1 + err = result.errors[0] + assert "dag_id=" not in err.source_reference # keyed by the file, not the document + assert "dag_id='bad'" in err.message # the failing document is named, not mashed + + def test_import_dependency_cycle_errors(mock_bundle): - # The parser allows a cycle (full cycle detection is the builder's job); the build rejects it. definition = _def( mock_bundle, "dag.yaml", From d83a145b9812a60bfd9b9eeeabcfb355bd3657ce Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 16:00:48 +0100 Subject: [PATCH 13/25] Reject duplicate dag_id in a file --- .../sdk/importers/yaml_importer/importer.py | 14 +++++++++----- .../importers/yaml_importer/test_importer.py | 9 +++++++++ 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py index 8813eac46149b..ebc1dcf9c31a7 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py @@ -41,7 +41,7 @@ find_file_dag_definitions, get_file_suffix, ) -from airflow.sdk.importers.yaml_importer.builder import DagBuilder +from airflow.sdk.importers.yaml_importer.builder import DagBuilder, YamlDagBuildError from airflow.sdk.importers.yaml_importer.parser import YamlDagParseError, parse_documents if TYPE_CHECKING: @@ -109,15 +109,19 @@ def import_definition( ) return result - dags: list[DAG] = [] + dags: dict[str, DAG] = {} build_errors: list[Exception] = [] with definition.import_context(): for document in documents: try: - dags.append(DagBuilder(definition=definition, bundle=bundle).build(document)) + dag = DagBuilder(definition=definition, bundle=bundle).build(document) + if dag.dag_id in dags: + raise YamlDagBuildError(f"duplicate dag_id {dag.dag_id!r} across documents") except Exception as exc: exc.add_note(f"dag_id={document.dag_id!r}") build_errors.append(exc) + else: + dags[dag.dag_id] = dag if build_errors: summary = f"{len(build_errors)} of {len(documents)} document(s) failed to build" @@ -137,8 +141,8 @@ def import_definition( ) return result - result.dags.extend(dags) - for dag in dags: + result.dags.extend(dags.values()) + for dag in dags.values(): log.debug("Built Dag %s from %r", dag.dag_id, definition) return result diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py index 25fa9f0938ca8..5d4a6cfe1b59e 100644 --- a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py @@ -453,6 +453,15 @@ def test_import_one_bad_document_fails_whole_file(mock_bundle): assert "dag_id='bad'" in err.message # the failing document is named, not mashed +def test_import_duplicate_dag_id_across_documents_errors(mock_bundle): + one = _doc(f"- {{id: a, uses: {OP}, with: {{payload: x}}}}", dag_id="dup") + two = _doc(f"- {{id: b, uses: {OP}, with: {{payload: y}}}}", dag_id="dup") + definition = _def(mock_bundle, "dag.yaml", f"{one}\n---\n{two}") + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert any("duplicate dag_id 'dup'" in e.message for e in result.errors) + + def test_import_dependency_cycle_errors(mock_bundle): definition = _def( mock_bundle, From 1d3963047a843846690066c12e69026de3504078 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 16:09:06 +0100 Subject: [PATCH 14/25] Tighten might_contain_dag heuristic --- .../sdk/importers/yaml_importer/importer.py | 17 +++++++++++++++-- .../importers/yaml_importer/test_importer.py | 6 ++++++ 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py index ebc1dcf9c31a7..97d1d5887a898 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py @@ -42,6 +42,7 @@ get_file_suffix, ) from airflow.sdk.importers.yaml_importer.builder import DagBuilder, YamlDagBuildError +from airflow.sdk.importers.yaml_importer.migrator import _SCHEMA_URL from airflow.sdk.importers.yaml_importer.parser import YamlDagParseError, parse_documents if TYPE_CHECKING: @@ -54,6 +55,11 @@ log = logging.getLogger(__name__) +# The sniff matches the canonical Dag ``$schema`` URL, not a bare ``$schema``, so a JSON-Schema +# file (or any other YAML declaring a $schema) in the bundle is not queued and then stuck as a +# parse error. Derived from the migrator's URL template to stay in sync. +_DAG_SCHEMA_PREFIX = _SCHEMA_URL.partition("{date}")[0].encode() + class YamlDagImporter(AbstractDagImporter[FileDagDefinition]): """Importer for YAML/JSON DAG sources. Opt-in via ``dag_importer_configs`` (not a default).""" @@ -150,10 +156,17 @@ def get_source_code(self, definition: DagDefinition, dag_id: str | None = None) return DagSourceCode(source_code=definition.read_text(encoding="utf-8"), language="yaml") def might_contain_dag(self, definition: DagDefinition, safe_mode: bool) -> bool: - """Cheap sniff: a DAG document declares a ``$schema``. Keep the file if the read fails.""" + """ + Cheap sniff: the file declares a ``$schema`` set to the canonical Dag URL. + + Both are required so a file that merely mentions the URL (without a ``$schema`` declaration) + is skipped. The sniff stays lenient past that -- a malformed Dag still reaches the parser and + gets a clear error rather than silently disappearing. Keep the file if the read fails. + """ if not safe_mode: return True try: - return b"$schema" in definition.read_bytes() + data = definition.read_bytes() except OSError: return True + return b"$schema" in data and _DAG_SCHEMA_PREFIX in data diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py index 5d4a6cfe1b59e..b9a37e037f361 100644 --- a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py @@ -103,8 +103,14 @@ 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 From 8ae0b9c935b2ff5e82ffd69f9b22f0dccdccb925 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 16:21:12 +0100 Subject: [PATCH 15/25] Coerce YAML input temporal values to datetimes --- .../sdk/importers/yaml_importer/builder.py | 36 ++++++++++++++++--- .../importers/yaml_importer/test_importer.py | 21 +++++++++++ 2 files changed, 53 insertions(+), 4 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py index 18c65100f5ce7..5dfea477ed73f 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py @@ -33,6 +33,7 @@ from __future__ import annotations +import datetime import keyword import linecache from typing import TYPE_CHECKING, Any, cast @@ -81,7 +82,8 @@ class DagBuilder: def build(self, doc: DagDocument) -> DAG: tasks = {t.id_: t for t in doc.tasks} - with DAG(dag_id=doc.dag_id, schedule=_build_schedule(doc.schedule), **(doc.model_extra or {})) as dag: + extras = _coerce_temporal(doc.model_extra or {}) + with DAG(dag_id=doc.dag_id, schedule=_build_schedule(doc.schedule), **extras) as dag: # A task depended by another task must be built first. for task_id in _topological_order(doc.tasks): self._operators[task_id] = self._build_operator(tasks[task_id]) @@ -113,7 +115,7 @@ def _build_stub_task(self, task: Task, args: dict[str, Value]) -> BaseOperator: 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_, **(task.model_extra or {})) + decorator = task_decorators.stub(task_id=task.id_, **_coerce_temporal(task.model_extra or {})) return decorator(self._stub_callable(tuple(run_kwargs)))(**run_kwargs).operator def _build_operator(self, task: Task) -> BaseOperator: @@ -151,8 +153,8 @@ def _build_operator(self, task: Task) -> BaseOperator: render_fields.append(name) if off_template: raise YamlDagBuildError(f"task {task.id_!r}: value grammar only allowed on template_fields") - if extras := task.model_extra: # Merge in task-level arguments (raw, never templated). - kwargs.update(extras) + if task.model_extra: # Merge in task-level arguments (raw, never templated). + kwargs.update(_coerce_temporal(task.model_extra)) return _restrict_template_fields(op_cls(task_id=task.id_, **kwargs), task.with_, render_fields) @@ -279,6 +281,32 @@ def _reject_unbindable_run_marker(task_id: str, value: Any) -> None: _reject_unbindable_run_marker(task_id, item) +def _coerce_temporal(mapping: dict[str, Any]) -> dict[str, Any]: + """ + Promote YAML-native dates and ISO date/datetime strings to datetimes. + + Dag/operator arguments such as ``start_date`` are datetimes, but YAML has + distinct date and date-time formats. Convert dates to midnight times. + + Strings are also parsed automatically to datetimes if they LOOK like the ISO + 8601 format. This helps with JSON compatibility. + """ + return {key: _to_datetime(value) for key, value in mapping.items()} + + +def _to_datetime(value: Any) -> Any: + if isinstance(value, datetime.datetime): + return value + if isinstance(value, datetime.date): + return datetime.datetime.combine(value, datetime.time.min) + if isinstance(value, str): + try: + return datetime.datetime.fromisoformat(value) + except ValueError: + return value + return value + + def _build_schedule(schedule: Any) -> Any: """Build *schedule* into a schedule construct for :class:`DAG`.""" if isinstance(schedule, TimetableSchedule): diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py index b9a37e037f361..9a960201396a0 100644 --- a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py @@ -17,6 +17,7 @@ from __future__ import annotations +import datetime import inspect import textwrap import zipfile @@ -189,6 +190,26 @@ def test_import_schedule_cron(mock_bundle): assert result.dags[0].schedule == "0 3 * * *" +@pytest.mark.parametrize( + ("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): + # A YAML date or ISO string for a temporal arg is promoted to a (localized) datetime so DAG() + # accepts it, rather than failing with an unhelpful cast error (format section 2d). + 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_code_task_builds_stub_operator(mock_bundle): # A code task (run:) builds the stub operator Airflow understands (@task.stub). Its body is a # registered SDK task handler; the run mapping binds as that handler's arguments (op_kwargs), From 73f6c96aed43f9d0aa7a7594453bb1d36533348f Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 16:28:57 +0100 Subject: [PATCH 16/25] Rely on stdlib graphlib for ordering --- .../sdk/importers/yaml_importer/builder.py | 18 ++++++------------ .../importers/yaml_importer/test_importer.py | 7 ++++++- 2 files changed, 12 insertions(+), 13 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py index 5dfea477ed73f..b92c8a5418d1b 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py @@ -34,6 +34,7 @@ from __future__ import annotations import datetime +import graphlib import keyword import linecache from typing import TYPE_CHECKING, Any, cast @@ -343,15 +344,8 @@ def _field_markers(value: Any) -> tuple[bool, bool]: def _topological_order(tasks: list[Task]) -> list[str]: """Return task ids ordered so every task follows the tasks it depends on.""" - upstreams = {task.id_: task.get_dependency_ids() for task in tasks} - order: list[str] = [] - done: set[str] = set() - remaining = [task.id_ for task in tasks] - while remaining: - ready = [task_id for task_id in remaining if upstreams[task_id] <= done] - if not ready: - raise YamlDagBuildError(f"dependency cycle among tasks: {sorted(remaining)}") - order.extend(ready) - done.update(ready) - remaining = [task_id for task_id in remaining if task_id not in done] - return order + graph = {task.id_: task.get_dependency_ids() for task in tasks} + try: + return list(graphlib.TopologicalSorter(graph).static_order()) + except graphlib.CycleError as exc: + raise YamlDagBuildError(f"dependency cycle among tasks: {exc.args[1]}") from exc diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py index 9a960201396a0..469f591963d5f 100644 --- a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py @@ -490,14 +490,19 @@ def test_import_duplicate_dag_id_across_documents_errors(mock_bundle): def test_import_dependency_cycle_errors(mock_bundle): + # a <-> b form the cycle; c only depends on it and must not be reported as part of it. definition = _def( mock_bundle, "dag.yaml", _doc(f""" - {{id: a, uses: {OP}, with: {{payload: x}}, needs: [b]}} - {{id: b, uses: {OP}, with: {{payload: y}}, needs: [a]}} + - {{id: c, uses: {OP}, with: {{payload: z}}, needs: [a]}} """), ) result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) assert not result.dags - assert any("cycle" in e.message for e in result.errors) + (error,) = [e for e in result.errors if "cycle" in e.message] + assert "'a'" in error.message + assert "'b'" in error.message + assert "'c'" not in error.message From 2cdc51736e903d363ab203b6820c9524f1798536 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 16:31:07 +0100 Subject: [PATCH 17/25] Move unnecessary local import up --- task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py index b92c8a5418d1b..20291956b302b 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py @@ -41,7 +41,7 @@ import attrs -from airflow.sdk import DAG, BaseOperator, XComArg, literal +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, @@ -108,8 +108,6 @@ def _build_stub_task(self, task: Task, args: dict[str, Value]) -> BaseOperator: run argument names, so serialization can build and later validate (AIP-108) the handler's argument spec. """ - from airflow.sdk import task as task_decorators # Lazy to avoid provider registration. - run_kwargs: dict[str, Any] = {} for name, value in args.items(): if not name.isidentifier() or keyword.iskeyword(name): From 186877395a95edfc8635f2c61bc4e313a76482c0 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 16:39:26 +0100 Subject: [PATCH 18/25] Tighten tests --- .../importers/yaml_importer/test_importer.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py index 469f591963d5f..8044fd43258b7 100644 --- a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py @@ -38,6 +38,7 @@ 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) @@ -147,6 +148,8 @@ def test_import_builds_operator_dag(mock_bundle): 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): @@ -176,7 +179,10 @@ 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 result.errors + 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): @@ -338,7 +344,7 @@ def test_import_literal_by_default(mock_bundle): a = result.dags[0].get_task("a") assert a.payload == "{{ ds }}" # kept literal assert tuple(a.template_fields) == () # not rendered - assert tuple(a.template_ext) == () # and the suffix-based file-load foot-gun is closed + assert tuple(a.template_ext) == () def test_import_template_marker_sets_template_fields(mock_bundle): @@ -465,7 +471,7 @@ def test_import_unknown_operator_errors(mock_bundle): definition = _def(mock_bundle, "dag.yaml", _doc("- {id: a, uses: does.not.Exist, with: {}}")) result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) assert not result.dags - assert result.errors + assert any("No module named" in e.message for e in result.errors) def test_import_one_bad_document_fails_whole_file(mock_bundle): From 0d4479d96d19be096a1b7534833ccc795bb71849 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 17:06:25 +0100 Subject: [PATCH 19/25] Try to smartly convert datetimes --- .../sdk/importers/yaml_importer/builder.py | 98 +++++++++++++------ .../sdk/importers/yaml_importer/importer.py | 11 +-- .../importers/yaml_importer/test_importer.py | 13 +++ 3 files changed, 88 insertions(+), 34 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py index 20291956b302b..e8f9552d85f7d 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py @@ -34,10 +34,12 @@ from __future__ import annotations import datetime +import functools import graphlib import keyword import linecache -from typing import TYPE_CHECKING, Any, cast +import re +from typing import TYPE_CHECKING, Any, NamedTuple, cast import attrs @@ -83,7 +85,7 @@ class DagBuilder: def build(self, doc: DagDocument) -> DAG: tasks = {t.id_: t for t in doc.tasks} - extras = _coerce_temporal(doc.model_extra or {}) + extras = _coerce_temporal(doc.model_extra or {}, DAG) with DAG(dag_id=doc.dag_id, schedule=_build_schedule(doc.schedule), **extras) as dag: # A task depended by another task must be built first. for task_id in _topological_order(doc.tasks): @@ -114,7 +116,9 @@ def _build_stub_task(self, task: Task, args: dict[str, Value]) -> BaseOperator: 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 {})) + 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: @@ -153,7 +157,7 @@ def _build_operator(self, task: Task) -> BaseOperator: if off_template: raise YamlDagBuildError(f"task {task.id_!r}: value grammar only allowed on template_fields") if task.model_extra: # Merge in task-level arguments (raw, never templated). - kwargs.update(_coerce_temporal(task.model_extra)) + kwargs.update(_coerce_temporal(task.model_extra, BaseOperator)) return _restrict_template_fields(op_cls(task_id=task.id_, **kwargs), task.with_, render_fields) @@ -280,17 +284,21 @@ def _reject_unbindable_run_marker(task_id: str, value: Any) -> None: _reject_unbindable_run_marker(task_id, item) -def _coerce_temporal(mapping: dict[str, Any]) -> dict[str, Any]: +def _coerce_temporal(mapping: dict[str, Any], target: type) -> dict[str, Any]: """ Promote YAML-native dates and ISO date/datetime strings to datetimes. - Dag/operator arguments such as ``start_date`` are datetimes, but YAML has - distinct date and date-time formats. Convert dates to midnight times. + Some fields on a Dag or operator are datetimes, but YAML has distinct date + and date-time formats, so a bare ``2024-01-01`` or a quoted ISO string would + otherwise reach the datetime field unconverted. A date becomes midnight; an + ISO string is parsed (which also helps JSON compatibility). - Strings are also parsed automatically to datetimes if they LOOK like the ISO - 8601 format. This helps with JSON compatibility. + Only values bound to a *target* argument annotated as ``datetime`` are + coerced. A non-temporal argument is left untouched, so a string that merely + looks like a date keeps its type. """ - return {key: _to_datetime(value) for key, value in mapping.items()} + temporal = _temporal_fields(target) + return {key: _to_datetime(value) if key in temporal else value for key, value in mapping.items()} def _to_datetime(value: Any) -> Any: @@ -306,6 +314,36 @@ def _to_datetime(value: Any) -> Any: return value +# We annotate arguments with either stdlib datetime or pendulum DateTime. +# Matched by name (see _temporal_fields) rather than resolved to avoid having +# to evaluate dynamically. +_TEMPORAL_TYPE_NAMES = frozenset({"datetime", "DateTime"}) +_TYPE_NAME_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_.]*") + + +@functools.cache +def _temporal_fields(target: type) -> frozenset[str]: + """ + Names of *target*'s arguments annotated as a datetime type. + + The annotation is read as text and matched against the datetime type names + we use; it is never resolved, so one referring to a ``TYPE_CHECKING``-only + name does not need importing. This is not perfect; we'll try to fix edge + cases as we go. + """ + found: set[str] = set() + for klass in getattr(target, "__mro__", (target,)): + for name, annotation in getattr(klass, "__annotations__", {}).items(): + if name not in found and _is_temporal_annotation(str(annotation)): + found.add(name) + return frozenset(found) + + +def _is_temporal_annotation(annotation: str) -> bool: + """Return True if *annotation* names a datetime type, ignoring any module prefix.""" + return any(ref.rsplit(".", 1)[-1] in _TEMPORAL_TYPE_NAMES for ref in _TYPE_NAME_RE.findall(annotation)) + + def _build_schedule(schedule: Any) -> Any: """Build *schedule* into a schedule construct for :class:`DAG`.""" if isinstance(schedule, TimetableSchedule): @@ -313,31 +351,35 @@ def _build_schedule(schedule: Any) -> Any: return schedule -def _field_markers(value: Any) -> tuple[bool, bool]: +class _FieldMarkers(NamedTuple): + """Flags describing a value-grammar node's templating needs.""" + + needs_render: bool # carries $t or $x anywhere -- the field must render + has_marker: bool # carries $x/$t/$f -- only allowed on a templated field + + +def _field_markers(value: Any) -> _FieldMarkers: """ - Scan a value-grammar node for ``(needs_render, has_marker)``. + Scan a value-grammar node for its :class:`_FieldMarkers`. - ``needs_render`` is True when it carries a ``$t`` or ``$x`` anywhere (the field must render); - ``has_marker`` is True when it carries any of ``$x``/``$t``/``$f`` (only allowed on a templated - field). ``$const`` is a force-literal escape, not a marker. + ``$const`` is a force-literal escape, not a marker. """ - if isinstance(value, XComRef): - return True, True - if isinstance(value, TemplateRef): - return True, True + if isinstance(value, (XComRef, TemplateRef)): + return _FieldMarkers(needs_render=True, has_marker=True) if isinstance(value, FileRef): - return False, True + return _FieldMarkers(needs_render=False, has_marker=True) if isinstance(value, ConstRef): - return False, False + return _FieldMarkers(needs_render=False, has_marker=False) if isinstance(value, Literal): return _field_markers(value.root) - if isinstance(value, dict): - flags = [_field_markers(v) for v in value.values()] - return any(r for r, _ in flags), any(m for _, m in flags) - if isinstance(value, list): - flags = [_field_markers(v) for v in value] - return any(r for r, _ in flags), any(m for _, m in flags) - return False, False + if isinstance(value, (dict, list)): + children = value.values() if isinstance(value, dict) else value + flags = [_field_markers(v) for v in children] + return _FieldMarkers( + needs_render=any(f.needs_render for f in flags), + has_marker=any(f.has_marker for f in flags), + ) + return _FieldMarkers(needs_render=False, has_marker=False) def _topological_order(tasks: list[Task]) -> list[str]: diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py index 97d1d5887a898..16dc9144077fa 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py @@ -92,12 +92,13 @@ def import_definition( bundle: BaseDagBundle, ) -> DagImportResult: result = DagImportResult(definition=definition) + ref = repr(definition) try: source = definition.read_text(encoding="utf-8") except Exception as exc: result.errors.append( DagImportError( - source_reference=repr(definition), + source_reference=ref, message=str(exc), error_type="import", stacktrace=traceback.format_exc(), @@ -108,11 +109,9 @@ def import_definition( # A parse failure fails the whole file (like a Python module that won't # import). All documents' failures are reported together. try: - documents = list(parse_documents(source, source=repr(definition))) + documents = list(parse_documents(source, source=ref)) except YamlDagParseError as exc: - result.errors.append( - DagImportError(source_reference=repr(definition), message=str(exc), error_type="parse") - ) + result.errors.append(DagImportError(source_reference=ref, message=str(exc), error_type="parse")) return result dags: dict[str, DAG] = {} @@ -139,7 +138,7 @@ def import_definition( ) result.errors.append( DagImportError( - source_reference=repr(definition), + source_reference=ref, message=f"{summary}:\n{detail}", error_type="import", stacktrace=stacktrace, diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py index 8044fd43258b7..ca1ee2e107c94 100644 --- a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py @@ -216,6 +216,19 @@ def test_import_coerces_temporal_start_date(mock_bundle, value, hour): assert (start.year, start.month, start.day, start.hour) == (2024, 1, 1, hour) +def test_import_does_not_coerce_nontemporal_arg(mock_bundle): + # _coerce_temporal introspects the target's argument types and only promotes datetime-typed ones. + # A date-looking string on a non-temporal arg (owner: str) keeps its type instead of becoming a datetime. + 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_code_task_builds_stub_operator(mock_bundle): # A code task (run:) builds the stub operator Airflow understands (@task.stub). Its body is a # registered SDK task handler; the run mapping binds as that handler's arguments (op_kwargs), From 88318c6322ff9cf7e91f03019f32726cf9c0c5ed Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 17:19:44 +0100 Subject: [PATCH 20/25] Tighten --- task-sdk/src/airflow/sdk/importers/base.py | 2 +- .../sdk/importers/yaml_importer/builder.py | 15 +++++--- .../sdk/importers/yaml_importer/importer.py | 14 ++++---- .../sdk/importers/yaml_importer/migrator.py | 3 +- .../importers/yaml_importer/test_importer.py | 35 +++++++++++++++++++ 5 files changed, 55 insertions(+), 14 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/base.py b/task-sdk/src/airflow/sdk/importers/base.py index b1c04842c187d..caef6070e3cd8 100644 --- a/task-sdk/src/airflow/sdk/importers/base.py +++ b/task-sdk/src/airflow/sdk/importers/base.py @@ -81,7 +81,7 @@ def sibling(self, relpath: str) -> _SiblingResource: The base cannot locate siblings; a definition type that can hold them overrides this. """ - raise ValueError(f"{type(self).__name__} cannot resolve sibling resources") + raise NotImplementedError(f"{type(self).__name__} cannot resolve sibling resources") @abstractmethod def as_file(self) -> contextlib.AbstractContextManager[Path]: diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py index e8f9552d85f7d..3870eca102867 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py @@ -67,6 +67,7 @@ class YamlDagBuildError(ValueError): _BASE_TASK_ARGS = frozenset(BaseOperator.__init__._BaseOperatorMeta__param_names) # type: ignore[attr-defined] +_UNSUPPORTED_TASK_ARGS = frozenset({"task_id", "dag", "task_group", "default_args"}) @attrs.define(kw_only=True) @@ -87,11 +88,9 @@ 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: - # A task depended by another task must be built first. for task_id in _topological_order(doc.tasks): self._operators[task_id] = self._build_operator(tasks[task_id]) - # Wire dependency edges explicitly. for task in doc.tasks: for upstream in sorted(task.get_dependency_ids()): self._operators[upstream] >> self._operators[task.id_] @@ -123,6 +122,13 @@ def _build_stub_task(self, task: Task, args: dict[str, Value]) -> BaseOperator: 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'" @@ -142,7 +148,6 @@ def _build_operator(self, task: Task) -> BaseOperator: else: op_cls = import_string(task.uses) - # Validate template fields. template_fields = set(getattr(op_cls, "template_fields", ())) kwargs: dict[str, Any] = {} render_fields: list[str] = [] @@ -155,7 +160,9 @@ def _build_operator(self, task: Task) -> BaseOperator: if needs_render: render_fields.append(name) if off_template: - raise YamlDagBuildError(f"task {task.id_!r}: value grammar only allowed on template_fields") + 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)) diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py index 16dc9144077fa..080e0fc97afb3 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py @@ -16,7 +16,7 @@ # under the License. """ -AIP-85 YAML/JSON Dag importer. +AIP-85 YAML Dag importer. This discovers ``.yaml`` and ``.yml`` Dag files in a bundle, and builds the Dag objects they describe. It is not turned on by default; enable with:: @@ -42,7 +42,7 @@ get_file_suffix, ) from airflow.sdk.importers.yaml_importer.builder import DagBuilder, YamlDagBuildError -from airflow.sdk.importers.yaml_importer.migrator import _SCHEMA_URL +from airflow.sdk.importers.yaml_importer.migrator import DAG_SCHEMA_PREFIX from airflow.sdk.importers.yaml_importer.parser import YamlDagParseError, parse_documents if TYPE_CHECKING: @@ -57,12 +57,12 @@ # The sniff matches the canonical Dag ``$schema`` URL, not a bare ``$schema``, so a JSON-Schema # file (or any other YAML declaring a $schema) in the bundle is not queued and then stuck as a -# parse error. Derived from the migrator's URL template to stay in sync. -_DAG_SCHEMA_PREFIX = _SCHEMA_URL.partition("{date}")[0].encode() +# parse error. Reuses the migrator's public prefix to stay in sync. +_DAG_SCHEMA_PREFIX = DAG_SCHEMA_PREFIX.encode() class YamlDagImporter(AbstractDagImporter[FileDagDefinition]): - """Importer for YAML/JSON DAG sources. Opt-in via ``dag_importer_configs`` (not a default).""" + """Importer for YAML DAG sources. Opt-in via ``dag_importer_configs`` (not a default).""" supported_extensions = [".yaml", ".yml"] @@ -132,9 +132,7 @@ def import_definition( summary = f"{len(build_errors)} of {len(documents)} document(s) failed to build" detail = "\n".join(f"{err.__notes__[-1]}: {err}" for err in build_errors) stacktrace = "\n".join( - f"{err.__notes__[-1]}:\n" - + "".join(traceback.format_exception(type(err), err, err.__traceback__)) - for err in build_errors + "".join(traceback.format_exception(type(err), err, err.__traceback__)) for err in build_errors ) result.errors.append( DagImportError( diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/migrator.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/migrator.py index 9a27c07e4b99c..90f6ebed807d9 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/migrator.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/migrator.py @@ -44,9 +44,10 @@ from cadwyn import VersionBundle _SCHEMA_URL = "https://airflow.apache.org/schemas/dag/{date}.json" -# The canonical ``$schema`` URL with the version captured; keep in sync with ``_SCHEMA_URL``. _SCHEMA_URL_RE = re.compile(r"https://airflow\.apache\.org/schemas/dag/(\d{4}-\d{2}-\d{2})\.json") +DAG_SCHEMA_PREFIX = _SCHEMA_URL.partition("{date}")[0] + def schema_from_version(date: str) -> str: """Build the canonical ``$schema`` URL for a published version *date*.""" diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py index ca1ee2e107c94..e08051db441e3 100644 --- a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py @@ -27,6 +27,7 @@ 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 @@ -196,6 +197,40 @@ def test_import_schedule_cron(mock_bundle): assert result.dags[0].schedule == "0 3 * * *" +def test_import_schedule_timetable(mock_bundle): + # A mapping schedule (uses/with) is built into a timetable via its import path (_build_schedule). + 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): + # $x with {task, key} on a uses: task resolves to a keyed XComArg (the _xcom_arg key branch); + # keyed XCom is rejected only in run: tasks. + 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"} + + @pytest.mark.parametrize( ("value", "hour"), [("2024-01-01", 0), ("'2024-01-01'", 0), ("2024-01-01T03:30:00", 3)], # YAML date, ISO string, datetime From 86e348f8deab67121bff85fa1059041502b93dca Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 17:28:48 +0100 Subject: [PATCH 21/25] Fix leaked rst in schema --- .../sdk/importers/yaml_importer/models.py | 27 ++++++++++++++----- .../sdk/importers/yaml_importer/schema.json | 4 +-- 2 files changed, 23 insertions(+), 8 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/models.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/models.py index 4fc628e1f5626..35628dfed7131 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/models.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/models.py @@ -147,11 +147,22 @@ class XComRef(_MarkerRef): class TemplateRef(_MarkerRef): """ - A Jinja template, rendered server-side. + A Jinja template, rendered when the task runs. - ``source`` is a Jinja string, or a :class:`FileRef` to load a template file. + The template is a Jinja string, or a file reference (``$f``) to load a + template file. """ + model_config = ConfigDict( + extra="forbid", + json_schema_extra={ + "description": ( + "A Jinja template, rendered when the task runs. This can be an " + "inline string, or a file reference ($f) to load a template file." + ), + }, + ) + _marker_keys = TEMPLATE_KEYS source: str | FileRef = Field(serialization_alias="$t", validation_alias=AliasChoices(*TEMPLATE_KEYS)) @@ -159,13 +170,17 @@ class TemplateRef(_MarkerRef): class FileRef(_MarkerRef): """ - A file input (AIP-80's FileTemplate). + A file input (AIP-80 FileTemplate). - ``path`` is a static path. The inverse nesting (templated filename) is - currently reserved and rejected since it does not work with current Airflow - internals. + The path is a static string. A templated filename (a ``$f`` wrapping a + ``$t``) is reserved for a future release and rejected for now. """ + model_config = ConfigDict( + extra="forbid", + json_schema_extra={"description": "A file input."}, + ) + _marker_keys = FILE_KEYS path: str | TemplateRef = Field(serialization_alias="$f", validation_alias=AliasChoices(*FILE_KEYS)) diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/schema.json b/task-sdk/src/airflow/sdk/importers/yaml_importer/schema.json index 10f5a0f559c1f..67686d202ae1c 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/schema.json +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/schema.json @@ -16,7 +16,7 @@ }, "FileRef": { "additionalProperties": false, - "description": "A file input (AIP-80's FileTemplate).\n\n``path`` is a static path. The inverse nesting (templated filename) is\ncurrently reserved and rejected since it does not work with current Airflow\ninternals.", + "description": "A file input.", "oneOf": [ { "required": [ @@ -234,7 +234,7 @@ }, "TemplateRef": { "additionalProperties": false, - "description": "A Jinja template, rendered server-side.\n\n``source`` is a Jinja string, or a :class:`FileRef` to load a template file.", + "description": "A Jinja template, rendered when the task runs. This can be an inline string, or a file reference ($f) to load a template file.", "oneOf": [ { "required": [ From 0d195121918fe80e45eda2c8d3dff9aeb97fb0ba Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 17:40:38 +0100 Subject: [PATCH 22/25] Clean up in-test comments --- .../importers/yaml_importer/test_importer.py | 58 ++++++++----------- 1 file changed, 24 insertions(+), 34 deletions(-) diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py index e08051db441e3..515709aec984f 100644 --- a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py @@ -70,6 +70,17 @@ def execute(self, context): 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__}" + + @pytest.fixture def mock_bundle(tmp_path): bundle_dir = tmp_path / "bundle" @@ -175,8 +186,6 @@ def test_import_bodyless_task_becomes_empty_operator(mock_bundle): def test_import_bodyless_task_with_args_errors(mock_bundle): - # A body-less task becomes an EmptyOperator, which rejects the extra 'with' arguments - # (BaseOperator raises on unknown kwargs) -- so it fails to build rather than silently dropping. definition = _def(mock_bundle, "dag.yaml", _doc("- {id: a, with: {x: 1}}")) result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) assert not result.dags @@ -198,7 +207,6 @@ def test_import_schedule_cron(mock_bundle): def test_import_schedule_timetable(mock_bundle): - # A mapping schedule (uses/with) is built into a timetable via its import path (_build_schedule). definition = _def( mock_bundle, "dag.yaml", @@ -213,8 +221,6 @@ def test_import_schedule_timetable(mock_bundle): def test_import_keyed_xcom_on_uses_task(mock_bundle): - # $x with {task, key} on a uses: task resolves to a keyed XComArg (the _xcom_arg key branch); - # keyed XCom is rejected only in run: tasks. definition = _def( mock_bundle, "dag.yaml", @@ -236,8 +242,6 @@ def test_import_keyed_xcom_on_uses_task(mock_bundle): [("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): - # A YAML date or ISO string for a temporal arg is promoted to a (localized) datetime so DAG() - # accepts it, rather than failing with an unhelpful cast error (format section 2d). definition = _def( mock_bundle, "dag.yaml", @@ -252,8 +256,6 @@ def test_import_coerces_temporal_start_date(mock_bundle, value, hour): def test_import_does_not_coerce_nontemporal_arg(mock_bundle): - # _coerce_temporal introspects the target's argument types and only promotes datetime-typed ones. - # A date-looking string on a non-temporal arg (owner: str) keeps its type instead of becoming a datetime. definition = _def( mock_bundle, "dag.yaml", @@ -264,10 +266,20 @@ def test_import_does_not_coerce_nontemporal_arg(mock_bundle): 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): - # A code task (run:) builds the stub operator Airflow understands (@task.stub). Its body is a - # registered SDK task handler; the run mapping binds as that handler's arguments (op_kwargs), - # and $x inside it wires the dependency natively. definition = _def( mock_bundle, "dag.yaml", @@ -291,8 +303,6 @@ def test_import_code_task_builds_stub_operator(mock_bundle): def test_import_mixed_container_keeps_literal_siblings_literal(mock_bundle): - # When a container field renders (it holds a $x), its literal siblings must stay literal rather - # than being rendered as Jinja. The importer wraps them in LiteralValue; only $x/$t render. definition = _def( mock_bundle, "dag.yaml", @@ -317,7 +327,6 @@ def test_import_code_task_rejects_non_identifier_arg(mock_bundle): def test_import_run_rejects_template(mock_bundle): - # A $t in run: cannot be carried by the handler binding (it would arrive unrendered); reject it. 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 @@ -325,7 +334,6 @@ def test_import_run_rejects_template(mock_bundle): def test_import_run_rejects_keyed_xcom(mock_bundle): - # Only a task's return value crosses the handler boundary; a keyed $x is rejected at build time. definition = _def( mock_bundle, "dag.yaml", @@ -341,9 +349,6 @@ def test_import_run_rejects_keyed_xcom(mock_bundle): @pytest.mark.parametrize("marker", ["{$x: a}", "{$t: 'v'}", "{$f: q.sql}"]) def test_import_marker_on_non_template_field_errors(mock_bundle, marker): - # A marker only works where the operator class already templates. ``note`` is not in - # _RecordingOperator.template_fields, so any $f/$t/$x on it is an error -- the importer never - # adds a field to template_fields to make a marker work. (mock_bundle.path / "q.sql").write_text("SELECT 1") # exists, so this is not a missing-file error definition = _def( mock_bundle, @@ -359,8 +364,6 @@ def test_import_marker_on_non_template_field_errors(mock_bundle, marker): def test_import_base_operator_arg_in_with_errors(mock_bundle): - # `with:` is the operator-specific namespace; a BaseOperator (task-level) arg like `retries` - # there is rejected rather than silently accepted. definition = _def( mock_bundle, "dag.yaml", _doc(f"- {{id: a, uses: {OP}, with: {{payload: x, retries: 2}}}}") ) @@ -373,8 +376,6 @@ def test_import_base_operator_arg_in_with_errors(mock_bundle): "body", ["- {{id: a, uses: {op}, with: {{payload: x}}, note: 1}}", "- {{id: a, run: {{}}, note: 1}}"] ) def test_import_operator_arg_at_task_level_errors(mock_bundle, body): - # Task-level keys must be BaseOperator args; an operator-specific key (`note`) there is rejected - # (for both `uses:` and `run:` tasks) instead of being passed through raw. definition = _def(mock_bundle, "dag.yaml", _doc(body.format(op=OP))) result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) assert not result.dags @@ -382,8 +383,6 @@ def test_import_operator_arg_at_task_level_errors(mock_bundle, body): def test_import_literal_by_default(mock_bundle): - # A bare Jinja-looking string on a field the operator class *does* template is NOT rendered: - # literal by default, and the instance's template_fields is emptied of it. definition = _def( mock_bundle, "dag.yaml", _doc(f"- {{id: a, uses: {OP}, with: {{payload: '{{{{ ds }}}}'}}}}") ) @@ -407,8 +406,6 @@ def test_import_template_marker_sets_template_fields(mock_bundle): def test_import_keeps_operator_default_template_fields(mock_bundle): - # The operator templates ``preset`` by default; the author only sets ``payload`` (unmarked). - # Literal-by-default drops the author's unmarked field but keeps the operator's own default. definition = _def( mock_bundle, "dag.yaml", _doc(f"- {{id: a, uses: {OP_DEFAULT_TEMPLATE}, with: {{payload: hi}}}}") ) @@ -433,8 +430,6 @@ def test_import_file_ref_reads_verbatim(mock_bundle): @pytest.mark.parametrize("bad_path", ["../../etc/passwd", "/etc/passwd"]) def test_import_file_ref_outside_dag_dir_is_rejected(mock_bundle, bad_path): - # $f resolves relative to the Dag file's directory, like a Python Dag's template search path; - # a path that escapes it (traversal or absolute) is rejected rather than read. definition = _def( mock_bundle, "dag.yaml", _doc(f"- {{id: a, uses: {OP}, with: {{payload: {{$f: '{bad_path}'}}}}}}") ) @@ -460,8 +455,6 @@ def execute(self, context): def test_import_yaml_inside_zip_bundle(tmp_path): - # A YAML Dag inside a zip: `uses:` imports a module next to it (import_context puts the archive on - # sys.path), and `$f` reads a sibling member from the same archive. zip_path = tmp_path / "bundle.zip" with zipfile.ZipFile(zip_path, "w") as z: z.writestr("_ziptest_ops.py", _ZIP_OP_SOURCE) @@ -494,8 +487,6 @@ def test_import_template_file_reads_then_templates(mock_bundle): def test_import_templated_filename_is_reserved(mock_bundle): - # $f wrapping $t (a templated filename) needs run-time file resolution; the parser reserves it, - # so the whole file fails to parse. definition = _def( mock_bundle, "dag.yaml", @@ -544,7 +535,6 @@ def test_import_duplicate_dag_id_across_documents_errors(mock_bundle): def test_import_dependency_cycle_errors(mock_bundle): - # a <-> b form the cycle; c only depends on it and must not be reported as part of it. definition = _def( mock_bundle, "dag.yaml", From d9cd66128fe611129acc77c490367f3122519822 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 17:46:20 +0100 Subject: [PATCH 23/25] Improve temporal coersion --- task-sdk/src/airflow/sdk/importers/base.py | 15 ++++---- .../sdk/importers/yaml_importer/builder.py | 34 +++++++++++-------- 2 files changed, 29 insertions(+), 20 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/base.py b/task-sdk/src/airflow/sdk/importers/base.py index caef6070e3cd8..390c28bfc2fe5 100644 --- a/task-sdk/src/airflow/sdk/importers/base.py +++ b/task-sdk/src/airflow/sdk/importers/base.py @@ -74,12 +74,15 @@ def sibling(self, relpath: str) -> _SiblingResource: """ Return a readable handle to a resource stored next to this definition. - The handle supports ``read_bytes()``. Resolution matches a Python Dag's - default Jinja search path (relative to the directory of the definition's - own file) and raises if *relpath* escapes it. - - The base cannot locate siblings; a definition type that can - hold them overrides this. + The handle supports ``read_bytes()``. Resolution is relative to the + directory of the definition's own file and raises if *relpath* escapes + it. This is close to a Python Dag's default Jinja search path, but + stricter: a path escaping via a symlink is rejected (Jinja's + ``FileSystemLoader`` would allow it), and a Dag-level + ``template_searchpath`` is not consulted. + + The base cannot locate siblings; a definition type that can hold them + overrides this. """ raise NotImplementedError(f"{type(self).__name__} cannot resolve sibling resources") diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py index 3870eca102867..60bd857c76f6f 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py @@ -36,6 +36,7 @@ import datetime import functools import graphlib +import inspect import keyword import linecache import re @@ -149,12 +150,16 @@ def _build_operator(self, task: Task) -> BaseOperator: 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) - kwargs[name] = self._lower(value, protect=needs_render) + 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: @@ -321,7 +326,7 @@ def _to_datetime(value: Any) -> Any: return value -# We annotate arguments with either stdlib datetime or pendulum DateTime. +# Airflow annotates arguments with either stdlib datetime or pendulum DateTime. # Matched by name (see _temporal_fields) rather than resolved to avoid having # to evaluate dynamically. _TEMPORAL_TYPE_NAMES = frozenset({"datetime", "DateTime"}) @@ -331,19 +336,20 @@ def _to_datetime(value: Any) -> Any: @functools.cache def _temporal_fields(target: type) -> frozenset[str]: """ - Names of *target*'s arguments annotated as a datetime type. - - The annotation is read as text and matched against the datetime type names - we use; it is never resolved, so one referring to a ``TYPE_CHECKING``-only - name does not need importing. This is not perfect; we'll try to fix edge - cases as we go. + Names of *target*'s constructor arguments annotated as a datetime type. + + The annotation text is matched against the datetime type names Airflow uses; + it is never resolved, so one referring to a ``TYPE_CHECKING``-only name does + not need importing. The constructor signature is read -- not class + attributes -- so attrs field aliases and operator-specific arguments are + covered. This is not perfect (a datetime type reached under another alias is + missed); we'll fix edge cases as we go. """ - found: set[str] = set() - for klass in getattr(target, "__mro__", (target,)): - for name, annotation in getattr(klass, "__annotations__", {}).items(): - if name not in found and _is_temporal_annotation(str(annotation)): - found.add(name) - return frozenset(found) + try: + params = inspect.signature(target).parameters + except (TypeError, ValueError): + return frozenset() + return frozenset(name for name, p in params.items() if _is_temporal_annotation(str(p.annotation))) def _is_temporal_annotation(annotation: str) -> bool: From 7915e617f63cfce0e14902862f103c28656dac55 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 22:05:28 +0100 Subject: [PATCH 24/25] Harden --- .../sdk/importers/yaml_importer/builder.py | 79 +++++++++++++++---- .../importers/yaml_importer/test_importer.py | 59 +++++++++++++- 2 files changed, 121 insertions(+), 17 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py index 60bd857c76f6f..02d49c6a8ffc3 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py @@ -37,6 +37,7 @@ import functools import graphlib import inspect +import json import keyword import linecache import re @@ -46,7 +47,12 @@ from airflow.sdk import DAG, BaseOperator, XComArg, literal, task as task_decorators from airflow.sdk._shared.module_loading import import_string +from airflow.sdk.definitions.context import KNOWN_CONTEXT_KEYS from airflow.sdk.importers.yaml_importer.models import ( + CONST_KEY, + FILE_KEYS, + TEMPLATE_KEYS, + XCOM_KEYS, ConstRef, FileRef, Literal, @@ -69,6 +75,7 @@ class YamlDagBuildError(ValueError): _BASE_TASK_ARGS = frozenset(BaseOperator.__init__._BaseOperatorMeta__param_names) # type: ignore[attr-defined] _UNSUPPORTED_TASK_ARGS = frozenset({"task_id", "dag", "task_group", "default_args"}) +_MARKER_KEYS = frozenset({*XCOM_KEYS, *TEMPLATE_KEYS, *FILE_KEYS, CONST_KEY}) @attrs.define(kw_only=True) @@ -87,6 +94,7 @@ class DagBuilder: def build(self, doc: DagDocument) -> DAG: tasks = {t.id_: t for t in doc.tasks} + _reject_reserved_extra_keys(doc.model_extra or {}) 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): @@ -104,18 +112,35 @@ 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. + 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. + + The "real" task body should be a registered SDK task handler resolved at + run time; the run mapping binds as that handler's arguments. Only a + top-level ``$x`` (an upstream return value) crosses the binding; every + other argument must be a JSON-serializable literal. """ + if task.with_: + raise YamlDagBuildError( + f"task {task.id_!r}: a run: task takes its inputs under 'run:', not 'with:'" + ) 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") + if name in KNOWN_CONTEXT_KEYS: + raise YamlDagBuildError(f"task {task.id_!r}: run argument cannot be named {name!r}") _reject_unbindable_run_marker(task.id_, value) - run_kwargs[name] = self._lower(value, protect=False) + if not isinstance(lowered := self._lower(value, protect=False), XComArg): + try: + json.dumps(lowered, allow_nan=False) + except (TypeError, ValueError): + raise YamlDagBuildError( + f"task {task.id_!r}: run argument {name!r} must be JSON-serializable, " + f"not a {type(lowered).__name__}" + ) from None + run_kwargs[name] = lowered decorator = task_decorators.stub( task_id=task.id_, **_coerce_temporal(task.model_extra or {}, BaseOperator) ) @@ -139,6 +164,7 @@ def _build_operator(self, task: Task) -> BaseOperator: f"task {task.id_!r}: {unknown} are not task-level (BaseOperator) arguments; " "operator-specific arguments belong under 'with'" ) + _reject_reserved_extra_keys(task.model_extra or {}) if (args := task.run) is not None: return self._build_stub_task(task, args) @@ -264,14 +290,14 @@ def _restrict_template_fields( return op -def _reject_unbindable_run_marker(task_id: str, value: Any) -> None: +def _reject_unbindable_run_marker(task_id: str, value: Any, *, top_level: bool = True) -> 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. + The binding carries a JSON literal or a single top-level ``$x`` (an upstream + task's return value). So a template (``$t``) or a keyed ``$x`` anywhere, or a + plain ``$x`` nested inside a container instead of passed as its own argument, + is an error. ``$f`` and ``$const`` lower to literals and are fine. """ if isinstance(value, TemplateRef): raise YamlDagBuildError( @@ -285,15 +311,40 @@ def _reject_unbindable_run_marker(task_id: str, value: Any) -> None: f"task {task_id!r}: a keyed XCom ($x with 'key') is not supported in a run: task; " "only a task's return value crosses the handler boundary" ) + if not top_level: + raise YamlDagBuildError( + f"task {task_id!r}: an XCom ($x) nested inside a run argument is not supported; pass " + "the upstream output as its own top-level run argument" + ) return if isinstance(value, Literal): - _reject_unbindable_run_marker(task_id, value.root) + _reject_unbindable_run_marker(task_id, value.root, top_level=top_level) elif isinstance(value, dict): for item in value.values(): - _reject_unbindable_run_marker(task_id, item) + _reject_unbindable_run_marker(task_id, item, top_level=False) elif isinstance(value, list): for item in value: - _reject_unbindable_run_marker(task_id, item) + _reject_unbindable_run_marker(task_id, item, top_level=False) + + +def _reject_reserved_extra_keys(extras: Any) -> None: + """ + Reject a marker key ($x/$t/$f/$const) nested in Dag- or task-level arguments. + + These arguments pass straight to ``DAG`` or the operator, never through the + value grammar, so a marker written there (e.g. ``params: {p: {$f: x}}``) + would be taken as a literal dict rather than flagged. Only marker keys are + rejected, so a genuine ``$`` key such as a JSON Schema ``$ref`` in a param + default is kept. + """ + if isinstance(extras, dict): + for key, val in extras.items(): + if key in _MARKER_KEYS: + raise YamlDagBuildError(f"marker key {key!r} may only be used in 'with:' or 'run:'") + _reject_reserved_extra_keys(val) + elif isinstance(extras, list): + for item in extras: + _reject_reserved_extra_keys(item) def _coerce_temporal(mapping: dict[str, Any], target: type) -> dict[str, Any]: diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py index 515709aec984f..84b7c53511a57 100644 --- a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py @@ -303,20 +303,26 @@ def test_import_code_task_builds_stub_operator(mock_bundle): def test_import_mixed_container_keeps_literal_siblings_literal(mock_bundle): + (mock_bundle.path / "note.txt").write_text("hello") 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 }}}}'}}}}}} + - {{id: t, uses: {OP}, with: {{payload: {{a: {{$x: up}}, b: '{{{{ not a template }}}}', c: {{$const: 7}}, d: {{$f: note.txt}}}}}}}} """), ) 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 + # the container renders ($x present), so every literal sibling -- plain, $const, $f -- stays wrapped + assert isinstance(payload["b"], LiteralValue) assert payload["b"].value == "{{ not a template }}" + assert isinstance(payload["c"], LiteralValue) + assert payload["c"].value == 7 + assert isinstance(payload["d"], LiteralValue) + assert payload["d"].value == "hello" def test_import_code_task_rejects_non_identifier_arg(mock_bundle): @@ -347,6 +353,53 @@ def test_import_run_rejects_keyed_xcom(mock_bundle): assert any("keyed XCom" in e.message for e in result.errors) +def test_import_run_task_rejects_with(mock_bundle): + definition = _def(mock_bundle, "dag.yaml", _doc("- {id: t, run: {n: 1}, with: {region: eu}}")) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert any("takes its inputs under 'run:'" in e.message for e in result.errors) + + +def test_import_run_rejects_nested_xcom(mock_bundle): + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f""" + - {{id: a, uses: {OP}, with: {{payload: x}}}} + - {{id: t, run: {{items: [{{$x: a}}]}}}} + """), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert any("nested inside a run argument" in e.message for e in result.errors) + + +def test_import_run_rejects_non_json_literal(mock_bundle): + definition = _def(mock_bundle, "dag.yaml", _doc("- {id: t, run: {day: 2024-01-01}}")) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert any("must be JSON-serializable" in e.message for e in result.errors) + + +def test_import_run_rejects_context_key_arg(mock_bundle): + definition = _def(mock_bundle, "dag.yaml", _doc("- {id: t, run: {ds: 1}}")) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert any("cannot be named" in e.message for e in result.errors) + + +def test_import_rejects_marker_key_in_extras(mock_bundle): + # Dag-level extras bypass the value grammar; a marker nested there is a typo, not a literal dict. + definition = _def( + mock_bundle, + "dag.yaml", + _doc(f"- {{id: a, uses: {OP}, with: {{payload: x}}}}", extra="params: {p: {$f: README.md}}\n"), + ) + result = YamlDagImporter().import_definition(definition, bundle=mock_bundle) + assert not result.dags + assert any("marker key '$f'" in e.message for e in result.errors) + + @pytest.mark.parametrize("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 @@ -360,7 +413,7 @@ def test_import_marker_on_non_template_field_errors(mock_bundle, 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) + assert any("only allowed on templated fields" in e.message for e in result.errors) def test_import_base_operator_arg_in_with_errors(mock_bundle): From 79c433f135e7b54d4772c7098810461984df3e61 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 12 Oct 2026 22:19:34 +0100 Subject: [PATCH 25/25] Fix zip import function arg name --- task-sdk/src/airflow/sdk/importers/zip_importer.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/task-sdk/src/airflow/sdk/importers/zip_importer.py b/task-sdk/src/airflow/sdk/importers/zip_importer.py index ee51f24ff3b8c..ff40089f75070 100644 --- a/task-sdk/src/airflow/sdk/importers/zip_importer.py +++ b/task-sdk/src/airflow/sdk/importers/zip_importer.py @@ -102,16 +102,16 @@ def read_bytes(self) -> bytes: self._content = z.read(self.file_path) return self._content - def sibling(self, rel_path: str) -> zipfile.Path: + def sibling(self, relpath: str) -> zipfile.Path: base_dir = posixpath.dirname(self.file_path) - member = posixpath.normpath(posixpath.join(base_dir, rel_path)) + member = posixpath.normpath(posixpath.join(base_dir, relpath)) within = ( (member == base_dir or member.startswith(f"{base_dir}/")) if base_dir else not (member == ".." or member.startswith("../")) ) if posixpath.isabs(member) or not within: - raise ValueError(f"{rel_path!r} resolves outside the Dag file's directory") + raise ValueError(f"{relpath!r} resolves outside the Dag file's directory") return zipfile.Path(self.zip_path, member) @contextlib.contextmanager