kaxil commented on code in PR #74084:
URL: https://github.com/apache/airflow/pull/74084#discussion_r4237568329
##########
task-sdk/pyproject.toml:
##########
@@ -53,6 +53,7 @@ dependencies = [
"asgiref>=3.11.1; python_version >= '3.14'",
"attrs>=24.2.0, !=25.2.0",
"babel>=2.17.0",
+ "cadwyn>=6.1.1",
Review Comment:
`uv.lock` wasn't refreshed for this: the task-sdk entry lists cadwyn in
neither `dependencies` nor `requires-dist`, so `uv lock --check` fails on this
branch. The `update-uv-lock` hook runs with `--frozen` and won't catch it, so
running `uv lock` and committing the result should be enough.
##########
task-sdk/src/airflow/sdk/importers/yaml_importer/parser.py:
##########
@@ -0,0 +1,101 @@
+# 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.
+
+"""Parse a YAML/JSON DAG file into validated :class:`~.models.DagDocument`
objects."""
+
+from __future__ import annotations
+
+import collections.abc
+import datetime
+from typing import TYPE_CHECKING, Any
+
+import yaml
+from pydantic import ValidationError
+
+from airflow.sdk.importers.yaml_importer import migrator
+from airflow.sdk.importers.yaml_importer.models import DagDocument
+
+if TYPE_CHECKING:
+ from collections.abc import Hashable, Iterator
+ from typing import IO
+
+ from yaml.nodes import MappingNode
+
+
+class _UniqueKeySafeLoader(yaml.SafeLoader):
+ """
+ A ``SafeLoader`` that rejects duplicate mapping keys.
+
+ PyYAML keeps the last value for a repeated key, so a second ``tasks:``
block
+ or a second ``with:`` or ``needs:`` on a task would silently drop the
first.
+ This can be a footgun for a hand-edited format, so we raise a
+ ``ConstructorError`` instead of silently shadowing.
+ """
+
+ def construct_mapping(self, node: MappingNode, deep: bool = False) ->
dict[Hashable, Any]:
+ seen: set[Hashable] = set()
+ for key_node, _ in node.value:
+ key = self.construct_object(key_node, deep=deep)
Review Comment:
This constructs every key before `super().construct_mapping` gets to
`flatten_mapping`, so YAML merge keys break. `<<` is tagged
`tag:yaml.org,2002:merge` and there's no constructor for that tag. `tasks:
[{<<: *defaults, id: a, run: {}}]` fails with `invalid YAML: could not
determine a constructor for the tag 'tag:yaml.org,2002:merge'`, but
`yaml.safe_load` handles it.
Skipping merge-tagged nodes in the scan should fix it:
```python
if key_node.tag == "tag:yaml.org,2002:merge":
continue
```
Explicit keys still get the duplicate check and can still override merged
ones. If merge keys are deliberately unsupported (use `extends` instead), I'd
raise an error that says that rather than "invalid YAML". Can we add a test for
whichever way you go?
##########
task-sdk/src/airflow/sdk/importers/yaml_importer/models.py:
##########
@@ -0,0 +1,409 @@
+# 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.
+"""
+Pydantic models for the YAML/JSON DAG format.
+
+These are the single source of truth for the format; they validate an authored
+document (rich Pydantic errors) and back the published JSON Schema. Only the
+format skeleton and value grammar only is described; operator ``with:`` args,
+Dag- and task-level attributes pass through untyped because the Python
+constructors validate those at import time.
+"""
+
+from __future__ import annotations
+
+import collections
+import itertools
+from typing import TYPE_CHECKING, Annotated, Any, ClassVar
+
+from pydantic import (
+ AliasChoices,
+ BaseModel,
+ ConfigDict,
+ Discriminator,
+ Field,
+ RootModel,
+ Tag,
+ model_validator,
+)
+
+if TYPE_CHECKING:
+ from pydantic import GetJsonSchemaHandler
+ from pydantic.json_schema import JsonSchemaValue
+ from pydantic_core import CoreSchema
+
+XCOM_KEYS = ("$x", "$xcom")
+TEMPLATE_KEYS = ("$t", "$template")
+CONST_KEY = "$const"
+
+# Keys deferred to a later edition...
+DAG_CALLBACK_KEYS = {"on_success_callback", "on_failure_callback"}
+TASK_CALLBACK_KEYS = {
+ "on_success_callback",
+ "on_failure_callback",
+ "on_retry_callback",
+ "on_execute_callback",
+ "on_skipped_callback",
+}
+TASK_ASSET_IO_KEYS = {"inlets", "outlets"}
+
+
+class XComTarget(BaseModel):
+ """The object form of an XCom reference."""
+
+ task: str
+ key: str | None = None
+
+ model_config = ConfigDict(extra="forbid")
+
+
+class _MarkerRef(BaseModel):
+ """Base for a reserved ``$``-marker object."""
+
+ model_config = ConfigDict(extra="forbid")
+
+ _marker_keys: ClassVar[tuple[str, ...]]
+
+ @classmethod
+ def __get_pydantic_json_schema__(
+ cls, core_schema: CoreSchema, handler: GetJsonSchemaHandler
+ ) -> JsonSchemaValue:
+ # Emit the marker as a single key in every accepted spelling, nothing
else.
+ schema = handler(core_schema)
+ (value_schema,) = schema.get("properties", {}).values()
+ schema["properties"] = {key: value_schema for key in cls._marker_keys}
+ schema.pop("required", None)
+ schema["oneOf"] = [{"required": [key]} for key in cls._marker_keys]
+ schema["additionalProperties"] = False
+ return schema
+
+
+class XComRef(_MarkerRef):
+ """An upstream task's XCom output."""
+
+ _marker_keys = XCOM_KEYS
+
+ target: str | XComTarget = Field(
+ serialization_alias="$x",
+ validation_alias=AliasChoices(*XCOM_KEYS),
+ )
+
+
+class TemplateRef(_MarkerRef):
+ """A Jinja template."""
+
+ _marker_keys = TEMPLATE_KEYS
+
+ source: str = Field(serialization_alias="$t",
validation_alias=AliasChoices(*TEMPLATE_KEYS))
+
+
+class ConstRef(_MarkerRef):
+ """Force a literal; the value is taken verbatim."""
+
+ _marker_keys = (CONST_KEY,)
+
+ value: Any = Field(serialization_alias="$const",
validation_alias=AliasChoices(CONST_KEY))
+
+
+def _value_discriminator(v: Any) -> str:
+ """
+ Route a raw value to a marker branch, or to ``literal``.
+
+ The prefix ``$`` is reserved to express Airflow constructs. To prevent
+ typos, any unknown ``$``-prefixed strings are routed here and rejected.
+
+ If such a key is needed, it can be wrapped in ``$const``, which would
+ avoid this discrimator (only used by :class:`_Literal`).
+ """
+ if isinstance(v, dict) and len(v) == 1:
+ (key,) = v
+ if key in XCOM_KEYS:
+ return "xcom"
+ if key in TEMPLATE_KEYS:
+ return "template"
+ if key == CONST_KEY:
+ return "const"
+ if isinstance(v, XComRef):
+ return "xcom"
+ if isinstance(v, TemplateRef):
+ return "template"
+ if isinstance(v, ConstRef):
+ return "const"
+ return "literal"
+
+
+class _Literal(RootModel):
+ """
+ A literal value.
+
+ Containers are validated recursively so nested markers resolve; scalars
pass
+ through. A ``$``-prefixed key here is a misused marker, so it is rejected.
+ """
+
+ root: dict[str, Value] | list[Value] | str | int | float | bool | None
Review Comment:
PyYAML loads an unquoted `2026-01-01` as `datetime.date`, and nothing in
this union accepts a date. So `with: {target_time: 2026-01-01}` fails with a
long list of union-branch errors. The same value as a task-level or Dag-level
extra goes through fine as a date.
I'd add `datetime.date | datetime.datetime` to the union. A single "quote
this date" error would also work. The parser already handles an unquoted date
for `$schema`.
##########
task-sdk/tests/task_sdk/importers/yaml_importer/test_parser.py:
##########
@@ -0,0 +1,433 @@
+# 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.
+"""Unit tests for the YAML DAG format parser
(airflow.sdk.importers.yaml_importer)."""
+
+from __future__ import annotations
+
+import io
+import json
+import pathlib
+import textwrap
+
+import jsonschema
+import pytest
+import yaml
+
+from airflow.sdk.importers import yaml_importer
+from airflow.sdk.importers.yaml_importer import YamlDagParseError,
parse_documents
+from airflow.sdk.importers.yaml_importer.models import (
+ ConstRef,
+ DagDocument,
+ Task,
+ TemplateRef,
+ TimetableSchedule,
+ XComRef,
+ XComTarget,
+ _Literal,
+)
+
+SCHEMA_PATH = pathlib.Path(yaml_importer.__file__).with_name("schema.json")
+
+HEAD = "2026-10-30"
+HEAD_SCHEMA = f"https://airflow.apache.org/schemas/dag/{HEAD}.json";
+
+
+def _one(body: str):
+ return next(iter(parse_documents(f"$schema: {HEAD_SCHEMA}\ndag_id:
d\n{textwrap.dedent(body)}")))
+
+
+def _task(task_yaml: str):
+ return _one("tasks:\n" + textwrap.indent(textwrap.dedent(task_yaml), "
")).tasks[0]
+
+
+def test_literal_by_default():
+ t = _task("- {id: t, run: {region: 'eu', n: 5, flag: true}}")
+ assert isinstance(t.run["region"], _Literal)
+ assert t.run["region"].root == "eu"
+ assert t.run["n"].root == 5
+ assert t.run["flag"].root is True
+
+
+def test_xcom_short_long_and_keyed():
+ doc = _one(
+ "tasks:\n - {id: up, run: {}}\n - {id: t, run: {a: {$x: up}, b:
{$xcom: up}, c: {$x: {task: up, key: rows}}}}"
+ )
+ t = {x.id_: x for x in doc.tasks}["t"]
+ assert isinstance(t.run["a"], XComRef)
+ assert t.run["a"].target == "up"
+ assert isinstance(t.run["b"], XComRef)
+ assert t.run["b"].target == "up"
+ assert isinstance(t.run["c"].target, XComTarget)
+ assert t.run["c"].target.key == "rows"
+
+
+def test_template_marker():
+ t = _task("- {id: t, uses: X, with: {k: {$t: 'v-{{ ds }}'}, k2:
{$template: 'x'}}}")
+ assert isinstance(t.with_["k"], TemplateRef)
+ assert "{{ ds }}" in t.with_["k"].source
+ assert isinstance(t.with_["k2"], TemplateRef)
+
+
+def test_const_is_verbatim_not_recursed():
+ t = _task("- {id: t, run: {a: {$const: {$x: not-a-ref}}}}")
+ assert isinstance(t.run["a"], ConstRef)
+ assert t.run["a"].value == {"$x": "not-a-ref"} # inner marker NOT
interpreted
+
+
[email protected](
+ "value",
+ [
+ pytest.param("{$x: up, extra: 1}", id="marker-key-plus-extra"),
+ pytest.param("{$xcomm: up}", id="unknown-marker-typo"),
+ pytest.param("{$tempalte: x}", id="unknown-template-typo"), #
codespell:ignore tempalte
+ pytest.param("{cfg: {$weird: 1}}", id="dollar-key-nested-in-literal"),
+ ],
+)
+def test_reserved_dollar_keys_rejected(value):
+ # A '$'-prefixed key that is not a recognised sole-key marker is rejected,
not silently literal.
+ with pytest.raises(YamlDagParseError, match="reserved"):
+ _task(f"- {{id: t, run: {{a: {value}}}}}")
+
+
+def test_nested_marker_inside_literal_resolved():
+ t = _task("- {id: t, uses: X, with: {cfg: {url: {$t: '{{ ds }}'}, n: 1}}}")
+ inner = t.with_["cfg"].root
+ assert isinstance(inner["url"], TemplateRef)
+ assert inner["n"].root == 1
+
+
+def test_operator_discrimination():
+ op = _task("- {id: t, uses: a.b.C, with: {x: 1}}")
+ assert isinstance(op, Task)
+ assert op.uses == "a.b.C"
+ assert op.run is None
+
+
+def test_code_discrimination():
+ code = _task("- {id: t, run: {}}")
+ assert isinstance(code, Task)
+ assert code.run == {}
+ assert code.uses is None
+
+
+def test_no_body_is_error():
+ with pytest.raises(YamlDagParseError, match="exactly one"):
+ _task("- {id: t}")
+
+
+def test_both_bodies_is_error():
+ with pytest.raises(YamlDagParseError, match="exactly one"):
+ _task("- {id: t, uses: X, run: {}}")
+
+
[email protected](
+ "task",
+ [
+ pytest.param("{id: t, uses: null}", id="uses-null"),
+ pytest.param('{id: t, uses: ""}', id="uses-empty"),
+ pytest.param("{id: t, run: null}", id="run-null"),
+ ],
+)
+def test_empty_body_is_error(task):
+ # A present-but-empty body (null / "") is not a usable body, even though
the key exists.
+ with pytest.raises(YamlDagParseError, match="must not be empty"):
+ _task(f"- {task}")
+
+
+def test_python_field_names_not_accepted_as_aliases():
+ with pytest.raises(YamlDagParseError):
+ _one("tasks:\n - {id_: t, run: {}}") # 'id_' is not 'id' -> id
missing
+ t = _task("- {id: t, uses: X, with_: {k: 1}}")
+ assert t.with_ == {} # 'with_' did not populate the 'with' field
+ assert t.model_extra.get("with_") == {"k": 1} # it is just a pass-through
extra
+
+
+def test_default_args_rejected():
+ with pytest.raises(YamlDagParseError, match="default_args"):
+ _one("default_args: {retries: 1}\ntasks: []")
+
+
+def test_callbacks_rejected():
+ with pytest.raises(YamlDagParseError, match="callback"):
+ _one("on_failure_callback: cb\ntasks: []")
+ with pytest.raises(YamlDagParseError, match="callback"):
+ _task("- {id: t, run: {}, on_success_callback: cb}")
+
+
+def test_inlets_outlets_rejected():
+ with pytest.raises(YamlDagParseError, match="not yet implemented"):
+ _task("- {id: t, run: {}, outlets: [a]}")
+
+
+def test_extends_merges_shallow_task_wins():
+ doc = _one(
+ """
+ templates:
+ retryable: {retries: 2, retry_delay: '5m'}
+ load:
+ extends: [retryable]
+ uses: a.b.S3ToRedshiftOperator
+ with: {schema: public, conn_id: warehouse}
+ tasks:
+ - id: t
+ extends: [load]
+ with: {table: sales, conn_id: override}
+ """
+ )
+ t = doc.tasks[0]
+ assert isinstance(t, Task)
+ assert t.uses.endswith("S3ToRedshiftOperator")
+ assert t.model_extra["retries"] == 2 # from retryable
+ assert isinstance(t.with_["schema"], _Literal) # inherited
+ assert t.with_["table"].root == "sales" # own
+ assert t.with_["conn_id"].root == "override" # one-level with-merge, task
wins
+
+
+def test_unknown_template_errors():
+ with pytest.raises(YamlDagParseError, match="unknown template"):
+ _one("tasks:\n - {id: t, extends: [ghost], run: {}}")
+
+
+def test_template_cycle_errors():
+ with pytest.raises(YamlDagParseError, match="cycle"):
+ _one("templates: {a: {extends: [b]}, b: {extends: [a]}}\ntasks: [{id:
t, extends: [a], run: {}}]")
+
+
[email protected](
+ "body",
+ [
+ pytest.param("tasks:\n - {id: t, run: {}, extends: [[base]]}",
id="task-extends-nested-list"),
+ pytest.param("tasks:\n - {id: t, run: {}, extends: [{base: 1}]}",
id="task-extends-dict-element"),
+ pytest.param("tasks:\n - {id: t, run: {}, extends: 1}",
id="task-extends-non-string"),
+ ],
+)
+def test_task_extends_must_be_a_name_or_list_of_names(body):
+ # A non-string entry (or a non-string/non-list value) is a parse error,
not a TypeError.
+ with pytest.raises(YamlDagParseError, match="extends"):
+ _one(body)
+
+
+def test_template_extends_must_be_a_name_or_list_of_names():
+ with pytest.raises(YamlDagParseError, match="extends"):
+ _one(
+ "templates: {base: {retries: 1}, child: {extends:
[[base]]}}\ntasks: [{id: t, extends: [child], run: {}}]"
+ )
+
+
+def test_extends_accepts_a_bare_string():
+ # `extends: base` is shorthand for `extends: [base]`, at both task and
template level.
+ doc = _one(
+ "templates:\n gp: {retries: 3}\n base: {extends: gp}\ntasks:\n -
{id: t, extends: base, uses: X}"
+ )
+ assert doc.tasks[0].model_extra["retries"] == 3 # merged gp -> base ->
task
+
+
+def test_unknown_needs_rejected():
+ with pytest.raises(YamlDagParseError, match="unknown task"):
+ _one("tasks:\n - {id: t, run: {}, needs: [ghost]}")
+
+
[email protected](
+ "task",
+ [
+ pytest.param("{id: t, run: {e: {$x: ghost}}}", id="run-top-level"),
+ pytest.param("{id: t, run: {e: [{$x: ghost}]}}", id="run-in-list"),
+ pytest.param("{id: t, run: {e: {nested: {$x: ghost}}}}",
id="run-nested-dict"),
+ pytest.param("{id: t, uses: X, with: {a: {$x: ghost}}}",
id="with-on-operator"),
+ ],
+)
+def test_unknown_xcom_target_rejected(task):
+ # XCom refs must be found at any depth inside with/run (the list and
nested-dict branches too).
+ with pytest.raises(YamlDagParseError, match="unknown task"):
+ _one(f"tasks:\n - {task}")
+
+
+def test_duplicate_task_ids_rejected():
+ with pytest.raises(YamlDagParseError, match="duplicate task id"):
+ _one("tasks:\n - {id: a, run: {}}\n - {id: a, uses: x.Y}")
+
+
[email protected](
+ "body",
+ [
+ pytest.param("tasks:\n - {id: a, run: {}, needs: [a]}",
id="self-needs"),
+ pytest.param("tasks:\n - {id: a, run: {x: {$x: a}}}", id="self-xcom"),
+ ],
+)
+def test_self_dependency_rejected(body):
+ with pytest.raises(YamlDagParseError, match="depends on itself"):
+ _one(body)
+
+
+def test_schedule_forms():
+ assert _one("schedule: '0 3 * * *'\ntasks: []").schedule == "0 3 * * *"
+ assert _one("schedule: null\ntasks: []").schedule is None
+ tt = _one("schedule: {uses: a.b.MyTimetable, with: {n: 1}}\ntasks:
[]").schedule
+ assert isinstance(tt, TimetableSchedule)
+ assert tt.uses.endswith("MyTimetable")
+
+
+def test_extra_top_level_keys_pass_through():
+ # Unknown top-level keys are Dag-level arguments: accepted and surfaced
via pydantic extras.
+ doc = _one("catchup: false\nmax_active_runs: 3\ntags: [etl]\ntasks: []")
+ assert doc.model_extra == {"catchup": False, "max_active_runs": 3, "tags":
["etl"]}
+
+
+def test_schema_required():
+ with pytest.raises(YamlDagParseError, match=r"\$schema"):
+ list(parse_documents("dag_id: d\ntasks: []"))
+
+
+def test_unquoted_date_schema_is_coerced_to_url():
+ # An unquoted date loads as datetime.date; accept it and normalise to the
canonical schema URL.
+ doc = next(iter(parse_documents("$schema: 2026-10-30\ndag_id: d\ntasks:
[]")))
+ assert doc.schema_ == HEAD_SCHEMA
+
+
[email protected](
+ "schema_value",
+ [
+ pytest.param("null", id="null"),
+ pytest.param('""', id="empty-string"),
+ pytest.param("false", id="false"),
+ pytest.param("123", id="int"),
+ pytest.param("[a, b]", id="list"),
+ ],
+)
+def test_schema_must_be_a_non_empty_string_or_date(schema_value):
+ # A present-but-falsy or wrong-typed $schema is reported as such, not as
"missing".
+ with pytest.raises(YamlDagParseError, match="must be a non-empty string or
a date"):
+ list(parse_documents(f"$schema: {schema_value}\ndag_id: d\ntasks: []"))
+
+
[email protected](
+ "body",
+ [
+ pytest.param("tasks: []\ntasks: []", id="dup-tasks-top-level"),
+ pytest.param("tasks:\n - {id: a, uses: X, with: {x: 1}, with: {y:
2}}", id="dup-with-on-task"),
+ pytest.param("tasks:\n - {id: a, run: {}, needs: [b], needs: [c]}",
id="dup-needs-on-task"),
+ ],
+)
+def test_duplicate_keys_rejected(body):
+ # PyYAML keeps the last value for a repeated key; a hand-edited format
should reject it instead
+ # of silently dropping the first. The loader raises a ConstructorError (a
YAMLError).
+ with pytest.raises(YamlDagParseError, match="duplicate key"):
+ _one(body)
+
+
+def test_unknown_schema_version_is_rejected():
+ # Strict exact match (mirroring the supervisor migrator): any version that
is not published --
+ # newer, ancient, or unreadable -- is rejected rather than guessed.
+ base = "https://airflow.apache.org/schemas/dag/{}.json";
+ for url in (base.format("2099-01-01"), base.format("2000-01-01"),
"nonsense"):
+ with pytest.raises(YamlDagParseError, match="not valid"):
+ list(parse_documents(f"$schema: {url}\ndag_id: d\ntasks: []"))
+
+
+def test_accepts_a_file_like_stream():
+ stream = io.StringIO(f"$schema: {HEAD_SCHEMA}\ndag_id: d\ntasks: [{{id: t,
run: {{}}}}]")
+ docs = list(parse_documents(stream))
+ assert [d.dag_id for d in docs] == ["d"]
+
+
+def test_parse_is_lazy_errors_surface_on_iteration():
+ # Building the iterator does not parse; the error only surfaces when
consumed.
+ gen = parse_documents("dag_id: d\ntasks: []") # missing $schema
+ with pytest.raises(YamlDagParseError, match=r"\$schema"):
+ next(iter(gen))
+
+
+def test_multiple_documents():
+ docs = parse_documents(
+ f"$schema: {HEAD_SCHEMA}\ndag_id: a\ntasks: []\n---\n$schema:
{HEAD_SCHEMA}\ndag_id: b\ntasks: []\n"
+ )
+ assert [d.dag_id for d in docs] == ["a", "b"]
+
+
+def test_non_mapping_document_rejected():
+ with pytest.raises(YamlDagParseError, match="must be a mapping"):
+ list(parse_documents("- a\n- b"))
+
+
+def test_invalid_yaml_rejected():
+ with pytest.raises(YamlDagParseError, match="invalid YAML"):
+ list(parse_documents("key: [unclosed"))
+
+
+def test_error_label_points_at_offending_document():
+ # A failure in a later document of a multi-doc stream is labelled with its
position.
+ good = f"$schema: {HEAD_SCHEMA}\ndag_id: a\ntasks: []"
+ bad = "dag_id: b\ntasks: []" # missing $schema -> fails, and it is doc
index 1
+ with pytest.raises(YamlDagParseError, match=r"\[doc 1\]"):
+ list(parse_documents(f"{good}\n---\n{bad}"))
+
+
+def test_model_json_schema_describes_markers():
+ # The models are the schema source of truth (the prek dump script
decorates + snapshots
+ # them). Assert the generated schema shape here; the snapshot itself is
enforced by the
+ # generate-yaml-importer-schema-snapshot prek hook.
+ sch = DagDocument.model_json_schema(by_alias=True)
+ assert set(sch["properties"]) >= {"$schema", "dag_id", "schedule",
"tasks", "templates"}
+ assert {"XComRef", "TemplateRef", "ConstRef", "Task"} <= set(sch["$defs"])
+
+
+def test_parses_example_fixture():
+ fixture = pathlib.Path(__file__).parent / "example_dag.yaml"
+ with fixture.open(encoding="utf-8") as fh:
+ docs = list(parse_documents(fh, source=str(fixture)))
+ assert [d.dag_id for d in docs] == ["retail_daily_sales",
"events_pipeline"]
+
+ d1 = {t.id_: t for t in docs[0].tasks}
+ load = d1["load_sales"]
+ assert isinstance(load, Task)
+ assert load.uses.endswith("S3ToRedshiftOperator")
+ assert load.model_extra["retries"] == 2 # via redshift_load -> retryable
+ assert isinstance(load.with_["schema"], TemplateRef) # inherited $t
+ assert isinstance(load.with_["s3_key"], TemplateRef) # own $t
+ assert load.needs == ["wait_for_export"]
+ assert docs[0].model_extra["catchup"] is False
+
+ d2 = {t.id_: t for t in docs[1].tasks}
+ assert isinstance(d2["extract"], Task)
+ assert d2["extract"].run == {}
+ assert d2["extract"].queue == "extract-workers"
+ assert isinstance(d2["transform"].run["events"], XComRef) # {$x: extract}
edge
+ assert d2["load"].needs == ["transform"]
+
+
+def test_example_fixture_conforms_to_published_schema():
Review Comment:
The fixture only uses `$x`, `$t` and list-form `extends`, so the schema-side
fixes aren't pinned. Emitting only the short marker alias, dropping the marker
`oneOf`, or dropping the bare-string `extends` form each regenerates a
schema.json that rejects documents the parser accepts, and the suite stays
green. A parametrized test that runs each valid shape (`$xcom`, `$template`,
keyed `{$x: {task, key}}`, `$const` around a `$` key, `extends: base`, an
empty-dict value) through both `parse_documents` and the validator would catch
that.
##########
task-sdk/src/airflow/sdk/importers/yaml_importer/migrator.py:
##########
@@ -0,0 +1,163 @@
+# 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.
+"""
+Schema version resolution and forward migration to the head shape.
+
+A document with an older ``$schema`` is migrated to the head shape by walking
+the Cadwyn version bundle and applying each ``VersionChange``'s forward request
+converters (`@convert_request_to_next_version_for(DagDocument)`), oldest to
+newest.
+
+Only whole-document request converters keyed on `DagDocument` are applied; a
+bundle carrying any other cadwyn instruction is rejected at construction rather
+than migrating a document with half its changes applied. (See documentation on
+:meth:`DagDocumentMigrator._reject_unsupported_instructions` for rationale.)
+"""
+
+from __future__ import annotations
+
+import copy
+import functools
+import re
+from typing import TYPE_CHECKING, Any
+
+import attrs
+
+from airflow.sdk.importers.yaml_importer.models import DagDocument
+from airflow.sdk.importers.yaml_importer.versions import get_bundle
+
+if TYPE_CHECKING:
+ from cadwyn import VersionBundle
+
+_SCHEMA_URL = "https://airflow.apache.org/schemas/dag/{date}.json";
+_VERSION_RE = re.compile(r"\d{4}-\d{2}-\d{2}")
+
+
+def schema_url(date: str) -> str:
+ """Build the canonical ``$schema`` URL for a published version *date*."""
+ return _SCHEMA_URL.format(date=date)
+
+
+def version_from_schema(schema: str) -> str | None:
+ """Read the version (``YYYY-MM-DD``) token out of a ``$schema`` URL,
ignoring the host."""
+ match = _VERSION_RE.search(schema or "")
+ return match.group(0) if match else None
+
+
[email protected]
+class _RequestInfo:
+ """
+ Duck-type stand-in for Cadwyn's ``RequestInfo``.
+
+ ``cadwyn.structure.data._AlterDataInstruction.__call__`` only reads
+ and writes ``info.body``; the by-schema transformers we drive never
+ touch FastAPI's Request/Response. Passing this minimal object lets
+ us run cadwyn's migrations from a pure in-process code path with no
+ HTTP stack.
+ """
+
+ body: dict[str, Any]
+
+
+def _calculate_version_values(migrator: DagDocumentMigrator) -> frozenset[str]:
+ return frozenset(v.value for v in migrator._bundle.versions)
+
+
[email protected]
+class DagDocumentMigrator:
+ """YAML Dag document migrator; pins each document to its ``$schema``
version."""
+
+ _bundle: VersionBundle
+ _version_values: frozenset[str] = attrs.field(
+ init=False,
+ default=attrs.Factory(_calculate_version_values, takes_self=True),
+ )
+
+ def __attrs_post_init__(self) -> None:
+ self._reject_unsupported_instructions()
+
+ def _reject_unsupported_instructions(self) -> None:
+ """
+ Fail if the bundle carries unsupported migrations.
+
+ Only forward request converters keyed on :class:`DagDocument` may run.
+ This is due to difficulties posed by things like templates and XCom,
+ which can present in nested locations inside task arguments.
Technically
+ we could just ban model-level migration only for them, but that's
+ considered unnecessarily complicated since you can already do all
+ per-model-level changes on :class:`DagDocument`. We can relax this
+ restriction in the future if the ergonomics prove to be suboptimal.
+ """
+ unsupported = (
+ "alter_endpoint_instructions",
+ "alter_enum_instructions",
+ "alter_request_by_path_instructions",
+ "alter_response_by_path_instructions",
+ "alter_response_by_schema_instructions",
+ "alter_schema_instructions",
+ )
+ for version in self._bundle.versions:
+ for change in version.changes:
+ name = type(change).__name__
Review Comment:
`Version.changes` holds the `VersionChange` classes, not instances, so
`type(change).__name__` is `"type"` and the error reads `type: unsupported
cadwyn instructions [...]`. `change.__name__` would name the offending change.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
