This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 72568207e57 Implement AIP-85 YAML importer based on the parser (#74586)
72568207e57 is described below
commit 72568207e57d76847eff22b09308648c3e7102c5
Author: Tzu-ping Chung <[email protected]>
AuthorDate: Tue Oct 13 13:59:39 2026 +0800
Implement AIP-85 YAML importer based on the parser (#74586)
* Add missing PyYAML dependency to Task SDK
* Implement file-loading syntax $f
* Implement the YAML importer shell
* Fix double-rendering
* Protect file read using the same logic as Python templates
* Re-export to reduce import levels
* Correctly import in an import context
* Refactor DagBuilder interface
* Stricter top-level and with argument checks
* Tighten template field patching
* Tighten literal handling in run
* Consolidate import-time errors from a YAML file
* Reject duplicate dag_id in a file
* Tighten might_contain_dag heuristic
* Coerce YAML input temporal values to datetimes
* Rely on stdlib graphlib for ordering
* Move unnecessary local import up
* Tighten tests
* Try to smartly convert datetimes
* Tighten
* Fix leaked rst in schema
* Clean up in-test comments
* Improve temporal coersion
* Harden
* Fix zip import function arg name
---
task-sdk/pyproject.toml | 5 +-
task-sdk/src/airflow/sdk/importers/base.py | 31 +-
.../sdk/importers/yaml_importer/__init__.py | 9 +-
.../airflow/sdk/importers/yaml_importer/builder.py | 455 ++++++++++++++++
.../sdk/importers/yaml_importer/importer.py | 169 ++++++
.../sdk/importers/yaml_importer/migrator.py | 3 +-
.../airflow/sdk/importers/yaml_importer/models.py | 84 ++-
.../sdk/importers/yaml_importer/schema.json | 62 ++-
task-sdk/src/airflow/sdk/importers/zip_importer.py | 13 +
.../importers/yaml_importer/test_importer.py | 605 +++++++++++++++++++++
.../importers/yaml_importer/test_parser.py | 10 +
11 files changed, 1425 insertions(+), 21 deletions(-)
diff --git a/task-sdk/pyproject.toml b/task-sdk/pyproject.toml
index 0ce16d4b713..9a4c9d44ea9 100644
--- a/task-sdk/pyproject.toml
+++ b/task-sdk/pyproject.toml
@@ -54,14 +54,15 @@ dependencies = [
"babel>=2.17.0",
"cadwyn>=6.1.1",
"fsspec>=2023.10.0",
+ "greenback>=1.2.1",
"httpx>=0.27.0",
"jinja2>=3.1.5",
"methodtools>=0.4.7",
"msgspec>=0.19.0",
- "python-dateutil>=2.7.0",
"psutil>=6.1.0",
+ "python-dateutil>=2.7.0",
+ "pyyaml>=6.0.3",
"structlog>=26.1.0",
- "greenback>=1.2.1",
"tenacity>=8.3.0",
# Start of shared timezones dependencies
"pendulum>=3.1.0",
diff --git a/task-sdk/src/airflow/sdk/importers/base.py
b/task-sdk/src/airflow/sdk/importers/base.py
index e594e15ad24..390c28bfc2f 100644
--- a/task-sdk/src/airflow/sdk/importers/base.py
+++ b/task-sdk/src/airflow/sdk/importers/base.py
@@ -34,11 +34,17 @@ from airflow.sdk.execution_time.coordinator import
get_coordinator_manager
if TYPE_CHECKING:
from collections.abc import Generator, Iterable, Iterator
- from typing import Self
+ from typing import Protocol, Self
from airflow.dag_processing.bundles.base import BaseDagBundle # noqa:
SDK002
from airflow.sdk import DAG
+ class _SiblingResource(Protocol):
+ """A readable handle to a resource, e.g. a :class:`pathlib.Path` or
:class:`zipfile.Path`."""
+
+ def read_bytes(self) -> bytes: ...
+
+
log = logging.getLogger(__name__)
DefT = TypeVar("DefT", bound="DagDefinition")
@@ -64,6 +70,22 @@ class DagDefinition(ABC):
"""Read and return the content of the resource as a string."""
return self.read_bytes().decode(encoding)
+ def sibling(self, relpath: str) -> _SiblingResource:
+ """
+ Return a readable handle to a resource stored next to this definition.
+
+ The handle supports ``read_bytes()``. Resolution is relative to the
+ directory of the definition's own file and raises if *relpath* escapes
+ it. This is close to a Python Dag's default Jinja search path, but
+ stricter: a path escaping via a symlink is rejected (Jinja's
+ ``FileSystemLoader`` would allow it), and a Dag-level
+ ``template_searchpath`` is not consulted.
+
+ The base cannot locate siblings; a definition type that can hold them
+ overrides this.
+ """
+ raise NotImplementedError(f"{type(self).__name__} cannot resolve
sibling resources")
+
@abstractmethod
def as_file(self) -> contextlib.AbstractContextManager[Path]:
"""
@@ -133,6 +155,13 @@ class FilesystemDagDefinition(FileDagDefinition):
def read_bytes(self) -> bytes:
return self.path.read_bytes()
+ def sibling(self, relpath: str) -> Path:
+ base = self.path.parent.resolve()
+ target = base.joinpath(relpath).resolve()
+ if not target.is_relative_to(base):
+ raise ValueError(f"{relpath!r} resolves outside the Dag file's
directory {str(base)!r}")
+ return target
+
@contextlib.contextmanager
def as_file(self) -> Generator[Path, None, None]:
yield self.path
diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/__init__.py
b/task-sdk/src/airflow/sdk/importers/yaml_importer/__init__.py
index 2f229838802..a08ed337b70 100644
--- a/task-sdk/src/airflow/sdk/importers/yaml_importer/__init__.py
+++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/__init__.py
@@ -19,14 +19,15 @@
from __future__ import annotations
+from airflow.sdk.importers.yaml_importer.builder import YamlDagBuildError
+from airflow.sdk.importers.yaml_importer.importer import YamlDagImporter
from airflow.sdk.importers.yaml_importer.models import DagDocument
-from airflow.sdk.importers.yaml_importer.parser import (
- YamlDagParseError,
- parse_documents,
-)
+from airflow.sdk.importers.yaml_importer.parser import YamlDagParseError,
parse_documents
__all__ = [
"DagDocument",
+ "YamlDagBuildError",
+ "YamlDagImporter",
"YamlDagParseError",
"parse_documents",
]
diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py
b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py
new file mode 100644
index 00000000000..02d49c6a8ff
--- /dev/null
+++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py
@@ -0,0 +1,455 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Build a validated :class:`~.models.DagDocument` into a
:class:`~airflow.sdk.DAG`.
+
+The parser produces a validated document; this turns it into real DAG +
operator
+objects. Ready-made operators are instantiated from their ``uses`` import path;
+tasks with ``run`` are built into stub operators and handled by AIP-108 task
+handlers at execution time.
+
+Templating is explicit, realized importer-side (a la AIP-80). The format is
+literal by default; only values the author marked render as templates. Value
+grammar markers are only accepted on fields declared in ``template_fields``.
+
+Tasks are built in topological order so a task referenced by an upstream's XCom
+already exists when the downstream operator is constructed.
+"""
+
+from __future__ import annotations
+
+import datetime
+import functools
+import graphlib
+import inspect
+import json
+import keyword
+import linecache
+import re
+from typing import TYPE_CHECKING, Any, NamedTuple, cast
+
+import attrs
+
+from airflow.sdk import DAG, BaseOperator, XComArg, literal, task as
task_decorators
+from airflow.sdk._shared.module_loading import import_string
+from airflow.sdk.definitions.context import KNOWN_CONTEXT_KEYS
+from airflow.sdk.importers.yaml_importer.models import (
+ CONST_KEY,
+ FILE_KEYS,
+ TEMPLATE_KEYS,
+ XCOM_KEYS,
+ ConstRef,
+ FileRef,
+ Literal,
+ TemplateRef,
+ TimetableSchedule,
+ XComRef,
+)
+
+if TYPE_CHECKING:
+ from collections.abc import Callable
+
+ from airflow.dag_processing.bundles.base import BaseDagBundle # noqa:
SDK002
+ from airflow.sdk.importers.base import DagDefinition
+ from airflow.sdk.importers.yaml_importer.models import DagDocument, Task,
Value
+
+
+class YamlDagBuildError(ValueError):
+ """A validated document could not be built into a DAG."""
+
+
+_BASE_TASK_ARGS =
frozenset(BaseOperator.__init__._BaseOperatorMeta__param_names) # type:
ignore[attr-defined]
+_UNSUPPORTED_TASK_ARGS = frozenset({"task_id", "dag", "task_group",
"default_args"})
+_MARKER_KEYS = frozenset({*XCOM_KEYS, *TEMPLATE_KEYS, *FILE_KEYS, CONST_KEY})
+
+
[email protected](kw_only=True)
+class DagBuilder:
+ """
+ Build a :class:`DagDocument` into a :class:`DAG`.
+
+ This holds the state a single build shares. The state belongs to one build
+ and is discarded with the instance.
+ """
+
+ _definition: DagDefinition
+ _bundle: BaseDagBundle
+ _operators: dict[str, BaseOperator] = attrs.field(factory=dict, init=False)
+ _stub_callables: dict[tuple[str, ...], Callable[..., None]] =
attrs.field(factory=dict, init=False)
+
+ def build(self, doc: DagDocument) -> DAG:
+ tasks = {t.id_: t for t in doc.tasks}
+ _reject_reserved_extra_keys(doc.model_extra or {})
+ extras = _coerce_temporal(doc.model_extra or {}, DAG)
+ with DAG(dag_id=doc.dag_id, schedule=_build_schedule(doc.schedule),
**extras) as dag:
+ for task_id in _topological_order(doc.tasks):
+ self._operators[task_id] = self._build_operator(tasks[task_id])
+
+ for task in doc.tasks:
+ for upstream in sorted(task.get_dependency_ids()):
+ self._operators[upstream] >> self._operators[task.id_]
+ dag.bundle_name = self._bundle.name
+ dag.fileloc = repr(self._definition)
+ dag.relative_fileloc =
self._definition.get_relative_loc(self._bundle.path)
+ return dag
+
+ def _build_stub_task(self, task: Task, args: dict[str, Value]) ->
BaseOperator:
+ """
+ Build a stub from a "run:" task.
+
+ The stub is fed a callable whose parameters are the run argument names,
+ so serialization can build and later validate (AIP-108) the handler's
+ argument spec.
+
+ The "real" task body should be a registered SDK task handler resolved
at
+ run time; the run mapping binds as that handler's arguments. Only a
+ top-level ``$x`` (an upstream return value) crosses the binding; every
+ other argument must be a JSON-serializable literal.
+ """
+ if task.with_:
+ raise YamlDagBuildError(
+ f"task {task.id_!r}: a run: task takes its inputs under
'run:', not 'with:'"
+ )
+ run_kwargs: dict[str, Any] = {}
+ for name, value in args.items():
+ if not name.isidentifier() or keyword.iskeyword(name):
+ raise YamlDagBuildError(f"task {task.id_!r}: run argument
{name!r} is not a valid identifier")
+ if name in KNOWN_CONTEXT_KEYS:
+ raise YamlDagBuildError(f"task {task.id_!r}: run argument
cannot be named {name!r}")
+ _reject_unbindable_run_marker(task.id_, value)
+ if not isinstance(lowered := self._lower(value, protect=False),
XComArg):
+ try:
+ json.dumps(lowered, allow_nan=False)
+ except (TypeError, ValueError):
+ raise YamlDagBuildError(
+ f"task {task.id_!r}: run argument {name!r} must be
JSON-serializable, "
+ f"not a {type(lowered).__name__}"
+ ) from None
+ run_kwargs[name] = lowered
+ decorator = task_decorators.stub(
+ task_id=task.id_, **_coerce_temporal(task.model_extra or {},
BaseOperator)
+ )
+ return
decorator(self._stub_callable(tuple(run_kwargs)))(**run_kwargs).operator
+
+ def _build_operator(self, task: Task) -> BaseOperator:
+ """Instantiate *task* as an operator."""
+ if reserved := sorted(
+ k for k in (*task.with_, *(task.model_extra or {})) if k in
_UNSUPPORTED_TASK_ARGS
+ ):
+ raise YamlDagBuildError(
+ f"task {task.id_!r}: {reserved} cannot be set on a task; the
task id comes from 'id', and "
+ "dag/task_group/default_args are managed by the importer"
+ )
+ if misplaced := sorted(k for k in task.with_ if k in _BASE_TASK_ARGS):
+ raise YamlDagBuildError(
+ f"task {task.id_!r}: {misplaced} are task-level arguments; set
them at task level, not under 'with'"
+ )
+ if unknown := sorted(k for k in (task.model_extra or {}) if k not in
_BASE_TASK_ARGS):
+ raise YamlDagBuildError(
+ f"task {task.id_!r}: {unknown} are not task-level
(BaseOperator) arguments; "
+ "operator-specific arguments belong under 'with'"
+ )
+ _reject_reserved_extra_keys(task.model_extra or {})
+
+ if (args := task.run) is not None:
+ return self._build_stub_task(task, args)
+
+ op_cls: Any
+ if task.uses is None: # uses: null (or omitted) becomes EmptyOperator.
+ from airflow.providers.standard.operators.empty import
EmptyOperator as op_cls # noqa: SDK002
+ else:
+ op_cls = import_string(task.uses)
+
+ template_fields = set(getattr(op_cls, "template_fields", ()))
+ temporal = _temporal_fields(op_cls)
+ kwargs: dict[str, Any] = {}
+ render_fields: list[str] = []
+ off_template: list[str] = []
+ for name, value in task.with_.items():
+ needs_render, has_marker = _field_markers(value)
+ lowered = self._lower(value, protect=needs_render)
+ if name in temporal and not has_marker: # a literal date/ISO
string for a datetime arg
+ lowered = _to_datetime(lowered)
+ kwargs[name] = lowered
+ if has_marker and name not in template_fields:
+ off_template.append(name)
+ if needs_render:
+ render_fields.append(name)
+ if off_template:
+ raise YamlDagBuildError(
+ f"task {task.id_!r}: value grammar markers are only allowed on
templated fields, not {off_template}"
+ )
+ if task.model_extra: # Merge in task-level arguments (raw, never
templated).
+ kwargs.update(_coerce_temporal(task.model_extra, BaseOperator))
+
+ return _restrict_template_fields(op_cls(task_id=task.id_, **kwargs),
task.with_, render_fields)
+
+ def _lower(self, value: Any, *, protect: bool) -> Any:
+ """
+ Lower a resolved value-grammar node into an operator argument.
+
+ :param protect: This should be set when the enclosing field renders to
+ avoid double-rendering. :class:`LiteralValue` is used to wrap
+ rendered values so they stay literal.
+ """
+ if isinstance(value, XComRef):
+ return self._xcom_arg(value) # resolves at render time; never
wrapped
+ if isinstance(value, TemplateRef):
+ # $t over a string renders that string; $t over $f is a template
file (read, then render).
+ if isinstance(value.source, FileRef):
+ # A constructed FileRef has a str path (its validator rejects
a templated filename).
+ return self._read_file(cast("str", value.source.path))
+ return value.source
+ if isinstance(value, FileRef):
+ return self._protect(self._read_file(cast("str", value.path)),
protect)
+ if isinstance(value, ConstRef):
+ return self._protect(value.value, protect)
+ if isinstance(value, Literal):
+ return self._lower(value.root, protect=protect)
+ if isinstance(value, dict):
+ return {k: self._lower(v, protect=protect) for k, v in
value.items()}
+ if isinstance(value, list):
+ return [self._lower(v, protect=protect) for v in value]
+ return self._protect(value, protect)
+
+ @staticmethod
+ def _protect(value: Any, protect: bool) -> Any:
+ return literal(value) if protect else value
+
+ def _xcom_arg(self, value: XComRef) -> XComArg:
+ """Resolve an :class:`XComRef` against the already-built operators."""
+ target = value.target
+ if isinstance(target, str):
+ return XComArg(self._operators[target])
+ if target.key is None:
+ return XComArg(self._operators[target.task])
+ return XComArg(self._operators[target.task], key=target.key)
+
+ def _read_file(self, rel_path: str) -> str:
+ """Read a ``$f`` at build time, relative to the document."""
+ try:
+ return
self._definition.sibling(rel_path).read_bytes().decode("utf-8")
+ except (OSError, ValueError) as exc:
+ raise YamlDagBuildError(f"cannot read $f {rel_path!r}: {exc}")
from exc
+
+ def _stub_callable(self, param_names: tuple[str, ...]) -> Callable[...,
None]:
+ """
+ Build a placeholder callable to feed to ``@task.stub``.
+
+ The stub needs a real callable with a meaningful signature; core
+ serialization reads it to build the TaskFlow arg-binding spec a
lang-SDK
+ handler is validated against (AIP-108), and ``@task.stub`` reads the
+ callable's source to check the body is empty.
+
+ The callable declares exactly the run arguments as a fixed, ordered
+ parameter list with an empty body, and its source is registered with
+ :mod:`linecache` so :func:`inspect.getsource` can read it back.
+ """
+ if (cached := self._stub_callables.get(param_names)) is not None:
+ return cached
+ source = f"def _stub({', '.join(param_names)}):\n ...\n"
+ filename = f"<yaml-dag stub ({', '.join(param_names)})>"
+ linecache.cache[filename] = (len(source), None,
source.splitlines(keepends=True), filename)
+ namespace: dict[str, Any] = {}
+ exec(compile(source, filename, "exec"), namespace)
+ self._stub_callables[param_names] = stub = namespace["_stub"]
+ return stub
+
+
+def _restrict_template_fields(
+ op: BaseOperator, with_: dict[str, Any], render_fields: list[str]
+) -> BaseOperator:
+ """
+ Make *op* render literal-by-default, per instance.
+
+ Only fields the author set in ``with:`` without a marker are dropped from
+ ``template_fields``; fields the operator templates by default or marked as
+ templates (e.g. uses ``$t``) are kept, so an operator's own defaults (e.g.
+ ``ExternalTaskMarker.logical_date``) still render.
+
+ ``template_ext`` is emptied: a file is read only when asked for explicitly
+ (``$f``), never by a rendered value's suffix.
+ """
+ op.template_fields = tuple(f for f in op.template_fields if f not in with_
or f in render_fields)
+ op.template_ext = ()
+ return op
+
+
+def _reject_unbindable_run_marker(task_id: str, value: Any, *, top_level: bool
= True) -> None:
+ """
+ Reject value-grammar a task handler's argument binding cannot carry.
+
+ The binding carries a JSON literal or a single top-level ``$x`` (an
upstream
+ task's return value). So a template (``$t``) or a keyed ``$x`` anywhere,
or a
+ plain ``$x`` nested inside a container instead of passed as its own
argument,
+ is an error. ``$f`` and ``$const`` lower to literals and are fine.
+ """
+ if isinstance(value, TemplateRef):
+ raise YamlDagBuildError(
+ f"task {task_id!r}: a template ($t) reference is not supported in
a run: task; the task "
+ "handler binding carries literal values and XCom return values,
not rendered templates"
+ )
+ if isinstance(value, XComRef):
+ target = value.target
+ if not isinstance(target, str) and target.key is not None:
+ raise YamlDagBuildError(
+ f"task {task_id!r}: a keyed XCom ($x with 'key') is not
supported in a run: task; "
+ "only a task's return value crosses the handler boundary"
+ )
+ if not top_level:
+ raise YamlDagBuildError(
+ f"task {task_id!r}: an XCom ($x) nested inside a run argument
is not supported; pass "
+ "the upstream output as its own top-level run argument"
+ )
+ return
+ if isinstance(value, Literal):
+ _reject_unbindable_run_marker(task_id, value.root, top_level=top_level)
+ elif isinstance(value, dict):
+ for item in value.values():
+ _reject_unbindable_run_marker(task_id, item, top_level=False)
+ elif isinstance(value, list):
+ for item in value:
+ _reject_unbindable_run_marker(task_id, item, top_level=False)
+
+
+def _reject_reserved_extra_keys(extras: Any) -> None:
+ """
+ Reject a marker key ($x/$t/$f/$const) nested in Dag- or task-level
arguments.
+
+ These arguments pass straight to ``DAG`` or the operator, never through the
+ value grammar, so a marker written there (e.g. ``params: {p: {$f: x}}``)
+ would be taken as a literal dict rather than flagged. Only marker keys are
+ rejected, so a genuine ``$`` key such as a JSON Schema ``$ref`` in a param
+ default is kept.
+ """
+ if isinstance(extras, dict):
+ for key, val in extras.items():
+ if key in _MARKER_KEYS:
+ raise YamlDagBuildError(f"marker key {key!r} may only be used
in 'with:' or 'run:'")
+ _reject_reserved_extra_keys(val)
+ elif isinstance(extras, list):
+ for item in extras:
+ _reject_reserved_extra_keys(item)
+
+
+def _coerce_temporal(mapping: dict[str, Any], target: type) -> dict[str, Any]:
+ """
+ Promote YAML-native dates and ISO date/datetime strings to datetimes.
+
+ Some fields on a Dag or operator are datetimes, but YAML has distinct date
+ and date-time formats, so a bare ``2024-01-01`` or a quoted ISO string
would
+ otherwise reach the datetime field unconverted. A date becomes midnight; an
+ ISO string is parsed (which also helps JSON compatibility).
+
+ Only values bound to a *target* argument annotated as ``datetime`` are
+ coerced. A non-temporal argument is left untouched, so a string that merely
+ looks like a date keeps its type.
+ """
+ temporal = _temporal_fields(target)
+ return {key: _to_datetime(value) if key in temporal else value for key,
value in mapping.items()}
+
+
+def _to_datetime(value: Any) -> Any:
+ if isinstance(value, datetime.datetime):
+ return value
+ if isinstance(value, datetime.date):
+ return datetime.datetime.combine(value, datetime.time.min)
+ if isinstance(value, str):
+ try:
+ return datetime.datetime.fromisoformat(value)
+ except ValueError:
+ return value
+ return value
+
+
+# Airflow annotates arguments with either stdlib datetime or pendulum DateTime.
+# Matched by name (see _temporal_fields) rather than resolved to avoid having
+# to evaluate dynamically.
+_TEMPORAL_TYPE_NAMES = frozenset({"datetime", "DateTime"})
+_TYPE_NAME_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_.]*")
+
+
[email protected]
+def _temporal_fields(target: type) -> frozenset[str]:
+ """
+ Names of *target*'s constructor arguments annotated as a datetime type.
+
+ The annotation text is matched against the datetime type names Airflow
uses;
+ it is never resolved, so one referring to a ``TYPE_CHECKING``-only name
does
+ not need importing. The constructor signature is read -- not class
+ attributes -- so attrs field aliases and operator-specific arguments are
+ covered. This is not perfect (a datetime type reached under another alias
is
+ missed); we'll fix edge cases as we go.
+ """
+ try:
+ params = inspect.signature(target).parameters
+ except (TypeError, ValueError):
+ return frozenset()
+ return frozenset(name for name, p in params.items() if
_is_temporal_annotation(str(p.annotation)))
+
+
+def _is_temporal_annotation(annotation: str) -> bool:
+ """Return True if *annotation* names a datetime type, ignoring any module
prefix."""
+ return any(ref.rsplit(".", 1)[-1] in _TEMPORAL_TYPE_NAMES for ref in
_TYPE_NAME_RE.findall(annotation))
+
+
+def _build_schedule(schedule: Any) -> Any:
+ """Build *schedule* into a schedule construct for :class:`DAG`."""
+ if isinstance(schedule, TimetableSchedule):
+ return import_string(schedule.uses)(**schedule.with_)
+ return schedule
+
+
+class _FieldMarkers(NamedTuple):
+ """Flags describing a value-grammar node's templating needs."""
+
+ needs_render: bool # carries $t or $x anywhere -- the field must render
+ has_marker: bool # carries $x/$t/$f -- only allowed on a templated field
+
+
+def _field_markers(value: Any) -> _FieldMarkers:
+ """
+ Scan a value-grammar node for its :class:`_FieldMarkers`.
+
+ ``$const`` is a force-literal escape, not a marker.
+ """
+ if isinstance(value, (XComRef, TemplateRef)):
+ return _FieldMarkers(needs_render=True, has_marker=True)
+ if isinstance(value, FileRef):
+ return _FieldMarkers(needs_render=False, has_marker=True)
+ if isinstance(value, ConstRef):
+ return _FieldMarkers(needs_render=False, has_marker=False)
+ if isinstance(value, Literal):
+ return _field_markers(value.root)
+ if isinstance(value, (dict, list)):
+ children = value.values() if isinstance(value, dict) else value
+ flags = [_field_markers(v) for v in children]
+ return _FieldMarkers(
+ needs_render=any(f.needs_render for f in flags),
+ has_marker=any(f.has_marker for f in flags),
+ )
+ return _FieldMarkers(needs_render=False, has_marker=False)
+
+
+def _topological_order(tasks: list[Task]) -> list[str]:
+ """Return task ids ordered so every task follows the tasks it depends
on."""
+ graph = {task.id_: task.get_dependency_ids() for task in tasks}
+ try:
+ return list(graphlib.TopologicalSorter(graph).static_order())
+ except graphlib.CycleError as exc:
+ raise YamlDagBuildError(f"dependency cycle among tasks:
{exc.args[1]}") from exc
diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py
b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py
new file mode 100644
index 00000000000..080e0fc97af
--- /dev/null
+++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/importer.py
@@ -0,0 +1,169 @@
+# 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 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.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,
YamlDagBuildError
+from airflow.sdk.importers.yaml_importer.migrator import DAG_SCHEMA_PREFIX
+from airflow.sdk.importers.yaml_importer.parser import YamlDagParseError,
parse_documents
+
+if TYPE_CHECKING:
+ from collections.abc import Iterator
+ from pathlib import Path
+
+ from airflow.dag_processing.bundles.base import BaseDagBundle # noqa:
SDK002
+ from airflow.sdk import DAG
+ from airflow.sdk.importers.base import DagDefinition
+
+log = logging.getLogger(__name__)
+
+# The sniff matches the canonical Dag ``$schema`` URL, not a bare ``$schema``,
so a JSON-Schema
+# file (or any other YAML declaring a $schema) in the bundle is not queued and
then stuck as a
+# parse error. Reuses the migrator's public prefix to stay in sync.
+_DAG_SCHEMA_PREFIX = DAG_SCHEMA_PREFIX.encode()
+
+
+class YamlDagImporter(AbstractDagImporter[FileDagDefinition]):
+ """Importer for YAML 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)
+ ref = repr(definition)
+ try:
+ source = definition.read_text(encoding="utf-8")
+ except Exception as exc:
+ result.errors.append(
+ DagImportError(
+ source_reference=ref,
+ 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). All documents' failures are reported together.
+ try:
+ documents = list(parse_documents(source, source=ref))
+ except YamlDagParseError as exc:
+ result.errors.append(DagImportError(source_reference=ref,
message=str(exc), error_type="parse"))
+ return result
+
+ dags: dict[str, DAG] = {}
+ build_errors: list[Exception] = []
+ with definition.import_context():
+ for document in documents:
+ try:
+ dag = DagBuilder(definition=definition,
bundle=bundle).build(document)
+ if dag.dag_id in dags:
+ raise YamlDagBuildError(f"duplicate dag_id
{dag.dag_id!r} across documents")
+ except Exception as exc:
+ exc.add_note(f"dag_id={document.dag_id!r}")
+ build_errors.append(exc)
+ else:
+ dags[dag.dag_id] = dag
+
+ if build_errors:
+ summary = f"{len(build_errors)} of {len(documents)} document(s)
failed to build"
+ detail = "\n".join(f"{err.__notes__[-1]}: {err}" for err in
build_errors)
+ stacktrace = "\n".join(
+ "".join(traceback.format_exception(type(err), err,
err.__traceback__)) for err in build_errors
+ )
+ result.errors.append(
+ DagImportError(
+ source_reference=ref,
+ message=f"{summary}:\n{detail}",
+ error_type="import",
+ stacktrace=stacktrace,
+ )
+ )
+ return result
+
+ result.dags.extend(dags.values())
+ for dag in dags.values():
+ 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: the file declares a ``$schema`` set to the canonical Dag
URL.
+
+ Both are required so a file that merely mentions the URL (without a
``$schema`` declaration)
+ is skipped. The sniff stays lenient past that -- a malformed Dag still
reaches the parser and
+ gets a clear error rather than silently disappearing. Keep the file if
the read fails.
+ """
+ if not safe_mode:
+ return True
+ try:
+ data = definition.read_bytes()
+ except OSError:
+ return True
+ return b"$schema" in data and _DAG_SCHEMA_PREFIX in data
diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/migrator.py
b/task-sdk/src/airflow/sdk/importers/yaml_importer/migrator.py
index 9a27c07e4b9..90f6ebed807 100644
--- a/task-sdk/src/airflow/sdk/importers/yaml_importer/migrator.py
+++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/migrator.py
@@ -44,9 +44,10 @@ if TYPE_CHECKING:
from cadwyn import VersionBundle
_SCHEMA_URL = "https://airflow.apache.org/schemas/dag/{date}.json";
-# The canonical ``$schema`` URL with the version captured; keep in sync with
``_SCHEMA_URL``.
_SCHEMA_URL_RE =
re.compile(r"https://airflow\.apache\.org/schemas/dag/(\d{4}-\d{2}-\d{2})\.json")
+DAG_SCHEMA_PREFIX = _SCHEMA_URL.partition("{date}")[0]
+
def schema_from_version(date: str) -> str:
"""Build the canonical ``$schema`` URL for a published version *date*."""
diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/models.py
b/task-sdk/src/airflow/sdk/importers/yaml_importer/models.py
index 94fdd50f852..35628dfed71 100644
--- a/task-sdk/src/airflow/sdk/importers/yaml_importer/models.py
+++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/models.py
@@ -50,6 +50,7 @@ if TYPE_CHECKING:
XCOM_KEYS = ("$x", "$xcom")
TEMPLATE_KEYS = ("$t", "$template")
+FILE_KEYS = ("$f", "$file")
CONST_KEY = "$const"
# Keys deferred to a later edition...
@@ -106,13 +107,26 @@ class _MarkerRef(BaseModel):
_marker_keys: ClassVar[tuple[str, ...]]
+ @classmethod
+ def _process_value_schema(cls, field_schema: JsonSchemaValue) ->
JsonSchemaValue:
+ """
+ Post-process published schema for each marker key.
+
+ The default publishes the field's schema as a disjoint union.
+ :class:`FileRef` overrides it to reject a templated filename in the
+ schema too.
+ """
+ return _as_exclusive_union(field_schema)
+
@classmethod
def __get_pydantic_json_schema__(
cls, core_schema: CoreSchema, handler: GetJsonSchemaHandler
) -> JsonSchemaValue:
schema = handler(core_schema)
- value_schema = _as_exclusive_union(next(iter(schema.get("properties",
{}).values())))
- schema["properties"] = {key: value_schema for key in cls._marker_keys}
+ schema["properties"] = {
+ key: cls._process_value_schema(next(iter(schema.get("properties",
{}).values())))
+ for key in cls._marker_keys
+ }
schema["additionalProperties"] = False
if len(cls._marker_keys) > 1: # Pydantic sets the correct 'required'
for the single-variant case.
schema.pop("required", None)
@@ -132,11 +146,54 @@ class XComRef(_MarkerRef):
class TemplateRef(_MarkerRef):
- """A Jinja template."""
+ """
+ A Jinja template, rendered when the task runs.
+
+ The template is a Jinja string, or a file reference (``$f``) to load a
+ template file.
+ """
+
+ model_config = ConfigDict(
+ extra="forbid",
+ json_schema_extra={
+ "description": (
+ "A Jinja template, rendered when the task runs. This can be an
"
+ "inline string, or a file reference ($f) to load a template
file."
+ ),
+ },
+ )
_marker_keys = TEMPLATE_KEYS
- source: str = Field(serialization_alias="$t",
validation_alias=AliasChoices(*TEMPLATE_KEYS))
+ source: str | FileRef = Field(serialization_alias="$t",
validation_alias=AliasChoices(*TEMPLATE_KEYS))
+
+
+class FileRef(_MarkerRef):
+ """
+ A file input (AIP-80 FileTemplate).
+
+ The path is a static string. A templated filename (a ``$f`` wrapping a
+ ``$t``) is reserved for a future release and rejected for now.
+ """
+
+ model_config = ConfigDict(
+ extra="forbid",
+ json_schema_extra={"description": "A file input."},
+ )
+
+ _marker_keys = FILE_KEYS
+
+ path: str | TemplateRef = Field(serialization_alias="$f",
validation_alias=AliasChoices(*FILE_KEYS))
+
+ @classmethod
+ def _process_value_schema(cls, field_schema: JsonSchemaValue) ->
JsonSchemaValue:
+ return {"type": "string"}
+
+ @model_validator(mode="after")
+ def _reject_templated_filename(self) -> FileRef:
+ if isinstance(self.path, TemplateRef):
+ raise ValueError("templated filename ($f wrapping $t) is not
supported yet")
+ return self
class ConstRef(_MarkerRef):
@@ -163,12 +220,16 @@ def _value_discriminator(v: Any) -> str:
return "xcom"
if key in TEMPLATE_KEYS:
return "template"
+ if key in FILE_KEYS:
+ return "file"
if key == CONST_KEY:
return "const"
if isinstance(v, XComRef):
return "xcom"
if isinstance(v, TemplateRef):
return "template"
+ if isinstance(v, FileRef):
+ return "file"
if isinstance(v, ConstRef):
return "const"
return "literal"
@@ -222,11 +283,14 @@ class Literal(RootModel):
Value = Annotated[
Annotated[XComRef, Tag("xcom")]
| Annotated[TemplateRef, Tag("template")]
+ | Annotated[FileRef, Tag("file")]
| Annotated[ConstRef, Tag("const")]
| Annotated[Literal, Tag("literal")],
Discriminator(_value_discriminator),
]
+TemplateRef.model_rebuild()
+FileRef.model_rebuild()
Literal.model_rebuild()
@@ -265,6 +329,13 @@ class Task(BaseModel):
raise ValueError(f"task {data.get('id')!r} may not set both
'uses' and 'run'")
return data
+ def get_dependency_ids(self) -> set[str]:
+ """Ids of the tasks this task depends on: explicit ``needs`` plus any
XCom references."""
+ ids = set(self.needs)
+ for value in itertools.chain(self.with_.values(), (self.run or
{}).values()):
+ ids.update(_iter_xcom_task_ids(value))
+ return ids
+
@classmethod
def __get_pydantic_json_schema__(
cls, core_schema: CoreSchema, handler: GetJsonSchemaHandler
@@ -449,10 +520,7 @@ class DagDocument(BaseModel):
raise ValueError(f"duplicate task id(s): {duplicates}")
validated_task_ids = set(task_id_counts)
for t in self.tasks:
- refs = set(t.needs)
- for v in itertools.chain(t.with_.values(), (t.run or {}).values()):
- refs.update(_iter_xcom_task_ids(v))
- if t.id_ in refs:
+ if t.id_ in (refs := t.get_dependency_ids()):
raise ValueError(f"task {t.id_!r} depends on itself")
if missing := (refs - validated_task_ids):
raise ValueError(f"task {t.id_!r} references unknown task(s):
{sorted(missing)}")
diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/schema.json
b/task-sdk/src/airflow/sdk/importers/yaml_importer/schema.json
index 953c61b6ccc..67686d202ae 100644
--- a/task-sdk/src/airflow/sdk/importers/yaml_importer/schema.json
+++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/schema.json
@@ -14,6 +14,32 @@
"title": "ConstRef",
"type": "object"
},
+ "FileRef": {
+ "additionalProperties": false,
+ "description": "A file input.",
+ "oneOf": [
+ {
+ "required": [
+ "$f"
+ ]
+ },
+ {
+ "required": [
+ "$file"
+ ]
+ }
+ ],
+ "properties": {
+ "$f": {
+ "type": "string"
+ },
+ "$file": {
+ "type": "string"
+ }
+ },
+ "title": "FileRef",
+ "type": "object"
+ },
"Literal": {
"description": "A literal value.",
"not": {
@@ -34,6 +60,9 @@
{
"$ref": "#/$defs/TemplateRef"
},
+ {
+ "$ref": "#/$defs/FileRef"
+ },
{
"$ref": "#/$defs/ConstRef"
},
@@ -53,6 +82,9 @@
{
"$ref": "#/$defs/TemplateRef"
},
+ {
+ "$ref": "#/$defs/FileRef"
+ },
{
"$ref": "#/$defs/ConstRef"
},
@@ -140,6 +172,9 @@
{
"$ref": "#/$defs/TemplateRef"
},
+ {
+ "$ref": "#/$defs/FileRef"
+ },
{
"$ref": "#/$defs/ConstRef"
},
@@ -167,6 +202,9 @@
{
"$ref": "#/$defs/TemplateRef"
},
+ {
+ "$ref": "#/$defs/FileRef"
+ },
{
"$ref": "#/$defs/ConstRef"
},
@@ -196,7 +234,7 @@
},
"TemplateRef": {
"additionalProperties": false,
- "description": "A Jinja template.",
+ "description": "A Jinja template, rendered when the task runs. This can
be an inline string, or a file reference ($f) to load a template file.",
"oneOf": [
{
"required": [
@@ -211,12 +249,26 @@
],
"properties": {
"$t": {
- "title": "$T",
- "type": "string"
+ "oneOf": [
+ {
+ "type": "string"
+ },
+ {
+ "$ref": "#/$defs/FileRef"
+ }
+ ],
+ "title": "$T"
},
"$template": {
- "title": "$T",
- "type": "string"
+ "oneOf": [
+ {
+ "type": "string"
+ },
+ {
+ "$ref": "#/$defs/FileRef"
+ }
+ ],
+ "title": "$T"
}
},
"title": "TemplateRef",
diff --git a/task-sdk/src/airflow/sdk/importers/zip_importer.py
b/task-sdk/src/airflow/sdk/importers/zip_importer.py
index 1b7c53ad104..ff40089f750 100644
--- a/task-sdk/src/airflow/sdk/importers/zip_importer.py
+++ b/task-sdk/src/airflow/sdk/importers/zip_importer.py
@@ -21,6 +21,7 @@ from __future__ import annotations
import contextlib
import logging
import os
+import posixpath
import sys
import tempfile
import threading
@@ -101,6 +102,18 @@ class ZipMemberDagDefinition(FileDagDefinition):
self._content = z.read(self.file_path)
return self._content
+ def sibling(self, relpath: str) -> zipfile.Path:
+ base_dir = posixpath.dirname(self.file_path)
+ member = posixpath.normpath(posixpath.join(base_dir, relpath))
+ within = (
+ (member == base_dir or member.startswith(f"{base_dir}/"))
+ if base_dir
+ else not (member == ".." or member.startswith("../"))
+ )
+ if posixpath.isabs(member) or not within:
+ raise ValueError(f"{relpath!r} resolves outside the Dag file's
directory")
+ return zipfile.Path(self.zip_path, member)
+
@contextlib.contextmanager
def as_file(self) -> Generator[Path, None, None]:
# Keep the member's file name: importers may derive identity from it,
such as a Dag id.
diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py
b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py
new file mode 100644
index 00000000000..84b7c53511a
--- /dev/null
+++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_importer.py
@@ -0,0 +1,605 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import datetime
+import inspect
+import textwrap
+import zipfile
+from types import SimpleNamespace
+
+import pytest
+
+from airflow.sdk import BaseOperator, XComArg
+from airflow.sdk.definitions._internal.templater import LiteralValue
+from airflow.sdk.definitions.timetables.simple import OnceTimetable
+from airflow.sdk.importers import ZipMemberDagDefinition
+from airflow.sdk.importers.base import FilesystemDagDefinition
+from airflow.sdk.importers.yaml_importer.importer import YamlDagImporter
+
+HEAD_SCHEMA = "https://airflow.apache.org/schemas/dag/2026-10-30.json";
+
+
+class _RecordingOperator(BaseOperator):
+ """A trivial operator with one templated field (``payload``) and one plain
field (``note``)."""
+
+ template_fields = ("payload",)
+ template_ext = (".sql",) # non-empty so the importer emptying it to () is
observable, not vacuous
+
+ def __init__(self, *, payload=None, note=None, **kwargs):
+ super().__init__(**kwargs)
+ self.payload = payload
+ self.note = note
+
+ def execute(self, context):
+ return self.payload
+
+
+OP = f"{__name__}.{_RecordingOperator.__name__}"
+
+
+class _DefaultTemplateOperator(BaseOperator):
+ """Templates ``preset`` by default (set in __init__, not by the author)
plus author-set ``payload``."""
+
+ template_fields = ("payload", "preset")
+
+ def __init__(self, *, payload=None, preset="{{ ds }}", **kwargs):
+ super().__init__(**kwargs)
+ self.payload = payload
+ self.preset = preset
+
+ def execute(self, context):
+ return self.payload
+
+
+OP_DEFAULT_TEMPLATE = f"{__name__}.{_DefaultTemplateOperator.__name__}"
+
+
+class _TemporalOperator(BaseOperator):
+ """An operator with a datetime-typed argument (``cutoff``)."""
+
+ def __init__(self, *, cutoff: datetime.datetime | None = None, **kwargs):
+ super().__init__(**kwargs)
+ self.cutoff = cutoff
+
+
+OP_TEMPORAL = f"{__name__}.{_TemporalOperator.__name__}"
+
+
[email protected]
+def mock_bundle(tmp_path):
+ bundle_dir = tmp_path / "bundle"
+ bundle_dir.mkdir(parents=True, exist_ok=True)
+ return SimpleNamespace(name="test_bundle", path=bundle_dir)
+
+
+def _def(bundle, name: str, body: str) -> FilesystemDagDefinition:
+ path = bundle.path / name
+ path.write_text(textwrap.dedent(body))
+ return FilesystemDagDefinition(path=path)
+
+
+def _doc(tasks: str, *, dag_id: str = "d", extra: str = "") -> str:
+ return f"$schema: {HEAD_SCHEMA}\ndag_id:
{dag_id}\n{extra}tasks:\n{textwrap.indent(textwrap.dedent(tasks), ' ')}"
+
+
+def test_can_handle():
+ importer = YamlDagImporter()
+ assert importer.can_handle("dag.yaml")
+ assert importer.can_handle("dag.yml")
+ assert not importer.can_handle("dag.py")
+ assert not importer.can_handle("dag.json")
+
+
+def test_get_source_code(mock_bundle):
+ definition = _def(mock_bundle, "d.yaml", f"$schema: {HEAD_SCHEMA}\ndag_id:
d\ntasks: []\n")
+ src = YamlDagImporter().get_source_code(definition)
+ assert src.language == "yaml"
+ assert "$schema" in src.source_code
+
+
+def test_might_contain_dag(mock_bundle):
+ importer = YamlDagImporter()
+ has = _def(mock_bundle, "has.yaml", f"$schema: {HEAD_SCHEMA}\ndag_id:
d\ntasks: []\n")
+ nope = _def(mock_bundle, "nope.yaml", "just: config\n")
+ # A $schema that is not the canonical Dag URL (e.g. a JSON-Schema file)
must not be picked up.
+ other = _def(mock_bundle, "schema.yaml", "$schema:
https://json-schema.org/draft/2020-12/schema\n";)
+ # The canonical URL merely mentioned, without a $schema declaration, must
not be picked up either.
+ mention = _def(mock_bundle, "mention.yaml", f"note: see {HEAD_SCHEMA}\n")
+ assert importer.might_contain_dag(has, safe_mode=True)
+ assert not importer.might_contain_dag(nope, safe_mode=True)
+ assert not importer.might_contain_dag(other, safe_mode=True)
+ assert not importer.might_contain_dag(mention, safe_mode=True)
+ assert importer.might_contain_dag(nope, safe_mode=False) # sniff skipped
+
+
+def test_list_dag_definitions_skips_non_dag_yaml(mock_bundle):
+ _def(mock_bundle, "dag.yaml", f"$schema: {HEAD_SCHEMA}\ndag_id: d\ntasks:
[]\n")
+ _def(mock_bundle, "config.yaml", "just: config\n") # no $schema -> skipped
+ _def(mock_bundle, "notes.txt", "hello") # wrong extension
+ names = {d.path.name for d in
YamlDagImporter().list_dag_definitions(mock_bundle)}
+ assert names == {"dag.yaml"}
+
+
+def test_import_builds_operator_dag(mock_bundle):
+ definition = _def(
+ mock_bundle,
+ "dag.yaml",
+ _doc(
+ f"""
+ - {{id: a, uses: {OP}, with: {{payload: hi}}}}
+ - {{id: b, uses: {OP}, with: {{payload: {{$x: a}}}}}}
+ """,
+ dag_id="pipeline",
+ extra="catchup: false\n",
+ ),
+ )
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert result.success, result.errors
+ (dag,) = result.dags
+ assert dag.dag_id == "pipeline"
+ assert set(dag.task_ids) == {"a", "b"}
+ assert dag.catchup is False # Dag-level attribute passed through
+ assert dag.bundle_name == "test_bundle"
+ a, b = dag.get_task("a"), dag.get_task("b")
+ assert a.payload == "hi" # literal
+ assert isinstance(b.payload, XComArg) # $x -> XComArg
+ assert b.upstream_task_ids == {"a"} # XCom wired the dependency
+ assert "payload" in b.template_fields # $x opted the field into rendering
+ assert "payload" not in a.template_fields # a's unmarked literal was
dropped from rendering
+
+
+def test_import_wires_needs(mock_bundle):
+ definition = _def(
+ mock_bundle,
+ "dag.yaml",
+ _doc(f"""
+ - {{id: a, uses: {OP}, with: {{payload: x}}}}
+ - {{id: b, uses: {OP}, with: {{payload: y}}, needs: [a]}}
+ """),
+ )
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert result.success, result.errors
+ assert result.dags[0].get_task("b").upstream_task_ids == {"a"}
+
+
+def test_import_bodyless_task_becomes_empty_operator(mock_bundle):
+ definition = _def(mock_bundle, "dag.yaml", _doc("- {id: a}"))
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert result.success, result.errors
+ assert type(result.dags[0].get_task("a")).__name__ == "EmptyOperator"
+
+
+def test_import_bodyless_task_with_args_errors(mock_bundle):
+ definition = _def(mock_bundle, "dag.yaml", _doc("- {id: a, with: {x: 1}}"))
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert not result.dags
+ assert any(
+ "Invalid arguments were passed to EmptyOperator" in e.message and
"'x'" in e.message
+ for e in result.errors
+ )
+
+
+def test_import_schedule_cron(mock_bundle):
+ definition = _def(
+ mock_bundle,
+ "dag.yaml",
+ _doc(f"- {{id: a, uses: {OP}, with: {{payload: x}}}}",
extra="schedule: '0 3 * * *'\n"),
+ )
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert result.success, result.errors
+ assert result.dags[0].schedule == "0 3 * * *"
+
+
+def test_import_schedule_timetable(mock_bundle):
+ definition = _def(
+ mock_bundle,
+ "dag.yaml",
+ _doc(
+ f"- {{id: a, uses: {OP}, with: {{payload: x}}}}",
+ extra="schedule: {uses:
airflow.sdk.definitions.timetables.simple.OnceTimetable}\n",
+ ),
+ )
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert result.success, result.errors
+ assert isinstance(result.dags[0].timetable, OnceTimetable)
+
+
+def test_import_keyed_xcom_on_uses_task(mock_bundle):
+ definition = _def(
+ mock_bundle,
+ "dag.yaml",
+ _doc(f"""
+ - {{id: a, uses: {OP}, with: {{payload: x}}}}
+ - {{id: b, uses: {OP}, with: {{payload: {{$x: {{task: a, key:
out}}}}}}}}
+ """),
+ )
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert result.success, result.errors
+ b = result.dags[0].get_task("b")
+ assert isinstance(b.payload, XComArg)
+ assert b.payload.key == "out"
+ assert b.upstream_task_ids == {"a"}
+
+
[email protected](
+ ("value", "hour"),
+ [("2024-01-01", 0), ("'2024-01-01'", 0), ("2024-01-01T03:30:00", 3)], #
YAML date, ISO string, datetime
+)
+def test_import_coerces_temporal_start_date(mock_bundle, value, hour):
+ definition = _def(
+ mock_bundle,
+ "dag.yaml",
+ _doc(f"- {{id: a, uses: {OP}, with: {{payload: x}}}}",
extra=f"start_date: {value}\n"),
+ )
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert result.success, result.errors
+ start = result.dags[0].start_date
+ assert isinstance(start, datetime.datetime) # promoted from date/string
+ assert start.tzinfo is not None # localized by DAG()
+ assert (start.year, start.month, start.day, start.hour) == (2024, 1, 1,
hour)
+
+
+def test_import_does_not_coerce_nontemporal_arg(mock_bundle):
+ definition = _def(
+ mock_bundle,
+ "dag.yaml",
+ _doc(f"- {{id: a, uses: {OP}, owner: '2024-01-01', with: {{payload:
x}}}}"),
+ )
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert result.success, result.errors
+ assert result.dags[0].get_task("a").owner == "2024-01-01"
+
+
+def test_import_coerces_temporal_with_arg(mock_bundle):
+ definition = _def(
+ mock_bundle,
+ "dag.yaml",
+ _doc(f"- {{id: a, uses: {OP_TEMPORAL}, with: {{cutoff:
2024-01-01}}}}"),
+ )
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert result.success, result.errors
+ cutoff = result.dags[0].get_task("a").cutoff
+ assert isinstance(cutoff, datetime.datetime)
+ assert (cutoff.year, cutoff.month, cutoff.day) == (2024, 1, 1)
+
+
+def test_import_code_task_builds_stub_operator(mock_bundle):
+ definition = _def(
+ mock_bundle,
+ "dag.yaml",
+ _doc(f"""
+ - {{id: a, uses: {OP}, with: {{payload: hi}}}}
+ - {{id: b, run: {{value: {{$x: a}}, label: done}}, queue:
handlers}}
+ """),
+ )
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert result.success, result.errors
+ b = result.dags[0].get_task("b")
+ assert type(b).__name__ == "_StubOperator"
+ assert getattr(b, "is_stub", False) is True # the real stub operator
Airflow serializes
+ assert b.queue == "handlers" # task-level attribute passed through
+ assert isinstance(b.op_kwargs["value"], XComArg) # $x -> XComArg, bound
as a handler argument
+ assert b.op_kwargs["label"] == "done" # literal bound verbatim (run args
are never rendered)
+ assert b.upstream_task_ids == {"a"} # $x wired the dependency
+ # The stub is fed a callable whose ordered parameters are the run argument
names, so core
+ # serialization can build the handler arg-binding spec (AIP-108) from a
real signature.
+ assert list(inspect.signature(b.python_callable).parameters) == ["value",
"label"]
+
+
+def test_import_mixed_container_keeps_literal_siblings_literal(mock_bundle):
+ (mock_bundle.path / "note.txt").write_text("hello")
+ definition = _def(
+ mock_bundle,
+ "dag.yaml",
+ _doc(f"""
+ - {{id: up, uses: {OP}, with: {{payload: x}}}}
+ - {{id: t, uses: {OP}, with: {{payload: {{a: {{$x: up}}, b: '{{{{
not a template }}}}', c: {{$const: 7}}, d: {{$f: note.txt}}}}}}}}
+ """),
+ )
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert result.success, result.errors
+ payload = result.dags[0].get_task("t").payload
+ assert isinstance(payload["a"], XComArg) # $x still resolves
+ # the container renders ($x present), so every literal sibling -- plain,
$const, $f -- stays wrapped
+ assert isinstance(payload["b"], LiteralValue)
+ assert payload["b"].value == "{{ not a template }}"
+ assert isinstance(payload["c"], LiteralValue)
+ assert payload["c"].value == 7
+ assert isinstance(payload["d"], LiteralValue)
+ assert payload["d"].value == "hello"
+
+
+def test_import_code_task_rejects_non_identifier_arg(mock_bundle):
+ definition = _def(mock_bundle, "dag.yaml", _doc("- {id: a, run:
{'bad-name': 1}}"))
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert not result.dags
+ assert any("not a valid identifier" in e.message for e in result.errors)
+
+
+def test_import_run_rejects_template(mock_bundle):
+ definition = _def(mock_bundle, "dag.yaml", _doc("- {id: a, run: {day: {$t:
'{{ ds }}'}}}"))
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert not result.dags
+ assert any("template ($t)" in e.message for e in result.errors)
+
+
+def test_import_run_rejects_keyed_xcom(mock_bundle):
+ definition = _def(
+ mock_bundle,
+ "dag.yaml",
+ _doc("""
+ - {id: a, run: {}}
+ - {id: b, run: {v: {$x: {task: a, key: k}}}}
+ """),
+ )
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert not result.dags
+ assert any("keyed XCom" in e.message for e in result.errors)
+
+
+def test_import_run_task_rejects_with(mock_bundle):
+ definition = _def(mock_bundle, "dag.yaml", _doc("- {id: t, run: {n: 1},
with: {region: eu}}"))
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert not result.dags
+ assert any("takes its inputs under 'run:'" in e.message for e in
result.errors)
+
+
+def test_import_run_rejects_nested_xcom(mock_bundle):
+ definition = _def(
+ mock_bundle,
+ "dag.yaml",
+ _doc(f"""
+ - {{id: a, uses: {OP}, with: {{payload: x}}}}
+ - {{id: t, run: {{items: [{{$x: a}}]}}}}
+ """),
+ )
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert not result.dags
+ assert any("nested inside a run argument" in e.message for e in
result.errors)
+
+
+def test_import_run_rejects_non_json_literal(mock_bundle):
+ definition = _def(mock_bundle, "dag.yaml", _doc("- {id: t, run: {day:
2024-01-01}}"))
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert not result.dags
+ assert any("must be JSON-serializable" in e.message for e in result.errors)
+
+
+def test_import_run_rejects_context_key_arg(mock_bundle):
+ definition = _def(mock_bundle, "dag.yaml", _doc("- {id: t, run: {ds: 1}}"))
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert not result.dags
+ assert any("cannot be named" in e.message for e in result.errors)
+
+
+def test_import_rejects_marker_key_in_extras(mock_bundle):
+ # Dag-level extras bypass the value grammar; a marker nested there is a
typo, not a literal dict.
+ definition = _def(
+ mock_bundle,
+ "dag.yaml",
+ _doc(f"- {{id: a, uses: {OP}, with: {{payload: x}}}}", extra="params:
{p: {$f: README.md}}\n"),
+ )
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert not result.dags
+ assert any("marker key '$f'" in e.message for e in result.errors)
+
+
[email protected]("marker", ["{$x: a}", "{$t: 'v'}", "{$f: q.sql}"])
+def test_import_marker_on_non_template_field_errors(mock_bundle, marker):
+ (mock_bundle.path / "q.sql").write_text("SELECT 1") # exists, so this is
not a missing-file error
+ definition = _def(
+ mock_bundle,
+ "dag.yaml",
+ _doc(f"""
+ - {{id: a, uses: {OP}, with: {{payload: hi}}}}
+ - {{id: b, uses: {OP}, with: {{payload: ok, note: {marker}}}}}
+ """),
+ )
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert not result.dags
+ assert any("only allowed on templated fields" in e.message for e in
result.errors)
+
+
+def test_import_base_operator_arg_in_with_errors(mock_bundle):
+ definition = _def(
+ mock_bundle, "dag.yaml", _doc(f"- {{id: a, uses: {OP}, with:
{{payload: x, retries: 2}}}}")
+ )
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert not result.dags
+ assert any("'retries'" in e.message and "not under 'with'" in e.message
for e in result.errors)
+
+
[email protected](
+ "body", ["- {{id: a, uses: {op}, with: {{payload: x}}, note: 1}}", "-
{{id: a, run: {{}}, note: 1}}"]
+)
+def test_import_operator_arg_at_task_level_errors(mock_bundle, body):
+ definition = _def(mock_bundle, "dag.yaml", _doc(body.format(op=OP)))
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert not result.dags
+ assert any("not task-level" in e.message and "'note'" in e.message for e
in result.errors)
+
+
+def test_import_literal_by_default(mock_bundle):
+ 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) == ()
+
+
+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_keeps_operator_default_template_fields(mock_bundle):
+ definition = _def(
+ mock_bundle, "dag.yaml", _doc(f"- {{id: a, uses:
{OP_DEFAULT_TEMPLATE}, with: {{payload: hi}}}}")
+ )
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert result.success, result.errors
+ a = result.dags[0].get_task("a")
+ assert set(a.template_fields) == {"preset"} # operator default kept;
author's unmarked payload dropped
+ assert a.payload == "hi" # author-set unmarked field no longer renders
+
+
+def test_import_file_ref_reads_verbatim(mock_bundle):
+ (mock_bundle.path / "query.sql").write_text("SELECT 1")
+ definition = _def(
+ 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
+
+
[email protected]("bad_path", ["../../etc/passwd", "/etc/passwd"])
+def test_import_file_ref_outside_dag_dir_is_rejected(mock_bundle, bad_path):
+ definition = _def(
+ mock_bundle, "dag.yaml", _doc(f"- {{id: a, uses: {OP}, with:
{{payload: {{$f: '{bad_path}'}}}}}}")
+ )
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert not result.dags
+ assert any("resolves outside" in e.message for e in result.errors)
+
+
+_ZIP_OP_SOURCE = """
+from airflow.sdk import BaseOperator
+
+
+class ZipOp(BaseOperator):
+ template_fields = ("payload",)
+
+ def __init__(self, *, payload=None, **kwargs):
+ super().__init__(**kwargs)
+ self.payload = payload
+
+ def execute(self, context):
+ return self.payload
+"""
+
+
+def test_import_yaml_inside_zip_bundle(tmp_path):
+ zip_path = tmp_path / "bundle.zip"
+ with zipfile.ZipFile(zip_path, "w") as z:
+ z.writestr("_ziptest_ops.py", _ZIP_OP_SOURCE)
+ z.writestr("q.sql", "SELECT 42")
+ z.writestr(
+ "dag.yaml",
+ _doc("- {id: a, uses: _ziptest_ops.ZipOp, with: {payload: {$f:
q.sql}}}", dag_id="z"),
+ )
+ definition = ZipMemberDagDefinition(zip_path=zip_path,
file_path="dag.yaml")
+ bundle = SimpleNamespace(name="zbundle", path=tmp_path)
+ result = YamlDagImporter().import_definition(definition, bundle=bundle)
+ assert result.success, result.errors
+ a = result.dags[0].get_task("a")
+ assert type(a).__name__ == "ZipOp" # imported from a module inside the
zip (import_context)
+ assert a.payload == "SELECT 42" # $f read a sibling member in the same
archive
+
+
+def test_import_template_file_reads_then_templates(mock_bundle):
+ (mock_bundle.path / "query.sql").write_text("SELECT {{ ds }}")
+ definition = _def(
+ 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):
+ 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 any("No module named" in e.message for e in result.errors)
+
+
+def test_import_one_bad_document_fails_whole_file(mock_bundle):
+ good = _doc(f"- {{id: a, uses: {OP}, with: {{payload: x}}}}",
dag_id="good")
+ bad = _doc("- {id: a, uses: does.not.Exist, with: {}}", dag_id="bad")
+ definition = _def(mock_bundle, "dag.yaml", f"{good}\n---\n{bad}")
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert not result.dags
+ assert len(result.errors) == 1
+ err = result.errors[0]
+ assert "dag_id=" not in err.source_reference # keyed by the file, not the
document
+ assert "dag_id='bad'" in err.message # the failing document is named, not
mashed
+
+
+def test_import_duplicate_dag_id_across_documents_errors(mock_bundle):
+ one = _doc(f"- {{id: a, uses: {OP}, with: {{payload: x}}}}", dag_id="dup")
+ two = _doc(f"- {{id: b, uses: {OP}, with: {{payload: y}}}}", dag_id="dup")
+ definition = _def(mock_bundle, "dag.yaml", f"{one}\n---\n{two}")
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert not result.dags
+ assert any("duplicate dag_id 'dup'" in e.message for e in result.errors)
+
+
+def test_import_dependency_cycle_errors(mock_bundle):
+ definition = _def(
+ mock_bundle,
+ "dag.yaml",
+ _doc(f"""
+ - {{id: a, uses: {OP}, with: {{payload: x}}, needs: [b]}}
+ - {{id: b, uses: {OP}, with: {{payload: y}}, needs: [a]}}
+ - {{id: c, uses: {OP}, with: {{payload: z}}, needs: [a]}}
+ """),
+ )
+ result = YamlDagImporter().import_definition(definition,
bundle=mock_bundle)
+ assert not result.dags
+ (error,) = [e for e in result.errors if "cycle" in e.message]
+ assert "'a'" in error.message
+ assert "'b'" in error.message
+ assert "'c'" not in error.message
diff --git a/task-sdk/tests/task_sdk/importers/yaml_importer/test_parser.py
b/task-sdk/tests/task_sdk/importers/yaml_importer/test_parser.py
index 22a69df2c75..91b0dd33f17 100644
--- a/task-sdk/tests/task_sdk/importers/yaml_importer/test_parser.py
+++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_parser.py
@@ -531,6 +531,14 @@ def
test_published_schema_accepts_null_and_empty_bodies(task):
"tasks:\n - {id: t, uses: X, with: {k: {$t: 'v-{{ ds }}'}}}",
id="template",
),
+ pytest.param(
+ "tasks:\n - {id: t, uses: X, with: {k: {$f:
'reports/daily.sql'}}}",
+ id="file",
+ ),
+ pytest.param(
+ "tasks:\n - {id: t, uses: X, with: {k: {$t: {$f:
'reports/daily.sql'}}}}",
+ id="template-file",
+ ),
pytest.param(
"templates: {base: {run: {r: 1}}}\ntasks:\n - {id: t, extends:
base}",
id="extends-bare",
@@ -564,6 +572,8 @@ def test_valid_shapes_round_trip(body):
pytest.param("{cfg: {$weird: 1}}", id="dollar-key-nested"),
pytest.param("{$x: up, $xcom: up}", id="two-marker-keys"),
pytest.param("{$x: {task: up, nope: 1}}", id="xcom-target-extra-key"),
+ pytest.param("{$f: {$t: 'q-{{ ds }}.sql'}}",
id="templated-filename-reserved"),
+ pytest.param("{$t: {$f: {$t: x}}}",
id="templated-filename-nested-reserved"),
],
)
def test_invalid_shapes_rejected_by_both(value):
