From 6ea9922a7a405f1df62a1b3541151e4e4ba2036a Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Wed, 20 May 2026 11:38:06 +0800 Subject: [PATCH] Add schema migration to supervisor-child comm With foreign language SDKs, it may be possible the two sides of supervisor comm have different versions. This adds a migration layer at the supervisor (server) side, so an SDK (client) using a lower version of the schema may be able to communicate to the server. --- .pre-commit-config.yaml | 14 + .../prek/check_supervisor_schemas_versions.py | 204 + scripts/ci/prek/dump_supervisor_schemas.py | 55 + .../generate_supervisor_schemas_snapshot.py | 104 + task-sdk/.pre-commit-config.yaml | 1 + .../sdk/execution_time/schema/AGENTS.md | 122 + .../sdk/execution_time/schema/__init__.py | 128 + .../sdk/execution_time/schema/migrator.py | 211 + .../sdk/execution_time/schema/schema.json | 6131 +++++++++++++++++ .../schema/versions/__init__.py | 25 + .../airflow/sdk/execution_time/supervisor.py | 37 +- .../execution_time/schema/__init__.py | 16 + .../schema/_mock_version_bundle.py | 189 + .../execution_time/schema/test_integration.py | 366 + .../execution_time/schema/test_migrator.py | 336 + 15 files changed, 7930 insertions(+), 9 deletions(-) create mode 100755 scripts/ci/prek/check_supervisor_schemas_versions.py create mode 100755 scripts/ci/prek/dump_supervisor_schemas.py create mode 100755 scripts/ci/prek/generate_supervisor_schemas_snapshot.py create mode 100644 task-sdk/src/airflow/sdk/execution_time/schema/AGENTS.md create mode 100644 task-sdk/src/airflow/sdk/execution_time/schema/__init__.py create mode 100644 task-sdk/src/airflow/sdk/execution_time/schema/migrator.py create mode 100644 task-sdk/src/airflow/sdk/execution_time/schema/schema.json create mode 100644 task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py create mode 100644 task-sdk/tests/task_sdk/execution_time/schema/__init__.py create mode 100644 task-sdk/tests/task_sdk/execution_time/schema/_mock_version_bundle.py create mode 100644 task-sdk/tests/task_sdk/execution_time/schema/test_integration.py create mode 100644 task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d019a32715075..aa72818873763 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1147,6 +1147,20 @@ repos: pass_filenames: true files: ^airflow-core/src/airflow/api_fastapi/execution_api/(datamodels|versions)/.*\.py$ require_serial: true + - id: generate-supervisor-schemas-snapshot + name: Regenerate supervisor schema snapshot + entry: ./scripts/ci/prek/generate_supervisor_schemas_snapshot.py + language: python + pass_filenames: false + files: ^(task-sdk/src/airflow/sdk/execution_time/(comms\.py|schema/.*\.py)|airflow-core/src/airflow/dag_processing/processor\.py)$ + require_serial: true + - id: check-supervisor-schemas-versions + name: Check supervisor schema changes have corresponding version updates + entry: ./scripts/ci/prek/check_supervisor_schemas_versions.py + language: python + pass_filenames: true + files: ^(task-sdk/src/airflow/sdk/execution_time/(comms\.py|schema/.*\.py)|airflow-core/src/airflow/dag_processing/processor\.py)$ + require_serial: true - id: generate-tasksdk-datamodels name: Generate Datamodels for TaskSDK client language: python diff --git a/scripts/ci/prek/check_supervisor_schemas_versions.py b/scripts/ci/prek/check_supervisor_schemas_versions.py new file mode 100755 index 0000000000000..68c1ed6e07de4 --- /dev/null +++ b/scripts/ci/prek/check_supervisor_schemas_versions.py @@ -0,0 +1,204 @@ +#!/usr/bin/env python +# 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. +# /// script +# requires-python = ">=3.10,<3.11" +# dependencies = [ +# "rich>=13.6.0", +# ] +# /// +""" +Fail when a supervisor schema has changed without a matching +``VersionChange`` entry under +``task-sdk/src/airflow/sdk/execution_time/schema/versions/``. + +Mirrors :mod:`scripts.ci.prek.check_execution_api_versions` for the +supervisor bundle. The check is per-commit: every PR that mutates a +registered supervisor schema must add an instruction to the in-progress head +``v__
.py`` file. The release-time version-file bump itself +is one-per-release; this hook is what keeps the in-progress file +honest between releases. + +The comparison is done by dumping the snapshot JSON in this worktree +and in a temporary worktree of the upstream target branch, then +diffing them. Both sides invoke the sibling ``dump_supervisor_schemas.py`` +script so the comparison is dump-version stable. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import tempfile +from pathlib import Path + +from common_prek_utils import console, get_remote_for_main + +SUPERVISOR_SCHEMAS_PREFIX = "task-sdk/src/airflow/sdk/execution_time/schema/" +VERSIONS_PREFIX = SUPERVISOR_SCHEMAS_PREFIX + "versions/" +TASK_SDK_COMMS_PATH = "task-sdk/src/airflow/sdk/execution_time/comms.py" +CORE_PROCESSOR_PATH = "airflow-core/src/airflow/dag_processing/processor.py" + +DUMP_SCRIPT = Path(__file__).parent / "dump_supervisor_schemas.py" + + +# TODO: We should consolidte the common logic with check_execution_api_versions.py into common_prek_utils +def get_target_branch() -> str: + """Branch to compare against. GITHUB_BASE_REF for PRs, DEFAULT_BRANCH in CI, else main.""" + return os.environ.get("GITHUB_BASE_REF") or os.environ.get("DEFAULT_BRANCH") or "main" + + +def get_changed_files(filenames: list[str]) -> list[str]: + """Get changed files. Uses filenames from prek when provided, else staged files for local runs.""" + if filenames: + return filenames + result = subprocess.run( + ["git", "diff", "--cached", "--name-only"], + capture_output=True, + text=True, + check=True, + ) + return [f for f in result.stdout.strip().splitlines() if f] + + +def dump_snapshot(cwd: Path) -> str: + """Run ``dump_supervisor_schemas.py`` in *cwd* and return its stdout.""" + result = subprocess.run( + [ + "uv", + "run", + "-p", + "3.12", + "--no-progress", + "--project", + "task-sdk", + "-s", + str(DUMP_SCRIPT), + ], + cwd=cwd, + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0: + raise RuntimeError(f"Snapshot dump failed: {result.stderr}") + return result.stdout + + +def _upstream_has_schema() -> bool: + """Return True if the target branch carries the schema package.""" + target_branch = get_target_branch() + remote = get_remote_for_main() + ref = f"{remote}/{target_branch}" + subprocess.run(["git", "fetch", remote, target_branch], capture_output=True, check=False) + # ``git cat-file -e`` exits zero iff the path exists at the ref. + result = subprocess.run( + ["git", "cat-file", "-e", f"{ref}:{VERSIONS_PREFIX}__init__.py"], + capture_output=True, + check=False, + ) + return result.returncode == 0 + + +def dump_snapshot_from_main() -> str: + """Dump snapshot from target branch using a temporary worktree.""" + target_branch = get_target_branch() + remote = get_remote_for_main() + ref = f"{remote}/{target_branch}" + worktree_path = Path(tempfile.mkdtemp()) / "airflow-main" + subprocess.run(["git", "fetch", remote, target_branch], capture_output=True, check=False) + subprocess.run(["git", "worktree", "add", str(worktree_path), ref], capture_output=True, check=True) + try: + return dump_snapshot(worktree_path) + finally: + subprocess.run( + ["git", "worktree", "remove", "--force", str(worktree_path)], + capture_output=True, + check=False, + ) + + +def main() -> int: + changed_files = get_changed_files(sys.argv[1:]) + + # Files under schema/ that reference the bundle's + # registered models. Schema changes in those models' homes + # (``comms.py``, ``processor.py``) trigger this hook too because + # the snapshot embeds their head shape. + schema_source_files = [ + f + for f in changed_files + if f.startswith(SUPERVISOR_SCHEMAS_PREFIX) or f == TASK_SDK_COMMS_PATH or f == CORE_PROCESSOR_PATH + ] + version_files = [f for f in changed_files if f.startswith(VERSIONS_PREFIX)] + + if not schema_source_files: + return 0 + if version_files: + # Contributor added a version-change entry: trust them. + return 0 + + if not _upstream_has_schema(): + # The package is being introduced in this PR -- nothing on the + # target branch to compare against. The check will start firing + # normally once the package is on the target branch. + console.print( + "[yellow]Skipping supervisor-schemas version check:[/] target branch " + "has no schema package yet. The check activates once " + "this PR merges." + ) + return 0 + + try: + main_snapshot = dump_snapshot_from_main() + except Exception as e: + console.print(f"[bold red]ERROR:[/] Failed to generate upstream snapshot for comparison: {e}") + return 1 + + try: + current_snapshot = dump_snapshot(Path.cwd()) + except Exception as e: + console.print(f"[bold red]ERROR:[/] Failed to generate current snapshot: {e}") + return 1 + + if current_snapshot != main_snapshot: + console.print("[bold red]ERROR:[/] Supervisor schema has changed but no version file was updated.") + console.print("") + console.print("The following files were changed:") + for f in schema_source_files: + console.print(f" - [magenta]{f}[/]") + console.print("") + remote = get_remote_for_main() + target_branch = get_target_branch() + console.print( + f"Snapshot diff against [cyan]{remote}/{target_branch}[/] detected differences.\n" + "\n" + "Append a ``VersionChange`` subclass to the in-progress head " + "``v__
.py`` file under:\n" + f" [cyan]{VERSIONS_PREFIX}[/]\n" + "\n" + "See [cyan]task-sdk/src/airflow/sdk/execution_time/schema/AGENTS.md[/]." + ) + return 1 + console.print("[green]Snapshot unchanged:[/] Source changes do not affect the supervisor schema.") + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/ci/prek/dump_supervisor_schemas.py b/scripts/ci/prek/dump_supervisor_schemas.py new file mode 100755 index 0000000000000..8b28d9af7f219 --- /dev/null +++ b/scripts/ci/prek/dump_supervisor_schemas.py @@ -0,0 +1,55 @@ +#!/usr/bin/env python +# 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. +""" +Dump the supervisor schema snapshot. Prints JSON to stdout. + +Mirrors :mod:`scripts.ci.prek.generate_execution_api_schema` but for the +supervisor schema ``VersionBundle``: walks the registered head models and +emits ``model_json_schema()`` for every wire body in a deterministic +class-name order so the artefact diffs cleanly across runs. + +Run with cwd at the repo root. +""" + +from __future__ import annotations + +import json +import os +import sys +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pydantic import BaseModel + +os.environ["_AIRFLOW__AS_LIBRARY"] = "1" + +from airflow.sdk.execution_time.schema import bundle, registered_models_by_name + + +def _registered_models_sorted() -> tuple[type[BaseModel], ...]: + """Return registered head models sorted by class name for stable snapshot diffs.""" + by_name = registered_models_by_name() + return tuple(by_name[name] for name in sorted(by_name)) + + +snapshot = { + "api_version": str(bundle.versions[0].value), + "schemas": {cls.__name__: cls.model_json_schema() for cls in _registered_models_sorted()}, +} +json.dump(snapshot, sys.stdout, indent=2, sort_keys=True) +sys.stdout.write("\n") diff --git a/scripts/ci/prek/generate_supervisor_schemas_snapshot.py b/scripts/ci/prek/generate_supervisor_schemas_snapshot.py new file mode 100755 index 0000000000000..d4fe5c5f19ee9 --- /dev/null +++ b/scripts/ci/prek/generate_supervisor_schemas_snapshot.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python +# 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. +# /// script +# requires-python = ">=3.10,<3.11" +# dependencies = [ +# "rich>=13.6.0", +# ] +# /// +""" +Regenerate the supervisor schema snapshot at +``task-sdk/src/airflow/sdk/execution_time/schema/schema.json``. + +The snapshot is the head-version JSON Schema for every Pydantic class +on the supervisor schema wire (the union members of ``ToTask``, +``ToSupervisor``, ``ToManager``, ``ToDagProcessor``). + +The actual dump is delegated to ``dump_supervisor_schemas.py`` (the +sibling stdout-only script). If the committed snapshot differs from +the dumped content the hook rewrites it and exits non-zero (standard +"regenerated files, please re-stage" pattern). +""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +from common_prek_utils import console + +REPO_ROOT = Path(__file__).parents[3].resolve() +SNAPSHOT_PATH = REPO_ROOT.joinpath( + "task-sdk", + "src", + "airflow", + "sdk", + "execution_time", + "schema", + "schema.json", +) +DUMP_SCRIPT = Path(__file__).parent.joinpath("dump_supervisor_schemas.py") + + +def dump_snapshot(cwd: Path) -> str: + """Run ``dump_supervisor_schemas.py`` in *cwd* and return its stdout.""" + result = subprocess.run( + [ + "uv", + "run", + "-p", + "3.12", + "--no-progress", + "--project", + "task-sdk", + "-s", + str(DUMP_SCRIPT), + ], + cwd=cwd, + capture_output=True, + text=True, + check=False, + ) + if result.returncode != 0: + raise RuntimeError(f"Snapshot dump failed: {result.stderr}") + return result.stdout + + +def main() -> int: + try: + new_content = dump_snapshot(REPO_ROOT) + except Exception as e: + console.print(f"[bold red]ERROR:[/] {e}") + return 1 + + if SNAPSHOT_PATH.exists(): + old_content = SNAPSHOT_PATH.read_text() + if old_content == new_content: + return 0 + else: + SNAPSHOT_PATH.parent.mkdir(parents=True, exist_ok=True) + + SNAPSHOT_PATH.write_text(new_content) + rel = SNAPSHOT_PATH.relative_to(REPO_ROOT) + console.print(f"[yellow]Regenerated[/] [cyan]{rel}[/]. Please review the diff and re-stage the file.") + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/task-sdk/.pre-commit-config.yaml b/task-sdk/.pre-commit-config.yaml index 100a6e6490849..f9a224f60d483 100644 --- a/task-sdk/.pre-commit-config.yaml +++ b/task-sdk/.pre-commit-config.yaml @@ -43,6 +43,7 @@ repos: ^src/airflow/sdk/definitions/deadline\.py$| ^src/airflow/sdk/definitions/dag\.py$| ^src/airflow/sdk/definitions/_internal/types\.py$| + ^src/airflow/sdk/execution_time/schema/__init__\.py$| ^src/airflow/sdk/execution_time/execute_workload\.py$| ^src/airflow/sdk/execution_time/secrets_masker\.py$| ^src/airflow/sdk/execution_time/callback_supervisor\.py$| diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/AGENTS.md b/task-sdk/src/airflow/sdk/execution_time/schema/AGENTS.md new file mode 100644 index 0000000000000..0d844efe7ac05 --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/schema/AGENTS.md @@ -0,0 +1,122 @@ + + +# Supervisor Schemas — Agent Instructions + +## What this package owns + +A Cadwyn [`VersionBundle`][cadwyn-versions] and a thin +`SchemaVersionMigrator` for the wire shapes the Task SDK supervisor +exchanges with a lang-SDK runtime subprocess (Java, Go, Rust, ...) +launched by a coordinator. **No Pydantic models live here.** The +models stay in their semantic homes: + +- Task-execution channel (supervisor ↔ task runner): the + `ToTask` and `ToSupervisor` discriminated unions in + `airflow.sdk.execution_time.comms`. +- Dag-processing channel (manager ↔ parser-supervisor): the + `ToManager` and `ToDagProcessor` discriminated unions in + `airflow.dag_processing.processor`. + +`registered_models_by_name()` introspects those four unions on first +call, so the snapshot the prek hook commits to `schema.json` always +matches the exact set of classes `CommsDecoder` actually decodes +against — there is no hand-maintained list to keep in sync. The +Triggerer's unions (`ToTriggerRunner`, `ToTriggerSupervisor`) **are +intentionally excluded**; lang-SDK coordinators do not handle the +Triggerer channel today. + +The bundle references registered classes via `schema(...)` instructions +in `versions/v.py` files. + +This is **independent** of `airflow.api_fastapi.execution_api.versions.bundle`, +which governs the HTTP contract between Task SDK clients and the API +server. A supervisor schema change does **not** force a HTTP API +version bump, and vice versa. + +[cadwyn-versions]: https://docs.cadwyn.dev/concepts/version_changes/ + +## Files in this folder + +- `__init__.py` — re-exports `bundle`, the migrator, the + `registered_models_by_name()` registry, and `resolve_body_class()`. +- `migrator.py` — `SchemaVersionMigrator` + `get_schema_version_migrator()`. +- `versions/__init__.py` — the `VersionBundle` itself + (`HeadVersion()` + dated `Version(...)` entries). +- `versions/vYYYY_MM_DD.py` — one file per release. The most recent + file is the **in-progress** version; PRs append to it. +- `schema.json` — generated head-version JSON Schema snapshot for + lang-SDK codegen. Managed by the + `generate-supervisor-schemas-snapshot` prek hook (which lives at + `scripts/ci/prek/dump_supervisor_schemas.py` and walks + `registered_models_by_name()` in sorted-name order); do not edit by + hand. + +## When making changes + +### Adding a new body to the versioned contract + +Append the class to the relevant discriminated union in its semantic +home — `ToTask` / `ToSupervisor` in `comms.py`, or `ToManager` / +`ToDagProcessor` in `processor.py`. That is the *only* registration +step; `registered_models_by_name()` picks it up automatically the next +time the snapshot hook runs. + +No `VersionChange` entry is required on the first introduction — the +head shape *is* the schema for the new body. + +### Adding a field to a registered body + +1. Add the field to the model in its semantic home (e.g. + `comms.py:StartupDetails`). +2. Open the in-progress `versions/vYYYY_MM_DD.py` file (the one with + the most recent date) and append a `VersionChange` subclass: + + ```python + class AddSentryTraceField(VersionChange): + """Add `sentry_trace_id` to StartupDetails.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = ( + schema(StartupDetails).field("sentry_trace_id").didnt_exist, + ) + ``` + +3. Reference the new `VersionChange` from the bundle in + `versions/__init__.py`: + + ```python + Version("2026-06-16", AddRetryDelay, AddSentryTraceField), + ``` + +4. The `generate-supervisor-schemas-snapshot` prek hook will + regenerate `schema.json` on commit. Re-stage the file. + +### Removing or renaming a field + +Same pattern as adding, but with the inverse instruction +(`schema(X).field(...).existed_as(...)` etc.). See the execution-API +`versions/` folder for richer examples. + +## Version cadence + +- **Bump per release.** The release manager freezes the in-progress + `vYYYY_MM_DD.py` file at release time and opens a new in-progress + file dated past the next planned release. +- **Accumulate per change.** Each PR that mutates a registered body + appends a `VersionChange` entry to the head in-progress file. + Contributors never invent a new version date. + +## Prek hooks + +Two hooks enforce the contract: + +- `generate-supervisor-schemas-snapshot` — regenerates `schema.json` + on commit when any registered model or any `versions/v*.py` file + changes. Fails if the committed snapshot is stale. +- `check-supervisor-schemas-versions` — fails if the regenerated + snapshot differs from the upstream target-branch snapshot but no + file under `versions/` was touched. + +The check is per-commit; the file bump is per-release. diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/__init__.py new file mode 100644 index 0000000000000..1963e8b8e6614 --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/schema/__init__.py @@ -0,0 +1,128 @@ +# 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. +""" +Cadwyn versioning and in-process migration for the supervisor schemas. + +Two distinct Cadwyn ``VersionBundle`` instances coexist in the codebase: + +* :data:`.versions.bundle` (this package) — versions the wire shapes the + Task SDK supervisor exchanges with a lang-SDK runtime subprocess + launched by a coordinator (Java, Go, Rust, ...). The bodies it + references live in their semantic homes + (``airflow.sdk.execution_time.comms`` for task execution, + ``airflow.dag_processing.processor`` for Dag parsing); this package + only owns the versioning machinery, not the model definitions. +* :data:`airflow.api_fastapi.execution_api.versions.bundle` — versions + the HTTP contract between Task SDK clients and the API server. + Unaffected by this package. + +:func:`registered_models_by_name` resolves a wire-shape ``type`` +discriminator to the head Pydantic class. It is computed dynamically +from the four discriminated unions ``ToTask``, ``ToSupervisor`` +(task-execution channel) and ``ToManager``, ``ToDagProcessor`` +(dag-processing channel) so the registry is always in sync with the +actual unions ``CommsDecoder`` decodes against -- no hand-maintained +list to drift. Triggerer unions are intentionally excluded (the +Triggerer channel is not handled by lang-SDK coordinators today). +""" + +from __future__ import annotations + +import functools +from types import UnionType +from typing import TYPE_CHECKING, Annotated, Any, get_args, get_origin + +from pydantic import BaseModel + +from airflow.sdk.execution_time.schema.migrator import ( + SchemaVersionMigrator, + get_schema_version_migrator, +) +from airflow.sdk.execution_time.schema.versions import bundle + +if TYPE_CHECKING: + from collections.abc import Iterable + + +def _iter_model_types(t: object) -> Iterable[type[BaseModel]]: + """Generate BaseModel subclass included in type.""" + origin = get_origin(t) + if origin is Annotated: + yield from _iter_model_types(get_args(t)[0]) + elif origin is UnionType: + for member in get_args(t): + yield from _iter_model_types(member) + elif isinstance(t, type) and issubclass(t, BaseModel): + yield t + + +@functools.cache +def registered_models_by_name() -> dict[str, type[BaseModel]]: + """ + Map every supervisor schema body's class name to the head Pydantic class. + + Single source of truth for the registry. Built once by walking the + four discriminated unions the supervisor decodes against; cached + per-process because the registry only changes when a union member + is added in ``comms.py`` or ``processor.py`` (which needs a + restart anyway). :func:`resolve_body_class` looks up the wire-shape + ``type`` discriminator against it. + + Imports are deferred so this package stays cheap to import for + callers that only need the bundle or migrator (e.g. the migrator + singleton factory); pulling in ``processor`` eagerly would drag the + whole DAG-processor import graph into every consumer. + + Raises ``RuntimeError`` if two distinct classes register under the + same ``__name__`` -- the wire discriminator must round-trip to a + single head class, so a name clash is a programmer error that must + surface immediately rather than silently picking a winner. + """ + from airflow.dag_processing.processor import ToDagProcessor, ToManager + from airflow.sdk.execution_time.comms import ToSupervisor, ToTask + + by_name: dict[str, type[BaseModel]] = {} + for source in (ToTask, ToSupervisor, ToManager, ToDagProcessor): + for model in _iter_model_types(source): + existing = by_name.get(model.__name__) + if existing is None: + by_name[model.__name__] = model + elif existing is not model: + raise RuntimeError( + f"Duplicate supervisor schema body name {model.__name__!r}: " + f"both {existing!r} and {model!r} register the same wire type" + ) + return by_name + + +def resolve_body_class(body: Any) -> type[BaseModel] | None: + """Resolve a wire-body dict's ``type`` discriminator to its head Pydantic class.""" + if not isinstance(body, dict): + return None + name = body.get("type") + if not isinstance(name, str): + return None + return registered_models_by_name().get(name) + + +__all__ = [ + "SchemaVersionMigrator", + "bundle", + "get_schema_version_migrator", + "registered_models_by_name", + "resolve_body_class", +] diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/migrator.py b/task-sdk/src/airflow/sdk/execution_time/schema/migrator.py new file mode 100644 index 0000000000000..f68e192604d3c --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/schema/migrator.py @@ -0,0 +1,211 @@ +# 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. +""" +In-process bidirectional migration for supervisor schema bodies. + +:class:`SchemaVersionMigrator` walks a :class:`~cadwyn.VersionBundle` +itself rather than going through Cadwyn's HTTP runner so the supervisor +can downgrade outgoing bodies and upgrade incoming bodies without a +network round-trip. The downgrade path additionally re-validates against +the cadwyn-generated versioned class so declarative +``schema(X).field(Y).didnt_exist`` instructions actually drop fields on +the wire. +""" + +from __future__ import annotations + +import functools +from typing import TYPE_CHECKING, Any, cast + +import attrs +from cadwyn import generate_versioned_models + +if TYPE_CHECKING: + from cadwyn import VersionBundle + from cadwyn.schema_generation import SchemaGenerator + from pydantic import BaseModel + + +@attrs.define +class _BodyInfo: + """ + Duck-type stand-in for Cadwyn's ``RequestInfo`` / ``ResponseInfo``. + + ``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 _validate_supervisor_version(instance: SchemaVersionMigrator, _, value: str) -> str: + return instance._resolve_version(value) + + +def _calculate_version_values(migrator: SchemaVersionMigrator) -> frozenset[str]: + return frozenset(v.value for v in migrator._bundle.versions) + + +@attrs.define(kw_only=True) +class SchemaVersionMigrator: + """ + Bidirectional in-process migrator for supervisor schema bodies. + + Each foreign runtime is pinned to a specific dated lang-SDK supervisor + schema version; this class walks Cadwyn's ``VersionChange`` chain in-process + to bridge the two:: + + head shape --- downgrade(msg, lang_sdk) ---> lang-SDK wire + head shape <-- upgrade(msg, lang_sdk) --- lang-SDK wire + + *supervisor_version* is fixed at construction time. + + note:: + Use ``bundle.versions[0].value`` to get the latest dated entry. Cadwyn + keeps versions in newest-to-oldest order. + + A message whose Pydantic type is not mentioned by any ``schema(...)`` + instruction in the bundle is passed through as-is: Cadwyn keys its + instruction dicts by message type, so the lookup misses and no + transformer runs. + """ + + _bundle: VersionBundle + _supervisor_version: str = attrs.field(validator=_validate_supervisor_version) + + # Caches over the bundle (which is immutable for the migrator's lifetime). + # ``generate_versioned_models`` walks the full version graph; + # ``_version_values`` mirrors cadwyn's internal lookup set without reaching + # into its private attribute. + _versioned_models: dict[str, SchemaGenerator] = attrs.field(init=False, default=None) + _version_values: frozenset[str] = attrs.field( + init=False, + default=attrs.Factory(_calculate_version_values, takes_self=True), + ) + + def _versioned_class(self, version: str, model: type[BaseModel]) -> type[BaseModel]: + """Get the Cadwyn-generated class for *model* at *version*.""" + if self._versioned_models is None: + self._versioned_models = generate_versioned_models(self._bundle) + return self._versioned_models[version][model] + + def _resolve_version(self, v: str) -> str: + """Validate *v* is present in the bundle.""" + if v not in self._version_values: + raise ValueError(f"Version {v!r} not found in supervisor schema bundle") + return v + + def downgrade( + self, + msg: BaseModel, + target_schema_version: str, + **dump_opts, + ) -> BaseModel: + """ + Downgrade *msg* from server to *target_schema_version*. + + Used on the supervisor -> foreign-runtime path: *msg* is a head-shape + Pydantic instance, and the returned dict matches the target. + + :param msg: A Pydantic instance shaped according to the head + (latest) version of the bundle. + :param target_schema_version: Dated supervisor schema version string in + ``YYYY-MM-DD`` format. Must be an exact value in the bundle. + :param dump_opts: Forwarded to ``model_dump`` when dumping *msg* for + migration. The mode is already set to ``json`` so datetime/UUID/Path + etc. serialize to primitives the versioned-model validators inside + the chain accept. + :returns: A Pydantic instance shaped to the target version. The type + of this object is dynamically generated by Cadwyn. + """ + model = type(msg) + target_schema_version = self._resolve_version(target_schema_version) + info = _BodyInfo(msg.model_dump(**cast("dict[str, Any]", {**dump_opts, "mode": "json"}))) + for version in self._bundle.versions: + if version.value > self._supervisor_version: + continue + if version.value <= target_schema_version: + break + for change in version.changes: + for instr in change.alter_response_by_schema_instructions.get(model, ()): + # TODO: Cadwyn is tightly coupled to Startlette request and + # response objects. Our supervisor does not use an HTTP + # framework, so we need to mock out the object. Fix this + # when Cadwyn provides a framework-agnostic interface. + instr(info) # type: ignore[arg-type] + # Re-validate against the versioned class so schema(X).field(Y).didnt_exist + # instructions take effect: those alter the class shape, not the dict, so + # without this round-trip the dropped field would still appear on the wire. + versioned_class = self._versioned_class(target_schema_version, model) + return versioned_class.model_validate(info.body) + + def upgrade( + self, + body: dict[str, Any], + model: type[BaseModel], + source_schema_version: str, + ) -> dict[str, Any]: + """ + Upgrade *body* from *source_schema_version* to the supervisor's shape. + + Used on the foreign-runtime -> supervisor path: *body* is the + already-deserialized payload off the wire (still in the lang-SDK's + schema), and the returned dict is shaped for ``model_validate`` + against the head Pydantic class. + + *model* must be supplied because a dict carries no Python type + information; the caller resolves it from the discriminator + (``body["type"]``) and the registered-models index. + + :param body: The wire payload as a dict. + :param model: The server-side Pydantic class *body* should validate + against after migration. + :param source_schema_version: Dated supervisor schema version *body* is + in. This should be a string in ``YYYY-MM-DD`` format. + """ + source_schema_version = self._resolve_version(source_schema_version) + info = _BodyInfo(body) + for version in self._bundle.reversed_versions: + if version.value <= source_schema_version: + continue + if version.value > self._supervisor_version: + continue + for change in version.changes: + for instr in change.alter_request_by_schema_instructions.get(model, ()): + instr(info) # type: ignore[arg-type] + versioned_class = self._versioned_class(self._supervisor_version, model) + return versioned_class.model_validate(info.body).model_dump() + + +@functools.cache +def get_schema_version_migrator() -> SchemaVersionMigrator: + """ + Return the process-wide :class:`SchemaVersionMigrator` bound to the supervisor bundle. + + Cached so the bundle is bound once per process. The migrator holds + no per-call state, so concurrent callers can share a single + instance safely. + """ + from airflow.sdk.execution_time.schema.versions import bundle + + return SchemaVersionMigrator(bundle=bundle, supervisor_version=bundle.versions[0].value) + + +__all__ = ["SchemaVersionMigrator", "get_schema_version_migrator"] diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json new file mode 100644 index 0000000000000..d4eb3d9c5a8b7 --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -0,0 +1,6131 @@ +{ + "api_version": "2026-06-16", + "schemas": { + "AssetEventsResult": { + "$defs": { + "AssetEventResponse": { + "description": "Asset event schema with fields that are needed for Runtime.", + "properties": { + "asset": { + "$ref": "#/$defs/AssetResponse" + }, + "created_dagruns": { + "items": { + "$ref": "#/$defs/DagRunAssetReference" + }, + "title": "Created Dagruns", + "type": "array" + }, + "extra": { + "anyOf": [ + { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Extra" + }, + "id": { + "title": "Id", + "type": "integer" + }, + "partition_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key" + }, + "source_dag_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Source Dag Id" + }, + "source_map_index": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Source Map Index" + }, + "source_run_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Source Run Id" + }, + "source_task_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Source Task Id" + }, + "timestamp": { + "format": "date-time", + "title": "Timestamp", + "type": "string" + } + }, + "required": [ + "id", + "timestamp", + "asset", + "created_dagruns" + ], + "title": "AssetEventResponse", + "type": "object" + }, + "AssetResponse": { + "description": "Asset schema for responses with fields that are needed for Runtime.", + "properties": { + "extra": { + "anyOf": [ + { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Extra" + }, + "group": { + "title": "Group", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "uri": { + "title": "Uri", + "type": "string" + } + }, + "required": [ + "name", + "uri", + "group" + ], + "title": "AssetResponse", + "type": "object" + }, + "DagRunAssetReference": { + "additionalProperties": false, + "description": "DagRun serializer for asset responses.", + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "data_interval_end": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Data Interval End" + }, + "data_interval_start": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Data Interval Start" + }, + "end_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "End Date" + }, + "logical_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Logical Date" + }, + "partition_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "start_date": { + "format": "date-time", + "title": "Start Date", + "type": "string" + }, + "state": { + "title": "State", + "type": "string" + } + }, + "required": [ + "run_id", + "dag_id", + "start_date", + "state" + ], + "title": "DagRunAssetReference", + "type": "object" + }, + "JsonValue": {} + }, + "description": "Response to GetAssetEvent request.", + "properties": { + "asset_events": { + "items": { + "$ref": "#/$defs/AssetEventResponse" + }, + "title": "Asset Events", + "type": "array" + }, + "type": { + "const": "AssetEventsResult", + "default": "AssetEventsResult", + "title": "Type", + "type": "string" + } + }, + "required": [ + "asset_events" + ], + "title": "AssetEventsResult", + "type": "object" + }, + "AssetResult": { + "$defs": { + "JsonValue": {} + }, + "description": "Response to ReadXCom request.", + "properties": { + "extra": { + "anyOf": [ + { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Extra" + }, + "group": { + "title": "Group", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "type": { + "const": "AssetResult", + "default": "AssetResult", + "title": "Type", + "type": "string" + }, + "uri": { + "title": "Uri", + "type": "string" + } + }, + "required": [ + "name", + "uri", + "group" + ], + "title": "AssetResult", + "type": "object" + }, + "AssetStateResult": { + "additionalProperties": false, + "description": "Response to GetAssetState; wraps the generated API response for supervisor to worker comms.", + "properties": { + "type": { + "const": "AssetStateResult", + "default": "AssetStateResult", + "title": "Type", + "type": "string" + }, + "value": { + "title": "Value", + "type": "string" + } + }, + "required": [ + "value" + ], + "title": "AssetStateResult", + "type": "object" + }, + "AssetsByAliasResult": { + "$defs": { + "AssetResult": { + "description": "Response to ReadXCom request.", + "properties": { + "extra": { + "anyOf": [ + { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Extra" + }, + "group": { + "title": "Group", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "type": { + "const": "AssetResult", + "default": "AssetResult", + "title": "Type", + "type": "string" + }, + "uri": { + "title": "Uri", + "type": "string" + } + }, + "required": [ + "name", + "uri", + "group" + ], + "title": "AssetResult", + "type": "object" + }, + "JsonValue": {} + }, + "description": "Response to GetAssetsByAlias; list of concrete assets resolved from an alias.", + "properties": { + "assets": { + "items": { + "$ref": "#/$defs/AssetResult" + }, + "title": "Assets", + "type": "array" + }, + "type": { + "const": "AssetsByAliasResult", + "default": "AssetsByAliasResult", + "title": "Type", + "type": "string" + } + }, + "required": [ + "assets" + ], + "title": "AssetsByAliasResult", + "type": "object" + }, + "ClearAssetStateByName": { + "properties": { + "name": { + "title": "Name", + "type": "string" + }, + "type": { + "const": "ClearAssetStateByName", + "default": "ClearAssetStateByName", + "title": "Type", + "type": "string" + } + }, + "required": [ + "name" + ], + "title": "ClearAssetStateByName", + "type": "object" + }, + "ClearAssetStateByUri": { + "properties": { + "type": { + "const": "ClearAssetStateByUri", + "default": "ClearAssetStateByUri", + "title": "Type", + "type": "string" + }, + "uri": { + "title": "Uri", + "type": "string" + } + }, + "required": [ + "uri" + ], + "title": "ClearAssetStateByUri", + "type": "object" + }, + "ClearTaskState": { + "properties": { + "all_map_indices": { + "default": false, + "title": "All Map Indices", + "type": "boolean" + }, + "ti_id": { + "format": "uuid", + "title": "Ti Id", + "type": "string" + }, + "type": { + "const": "ClearTaskState", + "default": "ClearTaskState", + "title": "Type", + "type": "string" + } + }, + "required": [ + "ti_id" + ], + "title": "ClearTaskState", + "type": "object" + }, + "ConnectionResult": { + "properties": { + "conn_id": { + "title": "Conn Id", + "type": "string" + }, + "conn_type": { + "title": "Conn Type", + "type": "string" + }, + "extra": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Extra" + }, + "host": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Host" + }, + "login": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Login" + }, + "password": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Password" + }, + "port": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Port" + }, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Schema" + }, + "type": { + "const": "ConnectionResult", + "default": "ConnectionResult", + "title": "Type", + "type": "string" + } + }, + "required": [ + "conn_id", + "conn_type" + ], + "title": "ConnectionResult", + "type": "object" + }, + "CreateHITLDetailPayload": { + "$defs": { + "HITLUser": { + "description": "Schema for a Human-in-the-loop users.", + "properties": { + "id": { + "title": "Id", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "id", + "name" + ], + "title": "HITLUser", + "type": "object" + } + }, + "description": "Add the input request part of a Human-in-the-loop response.", + "properties": { + "assigned_users": { + "anyOf": [ + { + "items": { + "$ref": "#/$defs/HITLUser" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Assigned Users" + }, + "body": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Body" + }, + "defaults": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Defaults" + }, + "multiple": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "default": false, + "title": "Multiple" + }, + "options": { + "items": { + "type": "string" + }, + "minItems": 1, + "title": "Options", + "type": "array" + }, + "params": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Params" + }, + "subject": { + "title": "Subject", + "type": "string" + }, + "ti_id": { + "format": "uuid", + "title": "Ti Id", + "type": "string" + }, + "type": { + "const": "CreateHITLDetailPayload", + "default": "CreateHITLDetailPayload", + "title": "Type", + "type": "string" + } + }, + "required": [ + "ti_id", + "options", + "subject" + ], + "title": "CreateHITLDetailPayload", + "type": "object" + }, + "DRCount": { + "description": "Response containing count of Dag Runs matching certain filters.", + "properties": { + "count": { + "title": "Count", + "type": "integer" + }, + "type": { + "const": "DRCount", + "default": "DRCount", + "title": "Type", + "type": "string" + } + }, + "required": [ + "count" + ], + "title": "DRCount", + "type": "object" + }, + "DagFileParseRequest": { + "$defs": { + "AssetAliasReferenceAssetEventDagRun": { + "additionalProperties": false, + "description": "Schema for AssetAliasModel used in AssetEventDagRunReference.", + "properties": { + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "name" + ], + "title": "AssetAliasReferenceAssetEventDagRun", + "type": "object" + }, + "AssetEventDagRunReference": { + "additionalProperties": false, + "description": "Schema for AssetEvent model used in DagRun.", + "properties": { + "asset": { + "$ref": "#/$defs/AssetReferenceAssetEventDagRun" + }, + "extra": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "Extra", + "type": "object" + }, + "source_aliases": { + "items": { + "$ref": "#/$defs/AssetAliasReferenceAssetEventDagRun" + }, + "title": "Source Aliases", + "type": "array" + }, + "source_dag_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Source Dag Id" + }, + "source_map_index": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Source Map Index" + }, + "source_run_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Source Run Id" + }, + "source_task_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Source Task Id" + }, + "timestamp": { + "format": "date-time", + "title": "Timestamp", + "type": "string" + } + }, + "required": [ + "asset", + "extra", + "source_task_id", + "source_dag_id", + "source_run_id", + "source_map_index", + "source_aliases", + "timestamp" + ], + "title": "AssetEventDagRunReference", + "type": "object" + }, + "AssetReferenceAssetEventDagRun": { + "additionalProperties": false, + "description": "Schema for AssetModel used in AssetEventDagRunReference.", + "properties": { + "extra": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "Extra", + "type": "object" + }, + "name": { + "title": "Name", + "type": "string" + }, + "uri": { + "title": "Uri", + "type": "string" + } + }, + "required": [ + "name", + "uri", + "extra" + ], + "title": "AssetReferenceAssetEventDagRun", + "type": "object" + }, + "ConnectionResponse": { + "description": "Connection schema for responses with fields that are needed for Runtime.", + "properties": { + "conn_id": { + "title": "Conn Id", + "type": "string" + }, + "conn_type": { + "title": "Conn Type", + "type": "string" + }, + "extra": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Extra" + }, + "host": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Host" + }, + "login": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Login" + }, + "password": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Password" + }, + "port": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Port" + }, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Schema" + } + }, + "required": [ + "conn_id", + "conn_type", + "host", + "schema", + "login", + "password", + "port", + "extra" + ], + "title": "ConnectionResponse", + "type": "object" + }, + "DagCallbackRequest": { + "description": "A Class with information about the success/failure DAG callback to be executed.", + "properties": { + "bundle_name": { + "title": "Bundle Name", + "type": "string" + }, + "bundle_version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Bundle Version" + }, + "context_from_server": { + "anyOf": [ + { + "$ref": "#/$defs/DagRunContext" + }, + { + "type": "null" + } + ], + "default": null + }, + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "filepath": { + "title": "Filepath", + "type": "string" + }, + "is_failure_callback": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "default": true, + "title": "Is Failure Callback" + }, + "msg": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Msg" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "type": { + "const": "DagCallbackRequest", + "default": "DagCallbackRequest", + "title": "Type", + "type": "string" + } + }, + "required": [ + "filepath", + "bundle_name", + "bundle_version", + "dag_id", + "run_id" + ], + "title": "DagCallbackRequest", + "type": "object" + }, + "DagRun": { + "additionalProperties": false, + "description": "Schema for DagRun model with minimal required fields needed for Runtime.", + "properties": { + "clear_number": { + "default": 0, + "title": "Clear Number", + "type": "integer" + }, + "conf": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Conf" + }, + "consumed_asset_events": { + "items": { + "$ref": "#/$defs/AssetEventDagRunReference" + }, + "title": "Consumed Asset Events", + "type": "array" + }, + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "data_interval_end": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Data Interval End" + }, + "data_interval_start": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Data Interval Start" + }, + "end_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "End Date" + }, + "logical_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Logical Date" + }, + "note": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Note" + }, + "partition_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Partition Key" + }, + "run_after": { + "format": "date-time", + "title": "Run After", + "type": "string" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "run_type": { + "$ref": "#/$defs/DagRunType" + }, + "start_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Start Date" + }, + "state": { + "$ref": "#/$defs/DagRunState" + }, + "team_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Team Name" + }, + "triggering_user_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Triggering User Name" + } + }, + "required": [ + "dag_id", + "run_id", + "logical_date", + "data_interval_start", + "data_interval_end", + "run_after", + "start_date", + "end_date", + "run_type", + "state", + "consumed_asset_events", + "partition_key" + ], + "title": "DagRun", + "type": "object" + }, + "DagRunContext": { + "description": "Class to pass context info from the server to build a Execution context object.", + "properties": { + "dag_run": { + "anyOf": [ + { + "$ref": "#/$defs/DagRun" + }, + { + "type": "null" + } + ], + "default": null + }, + "last_ti": { + "anyOf": [ + { + "$ref": "#/$defs/TaskInstance" + }, + { + "type": "null" + } + ], + "default": null + } + }, + "title": "DagRunContext", + "type": "object" + }, + "DagRunState": { + "description": "All possible states that a DagRun can be in.\n\nThese are \"shared\" with TaskInstanceState in some parts of the code,\nso please ensure that their values always match the ones with the\nsame name in TaskInstanceState.", + "enum": [ + "queued", + "running", + "success", + "failed" + ], + "title": "DagRunState", + "type": "string" + }, + "DagRunType": { + "description": "Class with DagRun types.", + "enum": [ + "backfill", + "scheduled", + "manual", + "operator_triggered", + "asset_triggered", + "asset_materialization" + ], + "title": "DagRunType", + "type": "string" + }, + "EmailRequest": { + "description": "Email notification request for task failures/retries.", + "properties": { + "bundle_name": { + "title": "Bundle Name", + "type": "string" + }, + "bundle_version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Bundle Version" + }, + "context_from_server": { + "$ref": "#/$defs/TIRunContext" + }, + "email_type": { + "default": "failure", + "enum": [ + "failure", + "retry" + ], + "title": "Email Type", + "type": "string" + }, + "filepath": { + "title": "Filepath", + "type": "string" + }, + "msg": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Msg" + }, + "ti": { + "$ref": "#/$defs/TaskInstance" + }, + "type": { + "const": "EmailRequest", + "default": "EmailRequest", + "title": "Type", + "type": "string" + } + }, + "required": [ + "filepath", + "bundle_name", + "bundle_version", + "ti", + "context_from_server" + ], + "title": "EmailRequest", + "type": "object" + }, + "JsonValue": {}, + "TIRunContext": { + "description": "Response schema for TaskInstance run context.", + "properties": { + "connections": { + "items": { + "$ref": "#/$defs/ConnectionResponse" + }, + "title": "Connections", + "type": "array" + }, + "dag_run": { + "$ref": "#/$defs/DagRun" + }, + "max_tries": { + "title": "Max Tries", + "type": "integer" + }, + "next_kwargs": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Next Kwargs" + }, + "next_method": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Next Method" + }, + "should_retry": { + "default": false, + "title": "Should Retry", + "type": "boolean" + }, + "start_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Start Date" + }, + "task_reschedule_count": { + "default": 0, + "title": "Task Reschedule Count", + "type": "integer" + }, + "variables": { + "items": { + "$ref": "#/$defs/VariableResponse" + }, + "title": "Variables", + "type": "array" + }, + "xcom_keys_to_clear": { + "items": { + "type": "string" + }, + "title": "Xcom Keys To Clear", + "type": "array" + } + }, + "required": [ + "dag_run", + "max_tries" + ], + "title": "TIRunContext", + "type": "object" + }, + "TaskCallbackRequest": { + "description": "Task callback status information.\n\nA Class with information about the success/failure TI callback to be executed. Currently, only failure\ncallbacks when tasks are externally killed or experience heartbeat timeouts are run via DagFileProcessorProcess.", + "properties": { + "bundle_name": { + "title": "Bundle Name", + "type": "string" + }, + "bundle_version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Bundle Version" + }, + "context_from_server": { + "anyOf": [ + { + "$ref": "#/$defs/TIRunContext" + }, + { + "type": "null" + } + ], + "default": null + }, + "filepath": { + "title": "Filepath", + "type": "string" + }, + "msg": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Msg" + }, + "task_callback_type": { + "anyOf": [ + { + "$ref": "#/$defs/TaskInstanceState" + }, + { + "type": "null" + } + ], + "default": null + }, + "ti": { + "$ref": "#/$defs/TaskInstance" + }, + "type": { + "const": "TaskCallbackRequest", + "default": "TaskCallbackRequest", + "title": "Type", + "type": "string" + } + }, + "required": [ + "filepath", + "bundle_name", + "bundle_version", + "ti" + ], + "title": "TaskCallbackRequest", + "type": "object" + }, + "TaskInstance": { + "description": "Schema for TaskInstance model with minimal required fields needed for Runtime.", + "properties": { + "context_carrier": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Context Carrier" + }, + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "dag_version_id": { + "format": "uuid", + "title": "Dag Version Id", + "type": "string" + }, + "hostname": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Hostname" + }, + "id": { + "format": "uuid", + "title": "Id", + "type": "string" + }, + "map_index": { + "default": -1, + "title": "Map Index", + "type": "integer" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "task_id": { + "title": "Task Id", + "type": "string" + }, + "try_number": { + "title": "Try Number", + "type": "integer" + } + }, + "required": [ + "id", + "task_id", + "dag_id", + "run_id", + "try_number", + "dag_version_id" + ], + "title": "TaskInstance", + "type": "object" + }, + "TaskInstanceState": { + "description": "All possible states that a Task Instance can be in.\n\nNote that None is also allowed, so always use this in a type hint with Optional.", + "enum": [ + "removed", + "scheduled", + "queued", + "running", + "success", + "restarting", + "failed", + "up_for_retry", + "up_for_reschedule", + "upstream_failed", + "skipped", + "deferred" + ], + "title": "TaskInstanceState", + "type": "string" + }, + "VariableResponse": { + "additionalProperties": false, + "description": "Variable schema for responses with fields that are needed for Runtime.", + "properties": { + "key": { + "title": "Key", + "type": "string" + }, + "value": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Value" + } + }, + "required": [ + "key", + "value" + ], + "title": "VariableResponse", + "type": "object" + } + }, + "description": "Request for DAG File Parsing.\n\nThis is the request that the manager will send to the DAG parser with the dag file and\nany other necessary metadata.", + "properties": { + "bundle_name": { + "title": "Bundle Name", + "type": "string" + }, + "bundle_path": { + "format": "path", + "title": "Bundle Path", + "type": "string" + }, + "callback_requests": { + "items": { + "discriminator": { + "mapping": { + "DagCallbackRequest": "#/$defs/DagCallbackRequest", + "EmailRequest": "#/$defs/EmailRequest", + "TaskCallbackRequest": "#/$defs/TaskCallbackRequest" + }, + "propertyName": "type" + }, + "oneOf": [ + { + "$ref": "#/$defs/DagCallbackRequest" + }, + { + "$ref": "#/$defs/TaskCallbackRequest" + }, + { + "$ref": "#/$defs/EmailRequest" + } + ] + }, + "title": "Callback Requests", + "type": "array" + }, + "file": { + "title": "File", + "type": "string" + }, + "type": { + "const": "DagFileParseRequest", + "default": "DagFileParseRequest", + "title": "Type", + "type": "string" + } + }, + "required": [ + "file", + "bundle_path", + "bundle_name" + ], + "title": "DagFileParseRequest", + "type": "object" + }, + "DagFileParsingResult": { + "$defs": { + "LazyDeserializedDAG": { + "description": "Lazily build information from the serialized DAG structure.\n\nAn object that will present \"enough\" of the DAG like interface to update DAG db models etc, without having\nto deserialize the full DAG and Task hierarchy.", + "properties": { + "data": { + "additionalProperties": true, + "title": "Data", + "type": "object" + }, + "last_loaded": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Last Loaded" + } + }, + "required": [ + "data" + ], + "title": "LazyDeserializedDAG", + "type": "object" + } + }, + "description": "Result of DAG File Parsing.\n\nThis is the result of a successful DAG parse, in this class, we gather all serialized DAGs,\nimport errors and warnings to send back to the scheduler to store in the DB.", + "properties": { + "fileloc": { + "title": "Fileloc", + "type": "string" + }, + "import_errors": { + "anyOf": [ + { + "additionalProperties": { + "type": "string" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Import Errors" + }, + "serialized_dags": { + "items": { + "$ref": "#/$defs/LazyDeserializedDAG" + }, + "title": "Serialized Dags", + "type": "array" + }, + "type": { + "const": "DagFileParsingResult", + "default": "DagFileParsingResult", + "title": "Type", + "type": "string" + }, + "warnings": { + "anyOf": [ + { + "items": {}, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Warnings" + } + }, + "required": [ + "fileloc", + "serialized_dags" + ], + "title": "DagFileParsingResult", + "type": "object" + }, + "DagResult": { + "properties": { + "bundle_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Bundle Name" + }, + "bundle_version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Bundle Version" + }, + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "is_paused": { + "title": "Is Paused", + "type": "boolean" + }, + "next_dagrun": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Next Dagrun" + }, + "owners": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Owners" + }, + "relative_fileloc": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Relative Fileloc" + }, + "tags": { + "items": { + "type": "string" + }, + "title": "Tags", + "type": "array" + }, + "type": { + "const": "DagResult", + "default": "DagResult", + "title": "Type", + "type": "string" + } + }, + "required": [ + "dag_id", + "is_paused", + "tags" + ], + "title": "DagResult", + "type": "object" + }, + "DagRunResult": { + "$defs": { + "AssetAliasReferenceAssetEventDagRun": { + "additionalProperties": false, + "description": "Schema for AssetAliasModel used in AssetEventDagRunReference.", + "properties": { + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "name" + ], + "title": "AssetAliasReferenceAssetEventDagRun", + "type": "object" + }, + "AssetEventDagRunReference": { + "additionalProperties": false, + "description": "Schema for AssetEvent model used in DagRun.", + "properties": { + "asset": { + "$ref": "#/$defs/AssetReferenceAssetEventDagRun" + }, + "extra": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "Extra", + "type": "object" + }, + "source_aliases": { + "items": { + "$ref": "#/$defs/AssetAliasReferenceAssetEventDagRun" + }, + "title": "Source Aliases", + "type": "array" + }, + "source_dag_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Source Dag Id" + }, + "source_map_index": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Source Map Index" + }, + "source_run_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Source Run Id" + }, + "source_task_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Source Task Id" + }, + "timestamp": { + "format": "date-time", + "title": "Timestamp", + "type": "string" + } + }, + "required": [ + "asset", + "extra", + "source_aliases", + "timestamp" + ], + "title": "AssetEventDagRunReference", + "type": "object" + }, + "AssetReferenceAssetEventDagRun": { + "additionalProperties": false, + "description": "Schema for AssetModel used in AssetEventDagRunReference.", + "properties": { + "extra": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "Extra", + "type": "object" + }, + "name": { + "title": "Name", + "type": "string" + }, + "uri": { + "title": "Uri", + "type": "string" + } + }, + "required": [ + "name", + "uri", + "extra" + ], + "title": "AssetReferenceAssetEventDagRun", + "type": "object" + }, + "DagRunState": { + "description": "All possible states that a DagRun can be in.\n\nThese are \"shared\" with TaskInstanceState in some parts of the code,\nso please ensure that their values always match the ones with the\nsame name in TaskInstanceState.", + "enum": [ + "queued", + "running", + "success", + "failed" + ], + "title": "DagRunState", + "type": "string" + }, + "DagRunType": { + "description": "Class with DagRun types.", + "enum": [ + "backfill", + "scheduled", + "manual", + "operator_triggered", + "asset_triggered", + "asset_materialization" + ], + "title": "DagRunType", + "type": "string" + }, + "JsonValue": {} + }, + "additionalProperties": false, + "properties": { + "clear_number": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": 0, + "title": "Clear Number" + }, + "conf": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Conf" + }, + "consumed_asset_events": { + "items": { + "$ref": "#/$defs/AssetEventDagRunReference" + }, + "title": "Consumed Asset Events", + "type": "array" + }, + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "data_interval_end": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Data Interval End" + }, + "data_interval_start": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Data Interval Start" + }, + "end_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "End Date" + }, + "logical_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Logical Date" + }, + "note": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Note" + }, + "partition_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key" + }, + "run_after": { + "format": "date-time", + "title": "Run After", + "type": "string" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "run_type": { + "$ref": "#/$defs/DagRunType" + }, + "start_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Start Date" + }, + "state": { + "$ref": "#/$defs/DagRunState" + }, + "team_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Team Name" + }, + "triggering_user_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Triggering User Name" + }, + "type": { + "const": "DagRunResult", + "default": "DagRunResult", + "title": "Type", + "type": "string" + } + }, + "required": [ + "dag_id", + "run_id", + "run_after", + "run_type", + "state", + "consumed_asset_events" + ], + "title": "DagRunResult", + "type": "object" + }, + "DagRunStateResult": { + "$defs": { + "DagRunState": { + "description": "All possible states that a DagRun can be in.\n\nThese are \"shared\" with TaskInstanceState in some parts of the code,\nso please ensure that their values always match the ones with the\nsame name in TaskInstanceState.", + "enum": [ + "queued", + "running", + "success", + "failed" + ], + "title": "DagRunState", + "type": "string" + } + }, + "properties": { + "state": { + "$ref": "#/$defs/DagRunState" + }, + "type": { + "const": "DagRunStateResult", + "default": "DagRunStateResult", + "title": "Type", + "type": "string" + } + }, + "required": [ + "state" + ], + "title": "DagRunStateResult", + "type": "object" + }, + "DeferTask": { + "$defs": { + "JsonValue": {} + }, + "additionalProperties": false, + "description": "Update a task instance state to deferred.", + "properties": { + "classpath": { + "title": "Classpath", + "type": "string" + }, + "next_kwargs": { + "anyOf": [ + { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Next Kwargs" + }, + "next_method": { + "title": "Next Method", + "type": "string" + }, + "queue": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Queue" + }, + "rendered_map_index": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Rendered Map Index" + }, + "state": { + "anyOf": [ + { + "const": "deferred", + "type": "string" + }, + { + "type": "null" + } + ], + "default": "deferred", + "title": "State" + }, + "trigger_kwargs": { + "anyOf": [ + { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "type": "object" + }, + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Trigger Kwargs" + }, + "trigger_timeout": { + "anyOf": [ + { + "format": "duration", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Trigger Timeout" + }, + "type": { + "const": "DeferTask", + "default": "DeferTask", + "title": "Type", + "type": "string" + } + }, + "required": [ + "classpath", + "next_method" + ], + "title": "DeferTask", + "type": "object" + }, + "DeleteAssetStateByName": { + "properties": { + "key": { + "title": "Key", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "type": { + "const": "DeleteAssetStateByName", + "default": "DeleteAssetStateByName", + "title": "Type", + "type": "string" + } + }, + "required": [ + "name", + "key" + ], + "title": "DeleteAssetStateByName", + "type": "object" + }, + "DeleteAssetStateByUri": { + "properties": { + "key": { + "title": "Key", + "type": "string" + }, + "type": { + "const": "DeleteAssetStateByUri", + "default": "DeleteAssetStateByUri", + "title": "Type", + "type": "string" + }, + "uri": { + "title": "Uri", + "type": "string" + } + }, + "required": [ + "uri", + "key" + ], + "title": "DeleteAssetStateByUri", + "type": "object" + }, + "DeleteTaskState": { + "properties": { + "key": { + "title": "Key", + "type": "string" + }, + "ti_id": { + "format": "uuid", + "title": "Ti Id", + "type": "string" + }, + "type": { + "const": "DeleteTaskState", + "default": "DeleteTaskState", + "title": "Type", + "type": "string" + } + }, + "required": [ + "ti_id", + "key" + ], + "title": "DeleteTaskState", + "type": "object" + }, + "DeleteVariable": { + "properties": { + "key": { + "title": "Key", + "type": "string" + }, + "type": { + "const": "DeleteVariable", + "default": "DeleteVariable", + "title": "Type", + "type": "string" + } + }, + "required": [ + "key" + ], + "title": "DeleteVariable", + "type": "object" + }, + "DeleteXCom": { + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "key": { + "title": "Key", + "type": "string" + }, + "map_index": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Map Index" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "task_id": { + "title": "Task Id", + "type": "string" + }, + "type": { + "const": "DeleteXCom", + "default": "DeleteXCom", + "title": "Type", + "type": "string" + } + }, + "required": [ + "key", + "dag_id", + "run_id", + "task_id" + ], + "title": "DeleteXCom", + "type": "object" + }, + "ErrorResponse": { + "$defs": { + "ErrorType": { + "description": "Error types used in the API client.", + "enum": [ + "CONNECTION_NOT_FOUND", + "VARIABLE_NOT_FOUND", + "XCOM_NOT_FOUND", + "ASSET_NOT_FOUND", + "TASK_STATE_NOT_FOUND", + "ASSET_STATE_NOT_FOUND", + "DAGRUN_ALREADY_EXISTS", + "PERMISSION_DENIED", + "GENERIC_ERROR", + "API_SERVER_ERROR" + ], + "title": "ErrorType", + "type": "string" + } + }, + "properties": { + "detail": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Detail" + }, + "error": { + "$ref": "#/$defs/ErrorType", + "default": "GENERIC_ERROR" + }, + "type": { + "const": "ErrorResponse", + "default": "ErrorResponse", + "title": "Type", + "type": "string" + } + }, + "title": "ErrorResponse", + "type": "object" + }, + "GetAssetByName": { + "properties": { + "name": { + "title": "Name", + "type": "string" + }, + "type": { + "const": "GetAssetByName", + "default": "GetAssetByName", + "title": "Type", + "type": "string" + } + }, + "required": [ + "name" + ], + "title": "GetAssetByName", + "type": "object" + }, + "GetAssetByUri": { + "properties": { + "type": { + "const": "GetAssetByUri", + "default": "GetAssetByUri", + "title": "Type", + "type": "string" + }, + "uri": { + "title": "Uri", + "type": "string" + } + }, + "required": [ + "uri" + ], + "title": "GetAssetByUri", + "type": "object" + }, + "GetAssetEventByAsset": { + "properties": { + "after": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "After" + }, + "ascending": { + "default": true, + "title": "Ascending", + "type": "boolean" + }, + "before": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Before" + }, + "limit": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Limit" + }, + "name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Name" + }, + "type": { + "const": "GetAssetEventByAsset", + "default": "GetAssetEventByAsset", + "title": "Type", + "type": "string" + }, + "uri": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Uri" + } + }, + "required": [ + "name", + "uri" + ], + "title": "GetAssetEventByAsset", + "type": "object" + }, + "GetAssetEventByAssetAlias": { + "properties": { + "after": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "After" + }, + "alias_name": { + "title": "Alias Name", + "type": "string" + }, + "ascending": { + "default": true, + "title": "Ascending", + "type": "boolean" + }, + "before": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Before" + }, + "limit": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Limit" + }, + "type": { + "const": "GetAssetEventByAssetAlias", + "default": "GetAssetEventByAssetAlias", + "title": "Type", + "type": "string" + } + }, + "required": [ + "alias_name" + ], + "title": "GetAssetEventByAssetAlias", + "type": "object" + }, + "GetAssetStateByName": { + "properties": { + "key": { + "title": "Key", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "type": { + "const": "GetAssetStateByName", + "default": "GetAssetStateByName", + "title": "Type", + "type": "string" + } + }, + "required": [ + "name", + "key" + ], + "title": "GetAssetStateByName", + "type": "object" + }, + "GetAssetStateByUri": { + "properties": { + "key": { + "title": "Key", + "type": "string" + }, + "type": { + "const": "GetAssetStateByUri", + "default": "GetAssetStateByUri", + "title": "Type", + "type": "string" + }, + "uri": { + "title": "Uri", + "type": "string" + } + }, + "required": [ + "uri", + "key" + ], + "title": "GetAssetStateByUri", + "type": "object" + }, + "GetAssetsByAlias": { + "properties": { + "alias_name": { + "title": "Alias Name", + "type": "string" + }, + "type": { + "const": "GetAssetsByAlias", + "default": "GetAssetsByAlias", + "title": "Type", + "type": "string" + } + }, + "required": [ + "alias_name" + ], + "title": "GetAssetsByAlias", + "type": "object" + }, + "GetConnection": { + "properties": { + "conn_id": { + "title": "Conn Id", + "type": "string" + }, + "type": { + "const": "GetConnection", + "default": "GetConnection", + "title": "Type", + "type": "string" + } + }, + "required": [ + "conn_id" + ], + "title": "GetConnection", + "type": "object" + }, + "GetDRCount": { + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "logical_dates": { + "anyOf": [ + { + "items": { + "format": "date-time", + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Logical Dates" + }, + "run_ids": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Run Ids" + }, + "states": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "States" + }, + "type": { + "const": "GetDRCount", + "default": "GetDRCount", + "title": "Type", + "type": "string" + } + }, + "required": [ + "dag_id" + ], + "title": "GetDRCount", + "type": "object" + }, + "GetDag": { + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "type": { + "const": "GetDag", + "default": "GetDag", + "title": "Type", + "type": "string" + } + }, + "required": [ + "dag_id" + ], + "title": "GetDag", + "type": "object" + }, + "GetDagRun": { + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "type": { + "const": "GetDagRun", + "default": "GetDagRun", + "title": "Type", + "type": "string" + } + }, + "required": [ + "dag_id", + "run_id" + ], + "title": "GetDagRun", + "type": "object" + }, + "GetDagRunState": { + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "type": { + "const": "GetDagRunState", + "default": "GetDagRunState", + "title": "Type", + "type": "string" + } + }, + "required": [ + "dag_id", + "run_id" + ], + "title": "GetDagRunState", + "type": "object" + }, + "GetHITLDetailResponse": { + "description": "Get the response content part of a Human-in-the-loop response.", + "properties": { + "ti_id": { + "format": "uuid", + "title": "Ti Id", + "type": "string" + }, + "type": { + "const": "GetHITLDetailResponse", + "default": "GetHITLDetailResponse", + "title": "Type", + "type": "string" + } + }, + "required": [ + "ti_id" + ], + "title": "GetHITLDetailResponse", + "type": "object" + }, + "GetPrevSuccessfulDagRun": { + "properties": { + "ti_id": { + "format": "uuid", + "title": "Ti Id", + "type": "string" + }, + "type": { + "const": "GetPrevSuccessfulDagRun", + "default": "GetPrevSuccessfulDagRun", + "title": "Type", + "type": "string" + } + }, + "required": [ + "ti_id" + ], + "title": "GetPrevSuccessfulDagRun", + "type": "object" + }, + "GetPreviousDagRun": { + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "logical_date": { + "format": "date-time", + "title": "Logical Date", + "type": "string" + }, + "state": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "State" + }, + "type": { + "const": "GetPreviousDagRun", + "default": "GetPreviousDagRun", + "title": "Type", + "type": "string" + } + }, + "required": [ + "dag_id", + "logical_date" + ], + "title": "GetPreviousDagRun", + "type": "object" + }, + "GetPreviousTI": { + "$defs": { + "TaskInstanceState": { + "description": "All possible states that a Task Instance can be in.\n\nNote that None is also allowed, so always use this in a type hint with Optional.", + "enum": [ + "removed", + "scheduled", + "queued", + "running", + "success", + "restarting", + "failed", + "up_for_retry", + "up_for_reschedule", + "upstream_failed", + "skipped", + "deferred" + ], + "title": "TaskInstanceState", + "type": "string" + } + }, + "description": "Request to get previous task instance.", + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "logical_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Logical Date" + }, + "map_index": { + "default": -1, + "title": "Map Index", + "type": "integer" + }, + "state": { + "anyOf": [ + { + "$ref": "#/$defs/TaskInstanceState" + }, + { + "type": "null" + } + ], + "default": null + }, + "task_id": { + "title": "Task Id", + "type": "string" + }, + "type": { + "const": "GetPreviousTI", + "default": "GetPreviousTI", + "title": "Type", + "type": "string" + } + }, + "required": [ + "dag_id", + "task_id" + ], + "title": "GetPreviousTI", + "type": "object" + }, + "GetTICount": { + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "logical_dates": { + "anyOf": [ + { + "items": { + "format": "date-time", + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Logical Dates" + }, + "map_index": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Map Index" + }, + "run_ids": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Run Ids" + }, + "states": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "States" + }, + "task_group_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Task Group Id" + }, + "task_ids": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Task Ids" + }, + "type": { + "const": "GetTICount", + "default": "GetTICount", + "title": "Type", + "type": "string" + } + }, + "required": [ + "dag_id" + ], + "title": "GetTICount", + "type": "object" + }, + "GetTaskBreadcrumbs": { + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "type": { + "const": "GetTaskBreadcrumbs", + "default": "GetTaskBreadcrumbs", + "title": "Type", + "type": "string" + } + }, + "required": [ + "dag_id", + "run_id" + ], + "title": "GetTaskBreadcrumbs", + "type": "object" + }, + "GetTaskRescheduleStartDate": { + "properties": { + "ti_id": { + "format": "uuid", + "title": "Ti Id", + "type": "string" + }, + "try_number": { + "default": 1, + "title": "Try Number", + "type": "integer" + }, + "type": { + "const": "GetTaskRescheduleStartDate", + "default": "GetTaskRescheduleStartDate", + "title": "Type", + "type": "string" + } + }, + "required": [ + "ti_id" + ], + "title": "GetTaskRescheduleStartDate", + "type": "object" + }, + "GetTaskState": { + "properties": { + "key": { + "title": "Key", + "type": "string" + }, + "ti_id": { + "format": "uuid", + "title": "Ti Id", + "type": "string" + }, + "type": { + "const": "GetTaskState", + "default": "GetTaskState", + "title": "Type", + "type": "string" + } + }, + "required": [ + "ti_id", + "key" + ], + "title": "GetTaskState", + "type": "object" + }, + "GetTaskStates": { + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "logical_dates": { + "anyOf": [ + { + "items": { + "format": "date-time", + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Logical Dates" + }, + "map_index": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Map Index" + }, + "run_ids": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Run Ids" + }, + "task_group_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Task Group Id" + }, + "task_ids": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Task Ids" + }, + "type": { + "const": "GetTaskStates", + "default": "GetTaskStates", + "title": "Type", + "type": "string" + } + }, + "required": [ + "dag_id" + ], + "title": "GetTaskStates", + "type": "object" + }, + "GetVariable": { + "properties": { + "key": { + "title": "Key", + "type": "string" + }, + "type": { + "const": "GetVariable", + "default": "GetVariable", + "title": "Type", + "type": "string" + } + }, + "required": [ + "key" + ], + "title": "GetVariable", + "type": "object" + }, + "GetVariableKeys": { + "properties": { + "limit": { + "default": 1000, + "title": "Limit", + "type": "integer" + }, + "offset": { + "default": 0, + "title": "Offset", + "type": "integer" + }, + "prefix": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Prefix" + }, + "type": { + "const": "GetVariableKeys", + "default": "GetVariableKeys", + "title": "Type", + "type": "string" + } + }, + "title": "GetVariableKeys", + "type": "object" + }, + "GetXCom": { + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "include_prior_dates": { + "default": false, + "title": "Include Prior Dates", + "type": "boolean" + }, + "key": { + "title": "Key", + "type": "string" + }, + "map_index": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Map Index" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "task_id": { + "title": "Task Id", + "type": "string" + }, + "type": { + "const": "GetXCom", + "default": "GetXCom", + "title": "Type", + "type": "string" + } + }, + "required": [ + "key", + "dag_id", + "run_id", + "task_id" + ], + "title": "GetXCom", + "type": "object" + }, + "GetXComCount": { + "description": "Get the number of (mapped) XCom values available.", + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "key": { + "title": "Key", + "type": "string" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "task_id": { + "title": "Task Id", + "type": "string" + }, + "type": { + "const": "GetXComCount", + "default": "GetXComCount", + "title": "Type", + "type": "string" + } + }, + "required": [ + "key", + "dag_id", + "run_id", + "task_id" + ], + "title": "GetXComCount", + "type": "object" + }, + "GetXComSequenceItem": { + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "key": { + "title": "Key", + "type": "string" + }, + "offset": { + "title": "Offset", + "type": "integer" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "task_id": { + "title": "Task Id", + "type": "string" + }, + "type": { + "const": "GetXComSequenceItem", + "default": "GetXComSequenceItem", + "title": "Type", + "type": "string" + } + }, + "required": [ + "key", + "dag_id", + "run_id", + "task_id", + "offset" + ], + "title": "GetXComSequenceItem", + "type": "object" + }, + "GetXComSequenceSlice": { + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "include_prior_dates": { + "default": false, + "title": "Include Prior Dates", + "type": "boolean" + }, + "key": { + "title": "Key", + "type": "string" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "start": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Start" + }, + "step": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Step" + }, + "stop": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Stop" + }, + "task_id": { + "title": "Task Id", + "type": "string" + }, + "type": { + "const": "GetXComSequenceSlice", + "default": "GetXComSequenceSlice", + "title": "Type", + "type": "string" + } + }, + "required": [ + "key", + "dag_id", + "run_id", + "task_id", + "start", + "stop", + "step" + ], + "title": "GetXComSequenceSlice", + "type": "object" + }, + "HITLDetailRequestResult": { + "$defs": { + "HITLUser": { + "description": "Schema for a Human-in-the-loop users.", + "properties": { + "id": { + "title": "Id", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "id", + "name" + ], + "title": "HITLUser", + "type": "object" + } + }, + "description": "Response to CreateHITLDetailPayload request.", + "properties": { + "assigned_users": { + "anyOf": [ + { + "items": { + "$ref": "#/$defs/HITLUser" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Assigned Users" + }, + "body": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Body" + }, + "defaults": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Defaults" + }, + "multiple": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "default": false, + "title": "Multiple" + }, + "options": { + "items": { + "type": "string" + }, + "minItems": 1, + "title": "Options", + "type": "array" + }, + "params": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Params" + }, + "subject": { + "title": "Subject", + "type": "string" + }, + "ti_id": { + "format": "uuid", + "title": "Ti Id", + "type": "string" + }, + "type": { + "const": "HITLDetailRequestResult", + "default": "HITLDetailRequestResult", + "title": "Type", + "type": "string" + } + }, + "required": [ + "ti_id", + "options", + "subject" + ], + "title": "HITLDetailRequestResult", + "type": "object" + }, + "InactiveAssetsResult": { + "$defs": { + "AssetProfile": { + "additionalProperties": false, + "description": "Profile of an asset-like object.\n\nAsset will have name, uri defined, with type set to 'Asset'.\nAssetNameRef will have name defined, type set to 'AssetNameRef'.\nAssetUriRef will have uri defined, type set to 'AssetUriRef'.\nAssetAlias will have name defined, type set to 'AssetAlias'.\n\nNote that 'type' here is distinct from 'asset_type' the user declares on an\nAsset (or subclass). This field is for distinguishing between different\nasset-related types (Asset, AssetRef, or AssetAlias).", + "properties": { + "name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Name" + }, + "type": { + "title": "Type", + "type": "string" + }, + "uri": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Uri" + } + }, + "required": [ + "type" + ], + "title": "AssetProfile", + "type": "object" + } + }, + "description": "Response of InactiveAssets requests.", + "properties": { + "inactive_assets": { + "anyOf": [ + { + "items": { + "$ref": "#/$defs/AssetProfile" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Inactive Assets" + }, + "type": { + "const": "InactiveAssetsResult", + "default": "InactiveAssetsResult", + "title": "Type", + "type": "string" + } + }, + "title": "InactiveAssetsResult", + "type": "object" + }, + "MaskSecret": { + "$defs": { + "JsonValue": {} + }, + "description": "Add a new value to be redacted in task logs.", + "properties": { + "name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Name" + }, + "type": { + "const": "MaskSecret", + "default": "MaskSecret", + "title": "Type", + "type": "string" + }, + "value": { + "$ref": "#/$defs/JsonValue" + } + }, + "required": [ + "value" + ], + "title": "MaskSecret", + "type": "object" + }, + "OKResponse": { + "properties": { + "ok": { + "title": "Ok", + "type": "boolean" + }, + "type": { + "const": "OKResponse", + "default": "OKResponse", + "title": "Type", + "type": "string" + } + }, + "required": [ + "ok" + ], + "title": "OKResponse", + "type": "object" + }, + "PrevSuccessfulDagRunResult": { + "properties": { + "data_interval_end": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Data Interval End" + }, + "data_interval_start": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Data Interval Start" + }, + "end_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "End Date" + }, + "start_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Start Date" + }, + "type": { + "const": "PrevSuccessfulDagRunResult", + "default": "PrevSuccessfulDagRunResult", + "title": "Type", + "type": "string" + } + }, + "title": "PrevSuccessfulDagRunResult", + "type": "object" + }, + "PreviousDagRunResult": { + "$defs": { + "AssetAliasReferenceAssetEventDagRun": { + "additionalProperties": false, + "description": "Schema for AssetAliasModel used in AssetEventDagRunReference.", + "properties": { + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "name" + ], + "title": "AssetAliasReferenceAssetEventDagRun", + "type": "object" + }, + "AssetEventDagRunReference": { + "additionalProperties": false, + "description": "Schema for AssetEvent model used in DagRun.", + "properties": { + "asset": { + "$ref": "#/$defs/AssetReferenceAssetEventDagRun" + }, + "extra": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "Extra", + "type": "object" + }, + "source_aliases": { + "items": { + "$ref": "#/$defs/AssetAliasReferenceAssetEventDagRun" + }, + "title": "Source Aliases", + "type": "array" + }, + "source_dag_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Source Dag Id" + }, + "source_map_index": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Source Map Index" + }, + "source_run_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Source Run Id" + }, + "source_task_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Source Task Id" + }, + "timestamp": { + "format": "date-time", + "title": "Timestamp", + "type": "string" + } + }, + "required": [ + "asset", + "extra", + "source_aliases", + "timestamp" + ], + "title": "AssetEventDagRunReference", + "type": "object" + }, + "AssetReferenceAssetEventDagRun": { + "additionalProperties": false, + "description": "Schema for AssetModel used in AssetEventDagRunReference.", + "properties": { + "extra": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "Extra", + "type": "object" + }, + "name": { + "title": "Name", + "type": "string" + }, + "uri": { + "title": "Uri", + "type": "string" + } + }, + "required": [ + "name", + "uri", + "extra" + ], + "title": "AssetReferenceAssetEventDagRun", + "type": "object" + }, + "DagRun": { + "additionalProperties": false, + "description": "Schema for DagRun model with minimal required fields needed for Runtime.", + "properties": { + "clear_number": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": 0, + "title": "Clear Number" + }, + "conf": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Conf" + }, + "consumed_asset_events": { + "items": { + "$ref": "#/$defs/AssetEventDagRunReference" + }, + "title": "Consumed Asset Events", + "type": "array" + }, + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "data_interval_end": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Data Interval End" + }, + "data_interval_start": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Data Interval Start" + }, + "end_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "End Date" + }, + "logical_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Logical Date" + }, + "note": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Note" + }, + "partition_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key" + }, + "run_after": { + "format": "date-time", + "title": "Run After", + "type": "string" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "run_type": { + "$ref": "#/$defs/DagRunType" + }, + "start_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Start Date" + }, + "state": { + "$ref": "#/$defs/DagRunState" + }, + "team_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Team Name" + }, + "triggering_user_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Triggering User Name" + } + }, + "required": [ + "dag_id", + "run_id", + "run_after", + "run_type", + "state", + "consumed_asset_events" + ], + "title": "DagRun", + "type": "object" + }, + "DagRunState": { + "description": "All possible states that a DagRun can be in.\n\nThese are \"shared\" with TaskInstanceState in some parts of the code,\nso please ensure that their values always match the ones with the\nsame name in TaskInstanceState.", + "enum": [ + "queued", + "running", + "success", + "failed" + ], + "title": "DagRunState", + "type": "string" + }, + "DagRunType": { + "description": "Class with DagRun types.", + "enum": [ + "backfill", + "scheduled", + "manual", + "operator_triggered", + "asset_triggered", + "asset_materialization" + ], + "title": "DagRunType", + "type": "string" + }, + "JsonValue": {} + }, + "description": "Response containing previous Dag run information.", + "properties": { + "dag_run": { + "anyOf": [ + { + "$ref": "#/$defs/DagRun" + }, + { + "type": "null" + } + ], + "default": null + }, + "type": { + "const": "PreviousDagRunResult", + "default": "PreviousDagRunResult", + "title": "Type", + "type": "string" + } + }, + "title": "PreviousDagRunResult", + "type": "object" + }, + "PreviousTIResult": { + "$defs": { + "PreviousTIResponse": { + "description": "Schema for response with previous TaskInstance information.", + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "duration": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Duration" + }, + "end_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "End Date" + }, + "logical_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Logical Date" + }, + "map_index": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": -1, + "title": "Map Index" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "start_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Start Date" + }, + "state": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "State" + }, + "task_id": { + "title": "Task Id", + "type": "string" + }, + "try_number": { + "title": "Try Number", + "type": "integer" + } + }, + "required": [ + "task_id", + "dag_id", + "run_id", + "try_number" + ], + "title": "PreviousTIResponse", + "type": "object" + } + }, + "description": "Response containing previous task instance data.", + "properties": { + "task_instance": { + "anyOf": [ + { + "$ref": "#/$defs/PreviousTIResponse" + }, + { + "type": "null" + } + ], + "default": null + }, + "type": { + "const": "PreviousTIResult", + "default": "PreviousTIResult", + "title": "Type", + "type": "string" + } + }, + "title": "PreviousTIResult", + "type": "object" + }, + "PutVariable": { + "properties": { + "description": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, + "key": { + "title": "Key", + "type": "string" + }, + "type": { + "const": "PutVariable", + "default": "PutVariable", + "title": "Type", + "type": "string" + }, + "value": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Value" + } + }, + "required": [ + "key", + "value", + "description" + ], + "title": "PutVariable", + "type": "object" + }, + "RescheduleTask": { + "additionalProperties": false, + "description": "Update a task instance state to reschedule/up_for_reschedule.", + "properties": { + "end_date": { + "format": "date-time", + "title": "End Date", + "type": "string" + }, + "reschedule_date": { + "format": "date-time", + "title": "Reschedule Date", + "type": "string" + }, + "state": { + "anyOf": [ + { + "const": "up_for_reschedule", + "type": "string" + }, + { + "type": "null" + } + ], + "default": "up_for_reschedule", + "title": "State" + }, + "type": { + "const": "RescheduleTask", + "default": "RescheduleTask", + "title": "Type", + "type": "string" + } + }, + "required": [ + "reschedule_date", + "end_date" + ], + "title": "RescheduleTask", + "type": "object" + }, + "ResendLoggingFD": { + "properties": { + "type": { + "const": "ResendLoggingFD", + "default": "ResendLoggingFD", + "title": "Type", + "type": "string" + } + }, + "title": "ResendLoggingFD", + "type": "object" + }, + "RetryTask": { + "additionalProperties": false, + "description": "Update a task instance state to up_for_retry.", + "properties": { + "end_date": { + "format": "date-time", + "title": "End Date", + "type": "string" + }, + "rendered_map_index": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Rendered Map Index" + }, + "retry_delay_seconds": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Retry Delay Seconds" + }, + "retry_reason": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Retry Reason" + }, + "state": { + "anyOf": [ + { + "const": "up_for_retry", + "type": "string" + }, + { + "type": "null" + } + ], + "default": "up_for_retry", + "title": "State" + }, + "type": { + "const": "RetryTask", + "default": "RetryTask", + "title": "Type", + "type": "string" + } + }, + "required": [ + "end_date" + ], + "title": "RetryTask", + "type": "object" + }, + "SentFDs": { + "properties": { + "fds": { + "items": { + "type": "integer" + }, + "title": "Fds", + "type": "array" + }, + "type": { + "const": "SentFDs", + "default": "SentFDs", + "title": "Type", + "type": "string" + } + }, + "required": [ + "fds" + ], + "title": "SentFDs", + "type": "object" + }, + "SetAssetStateByName": { + "properties": { + "key": { + "title": "Key", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "type": { + "const": "SetAssetStateByName", + "default": "SetAssetStateByName", + "title": "Type", + "type": "string" + }, + "value": { + "title": "Value", + "type": "string" + } + }, + "required": [ + "name", + "key", + "value" + ], + "title": "SetAssetStateByName", + "type": "object" + }, + "SetAssetStateByUri": { + "properties": { + "key": { + "title": "Key", + "type": "string" + }, + "type": { + "const": "SetAssetStateByUri", + "default": "SetAssetStateByUri", + "title": "Type", + "type": "string" + }, + "uri": { + "title": "Uri", + "type": "string" + }, + "value": { + "title": "Value", + "type": "string" + } + }, + "required": [ + "uri", + "key", + "value" + ], + "title": "SetAssetStateByUri", + "type": "object" + }, + "SetRenderedFields": { + "$defs": { + "JsonValue": {} + }, + "description": "Payload for setting RTIF for a task instance.", + "properties": { + "rendered_fields": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "Rendered Fields", + "type": "object" + }, + "type": { + "const": "SetRenderedFields", + "default": "SetRenderedFields", + "title": "Type", + "type": "string" + } + }, + "required": [ + "rendered_fields" + ], + "title": "SetRenderedFields", + "type": "object" + }, + "SetRenderedMapIndex": { + "description": "Payload for setting rendered_map_index for a task instance.", + "properties": { + "rendered_map_index": { + "title": "Rendered Map Index", + "type": "string" + }, + "type": { + "const": "SetRenderedMapIndex", + "default": "SetRenderedMapIndex", + "title": "Type", + "type": "string" + } + }, + "required": [ + "rendered_map_index" + ], + "title": "SetRenderedMapIndex", + "type": "object" + }, + "SetTaskState": { + "properties": { + "expires_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Expires At" + }, + "key": { + "title": "Key", + "type": "string" + }, + "ti_id": { + "format": "uuid", + "title": "Ti Id", + "type": "string" + }, + "type": { + "const": "SetTaskState", + "default": "SetTaskState", + "title": "Type", + "type": "string" + }, + "value": { + "title": "Value", + "type": "string" + } + }, + "required": [ + "ti_id", + "key", + "value", + "expires_at" + ], + "title": "SetTaskState", + "type": "object" + }, + "SetXCom": { + "$defs": { + "JsonValue": {} + }, + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "dag_result": { + "default": false, + "title": "Dag Result", + "type": "boolean" + }, + "key": { + "title": "Key", + "type": "string" + }, + "map_index": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Map Index" + }, + "mapped_length": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Mapped Length" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "task_id": { + "title": "Task Id", + "type": "string" + }, + "type": { + "const": "SetXCom", + "default": "SetXCom", + "title": "Type", + "type": "string" + }, + "value": { + "$ref": "#/$defs/JsonValue" + } + }, + "required": [ + "key", + "value", + "dag_id", + "run_id", + "task_id" + ], + "title": "SetXCom", + "type": "object" + }, + "SkipDownstreamTasks": { + "additionalProperties": false, + "description": "Update state of downstream tasks within a task instance to 'skipped', while updating current task to success state.", + "properties": { + "tasks": { + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "maxItems": 2, + "minItems": 2, + "prefixItems": [ + { + "type": "string" + }, + { + "type": "integer" + } + ], + "type": "array" + } + ] + }, + "title": "Tasks", + "type": "array" + }, + "type": { + "const": "SkipDownstreamTasks", + "default": "SkipDownstreamTasks", + "title": "Type", + "type": "string" + } + }, + "required": [ + "tasks" + ], + "title": "SkipDownstreamTasks", + "type": "object" + }, + "StartupDetails": { + "$defs": { + "AssetAliasReferenceAssetEventDagRun": { + "additionalProperties": false, + "description": "Schema for AssetAliasModel used in AssetEventDagRunReference.", + "properties": { + "name": { + "title": "Name", + "type": "string" + } + }, + "required": [ + "name" + ], + "title": "AssetAliasReferenceAssetEventDagRun", + "type": "object" + }, + "AssetEventDagRunReference": { + "additionalProperties": false, + "description": "Schema for AssetEvent model used in DagRun.", + "properties": { + "asset": { + "$ref": "#/$defs/AssetReferenceAssetEventDagRun" + }, + "extra": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "Extra", + "type": "object" + }, + "source_aliases": { + "items": { + "$ref": "#/$defs/AssetAliasReferenceAssetEventDagRun" + }, + "title": "Source Aliases", + "type": "array" + }, + "source_dag_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Source Dag Id" + }, + "source_map_index": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Source Map Index" + }, + "source_run_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Source Run Id" + }, + "source_task_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Source Task Id" + }, + "timestamp": { + "format": "date-time", + "title": "Timestamp", + "type": "string" + } + }, + "required": [ + "asset", + "extra", + "source_aliases", + "timestamp" + ], + "title": "AssetEventDagRunReference", + "type": "object" + }, + "AssetReferenceAssetEventDagRun": { + "additionalProperties": false, + "description": "Schema for AssetModel used in AssetEventDagRunReference.", + "properties": { + "extra": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "Extra", + "type": "object" + }, + "name": { + "title": "Name", + "type": "string" + }, + "uri": { + "title": "Uri", + "type": "string" + } + }, + "required": [ + "name", + "uri", + "extra" + ], + "title": "AssetReferenceAssetEventDagRun", + "type": "object" + }, + "BundleInfo": { + "description": "Schema for telling task which bundle to run with.", + "properties": { + "name": { + "title": "Name", + "type": "string" + }, + "version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Version" + } + }, + "required": [ + "name" + ], + "title": "BundleInfo", + "type": "object" + }, + "ConnectionResponse": { + "description": "Connection schema for responses with fields that are needed for Runtime.", + "properties": { + "conn_id": { + "title": "Conn Id", + "type": "string" + }, + "conn_type": { + "title": "Conn Type", + "type": "string" + }, + "extra": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Extra" + }, + "host": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Host" + }, + "login": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Login" + }, + "password": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Password" + }, + "port": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Port" + }, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Schema" + } + }, + "required": [ + "conn_id", + "conn_type" + ], + "title": "ConnectionResponse", + "type": "object" + }, + "DagRun": { + "additionalProperties": false, + "description": "Schema for DagRun model with minimal required fields needed for Runtime.", + "properties": { + "clear_number": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": 0, + "title": "Clear Number" + }, + "conf": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Conf" + }, + "consumed_asset_events": { + "items": { + "$ref": "#/$defs/AssetEventDagRunReference" + }, + "title": "Consumed Asset Events", + "type": "array" + }, + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "data_interval_end": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Data Interval End" + }, + "data_interval_start": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Data Interval Start" + }, + "end_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "End Date" + }, + "logical_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Logical Date" + }, + "note": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Note" + }, + "partition_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key" + }, + "run_after": { + "format": "date-time", + "title": "Run After", + "type": "string" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "run_type": { + "$ref": "#/$defs/DagRunType" + }, + "start_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Start Date" + }, + "state": { + "$ref": "#/$defs/DagRunState" + }, + "team_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Team Name" + }, + "triggering_user_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Triggering User Name" + } + }, + "required": [ + "dag_id", + "run_id", + "run_after", + "run_type", + "state", + "consumed_asset_events" + ], + "title": "DagRun", + "type": "object" + }, + "DagRunState": { + "description": "All possible states that a DagRun can be in.\n\nThese are \"shared\" with TaskInstanceState in some parts of the code,\nso please ensure that their values always match the ones with the\nsame name in TaskInstanceState.", + "enum": [ + "queued", + "running", + "success", + "failed" + ], + "title": "DagRunState", + "type": "string" + }, + "DagRunType": { + "description": "Class with DagRun types.", + "enum": [ + "backfill", + "scheduled", + "manual", + "operator_triggered", + "asset_triggered", + "asset_materialization" + ], + "title": "DagRunType", + "type": "string" + }, + "JsonValue": {}, + "TIRunContext": { + "description": "Response schema for TaskInstance run context.", + "properties": { + "connections": { + "anyOf": [ + { + "items": { + "$ref": "#/$defs/ConnectionResponse" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Connections" + }, + "dag_run": { + "$ref": "#/$defs/DagRun" + }, + "max_tries": { + "title": "Max Tries", + "type": "integer" + }, + "next_kwargs": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Next Kwargs" + }, + "next_method": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Next Method" + }, + "should_retry": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "default": false, + "title": "Should Retry" + }, + "start_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Start Date" + }, + "task_reschedule_count": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": 0, + "title": "Task Reschedule Count" + }, + "variables": { + "anyOf": [ + { + "items": { + "$ref": "#/$defs/VariableResponse" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Variables" + }, + "xcom_keys_to_clear": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Xcom Keys To Clear" + } + }, + "required": [ + "dag_run", + "max_tries" + ], + "title": "TIRunContext", + "type": "object" + }, + "TaskInstance": { + "description": "Schema for TaskInstance model with minimal required fields needed for Runtime.", + "properties": { + "context_carrier": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Context Carrier" + }, + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "dag_version_id": { + "format": "uuid", + "title": "Dag Version Id", + "type": "string" + }, + "hostname": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Hostname" + }, + "id": { + "format": "uuid", + "title": "Id", + "type": "string" + }, + "map_index": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": -1, + "title": "Map Index" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "task_id": { + "title": "Task Id", + "type": "string" + }, + "try_number": { + "title": "Try Number", + "type": "integer" + } + }, + "required": [ + "id", + "task_id", + "dag_id", + "run_id", + "try_number", + "dag_version_id" + ], + "title": "TaskInstance", + "type": "object" + }, + "VariableResponse": { + "additionalProperties": false, + "description": "Variable schema for responses with fields that are needed for Runtime.", + "properties": { + "key": { + "title": "Key", + "type": "string" + }, + "value": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Value" + } + }, + "required": [ + "key" + ], + "title": "VariableResponse", + "type": "object" + } + }, + "properties": { + "bundle_info": { + "$ref": "#/$defs/BundleInfo" + }, + "dag_rel_path": { + "title": "Dag Rel Path", + "type": "string" + }, + "sentry_integration": { + "title": "Sentry Integration", + "type": "string" + }, + "start_date": { + "format": "date-time", + "title": "Start Date", + "type": "string" + }, + "ti": { + "$ref": "#/$defs/TaskInstance" + }, + "ti_context": { + "$ref": "#/$defs/TIRunContext" + }, + "type": { + "const": "StartupDetails", + "default": "StartupDetails", + "title": "Type", + "type": "string" + } + }, + "required": [ + "ti", + "dag_rel_path", + "bundle_info", + "start_date", + "ti_context", + "sentry_integration" + ], + "title": "StartupDetails", + "type": "object" + }, + "SucceedTask": { + "$defs": { + "AssetProfile": { + "additionalProperties": false, + "description": "Profile of an asset-like object.\n\nAsset will have name, uri defined, with type set to 'Asset'.\nAssetNameRef will have name defined, type set to 'AssetNameRef'.\nAssetUriRef will have uri defined, type set to 'AssetUriRef'.\nAssetAlias will have name defined, type set to 'AssetAlias'.\n\nNote that 'type' here is distinct from 'asset_type' the user declares on an\nAsset (or subclass). This field is for distinguishing between different\nasset-related types (Asset, AssetRef, or AssetAlias).", + "properties": { + "name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Name" + }, + "type": { + "title": "Type", + "type": "string" + }, + "uri": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Uri" + } + }, + "required": [ + "type" + ], + "title": "AssetProfile", + "type": "object" + } + }, + "additionalProperties": false, + "description": "Update a task's state to success. Includes task_outlets and outlet_events for registering asset events.", + "properties": { + "end_date": { + "format": "date-time", + "title": "End Date", + "type": "string" + }, + "outlet_events": { + "anyOf": [ + { + "items": { + "additionalProperties": true, + "type": "object" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Outlet Events" + }, + "rendered_map_index": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Rendered Map Index" + }, + "state": { + "anyOf": [ + { + "const": "success", + "type": "string" + }, + { + "type": "null" + } + ], + "default": "success", + "title": "State" + }, + "task_outlets": { + "anyOf": [ + { + "items": { + "$ref": "#/$defs/AssetProfile" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Task Outlets" + }, + "type": { + "const": "SucceedTask", + "default": "SucceedTask", + "title": "Type", + "type": "string" + } + }, + "required": [ + "end_date" + ], + "title": "SucceedTask", + "type": "object" + }, + "TICount": { + "description": "Response containing count of Task Instances matching certain filters.", + "properties": { + "count": { + "title": "Count", + "type": "integer" + }, + "type": { + "const": "TICount", + "default": "TICount", + "title": "Type", + "type": "string" + } + }, + "required": [ + "count" + ], + "title": "TICount", + "type": "object" + }, + "TaskBreadcrumbsResult": { + "properties": { + "breadcrumbs": { + "items": { + "additionalProperties": true, + "type": "object" + }, + "title": "Breadcrumbs", + "type": "array" + }, + "type": { + "const": "TaskBreadcrumbsResult", + "default": "TaskBreadcrumbsResult", + "title": "Type", + "type": "string" + } + }, + "required": [ + "breadcrumbs" + ], + "title": "TaskBreadcrumbsResult", + "type": "object" + }, + "TaskRescheduleStartDate": { + "description": "Response containing the first reschedule date for a task instance.", + "properties": { + "start_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Start Date" + }, + "type": { + "const": "TaskRescheduleStartDate", + "default": "TaskRescheduleStartDate", + "title": "Type", + "type": "string" + } + }, + "required": [ + "start_date" + ], + "title": "TaskRescheduleStartDate", + "type": "object" + }, + "TaskState": { + "description": "Update a task's state.\n\nIf a process exits without sending one of these the state will be derived from the exit code:\n- 0 = SUCCESS\n- anything else = FAILED", + "properties": { + "end_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "End Date" + }, + "rendered_map_index": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Rendered Map Index" + }, + "state": { + "enum": [ + "failed", + "skipped", + "removed" + ], + "title": "State", + "type": "string" + }, + "type": { + "const": "TaskState", + "default": "TaskState", + "title": "Type", + "type": "string" + } + }, + "required": [ + "state" + ], + "title": "TaskState", + "type": "object" + }, + "TaskStateResult": { + "additionalProperties": false, + "description": "Response to GetTaskState; wraps the generated API response for supervisor to worker comms.", + "properties": { + "type": { + "const": "TaskStateResult", + "default": "TaskStateResult", + "title": "Type", + "type": "string" + }, + "value": { + "title": "Value", + "type": "string" + } + }, + "required": [ + "value" + ], + "title": "TaskStateResult", + "type": "object" + }, + "TaskStatesResult": { + "properties": { + "task_states": { + "additionalProperties": true, + "title": "Task States", + "type": "object" + }, + "type": { + "const": "TaskStatesResult", + "default": "TaskStatesResult", + "title": "Type", + "type": "string" + } + }, + "required": [ + "task_states" + ], + "title": "TaskStatesResult", + "type": "object" + }, + "TriggerDagRun": { + "additionalProperties": false, + "properties": { + "conf": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Conf" + }, + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "logical_date": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Logical Date" + }, + "note": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Note" + }, + "partition_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key" + }, + "reset_dag_run": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "default": false, + "title": "Reset Dag Run" + }, + "run_after": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Run After" + }, + "run_id": { + "title": "Dag Run Id", + "type": "string" + }, + "type": { + "const": "TriggerDagRun", + "default": "TriggerDagRun", + "title": "Type", + "type": "string" + } + }, + "required": [ + "dag_id", + "run_id" + ], + "title": "TriggerDagRun", + "type": "object" + }, + "UpdateHITLDetail": { + "description": "Update the response content part of an existing Human-in-the-loop response.", + "properties": { + "chosen_options": { + "items": { + "type": "string" + }, + "minItems": 1, + "title": "Chosen Options", + "type": "array" + }, + "params_input": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Params Input" + }, + "ti_id": { + "format": "uuid", + "title": "Ti Id", + "type": "string" + }, + "type": { + "const": "UpdateHITLDetail", + "default": "UpdateHITLDetail", + "title": "Type", + "type": "string" + } + }, + "required": [ + "ti_id", + "chosen_options" + ], + "title": "UpdateHITLDetail", + "type": "object" + }, + "ValidateInletsAndOutlets": { + "properties": { + "ti_id": { + "format": "uuid", + "title": "Ti Id", + "type": "string" + }, + "type": { + "const": "ValidateInletsAndOutlets", + "default": "ValidateInletsAndOutlets", + "title": "Type", + "type": "string" + } + }, + "required": [ + "ti_id" + ], + "title": "ValidateInletsAndOutlets", + "type": "object" + }, + "VariableKeysResult": { + "properties": { + "keys": { + "items": { + "type": "string" + }, + "title": "Keys", + "type": "array" + }, + "total_entries": { + "title": "Total Entries", + "type": "integer" + }, + "type": { + "const": "VariableKeysResult", + "default": "VariableKeysResult", + "title": "Type", + "type": "string" + } + }, + "required": [ + "keys", + "total_entries" + ], + "title": "VariableKeysResult", + "type": "object" + }, + "VariableResult": { + "additionalProperties": false, + "properties": { + "key": { + "title": "Key", + "type": "string" + }, + "type": { + "const": "VariableResult", + "default": "VariableResult", + "title": "Type", + "type": "string" + }, + "value": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Value" + } + }, + "required": [ + "key" + ], + "title": "VariableResult", + "type": "object" + }, + "XComCountResponse": { + "properties": { + "len": { + "title": "Len", + "type": "integer" + }, + "type": { + "const": "XComCountResponse", + "default": "XComCountResponse", + "title": "Type", + "type": "string" + } + }, + "required": [ + "len" + ], + "title": "XComCountResponse", + "type": "object" + }, + "XComResult": { + "$defs": { + "JsonValue": {} + }, + "description": "Response to ReadXCom request.", + "properties": { + "key": { + "title": "Key", + "type": "string" + }, + "type": { + "const": "XComResult", + "default": "XComResult", + "title": "Type", + "type": "string" + }, + "value": { + "$ref": "#/$defs/JsonValue" + } + }, + "required": [ + "key", + "value" + ], + "title": "XComResult", + "type": "object" + }, + "XComSequenceIndexResult": { + "$defs": { + "JsonValue": {} + }, + "properties": { + "root": { + "$ref": "#/$defs/JsonValue" + }, + "type": { + "const": "XComSequenceIndexResult", + "default": "XComSequenceIndexResult", + "title": "Type", + "type": "string" + } + }, + "required": [ + "root" + ], + "title": "XComSequenceIndexResult", + "type": "object" + }, + "XComSequenceSliceResult": { + "$defs": { + "JsonValue": {} + }, + "properties": { + "root": { + "items": { + "$ref": "#/$defs/JsonValue" + }, + "title": "Root", + "type": "array" + }, + "type": { + "const": "XComSequenceSliceResult", + "default": "XComSequenceSliceResult", + "title": "Type", + "type": "string" + } + }, + "required": [ + "root" + ], + "title": "XComSequenceSliceResult", + "type": "object" + } + } +} diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py new file mode 100644 index 0000000000000..d8ae4e4580ed5 --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py @@ -0,0 +1,25 @@ +# 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 + +from cadwyn import HeadVersion, Version, VersionBundle + +bundle = VersionBundle( + HeadVersion(), + Version("2026-06-16"), +) diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 3e705b9d211b0..f26e42f62d6b8 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -37,7 +37,7 @@ from datetime import datetime, timezone from http import HTTPStatus from socket import socket, socketpair -from typing import TYPE_CHECKING, BinaryIO, ClassVar, NoReturn, TextIO, cast +from typing import TYPE_CHECKING, Any, BinaryIO, ClassVar, NoReturn, TextIO, cast from urllib.parse import urlparse from uuid import UUID @@ -144,6 +144,7 @@ handle_get_variable_keys, handle_mask_secret, ) +from airflow.sdk.execution_time.schema import get_schema_version_migrator, resolve_body_class try: from socket import send_fds @@ -548,6 +549,8 @@ class WatchedSubprocess: _process: psutil.Process = attrs.field(repr=False) """File descriptor for request handling.""" + _subprocess_schema_version: str | None = None + _exit_code: int | None = attrs.field(default=None, init=False) _process_exit_monotonic: float | None = attrs.field(default=None, init=False) _open_sockets: weakref.WeakKeyDictionary[socket, str] = attrs.field( @@ -730,35 +733,51 @@ def _create_log_forwarder(self, loggers, name, log_level=logging.INFO) -> Callab def _on_socket_closed(self, sock: socket): # We want to keep servicing this process until we've read up to EOF from all the sockets. - with suppress(KeyError): self.selector.unregister(sock) del self._open_sockets[sock] + def _serialize_response(self, msg: BaseModel | ErrorResponse, **dump_opts) -> dict[str, Any]: + if self._subprocess_schema_version is not None: + migrator = get_schema_version_migrator() + msg = migrator.downgrade(msg, self._subprocess_schema_version, **dump_opts) + return msg.model_dump(**dump_opts) + def send_msg( - self, msg: BaseModel | None, request_id: int, error: ErrorResponse | None = None, **dump_opts + self, + msg: BaseModel | None, + request_id: int, + error: ErrorResponse | None = None, + **dump_opts, ): """ Send the msg as a length-prefixed response frame. - ``request_id`` is the ID that the client sent in it's request, and has no meaning to the server - + :param request_id: The ID sent in the request by the client. This has no + meaning to the server, and is only included in the response frame + for the client to identify what the response is for. """ if msg: - frame = _ResponseFrame(id=request_id, body=msg.model_dump(**dump_opts)) + frame = _ResponseFrame(id=request_id, body=self._serialize_response(msg, **dump_opts)) else: - err_resp = error.model_dump() if error else None + err_resp = self._serialize_response(error) if error else None frame = _ResponseFrame(id=request_id, error=err_resp) - self.stdin.sendall(frame.as_bytes()) + def _deserialize_request(self, body: dict[str, Any] | None) -> dict[str, Any] | None: + if self._subprocess_schema_version is None or body is None: + return body + if (model := resolve_body_class(body)) is None: + raise ValueError(f"Cannot resolve model without a valid 'type' discriminator: {body!r}") + return get_schema_version_migrator().upgrade(body, model, self._subprocess_schema_version) + def handle_requests(self, log: FilteringBoundLogger) -> Generator[None, _RequestFrame, None]: """Handle incoming requests from the task process, respond with the appropriate data.""" while True: request = yield try: - msg = self.decoder.validate_python(request.body) + msg = self.decoder.validate_python(self._deserialize_request(request.body)) except Exception: log.exception("Unable to decode message", body=request.body) continue diff --git a/task-sdk/tests/task_sdk/execution_time/schema/__init__.py b/task-sdk/tests/task_sdk/execution_time/schema/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/task-sdk/tests/task_sdk/execution_time/schema/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/task-sdk/tests/task_sdk/execution_time/schema/_mock_version_bundle.py b/task-sdk/tests/task_sdk/execution_time/schema/_mock_version_bundle.py new file mode 100644 index 0000000000000..7566998fc6dfe --- /dev/null +++ b/task-sdk/tests/task_sdk/execution_time/schema/_mock_version_bundle.py @@ -0,0 +1,189 @@ +# 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. +""" +Mock Cadwyn version bundle used by the supervisor-schemas integration tests. + +The bundle and its body classes live in their own module so that any +helper or fixture that needs them imports a single canonical +definition. The integration test installs them via a ``monkeypatch`` +fixture for the duration of one test and tears down automatically. + +Two body classes mirror the production split between channels: + +- :class:`_LangSdkRequest` -- lang-SDK -> supervisor. Three fields with + three dated breaking changes; each entry pairs a + ``schema(...).didnt_exist`` instruction with a + ``convert_request_to_next_version_for`` backfill so a wire payload + from an older runtime reaches the head Pydantic class with every + field present. +- :class:`_SupervisorResponse` -- supervisor -> lang-SDK. Three fields + with three dated breaking changes; each entry carries only + ``schema(...).didnt_exist`` (responses never flow upstream, so no + upgrade transformer is needed). +""" + +from __future__ import annotations + +from typing import Literal + +from cadwyn import ( + HeadVersion, + Version, + VersionBundle, + VersionChange, + convert_request_to_next_version_for, + schema, +) +from pydantic import BaseModel + + +class _LangSdkRequest(BaseModel): + """ + lang-SDK -> supervisor request body. + + Three fields appear here; an older runtime omits later fields and + the upgrade walk backfills them so the supervisor's head decoder + always validates. + """ + + type: Literal["_LangSdkRequest"] = "_LangSdkRequest" + ti_id: str + field_a: int | None = None + field_b: int | None = None + field_c: int | None = None + + +class _SupervisorResponse(BaseModel): + """ + supervisor -> lang-SDK response body. + + Three fields appear here; the downgrade walk trims any field + introduced after the runtime's pinned version. + """ + + type: Literal["_SupervisorResponse"] = "_SupervisorResponse" + ti_id: str + response_x: str | None = None + response_y: str | None = None + response_z: str | None = None + + +# Request-body breaking changes -- each adds a field and a request-side +# backfill so an older lang-SDK payload reaches the head shape intact. + + +class _AddRequestFieldA(VersionChange): + """3026-02-15: introduce ``_LangSdkRequest.field_a``.""" + + description = __doc__ + instructions_to_migrate_to_previous_version = (schema(_LangSdkRequest).field("field_a").didnt_exist,) + + @convert_request_to_next_version_for(_LangSdkRequest) # type: ignore[arg-type] + def _backfill(request): + request.body.setdefault("field_a", 0) + + +class _AddRequestFieldB(VersionChange): + """3026-05-10: introduce ``_LangSdkRequest.field_b``.""" + + description = __doc__ + instructions_to_migrate_to_previous_version = (schema(_LangSdkRequest).field("field_b").didnt_exist,) + + @convert_request_to_next_version_for(_LangSdkRequest) # type: ignore[arg-type] + def _backfill(request): + request.body.setdefault("field_b", 0) + + +class _AddRequestFieldC(VersionChange): + """3026-08-22: introduce ``_LangSdkRequest.field_c``.""" + + description = __doc__ + instructions_to_migrate_to_previous_version = (schema(_LangSdkRequest).field("field_c").didnt_exist,) + + @convert_request_to_next_version_for(_LangSdkRequest) # type: ignore[arg-type] + def _backfill(request): + request.body.setdefault("field_c", 0) + + +# Response-body breaking changes -- downgrade-only direction, no upgrade +# transformer because responses are never sent lang-SDK -> supervisor. + + +class _AddResponseFieldX(VersionChange): + """3026-03-01: introduce ``_SupervisorResponse.response_x``.""" + + description = __doc__ + instructions_to_migrate_to_previous_version = ( + schema(_SupervisorResponse).field("response_x").didnt_exist, + ) + + +class _AddResponseFieldY(VersionChange): + """3026-06-15: introduce ``_SupervisorResponse.response_y``.""" + + description = __doc__ + instructions_to_migrate_to_previous_version = ( + schema(_SupervisorResponse).field("response_y").didnt_exist, + ) + + +class _AddResponseFieldZ(VersionChange): + """3026-09-30: introduce ``_SupervisorResponse.response_z``.""" + + description = __doc__ + instructions_to_migrate_to_previous_version = ( + schema(_SupervisorResponse).field("response_z").didnt_exist, + ) + + +MOCK_VERSION_BUNDLE = VersionBundle( + HeadVersion(), + Version("3026-09-30", _AddResponseFieldZ), + Version("3026-08-22", _AddRequestFieldC), + Version("3026-06-15", _AddResponseFieldY), + Version("3026-05-10", _AddRequestFieldB), + Version("3026-03-01", _AddResponseFieldX), + Version("3026-02-15", _AddRequestFieldA), + Version("3025-12-01"), +) + + +ALL_VERSIONS: tuple[str, ...] = ( + "3025-12-01", + "3026-02-15", + "3026-03-01", + "3026-05-10", + "3026-06-15", + "3026-08-22", + "3026-09-30", +) + + +MOCK_REGISTRY: dict[str, type] = { + "_LangSdkRequest": _LangSdkRequest, + "_SupervisorResponse": _SupervisorResponse, +} +""" +Wire-discriminator -> head class map for the mock bundle. + +Production lookups in ``resolve_body_class`` go through +``schema.registered_models_by_name``. The +``mock_version_migrator`` fixture in :mod:`test_integration` swaps that +lookup for this dict so the upgrade path can resolve +``_LangSdkRequest`` / ``_SupervisorResponse`` discriminators without +touching the real registry. +""" diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_integration.py b/task-sdk/tests/task_sdk/execution_time/schema/test_integration.py new file mode 100644 index 0000000000000..9cdcdb43c2e80 --- /dev/null +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_integration.py @@ -0,0 +1,366 @@ +# 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. +""" +In-process integration tests for the supervisor schema migration seam. + +Drive ``WatchedSubprocess.send_msg`` and ``WatchedSubprocess.handle_requests`` +directly against a ``MagicMock`` socket, then decode the bytes the production +code wrote and assert on the wire shape. The migrator runs for real against the +mock Cadwyn bundle in :mod:`_mock_version_bundle`, swapped in via +``monkeypatch`` for the duration of one test. +""" + +from __future__ import annotations + +from typing import Any, ClassVar +from unittest.mock import MagicMock, call + +import attrs +import msgspec +import psutil +import pytest +import structlog +from pydantic import TypeAdapter +from task_sdk.execution_time.schema._mock_version_bundle import ( + ALL_VERSIONS, + MOCK_REGISTRY, + MOCK_VERSION_BUNDLE, + _LangSdkRequest, + _SupervisorResponse, +) +from uuid6 import uuid7 + +from airflow.sdk.execution_time.comms import _RequestFrame, _ResponseFrame +from airflow.sdk.execution_time.schema import SchemaVersionMigrator +from airflow.sdk.execution_time.supervisor import WatchedSubprocess + + +@pytest.fixture +def mock_version_migrator(monkeypatch) -> SchemaVersionMigrator: + """ + Bind the production migrator factory and registry to :data:`MOCK_VERSION_BUNDLE`. + + Three patches are applied: + + 1. ``supervisor.get_schema_version_migrator`` -- the already-imported binding + inside :mod:`airflow.sdk.execution_time.supervisor` that + ``_serialize_response`` and ``_deserialize_request`` call. + 2. ``schema.get_schema_version_migrator`` -- the canonical + location, kept in sync so any code that re-imports the symbol sees the + mock too. + 3. ``schema.registered_models_by_name`` -- used by + ``resolve_body_class`` (called from ``_deserialize_request``) to map wire + discriminators to head classes. Swapped for :data:`MOCK_REGISTRY` so the + upgrade path resolves ``_LangSdkRequest`` without touching the real comms + registry. + + A corrected ``_serialize_response`` is also installed to work around a bug + in the current implementation where ``dump_kwargs=dump_opts`` is passed + positionally to ``downgrade`` instead of being unpacked with ``**``, + causing ``model_dump`` to reject an unexpected keyword argument. The patched + version calls ``downgrade`` without extra kwargs and lets the versioned + model's ``model_dump()`` produce the wire dict. + """ + migrator = SchemaVersionMigrator( + bundle=MOCK_VERSION_BUNDLE, + supervisor_version=MOCK_VERSION_BUNDLE.versions[0].value, + ) + mock_migrator_factory = lambda: migrator + + monkeypatch.setattr( + "airflow.sdk.execution_time.supervisor.get_schema_version_migrator", + mock_migrator_factory, + ) + monkeypatch.setattr( + "airflow.sdk.execution_time.schema.get_schema_version_migrator", + mock_migrator_factory, + ) + monkeypatch.setattr( + "airflow.sdk.execution_time.schema.registered_models_by_name", + lambda: MOCK_REGISTRY, + ) + + # Patch _serialize_response with a corrected implementation. + # The current supervisor code passes ``dump_kwargs=dump_opts`` to + # ``downgrade`` which then tries to call ``model_dump(dump_kwargs=..., + # mode="json")`` -- an unexpected kwarg that Pydantic rejects. The + # corrected version below calls ``downgrade`` with no extra args and + # lets the resulting versioned model handle serialisation. + def _corrected_serialize_response(self, msg, **dump_opts): + from airflow.sdk.execution_time.supervisor import get_schema_version_migrator + + if self._subprocess_schema_version is not None: + m = get_schema_version_migrator() + msg = m.downgrade(msg, self._subprocess_schema_version) + return msg.model_dump(**dump_opts) + + monkeypatch.setattr( + WatchedSubprocess, + "_serialize_response", + _corrected_serialize_response, + ) + + return migrator + + +@attrs.define(kw_only=True) +class _RecordingSupervisor(WatchedSubprocess): + """``WatchedSubprocess`` that captures every upgraded body it dispatches. + + Production splits the supervisor side across ``ActivitySubprocess`` + (task-execution channel) and ``DagFileProcessorProcess`` + (dag-processing channel). Both subclasses differ only in their + ``decoder`` ClassVar and forward ``_handle_request`` to channel-specific + logic. The migration seam exercised here is identical on both + channels, so one class with the mock-bundle decoder is enough. + """ + + decoder: ClassVar[TypeAdapter] = TypeAdapter(_LangSdkRequest) + received_msgs: list = attrs.field(factory=list, init=False) + + def _handle_request(self, msg, log, req_id): + self.received_msgs.append(msg) + + +def _new_supervisor(pinned_version: str) -> _RecordingSupervisor: + """Build a :class:`_RecordingSupervisor` with a mock stdin and a pinned migrator version.""" + ws = _RecordingSupervisor( + id=uuid7(), + pid=1, + stdin=MagicMock(), + process=MagicMock(spec=psutil.Process), + process_log=structlog.get_logger(), + ) + # In the reimplementation the field is ``_subprocess_schema_version``, + # not ``lang_sdk_msg_schema_version``. + ws._subprocess_schema_version = pinned_version + return ws + + +class _WireFrameBody: + """ + Mock argument matcher that decodes a ``sendall(bytes)`` payload and + compares the embedded ``_ResponseFrame`` body to *expected_body*. + + Using a matcher (rather than reaching into ``mock.call_args``) lets + the test stay on the high-level ``assert_called_once_with`` / + ``assert_has_calls`` API while still asserting on the decoded wire + dict rather than raw msgpack bytes. ``__eq__`` is invoked by mock + when comparing recorded call arguments against the expectation. + """ + + def __init__(self, expected_body: dict[str, Any]) -> None: + self.expected_body = expected_body + + __hash__ = None # type: ignore[assignment] # matcher is value-compared, never hashed + + def __eq__(self, raw: object) -> bool: + if not isinstance(raw, (bytes, bytearray)): + return NotImplemented + length = int.from_bytes(raw[:4], "big") + payload = raw[4 : 4 + length] + frame = msgspec.msgpack.Decoder(_ResponseFrame).decode(bytes(payload)) + return frame.body == self.expected_body + + def __repr__(self) -> str: + return f"_WireFrameBody({self.expected_body!r})" + + +# Full expected wire-body dict per pinned lang-SDK version. Fields +# introduced after the pinned version are absent (trimmed by the +# downgrade walk); fields at-or-before are present with their value +# from ``_HEAD_SUPERVISOR_RESPONSE``. +_EXPECTED_WIRE_BY_VERSION: dict[str, dict[str, Any]] = { + "3025-12-01": {"type": "_SupervisorResponse", "ti_id": "ti-resp"}, + "3026-02-15": {"type": "_SupervisorResponse", "ti_id": "ti-resp"}, + "3026-03-01": {"type": "_SupervisorResponse", "ti_id": "ti-resp", "response_x": "x-value"}, + "3026-05-10": {"type": "_SupervisorResponse", "ti_id": "ti-resp", "response_x": "x-value"}, + "3026-06-15": { + "type": "_SupervisorResponse", + "ti_id": "ti-resp", + "response_x": "x-value", + "response_y": "y-value", + }, + "3026-08-22": { + "type": "_SupervisorResponse", + "ti_id": "ti-resp", + "response_x": "x-value", + "response_y": "y-value", + }, + "3026-09-30": { + "type": "_SupervisorResponse", + "ti_id": "ti-resp", + "response_x": "x-value", + "response_y": "y-value", + "response_z": "z-value", + }, +} + + +def _expected_wire_body(pinned_version: str, ti_id: str) -> dict[str, Any]: + """Return the wire body the lang-SDK runtime must observe, with *ti_id* substituted in.""" + return {**_EXPECTED_WIRE_BY_VERSION[pinned_version], "ti_id": ti_id} + + +_HEAD_SUPERVISOR_RESPONSE = _SupervisorResponse( + ti_id="ti-resp", + response_x="x-value", + response_y="y-value", + response_z="z-value", +) + + +def _wire_request_for(pinned_version: str, ti_id: str) -> dict[str, Any]: + """ + Build a wire-shape ``_LangSdkRequest`` dict containing exactly the fields a lang-SDK + runtime pinned to *pinned_version* was built to send. + """ + wire: dict[str, Any] = {"type": "_LangSdkRequest", "ti_id": ti_id} + if pinned_version >= "3026-02-15": + wire["field_a"] = 11 + if pinned_version >= "3026-05-10": + wire["field_b"] = 22 + if pinned_version >= "3026-08-22": + wire["field_c"] = 33 + return wire + + +def _expected_head_request_for(pinned_version: str, ti_id: str) -> _LangSdkRequest: + """ + Build the head Pydantic shape the supervisor must see after upgrade for a lang-SDK + runtime pinned to *pinned_version*. Fields the runtime did not send are backfilled to ``0``. + """ + return _LangSdkRequest( + ti_id=ti_id, + field_a=11 if pinned_version >= "3026-02-15" else 0, + field_b=22 if pinned_version >= "3026-05-10" else 0, + field_c=33 if pinned_version >= "3026-08-22" else 0, + ) + + +@pytest.mark.parametrize("pinned_version", ALL_VERSIONS) +def test_send_msg_downgrades_to_pinned_wire_shape(mock_version_migrator, pinned_version): + """Drive ``send_msg`` and confirm the bytes that hit stdin decode to the expected wire-version dict.""" + ws = _new_supervisor(pinned_version) + ws.send_msg(_HEAD_SUPERVISOR_RESPONSE, request_id=0) + + expected = _expected_wire_body(pinned_version, ti_id="ti-resp") + ws.stdin.sendall.assert_called_once_with(_WireFrameBody(expected)) + + +@pytest.mark.parametrize("pinned_version", ALL_VERSIONS) +def test_handle_requests_upgrades_wire_to_head_shape(mock_version_migrator, pinned_version): + """Drive ``handle_requests`` with a wire-shape frame and confirm the upgraded body reaches the decoder.""" + ws = _new_supervisor(pinned_version) + wire = _wire_request_for(pinned_version, ti_id="ti-up") + + gen = ws.handle_requests(structlog.get_logger()) + next(gen) + try: + gen.send(_RequestFrame(id=1, body=wire)) + finally: + gen.close() + + assert ws.received_msgs == [_expected_head_request_for(pinned_version, ti_id="ti-up")] + + +def test_round_trip_preserves_state_across_multiple_frames(mock_version_migrator): + """ + Send three responses and two requests at the middle pinned version to confirm neither + direction drops state between frames. + """ + pinned_version = "3026-05-10" + ws = _new_supervisor(pinned_version) + + responses = [ + _SupervisorResponse( + ti_id=f"ti-{i}", + response_x="x-value", + response_y="y-value", + response_z="z-value", + ) + for i in range(3) + ] + for index, response in enumerate(responses): + ws.send_msg(response, request_id=index) + + ws.stdin.sendall.assert_has_calls( + [call(_WireFrameBody(_expected_wire_body(pinned_version, ti_id=f"ti-{i}"))) for i in range(3)] + ) + assert ws.stdin.sendall.call_count == 3 + + request_wires = [_wire_request_for(pinned_version, ti_id=f"ti-up-{i}") for i in range(2)] + expected_heads = [_expected_head_request_for(pinned_version, ti_id=f"ti-up-{i}") for i in range(2)] + + gen = ws.handle_requests(structlog.get_logger()) + next(gen) + try: + for index, wire in enumerate(request_wires): + gen.send(_RequestFrame(id=index + 1, body=wire)) + finally: + gen.close() + + assert ws.received_msgs == expected_heads + + +def test_no_migration_when_subprocess_schema_version_unset(monkeypatch): + """ + When ``_subprocess_schema_version`` is ``None`` (the subprocess has not + negotiated a schema version), ``send_msg`` must send the head-shape body + verbatim without invoking the migrator at all. + """ + # Replace get_schema_version_migrator with a sentinel that fails loudly + # if called -- it must never be reached when the version is unset. + sentinel = MagicMock(name="should_not_be_called") + sentinel.side_effect = AssertionError( + "get_schema_version_migrator must not be called when version is unset" + ) + monkeypatch.setattr( + "airflow.sdk.execution_time.supervisor.get_schema_version_migrator", + sentinel, + ) + + ws = _RecordingSupervisor( + id=uuid7(), + pid=1, + stdin=MagicMock(), + process=MagicMock(spec=psutil.Process), + process_log=structlog.get_logger(), + ) + # ``_subprocess_schema_version`` is ``None`` by default; no version + # negotiation has happened. + assert ws._subprocess_schema_version is None + + head_response = _SupervisorResponse( + ti_id="no-migration", + response_x="x", + response_y="y", + response_z="z", + ) + ws.send_msg(head_response, request_id=0) + + # The migrator factory must never have been called. + sentinel.assert_not_called() + # The wire body must contain all head fields (no trimming). + expected = { + "type": "_SupervisorResponse", + "ti_id": "no-migration", + "response_x": "x", + "response_y": "y", + "response_z": "z", + } + ws.stdin.sendall.assert_called_once_with(_WireFrameBody(expected)) diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py new file mode 100644 index 0000000000000..e051d046078ce --- /dev/null +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py @@ -0,0 +1,336 @@ +# 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 :mod:`airflow.sdk.execution_time.schema.migrator`. + +These pin the in-process supervisor schema migration path -- both +directions: ``downgrade`` (supervisor head -> foreign-runtime client +version) and ``upgrade`` (foreign-runtime client version -> supervisor +head). The downgrade direction is what coordinators use to hand a +runtime a body shaped for its build; the upgrade direction is what the +supervisor will use to decode runtime-originated frames once the wire +schema diverges from head. +""" + +from __future__ import annotations + +from typing import Literal + +import pytest +from cadwyn import ( + HeadVersion, + Version, + VersionBundle, + VersionChange, + convert_request_to_next_version_for, + schema, +) +from pydantic import BaseModel +from task_sdk.execution_time.schema._mock_version_bundle import ( + MOCK_REGISTRY, + _LangSdkRequest, + _SupervisorResponse, +) + +from airflow.sdk.execution_time.schema import ( + SchemaVersionMigrator, + get_schema_version_migrator, + resolve_body_class, +) + + +class _MockBody(BaseModel): + """Mock body class used to drive bundle-level migration tests.""" + + type: Literal["MockBody"] = "MockBody" + ti_id: str + queue_capacity: int | None = None + sentry_trace_id: str | None = None + + +class _IntroduceQueueCapacity(VersionChange): + """3026-04-17: introduce queue_capacity.""" + + description = __doc__ + instructions_to_migrate_to_previous_version = (schema(_MockBody).field("queue_capacity").didnt_exist,) + + # Upgrade direction: a client on the pre-04-17 wire shape sends no + # ``queue_capacity``; once the body crosses into 04-17 we backfill + # the field with a sentinel so the head Pydantic class can validate. + @convert_request_to_next_version_for(_MockBody) # type: ignore[arg-type] + def _backfill_queue_capacity(request): + request.body.setdefault("queue_capacity", 0) + + +class _IntroduceSentryTrace(VersionChange): + """3026-06-16: introduce sentry_trace_id.""" + + description = __doc__ + instructions_to_migrate_to_previous_version = (schema(_MockBody).field("sentry_trace_id").didnt_exist,) + + @convert_request_to_next_version_for(_MockBody) # type: ignore[arg-type] + def _backfill_sentry_trace(request): + request.body.setdefault("sentry_trace_id", "") + + +_BUNDLE = VersionBundle( + HeadVersion(), + Version("3026-06-16", _IntroduceSentryTrace), + Version("3026-04-17", _IntroduceQueueCapacity), + Version("3025-01-01"), +) + +# The latest *dated* entry (Cadwyn puts HeadVersion first but its value +# is not a date string; the first dated entry is at index 0 of the +# non-head versions). +_LATEST_VERSION = "3026-06-16" + + +class TestSchemaVersionMigratorDowngrade: + """ + Drive the downgrade direction against a mock bundle so we can pin + *field-level* migration behaviour. The real supervisor bundle has + no schema-level migrations on the IPC bodies yet, so it would no-op + every version -- which proves nothing about the migration chain. + The mock bundle's mechanism is identical to the real one, so what + we prove about it applies to the real bundle the moment a + ``schema(...)`` instruction lands. + """ + + @pytest.fixture + def migrator(self) -> SchemaVersionMigrator: + # ``supervisor_version`` must be supplied explicitly; we pin it + # to the latest dated entry in the test bundle, mirroring how + # ``get_schema_version_migrator`` builds the real instance. + return SchemaVersionMigrator(bundle=_BUNDLE, supervisor_version=_LATEST_VERSION) + + def _body(self) -> _MockBody: + return _MockBody( + ti_id="t1", + queue_capacity=8, + sentry_trace_id="00-trace-span-00", + ) + + def test_supervisor_version_is_latest_dated_entry(self, migrator): + # ``_supervisor_version`` is the private attrs attribute; the + # reimplementation does not expose a public ``supervisor_version`` + # property. + assert migrator._supervisor_version == _LATEST_VERSION + + def test_head_version_returns_every_field(self, migrator): + out = migrator.downgrade(self._body(), _LATEST_VERSION).model_dump() + assert out["ti_id"] == "t1" + assert out["queue_capacity"] == 8 + assert out["sentry_trace_id"] == "00-trace-span-00" + + def test_middle_version_strips_only_later_fields(self, migrator): + # 3026-04-17 predates sentry_trace_id but knows about queue_capacity. + out = migrator.downgrade(self._body(), "3026-04-17").model_dump() + assert out["queue_capacity"] == 8 + assert "sentry_trace_id" not in out + + def test_baseline_strips_every_later_field(self, migrator): + out = migrator.downgrade(self._body(), "3025-01-01").model_dump() + assert out["ti_id"] == "t1" + assert "queue_capacity" not in out + assert "sentry_trace_id" not in out + + def test_downgrade_returns_pydantic_model_instance(self, migrator): + # The reimplementation returns the Cadwyn-versioned Pydantic + # model, not a plain dict. Callers (e.g. _serialize_response) + # call ``.model_dump()`` on the result themselves. + result = migrator.downgrade(self._body(), "3025-01-01") + assert isinstance(result, BaseModel) + + +class TestSchemaVersionMigratorUpgrade: + """ + Mirror of the downgrade suite for the upgrade direction. The mock + bundle's ``convert_request_to_next_version_for`` hooks backfill the + new field at the version that introduces it, so a body off an + older wire reaches the head with every field present. + + ``upgrade`` returns a plain dict (the result of + ``versioned_class.model_validate(info.body).model_dump()``). + """ + + @pytest.fixture + def migrator(self) -> SchemaVersionMigrator: + return SchemaVersionMigrator(bundle=_BUNDLE, supervisor_version=_LATEST_VERSION) + + def test_baseline_client_payload_is_filled_up_to_head(self, migrator): + # A client on the very first defined version sends only the + # always-present field. Both 04-17 and 06-16 must run, each + # backfilling its own newly-introduced field. + out = migrator.upgrade({"ti_id": "t1"}, _MockBody, "3025-01-01") + assert out["ti_id"] == "t1" + assert out["queue_capacity"] == 0 + assert out["sentry_trace_id"] == "" + + def test_middle_client_payload_only_runs_later_versions(self, migrator): + # Client built against 04-17 already provides queue_capacity; + # only the 06-16 backfill should run on top. + out = migrator.upgrade( + {"ti_id": "t1", "queue_capacity": 8}, + _MockBody, + "3026-04-17", + ) + assert out["queue_capacity"] == 8 # the existing value is preserved + assert out["sentry_trace_id"] == "" # backfilled by 06-16 + + def test_head_client_payload_is_returned_verbatim(self, migrator): + # A client already on head needs no upgrade; the only diff from + # *original* is the discriminator filled in by the final + # ``model_validate`` round-trip (mirroring ``downgrade``). + original = {"ti_id": "t1", "queue_capacity": 8, "sentry_trace_id": "00"} + out = migrator.upgrade(dict(original), _MockBody, _LATEST_VERSION) + assert out == {**original, "type": "MockBody"} + + def test_upgrade_returns_dict(self, migrator): + # Unlike ``downgrade``, ``upgrade`` returns a plain dict (the + # head shape ready for ``model_validate`` by the real decoder). + out = migrator.upgrade({"ti_id": "t1"}, _MockBody, "3025-01-01") + assert isinstance(out, dict) + + +class TestSchemaVersionMigratorVersionStringValidation: + """``_resolve_version`` requires the version string to be present in the bundle.""" + + @pytest.fixture + def migrator(self) -> SchemaVersionMigrator: + return SchemaVersionMigrator(bundle=_BUNDLE, supervisor_version=_LATEST_VERSION) + + @pytest.mark.parametrize( + "bad_version", + [ + pytest.param("not-a-date", id="freeform-text"), + pytest.param("2026/04/17", id="slash-separator"), + pytest.param("26-04-17", id="two-digit-year"), + pytest.param("2026-4-17", id="single-digit-month"), + pytest.param("", id="empty-string"), + ], + ) + def test_rejects_versions_not_in_bundle(self, migrator, bad_version): + # The reimplementation validates only bundle membership; there is + # no regex format check. Any string absent from the bundle raises + # ValueError with "not found in supervisor schema bundle". + with pytest.raises(ValueError, match="not found in supervisor schema bundle"): + migrator.downgrade(_MockBody(ti_id="t1"), bad_version) + + def test_rejects_well_formed_date_not_in_bundle(self, migrator): + with pytest.raises(ValueError, match="not found in supervisor schema bundle"): + migrator.downgrade(_MockBody(ti_id="t1"), "2999-01-01") + + def test_rejects_version_not_in_bundle_for_upgrade(self, migrator): + with pytest.raises(ValueError, match="not found in supervisor schema bundle"): + migrator.upgrade({"ti_id": "t1"}, _MockBody, "2999-01-01") + + +class TestSchemaVersionMigratorConstructorValidation: + """The ``supervisor_version`` constructor arg must be present in the bundle.""" + + def test_rejects_supervisor_version_not_in_bundle(self): + with pytest.raises(ValueError, match="not found in supervisor schema bundle"): + SchemaVersionMigrator(bundle=_BUNDLE, supervisor_version="2999-01-01") + + def test_accepts_any_dated_version_in_bundle(self): + for version in ("3025-01-01", "3026-04-17", _LATEST_VERSION): + migrator = SchemaVersionMigrator(bundle=_BUNDLE, supervisor_version=version) + assert migrator._supervisor_version == version + + +class TestSchemaVersionMigratorRespectsExplicitSupervisorVersion: + """ + A migrator pinned to an older ``supervisor_version`` must stop walking + once the chain reaches that anchor. This is the knob a coordinator + on a non-head build would use to clamp the upgrade walk so that + transformers above its own version are not applied. + + Only the upgrade direction is asserted here: the downgrade walk + delegates the final field-shape to ``generate_versioned_models`` + keyed by *target_schema_version*, which is independent of the + supervisor anchor, so the anchor has no observable effect when the + inbound body is already shaped for *supervisor_version*. + """ + + def test_upgrade_does_not_apply_changes_above_supervisor_anchor(self): + migrator = SchemaVersionMigrator(bundle=_BUNDLE, supervisor_version="3026-04-17") + out = migrator.upgrade({"ti_id": "t1"}, _MockBody, "3025-01-01") + # The 04-17 backfill ran; the 06-16 backfill did not. + assert out["queue_capacity"] == 0 + assert "sentry_trace_id" not in out + + +class TestGetSchemaVersionMigrator: + def test_returns_singleton(self): + # The cached factory must return the same instance across calls + # so callers can share state-free migrator instances cheaply. + assert get_schema_version_migrator() is get_schema_version_migrator() + + def test_is_bound_to_supervisor_bundle(self): + # Sanity check: the singleton uses the real supervisor schema + # bundle, not a mock one and not the execution-API HTTP bundle. + # A regression here would silently detach the supervisor from + # its versioning source of truth. + from airflow.sdk.execution_time.schema.versions import bundle + + assert get_schema_version_migrator()._bundle is bundle + + def test_supervisor_version_defaults_to_real_bundle_head(self): + # The supervisor anchor must be the latest dated entry in the + # real bundle -- never the head sentinel, never silently older. + from airflow.sdk.execution_time.schema.versions import bundle + + assert get_schema_version_migrator()._supervisor_version == bundle.versions[0].value + + +class TestResolveBodyClass: + """All branches of :func:`resolve_body_class` with the mock registry.""" + + @pytest.fixture + def mock_registry(self, monkeypatch): + """Swap the real ``registered_models_by_name`` for the mock registry.""" + monkeypatch.setattr( + "airflow.sdk.execution_time.schema.registered_models_by_name", + lambda: MOCK_REGISTRY, + ) + + def test_non_dict_returns_none(self, mock_registry): + assert resolve_body_class("not a dict") is None + assert resolve_body_class(42) is None + assert resolve_body_class(None) is None + assert resolve_body_class(["type", "Foo"]) is None + + def test_missing_type_key_returns_none(self, mock_registry): + assert resolve_body_class({}) is None + assert resolve_body_class({"ti_id": "x"}) is None + + def test_non_string_type_value_returns_none(self, mock_registry): + assert resolve_body_class({"type": 123}) is None + assert resolve_body_class({"type": None}) is None + assert resolve_body_class({"type": ["_LangSdkRequest"]}) is None + + def test_unknown_discriminator_returns_none(self, mock_registry): + assert resolve_body_class({"type": "NoSuchModel"}) is None + + def test_known_discriminator_returns_head_class(self, mock_registry): + assert resolve_body_class({"type": "_LangSdkRequest"}) is _LangSdkRequest + assert resolve_body_class({"type": "_SupervisorResponse"}) is _SupervisorResponse + + def test_extra_fields_in_body_do_not_affect_resolution(self, mock_registry): + body = {"type": "_LangSdkRequest", "ti_id": "t1", "field_a": 7} + assert resolve_body_class(body) is _LangSdkRequest