kaxil commented on code in PR #74586: URL: https://github.com/apache/airflow/pull/74586#discussion_r4248373279
########## 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"<yaml-dag stub ({', '.join(param_names)})>" + linecache.cache[filename] = (len(source), None, source.splitlines(keepends=True), filename) + namespace: dict[str, Any] = {} + exec(compile(source, filename, "exec"), namespace) + self._stub_callables[param_names] = stub = namespace["_stub"] + return stub + + +def _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 Review Comment: Replacing `template_fields` wholesale also switches off fields the author never set. `ExternalTaskMarker` defaults `logical_date` to `"{{ logical_date.isoformat() }}"`, and `SQLIntervalCheckOperator` builds `sql1`/`sql2` with `{{ ds }}` in `__init__`. Through this importer both reach `execute()` unrendered, while the same operator in a Python Dag renders them. Literal-by-default is about values the author wrote, so could this drop only the fields set in `with:` without a marker? Something like `tuple(f for f in op.template_fields if f not in task.with_ or f in render_fields)`. ########## 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})", Review Comment: Core treats `source_reference` as a bundle-relative path (`DagBag._record_import_error`), so this becomes an `import_error` row with filename `dag.yaml (dag_id='x')`, and nothing matches it later. `_update_import_errors` only clears names in `files_parsed` (`dag.yaml`), so the row outlives the fix. Meanwhile `clear_orphaned_import_errors` deletes it on every bundle refresh while the document is still broken, and the next parse recreates it, firing `on_new_dag_import_error` again. Keying it by the file isn't a drop-in fix either: `_update_import_errors` then marks every Dag with that `relative_fileloc` stale, including the good documents from the same file. Core has no per-Dag error slot today, so would it be simpler to fail the whole file on a build error, the way a parse error already does? Nothing tests the one-bad-document case yet. ########## 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) Review Comment: If two documents in one file share a `dag_id`, both get appended and `DagBag._add_to_bag` lets the second silently replace the first, since it only rejects duplicates across different filelocs. Worth checking for repeated `dag_id`s across `documents` here? ########## 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() Review Comment: Any YAML containing `$schema` passes this, so a JSON Schema file (or any other YAML that declares one) in the bundle gets queued and then sits there as a permanent parse error. Matching the `https://airflow.apache.org/schemas/dag/` prefix (`_SCHEMA_URL` in `migrator.py`) would keep the sniff cheap and skip those. ########## 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. Review Comment: Importing `task` only binds the `TaskDecoratorCollection`; provider registration happens on the `.stub` attribute access (`TaskDecoratorCollection.__getattr__`). So this can move to the top next to `DAG, XComArg`. ########## 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. Review Comment: Line 146 looks like a leftover fragment. The bullets also say `needs_render` is what gets validated against `template_fields`, but the code validates on `has_marker` (line 127) and uses `needs_render` to choose the instance's `template_fields`. ########## 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__}" + + [email protected] +def mock_bundle(tmp_path): + bundle_dir = tmp_path / "bundle" + bundle_dir.mkdir(parents=True, exist_ok=True) + return SimpleNamespace(name="test_bundle", path=bundle_dir) + + +def _def(bundle, name: str, body: str) -> FilesystemDagDefinition: + path = bundle.path / name + path.write_text(textwrap.dedent(body)) + return FilesystemDagDefinition(path=path) + + +def _doc(tasks: str, *, dag_id: str = "d", extra: str = "") -> str: + return f"$schema: {HEAD_SCHEMA}\ndag_id: {dag_id}\n{extra}tasks:\n{textwrap.indent(textwrap.dedent(tasks), ' ')}" + + +def test_can_handle(): + importer = YamlDagImporter() + assert importer.can_handle("dag.yaml") + assert importer.can_handle("dag.yml") + assert not importer.can_handle("dag.py") + assert not importer.can_handle("dag.json") + + +def test_get_source_code(mock_bundle): + definition = _def(mock_bundle, "d.yaml", f"$schema: {HEAD_SCHEMA}\ndag_id: d\ntasks: []\n") + src = YamlDagImporter().get_source_code(definition) + assert src.language == "yaml" + assert "$schema" in src.source_code + + +def test_might_contain_dag(mock_bundle): + importer = YamlDagImporter() + has = _def(mock_bundle, "has.yaml", f"$schema: {HEAD_SCHEMA}\ndag_id: d\ntasks: []\n") + nope = _def(mock_bundle, "nope.yaml", "just: config\n") + 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 Review Comment: Could this also assert `"payload" in b.template_fields`? No test checks that `$x` opts the field into rendering, and without that the `XComArg` never resolves at runtime. ########## 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"<yaml-dag stub ({', '.join(param_names)})>" + linecache.cache[filename] = (len(source), None, source.splitlines(keepends=True), filename) + namespace: dict[str, Any] = {} + exec(compile(source, filename, "exec"), namespace) + self._stub_callables[param_names] = stub = namespace["_stub"] + return stub + + +def _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)}") Review Comment: This is a second cycle check next to `detect_cycle` in `_shared/dagnode/cycle.py`, which #74034 moved there to be the one implementation. The message also lists every task downstream of the cycle, not only the cycle. `graphlib.TopologicalSorter(...).static_order()` would give the order and raise `CycleError` carrying the actual cycle. ########## 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__}" + + [email protected] +def mock_bundle(tmp_path): + bundle_dir = tmp_path / "bundle" + bundle_dir.mkdir(parents=True, exist_ok=True) + return SimpleNamespace(name="test_bundle", path=bundle_dir) + + +def _def(bundle, name: str, body: str) -> FilesystemDagDefinition: + path = bundle.path / name + path.write_text(textwrap.dedent(body)) + return FilesystemDagDefinition(path=path) + + +def _doc(tasks: str, *, dag_id: str = "d", extra: str = "") -> str: + return f"$schema: {HEAD_SCHEMA}\ndag_id: {dag_id}\n{extra}tasks:\n{textwrap.indent(textwrap.dedent(tasks), ' ')}" + + +def test_can_handle(): + importer = YamlDagImporter() + assert importer.can_handle("dag.yaml") + assert importer.can_handle("dag.yml") + assert not importer.can_handle("dag.py") + assert not importer.can_handle("dag.json") + + +def test_get_source_code(mock_bundle): + definition = _def(mock_bundle, "d.yaml", f"$schema: {HEAD_SCHEMA}\ndag_id: d\ntasks: []\n") + src = YamlDagImporter().get_source_code(definition) + assert src.language == "yaml" + assert "$schema" in src.source_code + + +def test_might_contain_dag(mock_bundle): + importer = YamlDagImporter() + has = _def(mock_bundle, "has.yaml", f"$schema: {HEAD_SCHEMA}\ndag_id: d\ntasks: []\n") + nope = _def(mock_bundle, "nope.yaml", "just: config\n") + 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) + + [email protected]("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 Review Comment: `_RecordingOperator` doesn't set `template_ext`, so `BaseOperator`'s default `()` satisfies this without the reset in `_apply_template_lowering`. Giving the test operator `template_ext = (".sql",)` would make this assertion check something. ########## 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__}" + + [email protected] +def mock_bundle(tmp_path): + bundle_dir = tmp_path / "bundle" + bundle_dir.mkdir(parents=True, exist_ok=True) + return SimpleNamespace(name="test_bundle", path=bundle_dir) + + +def _def(bundle, name: str, body: str) -> FilesystemDagDefinition: + path = bundle.path / name + path.write_text(textwrap.dedent(body)) + return FilesystemDagDefinition(path=path) + + +def _doc(tasks: str, *, dag_id: str = "d", extra: str = "") -> str: + return f"$schema: {HEAD_SCHEMA}\ndag_id: {dag_id}\n{extra}tasks:\n{textwrap.indent(textwrap.dedent(tasks), ' ')}" + + +def test_can_handle(): + importer = YamlDagImporter() + assert importer.can_handle("dag.yaml") + assert importer.can_handle("dag.yml") + assert not importer.can_handle("dag.py") + assert not importer.can_handle("dag.json") + + +def test_get_source_code(mock_bundle): + definition = _def(mock_bundle, "d.yaml", f"$schema: {HEAD_SCHEMA}\ndag_id: d\ntasks: []\n") + src = YamlDagImporter().get_source_code(definition) + assert src.language == "yaml" + assert "$schema" in src.source_code + + +def test_might_contain_dag(mock_bundle): + importer = YamlDagImporter() + has = _def(mock_bundle, "has.yaml", f"$schema: {HEAD_SCHEMA}\ndag_id: d\ntasks: []\n") + nope = _def(mock_bundle, "nope.yaml", "just: config\n") + 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) + + [email protected]("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 + Review Comment: This and `test_import_bodyless_task_with_args_errors` pass on any error. Matching on the message (`does.not.Exist` here, the unexpected `x` kwarg there) would stop them passing for the wrong reason. ########## 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] Review Comment: Only the lowered value survives here, so `run: {day: {$t: '{{ ds }}'}}` reaches the handler as the string `{{ ds }}`. The stub arg bindings only have `xcom` and `literal` kinds (`build_arg_bindings` in `serialization/stub_arg_bindings.py`), and nothing renders the value. The docstring above says template references resolve natively. Should `$t` raise a `YamlDagBuildError` in `run:` until the binding protocol can carry a template? `$x: {task, key}` has a smaller version of this: it builds fine and then fails in core serialization with a `@task.stub` message, since only `return_value` crosses the boundary. ########## 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) Review Comment: When a key is set both at task level and under `with:`, the task-level value silently wins, but `render_fields` still comes from the `with:` side. So `payload: '{{ ds }}'` next to `with: {payload: {$t: x}}` renders the task-level value, which the comment says is never templated. Should overlapping keys raise? ########## 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()), Review Comment: Rendering is per field, so one `$t` or `$x` anywhere in a dict or list makes everything else in that field render too. With `payload: {a: {$t: '{{ x }}'}, b: '{{ x }}', c: {$const: '{{ x }}'}}` all three render, so `$const` does nothing here, and a sibling `$f` gets its file content rendered as well. An `env` or `op_kwargs` mixing one XCom with literal strings hits this. Could the unmarked leaves (plus `$const` and `$f` values) be wrapped in `airflow.sdk.literal()` when the container needs rendering? Rejecting mixed containers would also work if that's simpler for now. ########## 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: Review Comment: `start_date: 2024-01-01` is probably the first thing people will write, and YAML hands it over as a `datetime.date`, which `DAG()` rejects with `'date' object cannot be cast as 'str'` without naming the field. A quoted `'2024-01-01'` fails too (`'str' object has no attribute 'utcoffset'`); only full timestamps work. Could the builder coerce dates for `start_date`/`end_date`, or at least wrap `DAG(...)` so the error names the key? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
