diff --git a/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md b/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md index 6f9a8455fec15..abbe3d6545e9f 100644 --- a/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md +++ b/airflow-core/adr/lang-sdk/0011-mixed-language-dag-processing.md @@ -155,15 +155,21 @@ DagFileProcessorProcess(etl.py) ← manager spawn │ └─────────────────────────────────────────────────────────────────────┘ │ │ ┌─────────────────────────────────────────────────────────────────────┐ - │ │ Step 5: Compare per dag_id against the Dag just serialized — │ - │ │ union task_handlers[dag_id] across every process first │ + │ │ Step 5: Check each stub task against the one answer for its │ + │ │ coordinator and its Dag │ │ │ │ │ │ Python Dag "etl" (stub tasks) TaskHandlerDeclaration │ │ │ ────────────────────────────── ────────────────────────────── │ - │ │ task_id ↔ task_id (sets must match) │ - │ │ arg_bindings[*].name ↔ params[*].name (in order) │ + │ │ (dag_id, task_id) ↔ its declaration in that answer │ + │ │ none → an error │ + │ │ a handler with no stub task → not an error │ + │ │ arg_bindings[*] ↔ params[*] (per binding) │ + │ │ by position, or by folded or exact name │ + │ │ defaulted arguments dropped if that makes the count match │ + │ │ unmatched by name → a warning, not an error │ │ │ arg_bindings[*].value_schema ↔ params[*].value_schema │ - │ │ compared only where neither side is null │ + │ │ top-level JSON types, where both sides have a schema │ + │ │ a mapped stub task, or params=None → the handler's presence only │ │ │ │ │ │ On mismatch → import_errors["etl.py"] │ │ └─────────────────────────────────────────────────────────────────────┘ @@ -196,11 +202,13 @@ There is no fourth row. A `TaskHandlerRef` has no Dag, so no `DagImporter` — a - No importer knows about coordinators. `PythonDagImporter` is unchanged by this ADR; the stub-to-handler comparison sits in `_parse_file`, above every importer. - A mixed-language `dag_id` never appears in Dag processing results. No `Dag` registration exists for a `dag_id` a Python file already owns, so everything downstream sees exactly one record per `dag_id`, with no flag to interpret. -- Stub/implementation mismatches — missing handler, extra handler, parameter name or order, incompatible schema — surface as import errors against the Python file at parse time, - alongside the errors the parse already reports. An unannotated stub argument is checked by name and position only. +- Stub/implementation mismatches, such as a missing handler, a `positional` argument count that does not match or an incompatible schema, + surface as import errors against the Python file at parse time, alongside the errors the parse already reports. + A `named` argument or parameter that matches nothing is only logged as a warning, because the runtime runs the task anyway. + An unannotated stub argument is checked only for how it binds, by position or by name. - The Python Dag and Lang-SDK artifact can live in different DagBundles. -- A single Dag can have stubs targeting different queues, some Java, some Go. Each resolves to its own coordinator instance, and validation unions their declarations per `dag_id` - before comparing task ids. +- A single Dag can have stubs targeting different queues, some Java, some Go. Each resolves to its own coordinator instance, + and each stub task is checked only against the answer for its own coordinator and Dag. - Validating a file costs one extra process per (coordinator, artifact) pair its stubs resolve to — one for the common case of a file whose stubs all target a single runtime, and none at all for a file with no stub tasks. - Mixed-language is Python-primary only. Lang-SDK runtimes cannot define stub operators; a native Dag cannot delegate tasks to Python. @@ -233,12 +241,15 @@ with one `bundle.serve()`. One artifact can hold both, so neither the file nor t ### Appendix B — What is compared -Handler declarations from every process spawned in Step 4 are unioned per `dag_id` before comparison, since one Dag's stubs can target several queues. +Each stub task is compared only against the answer for its own coordinator and Dag, since one Dag's stubs can target several queues. +An answer covers every Dag its artifact registers handlers for; only the parsed file's stub tasks are looked up in it. -- `task_id` sets must match exactly. A missing or extra handler is an error. -- `arg_bindings[*].name` against `params[*].name`, in order — both sides bind positionally. -- `arg_bindings[*].value_schema` against `params[*].value_schema`, compared only where neither side is null. An unannotated `@task.stub` parameter produces `null` today, so a - strict comparison would make every untyped stub argument a parse error. +- Every stub task needs its handler in that answer, the declaration of its `(dag_id, task_id)`. A stub task with no such declaration is an error. + A handler with no stub task is not an error: one artifact serves many Dag files, and a handler can outlive its stub task. +- `arg_bindings` are checked against `params` the way the declaration's `binding` binds them, and `arg_bindings[*].value_schema` against + `params[*].value_schema` where both sides give one; [ADR-0012](0012-lang-sdk-parse-protocol.md) Appendix B states those rules. + An argument count or value type a parameter does not accept is an error; under `named`, an argument or parameter that matches nothing is only a warning. + A mapped stub task, whose arguments are not captured at parse time, and a declaration whose `params` is `None` are checked only for the handler's presence. -Any mismatch is reported against the Python file, which is the definition the author can act on, and travels back on `DagFileParsingResult.import_errors` with everything else the +Every error is reported against the Python file, which is the definition the author can act on, and travels back on `DagFileParsingResult.import_errors` with everything else the parse found. diff --git a/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md b/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md index 7a6e91e1567d7..da7e0ed14af88 100644 --- a/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md +++ b/airflow-core/adr/lang-sdk/0012-lang-sdk-parse-protocol.md @@ -109,6 +109,7 @@ class TaskHandlerParam: The request names no Dags. The runtime answers with every task handler the artifact registers, keyed by `dag_id`, and with `{}` when it registers none. The answer must depend only on the artifact, never on the request, so one answer serves every Dag whose stub tasks resolve to that artifact. Its keys need not match the parsed file's Dags and can include Dags of other files. +Each stub task is checked only against the answer for its own coordinator and Dag ([ADR-0011](0011-mixed-language-dag-processing.md)). `value_schema` reuses the `ArgValueSchema` definition `arg_bindings` already carries ([ADR-0007](0007-taskflow-across-language-boundary.md)), so both sides of a comparison are the same type. Two properties matter to validation: the field is nullable on both sides, and each declaration names its binding mode. Appendix B says what that forces. diff --git a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md index d043ea927941b..7351f163976d3 100644 --- a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md +++ b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md @@ -394,11 +394,7 @@ DagFileProcessorProcess(etl.py) [no DB — client con │ ◀─TaskHandlerParsingResult(task_handlers={"etl": […]})── runtime │ Get* from the runtime is relayed up ToManager unchanged │ - ├─7─ VALIDATE per dag_id, unioned across coordinators - │ task_id sets must match exactly - │ arg_bindings[*].name ↔ handler_params[*].name, in order - │ arg_bindings[*].schema ↔ handler_params[*].value_schema, - │ compared only where neither is null + ├─7─ VALIDATE each stub task against its coordinator's answer, per ADR-0011 │ two candidates claiming one (dag_id, task_id) → import error │ naming both paths │ diff --git a/airflow-core/src/airflow/dag_processing/task_handler_validation.py b/airflow-core/src/airflow/dag_processing/task_handler_validation.py new file mode 100644 index 0000000000000..4903f0405f5f9 --- /dev/null +++ b/airflow-core/src/airflow/dag_processing/task_handler_validation.py @@ -0,0 +1,404 @@ +# 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. +"""Check a Dag file's stub tasks against the task handlers that a Lang-SDK artifact registers.""" + +from __future__ import annotations + +import json +from collections import defaultdict +from typing import TYPE_CHECKING + +import attrs + +from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import ( + LiteralArgBinding, + get_arg_bindings_adapter, +) +from airflow.sdk.definitions._internal.abstractoperator import DEFAULT_QUEUE # noqa: SDK001 +from airflow.serialization.enums import Encoding + +if TYPE_CHECKING: + from collections.abc import Iterable, Mapping, Sequence + + from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import ArgValueSchema, TaskArgBinding + from airflow.dag_processing.processor import TaskHandlerDeclaration, TaskHandlerParam + from airflow.sdk.definitions.dag import DAG # noqa: SDK001 + from airflow.serialization.serialized_objects import LazyDeserializedDAG + +# Every top-level JSON type, in the order messages list them. +_JSON_TYPES = ("string", "integer", "number", "boolean", "array", "object", "null") + + +@attrs.frozen(kw_only=True) +class StubTask: + """A ``@task.stub`` task of a Dag that serialized, as the worker that runs it will see it.""" + + dag_id: str + task_id: str + queue: str + relative_fileloc: str + """The Dag's file, which keys its import error.""" + + arg_bindings: list[TaskArgBinding] + """The arguments of the call; ``[]`` for an argless call.""" + + is_mapped: bool + """A mapped stub task has no arguments captured at parse time, so only its handler's presence is checked.""" + + +@attrs.frozen(kw_only=True) +class TaskHandlerAnswer: + """The task handlers a Lang-SDK artifact registers, as its runtime answered a probe.""" + + bundle_name: str + """The Dag bundle that holds the artifact.""" + + rel_path: str + """The artifact's path in that Dag bundle.""" + + task_handlers: dict[str, list[TaskHandlerDeclaration]] + """Every task handler the artifact registers, by Dag id.""" + + +@attrs.frozen(kw_only=True) +class ArgumentCheck: + """How a stub task's arguments bind to one task handler's parameters.""" + + errors: list[str] + """Each makes the Dag file fail to import.""" + + passed_not_declared: list[str] + """Under ``named`` binding, passed arguments no parameter takes; logged, never an error.""" + + declared_not_passed: list[str] + """Under ``named`` binding, parameters no argument fills; logged, never an error.""" + + +@attrs.frozen(kw_only=True) +class TaskHandlerProblem: + """One line of a Dag file's import error.""" + + relative_fileloc: str + message: str + dag_id: str | None = None + """``None`` for a problem of no single stub task, which is listed first.""" + + task_id: str | None = None + + +@attrs.frozen(kw_only=True) +class StubTaskWarning: + """A name mismatch between a stub task and its ``named`` task handler, for the parse log.""" + + dag_id: str + task_id: str + artifact_bundle_name: str + artifact_rel_path: str + passed_not_declared: list[str] + declared_not_passed: list[str] + + +@attrs.frozen(kw_only=True) +class TaskHandlerMatch: + """The outcome of checking stub tasks against the task handlers of one artifact.""" + + problems: list[TaskHandlerProblem] + warnings: list[StubTaskWarning] + + +def collect_stub_tasks(dags: Iterable[DAG], serialized_dags: Iterable[LazyDeserializedDAG]) -> list[StubTask]: + """ + Return the stub tasks of every Dag in *serialized_dags*. + + The arguments are read from the serialized Dag as JSON, so they are the ones the worker gets: a tuple + becomes a list, and a mapping key becomes a string. A stub task with no queue of its own runs on the + default queue. + """ + dags_by_id = {dag.dag_id: dag for dag in dags} + stub_tasks: list[StubTask] = [] + for serialized_dag in serialized_dags: + dag = dags_by_id[serialized_dag.dag_id] + encoded_tasks = { + task[Encoding.VAR]["task_id"]: task[Encoding.VAR] for task in serialized_dag.data["dag"]["tasks"] + } + stub_tasks.extend( + StubTask( + dag_id=dag.dag_id, + task_id=task.task_id, + queue=DEFAULT_QUEUE if task.queue is None else task.queue, + relative_fileloc=dag.relative_fileloc or dag.fileloc, + arg_bindings=get_arg_bindings_adapter().validate_json( + json.dumps(encoded_tasks[task.task_id].get("_arg_bindings") or []) + ), + is_mapped=task.is_mapped, + ) + for task in dag.task_dict.values() + if task.is_stub + ) + return stub_tasks + + +def check_task_handler_arguments( + arg_bindings: Sequence[TaskArgBinding], declaration: TaskHandlerDeclaration +) -> ArgumentCheck: + """ + Check how *arg_bindings* bind to the parameters of *declaration*, as its runtime binds them. + + A defaulted argument, one filled from the stub signature's default, is type-checked when it binds and is + never reported as unmatched. + + - ``positional``: the count matches when all arguments, or those left after dropping the defaulted + ones, number the parameters; any other count is an error. Each argument must have a value type the + parameter at its position accepts. + - ``named``: a parameter takes the argument of its exact name, else, unless ``exact_name`` is set, the one + whose name folds to its own (lower case, underscores removed); a folded name two arguments share + matches neither. Each bound argument must have a value type its parameter accepts. A passed argument + no parameter takes and a parameter no argument fills are reported for the parse log. They are not + when the argument may be the whole value: no parameter matched, exactly one argument was passed, the + handler has parameters and none of them has ``exact_name`` set, and the argument may be an object. + - ``params`` is ``None``: nothing is checked. + """ + params = declaration.params + if params is None: + return ArgumentCheck(errors=[], passed_not_declared=[], declared_not_passed=[]) + passed = [arg for arg in arg_bindings if not _is_defaulted(arg)] + if declaration.binding == "positional": + return _check_positional_arguments(arg_bindings, passed, params) + return _check_named_arguments(arg_bindings, passed, params) + + +def check_value_schema( + stub_schema: ArgValueSchema | None, handler_schema: ArgValueSchema | None +) -> str | None: + """ + Return why a value the stub's schema allows may be one the handler's schema rejects, or ``None``. + + Only the top-level JSON types are compared, as the Go runtime compares them. A stub type of ``null`` + needs a handler type of ``null``, and, where the stub has other types, at least one of them must be a + handler type, with ``integer`` accepted by ``number``. The types come from ``type``, the union over + ``anyOf`` or ``oneOf``, or the values of ``const`` or ``enum``. A side with no schema, or a schema of + any other shape such as ``$ref`` or ``allOf``, is not compared. ``format``, ranges and nested items are + never compared. + """ + if stub_schema is None or handler_schema is None: + return None + stub_types = _get_json_types(stub_schema) + handler_types = _get_json_types(handler_schema) + if stub_types is None or handler_types is None: + return None + accepted = (handler_types | {"integer"}) if "number" in handler_types else handler_types + non_null_types = stub_types - {"null"} + if ("null" not in stub_types or "null" in handler_types) and ( + not non_null_types or non_null_types & accepted + ): + return None + return f"{_join_types(stub_types)}, the task handler takes {_join_types(handler_types)}" + + +def match_task_handlers(stub_tasks: Sequence[StubTask], answer: TaskHandlerAnswer) -> TaskHandlerMatch: + """ + Check each of *stub_tasks* against the task handler that *answer* registers for it. + + *stub_tasks* are the stub tasks of one Dag file that route to the coordinator *answer* came from; a + stub task of another coordinator would be reported as having no handler. + + A stub task that *answer* registers no handler for is a problem. A handler with no stub task is not. + The arguments of a stub task are checked against its handler unless the stub task is mapped or the + handler's ``params`` are ``None``. + """ + declarations = { + (dag_id, declaration.task_id): declaration + for dag_id, dag_declarations in answer.task_handlers.items() + for declaration in dag_declarations + } + artifact = f"{answer.rel_path!r} in Dag bundle {answer.bundle_name!r}" + problems: list[TaskHandlerProblem] = [] + warnings: list[StubTaskWarning] = [] + for stub in stub_tasks: + prefix = f"Dag {stub.dag_id!r}, task {stub.task_id!r}" + declaration = declarations.get((stub.dag_id, stub.task_id)) + if declaration is None: + problems.append(_make_problem(stub, f"{prefix}: {artifact} registers no task handler for it")) + continue + if stub.is_mapped: + continue + check = check_task_handler_arguments(stub.arg_bindings, declaration) + problems.extend(_make_problem(stub, f"{prefix} ({artifact}): {error}") for error in check.errors) + if check.passed_not_declared or check.declared_not_passed: + warnings.append( + StubTaskWarning( + dag_id=stub.dag_id, + task_id=stub.task_id, + artifact_bundle_name=answer.bundle_name, + artifact_rel_path=answer.rel_path, + passed_not_declared=check.passed_not_declared, + declared_not_passed=check.declared_not_passed, + ) + ) + return TaskHandlerMatch(problems=problems, warnings=warnings) + + +def format_import_errors(problems: Iterable[TaskHandlerProblem]) -> dict[str, str]: + """ + Build one import error per Dag file from its *problems*. + + Problems of no single stub task come first, in the order given, then the others by Dag and task id. + """ + by_file: defaultdict[str, list[TaskHandlerProblem]] = defaultdict(list) + for problem in problems: + by_file[problem.relative_fileloc].append(problem) + import_errors: dict[str, str] = {} + for relative_fileloc in sorted(by_file): + ordered = sorted( + by_file[relative_fileloc], + key=lambda problem: (problem.dag_id is not None, problem.dag_id or "", problem.task_id or ""), + ) + lines = [f"Stub tasks in {relative_fileloc} do not match their task handlers:"] + lines.extend(f"- {problem.message}" for problem in ordered) + import_errors[relative_fileloc] = "\n".join(lines) + return import_errors + + +def _check_positional_arguments( + arg_bindings: Sequence[TaskArgBinding], + passed: Sequence[TaskArgBinding], + params: Sequence[TaskHandlerParam], +) -> ArgumentCheck: + bound = arg_bindings if len(arg_bindings) == len(params) else passed + if len(bound) == len(params): + errors = _check_values(zip(bound, params)) + else: + count = f"{len(arg_bindings)} argument{'' if len(arg_bindings) == 1 else 's'}" + if len(passed) != len(arg_bindings): + count = f"{count} ({len(passed)} without defaults)" + errors = [f"passes {count}, the task handler takes {len(params)}"] + return ArgumentCheck(errors=errors, passed_not_declared=[], declared_not_passed=[]) + + +def _check_named_arguments( + arg_bindings: Sequence[TaskArgBinding], + passed: Sequence[TaskArgBinding], + params: Sequence[TaskHandlerParam], +) -> ArgumentCheck: + by_name = {binding.name: binding for binding in arg_bindings} + by_folded_name: defaultdict[str, list[TaskArgBinding]] = defaultdict(list) + for binding in arg_bindings: + by_folded_name[_fold_name(binding.name)].append(binding) + matched: list[tuple[TaskArgBinding, TaskHandlerParam]] = [] + declared_not_passed: list[str] = [] + for index, param in enumerate(params): + if param.name is None: + declared_not_passed.append(f"#{index}") + continue + arg: TaskArgBinding | None = by_name.get(param.name) + if arg is None and not param.exact_name: + candidates = by_folded_name[_fold_name(param.name)] + arg = candidates[0] if len(candidates) == 1 else None + if arg is None: + declared_not_passed.append(param.name) + else: + matched.append((arg, param)) + if not matched and _may_be_whole_value(passed, params): + return ArgumentCheck(errors=[], passed_not_declared=[], declared_not_passed=[]) + claimed = {arg.name for arg, _ in matched} + return ArgumentCheck( + errors=_check_values(matched), + passed_not_declared=[arg.name for arg in passed if arg.name not in claimed], + declared_not_passed=declared_not_passed, + ) + + +def _may_be_whole_value(passed: Sequence[TaskArgBinding], params: Sequence[TaskHandlerParam]) -> bool: + if len(passed) != 1 or not params or any(param.exact_name for param in params): + return False + schema = passed[0].value_schema + types = None if schema is None else _get_json_types(schema) + return types is None or "object" in types + + +def _is_defaulted(arg: TaskArgBinding) -> bool: + return isinstance(arg, LiteralArgBinding) and arg.from_default + + +def _fold_name(name: str) -> str: + return name.replace("_", "").lower() + + +def _check_values(pairs: Iterable[tuple[TaskArgBinding, TaskHandlerParam]]) -> list[str]: + return [ + f"argument {arg.name!r} is {mismatch}" + for arg, param in pairs + if (mismatch := check_value_schema(arg.value_schema, param.value_schema)) is not None + ] + + +def _get_json_types(schema: Mapping[str, object]) -> frozenset[str] | None: + if "type" in schema: + names = schema["type"] + if isinstance(names, str): + names = [names] + if not isinstance(names, list) or not names or not all(name in _JSON_TYPES for name in names): + return None + return frozenset(names) + for keyword in ("anyOf", "oneOf"): + if keyword in schema: + branches = schema[keyword] + if not isinstance(branches, list) or not branches: + return None + union: set[str] = set() + for branch in branches: + if not isinstance(branch, dict) or (types := _get_json_types(branch)) is None: + return None + union |= types + return frozenset(union) + if "const" in schema: + return frozenset([_get_value_type(schema["const"])]) + if isinstance(values := schema.get("enum"), list) and values: + return frozenset(_get_value_type(value) for value in values) + return None + + +def _get_value_type(value: object) -> str: + if value is None: + return "null" + if isinstance(value, bool): + return "boolean" + if isinstance(value, int): + return "integer" + if isinstance(value, float): + return "number" + if isinstance(value, str): + return "string" + if isinstance(value, list): + return "array" + return "object" + + +def _join_types(types: frozenset[str]) -> str: + return _join([name for name in _JSON_TYPES if name in types], "or") + + +def _join(values: Sequence[str], conjunction: str) -> str: + if len(values) == 1: + return values[0] + return f"{', '.join(values[:-1])} {conjunction} {values[-1]}" + + +def _make_problem(stub: StubTask, message: str) -> TaskHandlerProblem: + return TaskHandlerProblem( + relative_fileloc=stub.relative_fileloc, message=message, dag_id=stub.dag_id, task_id=stub.task_id + ) diff --git a/airflow-core/tests/unit/dag_processing/test_task_handler_validation.py b/airflow-core/tests/unit/dag_processing/test_task_handler_validation.py new file mode 100644 index 0000000000000..89212a423a310 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_task_handler_validation.py @@ -0,0 +1,668 @@ +# 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 pytest + +from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import LiteralArgBinding, XComArgBinding +from airflow.dag_processing.processor import TaskHandlerDeclaration, TaskHandlerParam +from airflow.dag_processing.task_handler_validation import ( + ArgumentCheck, + StubTask, + StubTaskWarning, + TaskHandlerAnswer, + TaskHandlerMatch, + TaskHandlerProblem, + check_task_handler_arguments, + check_value_schema, + collect_stub_tasks, + format_import_errors, + match_task_handlers, +) +from airflow.sdk import DAG, task +from airflow.serialization.serialized_objects import LazyDeserializedDAG + +BUNDLE_NAME = "go-task-handlers" +DAG_FILE = "dags/etl.py" + +INTEGER = {"type": "integer"} +NUMBER = {"type": "number"} +STRING = {"type": "string"} +NULL = {"type": "null"} +NULLABLE_INTEGER = {"anyOf": [INTEGER, NULL]} +OBJECT = {"type": "object", "additionalProperties": True} + + +def _literal(name, value=None, *, schema=None, from_default=False): + return LiteralArgBinding( + kind="literal", name=name, value=value, value_schema=schema, from_default=from_default + ) + + +def _default(name, value=None, *, schema=None): + return _literal(name, value, schema=schema, from_default=True) + + +def _xcom(name, *, schema=None): + return XComArgBinding(kind="xcom", name=name, task_id="extract", value_schema=schema) + + +def _param(name=None, *, schema=None, exact_name=False): + return TaskHandlerParam(name=name, value_schema=schema, exact_name=exact_name) + + +def _positional(*params, task_id="transform"): + return TaskHandlerDeclaration(task_id=task_id, binding="positional", params=list(params)) + + +def _named(*params, task_id="transform"): + return TaskHandlerDeclaration(task_id=task_id, binding="named", params=list(params)) + + +def _check(*, errors=(), passed_not_declared=(), declared_not_passed=()): + return ArgumentCheck( + errors=list(errors), + passed_not_declared=list(passed_not_declared), + declared_not_passed=list(declared_not_passed), + ) + + +def _stub(task_id="transform", *, arg_bindings=(), is_mapped=False, dag_id="etl"): + return StubTask( + dag_id=dag_id, + task_id=task_id, + queue="go", + relative_fileloc=DAG_FILE, + arg_bindings=list(arg_bindings), + is_mapped=is_mapped, + ) + + +def _answer(**task_handlers): + return TaskHandlerAnswer(bundle_name=BUNDLE_NAME, rel_path="bin/etl", task_handlers=task_handlers) + + +@pytest.mark.parametrize( + ("arg_bindings", "declaration", "expected"), + [ + pytest.param( + [_literal("a"), _xcom("b")], + _positional(_param(), _param()), + _check(), + id="positional-count-matches", + ), + pytest.param( + [_literal("a"), _literal("b"), _literal("c")], + _positional(_param(), _param()), + _check(errors=["passes 3 arguments, the task handler takes 2"]), + id="positional-too-many", + ), + pytest.param( + [_literal("a")], + _positional(_param(), _param()), + _check(errors=["passes 1 argument, the task handler takes 2"]), + id="positional-too-few", + ), + pytest.param( + [_literal("a"), _default("b")], _positional(_param()), _check(), id="defaulted-dropped-to-match" + ), + pytest.param( + [_literal("a"), _default("b")], + _positional(_param(), _param()), + _check(), + id="defaulted-kept-when-the-full-count-matches", + ), + pytest.param( + [_literal("a"), _default("b"), _literal("c")], + _positional(_param()), + _check(errors=["passes 3 arguments (2 without defaults), the task handler takes 1"]), + id="dropping-defaulted-still-mismatches", + ), + pytest.param( + [_literal("x", schema=INTEGER), _literal("y", schema=STRING)], + _positional(_param("y", schema=INTEGER), _param("x", schema=STRING)), + _check(), + id="positional-names-ignored", + ), + pytest.param( + [_literal("a", schema=INTEGER), _xcom("b", schema=STRING)], + _positional(_param(schema=INTEGER), _param(schema=INTEGER)), + _check(errors=["argument 'b' is string, the task handler takes integer"]), + id="positional-type-mismatch", + ), + pytest.param( + [_literal("a", schema=STRING), _literal("b", schema=STRING)], + _positional(_param(schema=INTEGER), _param(schema=INTEGER)), + _check( + errors=[ + "argument 'a' is string, the task handler takes integer", + "argument 'b' is string, the task handler takes integer", + ] + ), + id="every-mismatching-argument-is-reported", + ), + pytest.param([], _positional(), _check(), id="argless-against-no-params"), + pytest.param( + [], + _positional(_param()), + _check(errors=["passes 0 arguments, the task handler takes 1"]), + id="argless-against-positional-params", + ), + pytest.param( + [_literal("user_id", schema=INTEGER)], + _named(_param("user_id", schema=INTEGER)), + _check(), + id="named-exact", + ), + pytest.param( + [_literal("user_id"), _literal("region")], + _named(_param("UserId"), _param("region")), + _check(), + id="named-folded", + ), + pytest.param( + [_literal("user_id"), _literal("region")], + _named(_param("UserId", exact_name=True), _param("region")), + _check(passed_not_declared=["user_id"], declared_not_passed=["UserId"]), + id="exact-name-blocks-folding", + ), + pytest.param( + [_literal("region"), _literal("extra")], + _named(_param("region")), + _check(passed_not_declared=["extra"]), + id="passed-argument-no-param-takes", + ), + pytest.param( + [_literal("region"), _default("extra")], + _named(_param("region")), + _check(), + id="defaulted-argument-no-param-takes", + ), + pytest.param( + [_literal("region")], + _named(_param("region"), _param("limit")), + _check(declared_not_passed=["limit"]), + id="param-no-argument-fills", + ), + pytest.param( + [], + _named(_param("region")), + _check(declared_not_passed=["region"]), + id="argless-against-named-params", + ), + pytest.param( + [_literal("config"), _default("extra")], + _named(_param("region"), _param("limit")), + _check(), + id="lone-argument-without-a-schema-may-be-the-whole-value", + ), + pytest.param( + [_literal("config", schema=OBJECT)], + _named(_param("region"), _param("limit")), + _check(), + id="lone-object-argument-may-be-the-whole-value", + ), + pytest.param( + [_literal("config", schema={"$ref": "#/$defs/Config"})], + _named(_param("region"), _param("limit")), + _check(), + id="lone-argument-of-unknown-type-may-be-the-whole-value", + ), + pytest.param( + [_literal("config", schema=STRING)], + _named(_param("region"), _param("limit")), + _check(passed_not_declared=["config"], declared_not_passed=["region", "limit"]), + id="lone-string-argument-is-not-the-whole-value", + ), + pytest.param( + [_literal("config")], + _named(_param("region", exact_name=True), _param("limit")), + _check(passed_not_declared=["config"], declared_not_passed=["region", "limit"]), + id="lone-argument-with-an-exact-name-param", + ), + pytest.param( + [_literal("config")], + _named(), + _check(passed_not_declared=["config"]), + id="lone-argument-against-no-params", + ), + pytest.param( + [_literal("a"), _literal("b")], + _named(_param("region")), + _check(passed_not_declared=["a", "b"], declared_not_passed=["region"]), + id="two-unmatched-arguments", + ), + pytest.param( + [_literal("user_id"), _literal("userId"), _literal("region")], + _named(_param("UserID"), _param("region")), + _check(passed_not_declared=["user_id", "userId"], declared_not_passed=["UserID"]), + id="folded-name-two-arguments-share", + ), + pytest.param( + [_literal("user_id"), _literal("userId")], + _named(_param("user_id")), + _check(passed_not_declared=["userId"]), + id="exact-name-despite-a-shared-fold", + ), + pytest.param( + [_literal("region"), _literal("limit")], + _named(_param("region"), _param()), + _check(passed_not_declared=["limit"], declared_not_passed=["#1"]), + id="nameless-named-param", + ), + pytest.param( + [_literal("region", schema=INTEGER), _literal("limit", schema=INTEGER)], + _named(_param("limit", schema=INTEGER), _param("region", schema=STRING)), + _check(errors=["argument 'region' is integer, the task handler takes string"]), + id="named-type-mismatch", + ), + pytest.param( + [_literal("a", schema=INTEGER), _default("b", schema=STRING)], + _positional(_param(schema=INTEGER), _param(schema=INTEGER)), + _check(errors=["argument 'b' is string, the task handler takes integer"]), + id="positional-defaulted-argument-type-mismatch", + ), + pytest.param( + [_literal("region", schema=STRING), _default("limit", schema=STRING)], + _named(_param("region", schema=STRING), _param("limit", schema=INTEGER)), + _check(errors=["argument 'limit' is string, the task handler takes integer"]), + id="named-defaulted-argument-type-mismatch", + ), + pytest.param( + [_literal("a", schema=STRING)], + TaskHandlerDeclaration(task_id="transform", binding="named", params=None), + _check(), + id="unlisted-params", + ), + ], +) +def test_check_task_handler_arguments(arg_bindings, declaration, expected): + assert check_task_handler_arguments(arg_bindings, declaration) == expected + + +@pytest.mark.parametrize( + ("stub_schema", "handler_schema", "expected"), + [ + pytest.param(INTEGER, INTEGER, None, id="equal"), + pytest.param(INTEGER, NUMBER, None, id="integer-into-number"), + pytest.param(NUMBER, INTEGER, "number, the task handler takes integer", id="number-into-integer"), + pytest.param(STRING, INTEGER, "string, the task handler takes integer", id="string-into-integer"), + pytest.param( + NULLABLE_INTEGER, + INTEGER, + "integer or null, the task handler takes integer", + id="nullable-into-non-nullable", + ), + pytest.param(INTEGER, NULLABLE_INTEGER, None, id="non-nullable-into-nullable"), + pytest.param(NULLABLE_INTEGER, NULLABLE_INTEGER, None, id="nullable-into-nullable"), + pytest.param(NULL, NULLABLE_INTEGER, None, id="null-only-into-nullable"), + pytest.param( + {"const": None}, STRING, "null, the task handler takes string", id="null-into-non-nullable" + ), + pytest.param({"anyOf": [INTEGER, STRING]}, INTEGER, None, id="union-with-one-accepted-type"), + pytest.param( + {"anyOf": [STRING, {"type": "boolean"}]}, + INTEGER, + "string or boolean, the task handler takes integer", + id="union-with-no-accepted-type", + ), + pytest.param( + {"anyOf": [INTEGER, STRING]}, {"anyOf": [STRING, INTEGER, NULL]}, None, id="any-of-subset" + ), + pytest.param( + {"anyOf": [INTEGER, STRING, NULL]}, + {"oneOf": [STRING, INTEGER]}, + "string, integer or null, the task handler takes string or integer", + id="any-of-superset", + ), + pytest.param( + {"type": ["integer", "null"]}, + INTEGER, + "integer or null, the task handler takes integer", + id="type-list", + ), + pytest.param( + {"type": ["string", "null"]}, {"anyOf": [STRING, NULL]}, None, id="type-list-into-any-of" + ), + pytest.param({"const": "eu"}, STRING, None, id="const"), + pytest.param({"const": 1}, STRING, "integer, the task handler takes string", id="const-mismatch"), + pytest.param({"enum": [1, 2.5]}, NUMBER, None, id="enum"), + pytest.param( + {"enum": ["eu", None]}, + {"enum": ["eu", "us"]}, + "string or null, the task handler takes string", + id="enum-mismatch", + ), + pytest.param( + {"enum": [True]}, INTEGER, "boolean, the task handler takes integer", id="boolean-into-integer" + ), + pytest.param({"$ref": "#/$defs/Config"}, INTEGER, None, id="ref-skipped"), + pytest.param({"allOf": [INTEGER]}, STRING, None, id="all-of-skipped"), + pytest.param( + {"anyOf": [{"$ref": "#/$defs/Config"}, NULL]}, INTEGER, None, id="any-of-with-a-ref-skipped" + ), + pytest.param({}, INTEGER, None, id="empty-stub-schema-skipped"), + pytest.param(STRING, {}, None, id="empty-handler-schema-skipped"), + pytest.param(None, INTEGER, None, id="no-stub-schema"), + pytest.param(STRING, None, None, id="no-handler-schema"), + pytest.param( + {"type": "integer", "format": "int64"}, + {"type": "integer", "format": "int32", "minimum": 0}, + None, + id="format-and-range-ignored", + ), + ], +) +def test_check_value_schema(stub_schema, handler_schema, expected): + assert check_value_schema(stub_schema, handler_schema) == expected + + +def test_each_stub_task_is_checked_against_the_handler_of_its_own_dag(): + stub_tasks = [ + _stub("extract"), + _stub("transform", arg_bindings=[_literal("a")]), + _stub("transform", dag_id="reporting"), + ] + # "audit" has no stub task here; it and "reporting" both register a "transform" that takes nothing, + # while "etl"'s takes one. + answer = _answer( + reporting=[_positional(task_id="transform")], + etl=[_positional(task_id="extract"), _positional(_param(), task_id="transform")], + audit=[_positional(task_id="transform")], + ) + + assert match_task_handlers(stub_tasks, answer) == TaskHandlerMatch(problems=[], warnings=[]) + + +def test_a_stub_task_without_a_handler_is_a_problem(): + answer = _answer(other=[_positional()], etl=[_positional(task_id="extract")]) + + assert match_task_handlers([_stub()], answer) == TaskHandlerMatch( + problems=[ + TaskHandlerProblem( + relative_fileloc=DAG_FILE, + message="Dag 'etl', task 'transform': 'bin/etl' in Dag bundle 'go-task-handlers' " + "registers no task handler for it", + dag_id="etl", + task_id="transform", + ) + ], + warnings=[], + ) + + +def test_a_handler_without_a_stub_task_is_not_a_problem(): + answer = _answer(etl=[_positional(), _positional(task_id="load")], reporting=[_named(task_id="publish")]) + + assert match_task_handlers([_stub()], answer) == TaskHandlerMatch(problems=[], warnings=[]) + + +def test_a_mapped_stub_task_is_checked_for_its_handler_only(): + stub_tasks = [_stub(is_mapped=True), _stub("load", is_mapped=True)] + + match = match_task_handlers(stub_tasks, _answer(etl=[_positional(_param(), _param())])) + + assert match.warnings == [] + assert [problem.task_id for problem in match.problems] == ["load"] + + +def test_unlisted_params_are_checked_for_the_handler_only(): + stub = _stub(arg_bindings=[_literal("a", schema=STRING)]) + answer = _answer(etl=[TaskHandlerDeclaration(task_id="transform", binding="named", params=None)]) + + assert match_task_handlers([stub], answer) == TaskHandlerMatch(problems=[], warnings=[]) + + +@pytest.mark.parametrize( + ("arg_bindings", "params", "passed_not_declared", "declared_not_passed"), + [ + pytest.param( + [_literal("region"), _literal("unused_label")], + [_param("region"), _param("limit")], + ["unused_label"], + ["limit"], + id="both-sides-mismatch", + ), + pytest.param( + [_literal("region")], + [_param("region"), _param("limit")], + [], + ["limit"], + id="only-a-declared-param-is-unfilled", + ), + pytest.param( + [_literal("region"), _literal("extra")], + [_param("region")], + ["extra"], + [], + id="only-a-passed-argument-is-unclaimed", + ), + ], +) +def test_a_name_mismatch_is_a_warning_not_a_problem( + arg_bindings, params, passed_not_declared, declared_not_passed +): + stub = _stub(arg_bindings=arg_bindings) + answer = _answer(etl=[_named(*params)]) + + assert match_task_handlers([stub], answer) == TaskHandlerMatch( + problems=[], + warnings=[ + StubTaskWarning( + dag_id="etl", + task_id="transform", + artifact_bundle_name=BUNDLE_NAME, + artifact_rel_path="bin/etl", + passed_not_declared=passed_not_declared, + declared_not_passed=declared_not_passed, + ) + ], + ) + + +@pytest.mark.parametrize( + ("params", "error"), + [ + pytest.param( + [_param(schema=INTEGER), _param(), _param()], + "passes 2 arguments, the task handler takes 3", + id="count", + ), + pytest.param( + [_param(schema=INTEGER), _param()], + "argument 'count' is integer or null, the task handler takes integer", + id="value-type", + ), + ], +) +def test_an_argument_problem_names_the_artifact(params, error): + stub = _stub(arg_bindings=[_literal("count", schema=NULLABLE_INTEGER), _literal("label")]) + + match = match_task_handlers([stub], _answer(etl=[_positional(*params)])) + + assert match == TaskHandlerMatch( + problems=[ + TaskHandlerProblem( + relative_fileloc=DAG_FILE, + message=f"Dag 'etl', task 'transform' ('bin/etl' in Dag bundle 'go-task-handlers'): {error}", + dag_id="etl", + task_id="transform", + ) + ], + warnings=[], + ) + + +def test_each_mismatching_argument_is_its_own_problem(): + stub = _stub(arg_bindings=[_literal("a", schema=STRING), _literal("b", schema=STRING)]) + answer = _answer(etl=[_positional(_param(schema=INTEGER), _param(schema=INTEGER))]) + + match = match_task_handlers([stub], answer) + + assert match == TaskHandlerMatch( + problems=[ + TaskHandlerProblem( + relative_fileloc=DAG_FILE, + message="Dag 'etl', task 'transform' ('bin/etl' in Dag bundle 'go-task-handlers'): " + "argument 'a' is string, the task handler takes integer", + dag_id="etl", + task_id="transform", + ), + TaskHandlerProblem( + relative_fileloc=DAG_FILE, + message="Dag 'etl', task 'transform' ('bin/etl' in Dag bundle 'go-task-handlers'): " + "argument 'b' is string, the task handler takes integer", + dag_id="etl", + task_id="transform", + ), + ], + warnings=[], + ) + + +def test_collect_stub_tasks(): + with DAG(dag_id="etl", schedule=None) as dag: + + @task.stub(queue="go") + def extract(): ... + + @task.stub(queue="go") + def transform(data: str, region: str = "eu"): ... + + @task.stub(queue="go") + def fan_out(item: int): ... + + @task.stub + def unrouted(): ... + + @task + def report(): ... + + transform(extract()) + fan_out.expand(item=[1, 2]) + unrouted() + report() + + with DAG(dag_id="unserialized", schedule=None) as unserialized_dag: + + @task.stub(queue="go") + def load(): ... + + load() + + dag.relative_fileloc = DAG_FILE + + stub_tasks = collect_stub_tasks([dag, unserialized_dag], [LazyDeserializedDAG.from_dag(dag)]) + + assert stub_tasks == [ + StubTask( + dag_id="etl", + task_id="extract", + queue="go", + relative_fileloc=DAG_FILE, + arg_bindings=[], + is_mapped=False, + ), + StubTask( + dag_id="etl", + task_id="transform", + queue="go", + relative_fileloc=DAG_FILE, + arg_bindings=[ + XComArgBinding(kind="xcom", name="data", task_id="extract", value_schema=STRING), + _default("region", "eu", schema=STRING), + ], + is_mapped=False, + ), + StubTask( + dag_id="etl", + task_id="fan_out", + queue="go", + relative_fileloc=DAG_FILE, + arg_bindings=[], + is_mapped=True, + ), + StubTask( + dag_id="etl", + task_id="unrouted", + queue="default", + relative_fileloc=DAG_FILE, + arg_bindings=[], + is_mapped=False, + ), + ] + + +def test_collect_stub_tasks_reads_the_arguments_the_worker_gets(): + with DAG(dag_id="etl", schedule=None) as dag: + + @task.stub(queue="go") + def load(table: str, columns: tuple = ("id", "name")): ... + + @task.stub(queue="go") + def merge(mapping: dict): ... + + load("users") + merge({1: "a"}) + + stub_tasks = collect_stub_tasks([dag], [LazyDeserializedDAG.from_dag(dag)]) + + assert [stub.arg_bindings for stub in stub_tasks] == [ + [ + _literal("table", "users", schema=STRING), + _default("columns", ["id", "name"], schema={"type": "array", "items": {}}), + ], + [_literal("mapping", {"1": "a"}, schema=OBJECT)], + ] + + +def test_format_import_errors_groups_by_file_in_a_stable_order(): + problems = [ + TaskHandlerProblem(relative_fileloc="dags/z.py", message="z: load", dag_id="z", task_id="load"), + TaskHandlerProblem( + relative_fileloc=DAG_FILE, message="etl: transform", dag_id="etl", task_id="transform" + ), + TaskHandlerProblem(relative_fileloc=DAG_FILE, message="Coordinator 'java'"), + TaskHandlerProblem( + relative_fileloc=DAG_FILE, message="etl: load, count", dag_id="etl", task_id="load" + ), + TaskHandlerProblem( + relative_fileloc=DAG_FILE, message="audit: check", dag_id="audit", task_id="check" + ), + TaskHandlerProblem(relative_fileloc=DAG_FILE, message="Coordinator 'go-sdk'"), + TaskHandlerProblem( + relative_fileloc=DAG_FILE, message="etl: load, type", dag_id="etl", task_id="load" + ), + ] + + import_errors = format_import_errors(problems) + + assert list(import_errors) == [DAG_FILE, "dags/z.py"] + assert import_errors == { + "dags/z.py": "Stub tasks in dags/z.py do not match their task handlers:\n- z: load", + DAG_FILE: "\n".join( + [ + "Stub tasks in dags/etl.py do not match their task handlers:", + "- Coordinator 'java'", + "- Coordinator 'go-sdk'", + "- audit: check", + "- etl: load, count", + "- etl: load, type", + "- etl: transform", + ] + ), + }