diff --git a/task-sdk/pyproject.toml b/task-sdk/pyproject.toml index 0ce16d4b713e9..9a4c9d44ea97b 100644 --- a/task-sdk/pyproject.toml +++ b/task-sdk/pyproject.toml @@ -54,14 +54,15 @@ dependencies = [ "babel>=2.17.0", "cadwyn>=6.1.1", "fsspec>=2023.10.0", + "greenback>=1.2.1", "httpx>=0.27.0", "jinja2>=3.1.5", "methodtools>=0.4.7", "msgspec>=0.19.0", - "python-dateutil>=2.7.0", "psutil>=6.1.0", + "python-dateutil>=2.7.0", + "pyyaml>=6.0.3", "structlog>=26.1.0", - "greenback>=1.2.1", "tenacity>=8.3.0", # Start of shared timezones dependencies "pendulum>=3.1.0", diff --git a/task-sdk/src/airflow/sdk/importers/base.py b/task-sdk/src/airflow/sdk/importers/base.py index e594e15ad2402..390c28bfc2fe5 100644 --- a/task-sdk/src/airflow/sdk/importers/base.py +++ b/task-sdk/src/airflow/sdk/importers/base.py @@ -34,11 +34,17 @@ if TYPE_CHECKING: from collections.abc import Generator, Iterable, Iterator - from typing import Self + from typing import Protocol, Self from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 from airflow.sdk import DAG + class _SiblingResource(Protocol): + """A readable handle to a resource, e.g. a :class:`pathlib.Path` or :class:`zipfile.Path`.""" + + def read_bytes(self) -> bytes: ... + + log = logging.getLogger(__name__) DefT = TypeVar("DefT", bound="DagDefinition") @@ -64,6 +70,22 @@ def read_text(self, encoding: str = "utf-8") -> str: """Read and return the content of the resource as a string.""" return self.read_bytes().decode(encoding) + def sibling(self, relpath: str) -> _SiblingResource: + """ + Return a readable handle to a resource stored next to this definition. + + The handle supports ``read_bytes()``. Resolution 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 @@ def suffix(self) -> str: def read_bytes(self) -> bytes: return self.path.read_bytes() + def sibling(self, relpath: str) -> Path: + base = self.path.parent.resolve() + target = base.joinpath(relpath).resolve() + if not target.is_relative_to(base): + raise ValueError(f"{relpath!r} resolves outside the Dag file's directory {str(base)!r}") + return target + @contextlib.contextmanager def as_file(self) -> Generator[Path, None, None]: yield self.path diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/__init__.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/__init__.py index 2f229838802e7..a08ed337b706d 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/__init__.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/__init__.py @@ -19,14 +19,15 @@ from __future__ import annotations +from airflow.sdk.importers.yaml_importer.builder import YamlDagBuildError +from airflow.sdk.importers.yaml_importer.importer import YamlDagImporter from airflow.sdk.importers.yaml_importer.models import DagDocument -from airflow.sdk.importers.yaml_importer.parser import ( - YamlDagParseError, - parse_documents, -) +from airflow.sdk.importers.yaml_importer.parser import YamlDagParseError, parse_documents __all__ = [ "DagDocument", + "YamlDagBuildError", + "YamlDagImporter", "YamlDagParseError", "parse_documents", ] diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py new file mode 100644 index 0000000000000..02d49c6a8ffc3 --- /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}) + + +@attrs.define(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"" + 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_.]*") + + +@functools.cache +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 0000000000000..080e0fc97afb3 --- /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 9a27c07e4b99c..90f6ebed807d9 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/migrator.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/migrator.py @@ -44,9 +44,10 @@ from cadwyn import VersionBundle _SCHEMA_URL = "https://airflow.apache.org/schemas/dag/{date}.json" -# The canonical ``$schema`` URL with the version captured; keep in sync with ``_SCHEMA_URL``. _SCHEMA_URL_RE = re.compile(r"https://airflow\.apache\.org/schemas/dag/(\d{4}-\d{2}-\d{2})\.json") +DAG_SCHEMA_PREFIX = _SCHEMA_URL.partition("{date}")[0] + def schema_from_version(date: str) -> str: """Build the canonical ``$schema`` URL for a published version *date*.""" diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/models.py b/task-sdk/src/airflow/sdk/importers/yaml_importer/models.py index 94fdd50f852f4..35628dfed7131 100644 --- a/task-sdk/src/airflow/sdk/importers/yaml_importer/models.py +++ b/task-sdk/src/airflow/sdk/importers/yaml_importer/models.py @@ -50,6 +50,7 @@ XCOM_KEYS = ("$x", "$xcom") TEMPLATE_KEYS = ("$t", "$template") +FILE_KEYS = ("$f", "$file") CONST_KEY = "$const" # Keys deferred to a later edition... @@ -106,13 +107,26 @@ class _MarkerRef(BaseModel): _marker_keys: ClassVar[tuple[str, ...]] + @classmethod + def _process_value_schema(cls, field_schema: JsonSchemaValue) -> JsonSchemaValue: + """ + Post-process published schema for each marker key. + + The default publishes the field's schema as a disjoint union. + :class:`FileRef` overrides it to reject a templated filename in the + schema too. + """ + return _as_exclusive_union(field_schema) + @classmethod def __get_pydantic_json_schema__( cls, core_schema: CoreSchema, handler: GetJsonSchemaHandler ) -> JsonSchemaValue: schema = handler(core_schema) - value_schema = _as_exclusive_union(next(iter(schema.get("properties", {}).values()))) - schema["properties"] = {key: value_schema for key in cls._marker_keys} + schema["properties"] = { + key: cls._process_value_schema(next(iter(schema.get("properties", {}).values()))) + for key in cls._marker_keys + } schema["additionalProperties"] = False if len(cls._marker_keys) > 1: # Pydantic sets the correct 'required' for the single-variant case. schema.pop("required", None) @@ -132,11 +146,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 @@ def __get_pydantic_json_schema__( Value = Annotated[ Annotated[XComRef, Tag("xcom")] | Annotated[TemplateRef, Tag("template")] + | Annotated[FileRef, Tag("file")] | Annotated[ConstRef, Tag("const")] | Annotated[Literal, Tag("literal")], Discriminator(_value_discriminator), ] +TemplateRef.model_rebuild() +FileRef.model_rebuild() Literal.model_rebuild() @@ -265,6 +329,13 @@ def _check_task(cls, data: Any) -> Any: raise ValueError(f"task {data.get('id')!r} may not set both 'uses' and 'run'") return data + def get_dependency_ids(self) -> set[str]: + """Ids of the tasks this task depends on: explicit ``needs`` plus any XCom references.""" + ids = set(self.needs) + for value in itertools.chain(self.with_.values(), (self.run or {}).values()): + ids.update(_iter_xcom_task_ids(value)) + return ids + @classmethod def __get_pydantic_json_schema__( cls, core_schema: CoreSchema, handler: GetJsonSchemaHandler @@ -449,10 +520,7 @@ def _check_refs(self): raise ValueError(f"duplicate task id(s): {duplicates}") validated_task_ids = set(task_id_counts) for t in self.tasks: - refs = set(t.needs) - for v in itertools.chain(t.with_.values(), (t.run or {}).values()): - refs.update(_iter_xcom_task_ids(v)) - if t.id_ in refs: + if t.id_ in (refs := t.get_dependency_ids()): raise ValueError(f"task {t.id_!r} depends on itself") if missing := (refs - validated_task_ids): raise ValueError(f"task {t.id_!r} references unknown task(s): {sorted(missing)}") diff --git a/task-sdk/src/airflow/sdk/importers/yaml_importer/schema.json b/task-sdk/src/airflow/sdk/importers/yaml_importer/schema.json index 953c61b6ccca4..67686d202ae1c 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 1b7c53ad104e6..ff40089f75070 100644 --- a/task-sdk/src/airflow/sdk/importers/zip_importer.py +++ b/task-sdk/src/airflow/sdk/importers/zip_importer.py @@ -21,6 +21,7 @@ import contextlib import logging import os +import posixpath import sys import tempfile import threading @@ -101,6 +102,18 @@ def read_bytes(self) -> bytes: self._content = z.read(self.file_path) return self._content + def sibling(self, 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 0000000000000..84b7c53511a57 --- /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__}" + + +@pytest.fixture +def mock_bundle(tmp_path): + bundle_dir = tmp_path / "bundle" + bundle_dir.mkdir(parents=True, exist_ok=True) + return SimpleNamespace(name="test_bundle", path=bundle_dir) + + +def _def(bundle, name: str, body: str) -> FilesystemDagDefinition: + path = bundle.path / name + path.write_text(textwrap.dedent(body)) + return FilesystemDagDefinition(path=path) + + +def _doc(tasks: str, *, dag_id: str = "d", extra: str = "") -> str: + return f"$schema: {HEAD_SCHEMA}\ndag_id: {dag_id}\n{extra}tasks:\n{textwrap.indent(textwrap.dedent(tasks), ' ')}" + + +def test_can_handle(): + importer = YamlDagImporter() + assert importer.can_handle("dag.yaml") + assert importer.can_handle("dag.yml") + assert not importer.can_handle("dag.py") + assert not importer.can_handle("dag.json") + + +def test_get_source_code(mock_bundle): + definition = _def(mock_bundle, "d.yaml", f"$schema: {HEAD_SCHEMA}\ndag_id: d\ntasks: []\n") + src = YamlDagImporter().get_source_code(definition) + assert src.language == "yaml" + assert "$schema" in src.source_code + + +def test_might_contain_dag(mock_bundle): + importer = YamlDagImporter() + has = _def(mock_bundle, "has.yaml", f"$schema: {HEAD_SCHEMA}\ndag_id: d\ntasks: []\n") + nope = _def(mock_bundle, "nope.yaml", "just: config\n") + # 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"} + + +@pytest.mark.parametrize( + ("value", "hour"), + [("2024-01-01", 0), ("'2024-01-01'", 0), ("2024-01-01T03:30:00", 3)], # YAML date, ISO string, datetime +) +def test_import_coerces_temporal_start_date(mock_bundle, value, hour): + 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) + + +@pytest.mark.parametrize("marker", ["{$x: a}", "{$t: 'v'}", "{$f: q.sql}"]) +def test_import_marker_on_non_template_field_errors(mock_bundle, marker): + (mock_bundle.path / "q.sql").write_text("SELECT 1") # exists, so this is not a missing-file error + 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) + + +@pytest.mark.parametrize( + "body", ["- {{id: a, uses: {op}, with: {{payload: x}}, note: 1}}", "- {{id: a, run: {{}}, note: 1}}"] +) +def test_import_operator_arg_at_task_level_errors(mock_bundle, body): + 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 + + +@pytest.mark.parametrize("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 22a69df2c753f..91b0dd33f17cd 100644 --- a/task-sdk/tests/task_sdk/importers/yaml_importer/test_parser.py +++ b/task-sdk/tests/task_sdk/importers/yaml_importer/test_parser.py @@ -531,6 +531,14 @@ def test_published_schema_accepts_null_and_empty_bodies(task): "tasks:\n - {id: t, uses: X, with: {k: {$t: 'v-{{ ds }}'}}}", id="template", ), + pytest.param( + "tasks:\n - {id: t, uses: X, with: {k: {$f: 'reports/daily.sql'}}}", + id="file", + ), + pytest.param( + "tasks:\n - {id: t, uses: X, with: {k: {$t: {$f: 'reports/daily.sql'}}}}", + id="template-file", + ), pytest.param( "templates: {base: {run: {r: 1}}}\ntasks:\n - {id: t, extends: base}", id="extends-bare", @@ -564,6 +572,8 @@ def test_valid_shapes_round_trip(body): pytest.param("{cfg: {$weird: 1}}", id="dollar-key-nested"), pytest.param("{$x: up, $xcom: up}", id="two-marker-keys"), pytest.param("{$x: {task: up, nope: 1}}", id="xcom-target-extra-key"), + pytest.param("{$f: {$t: 'q-{{ ds }}.sql'}}", id="templated-filename-reserved"), + pytest.param("{$t: {$f: {$t: x}}}", id="templated-filename-nested-reserved"), ], ) def test_invalid_shapes_rejected_by_both(value):