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/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index f7f958bb84b6f..4570141b18bb3 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -2004,6 +2004,43 @@ workers: type: integer example: ~ default: "60" +sdk: + description: Settings for non-Python SDK runtime coordination + options: + coordinators: + description: | + JSON object mapping of coordinator keys to coordinator definitions. + + Each value is an object with ``classpath`` and optional ``kwargs``. + ``classpath`` is resolved via ``import_string`` and constructed with + ``kwargs`` on first use. Entries are + independent instances, so the same ``classpath`` can be configured + multiple times under different names with different ``kwargs`` (for + example, two ``JavaCoordinator`` instances pinned to different JDK + versions). + version_added: 3.3.0 + type: string + example: | + { + "jdk-17": { + "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", + "kwargs": {"java_executable": "/usr/lib/jvm/java-17-openjdk/bin/java", "jvm_args": ["-Xmx1024m"]} + } + } + default: ~ + queue_to_coordinator: + description: | + JSON mapping of queue names to a coordinator key from + ``[sdk] coordinators``. + + When a task's ``language`` field is not set, this mapping is checked + to route the task to a configured coordinator instance based on its + queue. This is useful when queues are used as environment or + isolation identifiers (e.g. ``legacy-java``, ``modern-java``). + version_added: 3.3.0 + type: string + example: '{"legacy-java": "jdk-11", "modern-java": "jdk-17"}' + default: ~ api_auth: description: Settings relating to authentication on the Airflow APIs options: diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index eff6ff0771474..5080e75750a48 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/airflow-core/src/airflow/executors/base_executor.py @@ -652,11 +652,10 @@ def run_workload( if isinstance(workload, ExecuteTask): from airflow.sdk.execution_time.supervisor import supervise_task + from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO as SDKTaskInstanceDTO - # workload.ti is a TaskInstanceDTO which duck-types as TaskInstance. - # TODO: Create a protocol for this. return supervise_task( - ti=workload.ti, # type: ignore[arg-type] + ti=SDKTaskInstanceDTO.model_validate(workload.ti, from_attributes=True), bundle_info=workload.bundle_info, dag_rel_path=workload.dag_rel_path, token=workload.token, diff --git a/airflow-core/src/airflow/executors/workloads/task.py b/airflow-core/src/airflow/executors/workloads/task.py index 9af3f33c10efd..4c563f34351d8 100644 --- a/airflow-core/src/airflow/executors/workloads/task.py +++ b/airflow-core/src/airflow/executors/workloads/task.py @@ -28,6 +28,8 @@ from airflow.utils.state import TaskInstanceState if TYPE_CHECKING: + import uuid + from airflow.api_fastapi.auth.tokens import JWTGenerator from airflow.models.taskinstance import TaskInstance as TIModel from airflow.models.taskinstancekey import TaskInstanceKey diff --git a/devel-common/src/tests_common/pytest_plugin.py b/devel-common/src/tests_common/pytest_plugin.py index 6f47d98fd649d..defee35f0e30c 100644 --- a/devel-common/src/tests_common/pytest_plugin.py +++ b/devel-common/src/tests_common/pytest_plugin.py @@ -2519,7 +2519,6 @@ def execute(self, context): from uuid6 import uuid7 from airflow.sdk import DAG - from airflow.sdk.api.datamodels._generated import TaskInstance from airflow.sdk.execution_time.comms import BundleInfo, StartupDetails from airflow.timetables.base import TimeRestriction @@ -2547,6 +2546,15 @@ def _create_task_instance( should_retry: bool | None = None, max_tries: int | None = None, ) -> RuntimeTaskInstance: + from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS + + if AIRFLOW_V_3_3_PLUS: + from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO + else: + from airflow.sdk.api.datamodels._generated import ( # type: ignore[no-redef,assignment] + TaskInstance as TaskInstanceDTO, + ) + from airflow.sdk.api.datamodels._generated import DagRun, DagRunState, TIRunContext from airflow.utils.types import DagRunType @@ -2624,14 +2632,17 @@ def _create_task_instance( } startup_details = StartupDetails( - ti=TaskInstance( + ti=TaskInstanceDTO( id=ti_id, task_id=task.task_id, dag_id=dag_id, run_id=run_id, try_number=try_number, - map_index=map_index, + map_index=map_index, # type: ignore[arg-type] dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), dag_rel_path="", bundle_info=BundleInfo(name="anything", version="any"), diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index f98d6edde306e..4fa36ffa98279 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -860,6 +860,7 @@ iTerm iterm itertools Jarek +JavaCoordinator javascript jaydebeapi Jdbc @@ -897,6 +898,7 @@ jsonl juli Jupyter jupyter +jvm jwks JWT jwt @@ -1130,6 +1132,7 @@ openai openapi openfaas OpenID +openjdk openlineage OpenSearch opensearch @@ -1860,6 +1863,7 @@ XComs Xiaodong xlarge xml +Xmx xpath XSS xyz diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index f29d123c740c1..a5e63d925f795 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -85,7 +85,6 @@ def stub( Stub tasks exist in the Dag graph only, but the execution must happen in an external environment via the Task Execution Interface. - """ return task_decorator_factory( decorated_operator_class=_StubOperator, 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/scripts/in_container/java_sdk_setup.sh b/scripts/in_container/java_sdk_setup.sh new file mode 100644 index 0000000000000..b3437b7fc4200 --- /dev/null +++ b/scripts/in_container/java_sdk_setup.sh @@ -0,0 +1,73 @@ +#!/usr/bin/env bash +# 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. + + + +# 1. Check Java +check_java() { + local java_bin="/files/openjdk/bin/java" + local version_output + + # First check if the locally installed OpenJDK exists and works. + if [ -x "$java_bin" ] && version_output=$("$java_bin" -version 2>&1); then + echo "Found existing OpenJDK at $java_bin. OK." + return + fi + + # On macOS, /usr/bin/java exists as a shim even without a JDK installed, + # so we must test with `java -version` directly. + if ! version_output=$(java -version 2>&1); then + echo "Java is not installed." + install_java + return + fi + + local java_version + java_version=$(echo "$version_output" | head -n1 | sed -E 's/.*"([0-9]+)(\.[0-9]+)*.*/\1/') + + if ! [[ "$java_version" =~ ^[0-9]+$ ]]; then + echo "Could not determine Java version." + install_java + return + fi + + if [ "$java_version" -ge 11 ]; then + echo "Java $java_version detected. OK." + else + echo "Java version $java_version found, but >= 11 is required." + install_java + fi +} + + +install_java() { + echo "Installing OpenJDK 11 in Breeze..." + + curl -L -o /files/openjdk-11-aarch64.tar.gz \ + https://github.com/adoptium/temurin11-binaries/releases/download/jdk-11.0.30+7/OpenJDK11U-jdk_aarch64_linux_hotspot_11.0.30_7.tar.gz + + rm -rf /files/openjdk && mkdir -p /files/openjdk && \ + tar -xzf /files/openjdk-11-aarch64.tar.gz --strip-components=1 -C /files/openjdk + + /files/openjdk/bin/java -version + echo "" +} + +check_java +# Install Java Provider +pip install -e /opt/airflow/providers/languages/java/ diff --git a/task-sdk/.pre-commit-config.yaml b/task-sdk/.pre-commit-config.yaml index 100a6e6490849..c1a19ac4c1e29 100644 --- a/task-sdk/.pre-commit-config.yaml +++ b/task-sdk/.pre-commit-config.yaml @@ -43,6 +43,8 @@ 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/coordinator\.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/docs/airflow-metadata.schema.json b/task-sdk/docs/airflow-metadata.schema.json new file mode 100644 index 0000000000000..31e41c96433bc --- /dev/null +++ b/task-sdk/docs/airflow-metadata.schema.json @@ -0,0 +1,65 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://airflow.apache.org/schemas/sdk-executable/airflow-metadata-1.0.schema.json", + "title": "Airflow Executable SDK Bundle Metadata", + "description": "Build-time manifest declaring DAG and task identifiers exposed by an Airflow native-executable SDK bundle. See the Bundle Spec Format documentation in apache-airflow-task-sdk.", + "type": "object", + "required": ["format_version", "sdk", "source", "dags"], + "additionalProperties": true, + "properties": { + "format_version": { + "type": "string", + "description": "Bundle-spec version this manifest conforms to (currently '1.0').", + "pattern": "^[0-9]+\\.[0-9]+(\\.[0-9]+)?$" + }, + "sdk": { + "type": "object", + "description": "Identifies the SDK that produced the bundle.", + "required": ["language", "version"], + "additionalProperties": true, + "properties": { + "language": { + "type": "string", + "description": "Lower-case source-language identifier (e.g. 'go', 'rust', 'cpp', 'zig').", + "pattern": "^[a-z][a-z0-9_+.\\-]*$" + }, + "version": { + "type": "string", + "description": "SDK version used at build time.", + "minLength": 1 + } + } + }, + "source": { + "type": "string", + "description": "Original filename of the primary DAG source file (e.g. 'example.go'). The file's bytes are embedded in the bundle's source region; this field is a display name used by the Airflow UI.", + "minLength": 1 + }, + "dags": { + "type": "object", + "description": "Mapping of dag_id to DAG entry. Every dag_id the bundle exposes must appear here.", + "minProperties": 1, + "additionalProperties": { + "$ref": "#/$defs/dagEntry" + } + } + }, + "$defs": { + "dagEntry": { + "type": "object", + "description": "Static description of a single DAG declared in the bundle.", + "required": ["tasks"], + "additionalProperties": true, + "properties": { + "tasks": { + "type": "array", + "description": "Static list of task_ids declared in the DAG.", + "items": { + "type": "string", + "minLength": 1 + } + } + } + } + } +} diff --git a/task-sdk/docs/bundle-spec.rst b/task-sdk/docs/bundle-spec.rst new file mode 100644 index 0000000000000..d231f0cf0ec03 --- /dev/null +++ b/task-sdk/docs/bundle-spec.rst @@ -0,0 +1,252 @@ + .. 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. + +Bundle Spec Format +================== + +This document specifies the on-disk format of a build artifact produced by an +Airflow native-executable SDK (Go, Rust, C++, Zig, ...) and consumed by +:class:`~airflow.sdk.coordinators.executable.ExecutableCoordinator` +at deployment time. + +The goal is a single, language-agnostic *bundle* shape so that scheduler, +worker, and UI behave identically regardless of which compiled SDK produced +the DAG. + +Format version: ``1.0``. + +Container +--------- + +A bundle is **the compiled executable itself, with a fixed-format footer +appended after the binary's normal end-of-file**. The executable remains +directly runnable; the footer is data that follows the last byte the OS +loader cares about and is invisible to ``exec()``. There is no enclosing +archive. + +A bundle file therefore has three regions, in order from offset 0: + +1. The native executable (ELF / Mach-O / PE), including any code-signing + structures the platform appends. +2. The primary DAG source file, embedded verbatim (UTF-8). MAY have length 0. +3. The build-time manifest (``airflow-metadata.yaml`` content, UTF-8). + +The file ends with a fixed 32-byte trailer that locates regions (2) and (3) +and identifies the file as a bundle. See :ref:`bundle-trailer-layout`. + +Filenames follow OS conventions for executables: no extension on Linux/macOS, +``.exe`` on Windows. The scanner identifies bundles by the trailer's magic, +not by the filename. + +.. _bundle-trailer-layout: + +Trailer Layout +-------------- + +The last 32 bytes of a conforming bundle are the trailer. All multi-byte +integers are little-endian. + +:: + + bytes 0..3 source_len uint32 length of the source region in bytes + bytes 4..7 metadata_len uint32 length of the metadata region in bytes + bytes 8..11 footer_ver uint32 currently 1 + bytes 12..23 reserved 12 bytes, MUST be zero + bytes 24..31 magic 8 bytes ASCII "AFBNDL01" + +The magic is the byte sequence ``0x41 0x46 0x42 0x4E 0x44 0x4C 0x30 0x31`` +(``"AFBNDL01"``). The trailing ``01`` is the footer-format version repeated +in ASCII so a human can identify a bundle at a glance +(``tail -c 8 ./mybundle | xxd``); the binary ``footer_ver`` field is the +authoritative source of truth for parsing. + +Reader algorithm: + +1. Open the file. Seek to ``EOF - 32``. Read 32 bytes. +2. Compare bytes ``24..31`` against ``"AFBNDL01"``. If different, the file + is not a bundle; the scanner MUST ignore it. +3. Parse ``footer_ver``. If unknown, fail with a versioning error. +4. Compute ``metadata_start = filesize - 32 - metadata_len`` and + ``source_start = metadata_start - source_len``. +5. Read ``metadata_len`` bytes from ``metadata_start`` for the manifest. +6. Read ``source_len`` bytes from ``source_start`` for the source view. + If ``source_len == 0``, no source is embedded; the UI displays + "(source not available)". +7. Validate ``source_start >= 0`` and that the implied binary region + (``[0, source_start)``) is non-empty. + +Source comes *before* metadata so a future ``footer_ver`` MAY introduce +additional trailing blobs (e.g. signed checksums, compressed deps) by +extending the trailer rather than inserting between existing blobs. + +.. _bundle-metadata-schema: + +``airflow-metadata.yaml`` schema +-------------------------------- + +The metadata region carries the same YAML manifest documented previously, +produced at build time from a static scan of the DAG source. A +machine-readable JSON Schema is published at +:download:`airflow-metadata.schema.json` for use by build tooling, validators, +and editors. + +.. code-block:: yaml + + format_version: "1.0" + sdk: + language: go + version: "0.1.0" + source: example.go + dags: + example_dag: + tasks: + - extract + - transform + - load + another_dag: + tasks: + - run + +Top-level keys: + +``format_version`` (string, required) + The bundle-spec version this manifest conforms to. Currently ``"1.0"``. + +``sdk`` (mapping, required) + Identifies the SDK that produced the bundle. + + - ``language`` (string, required): lower-case source-language identifier + (e.g. ``go``, ``rust``, ``cpp``, ``zig``). + - ``version`` (string, required): SDK version used at build time. + +``source`` (string, required) + Original filename of the primary DAG source file (e.g. ``example.go``). + The file's bytes live in the source region of the bundle, not at this + path; this field is a display name the Airflow UI uses to label the + source-view panel and pick a syntax-highlighting mode from the + extension. + +``dags`` (mapping, required) + Mapping of ``dag_id`` to a *DAG entry*. Every ``dag_id`` the bundle + exposes MUST appear here. The scanner uses these keys to match a DAG + parsing or task-execution request to the bundle that owns it. + +DAG entry fields: + +``tasks`` (list of strings, required) + Static list of ``task_id``\ s declared in the DAG. Empty lists are + permitted but discouraged. + +Unrecognised top-level or DAG-entry keys MUST be ignored by the consumer so +that future SDK versions can extend the manifest without breaking older +runtimes. + +Examples +-------- + +Go bundle:: + + example + ├── ELF/Mach-O/PE executable + ├── source region: contents of example.go + ├── metadata region: airflow-metadata.yaml (source: example.go) + └── trailer (32 B): AFBNDL01 magic + lengths + +Rust bundle:: + + pipeline + ├── ELF/Mach-O/PE executable + ├── source region: contents of main.rs + ├── metadata region: airflow-metadata.yaml (source: main.rs) + └── trailer (32 B): AFBNDL01 magic + lengths + +The bundle is one file. ``./example`` runs the binary; the appended data +is invisible to ``exec()``. + +Build Pipeline Ordering +----------------------- + +The footer is appended after the executable is otherwise complete. Producers +that perform additional post-build steps MUST observe the following order: + +- **Strip** debug symbols *before* appending the footer. Strip + implementations operate on the binary's defined end and either leave + trailing data intact or truncate it; do not rely on either behaviour. +- **Code-sign** *after* appending the footer on platforms whose signature + covers the entire file (Authenticode, certain notarisation flows). The + signature then attests to the footer's contents along with the binary. +- **Compressors** such as UPX are NOT supported. They rewrite the file + end-to-end and destroy the trailer. + +Determinism: the trailer is byte-identical for byte-identical inputs, so a +deterministic build plus a canonical (sorted-key) manifest serialisation +yields a byte-identical bundle file. + +Deployment Layout +----------------- + +Bundle files are placed **as-is** in any of the directories configured as the +``executables_root`` kwarg on the +:class:`~airflow.sdk.coordinators.executable.ExecutableCoordinator` entry +under ``[sdk] coordinators``. The scanner enumerates regular files in each +directory, reads the last 32 bytes of each, and treats files whose magic +matches ``"AFBNDL01"`` as bundles. Files without the magic are silently +ignored, so non-bundle files (READMEs, dotfiles) MAY share the directory +without interfering with the scan. + +:: + + /opt/airflow/executable-bundles/ + ├── example + ├── pipeline + └── analytics + +At task-execution time the runtime execs the bundle file directly with the +coordinator arguments (``--comm=`` / ``--logs=``). No extraction, +no transient cache directory, no chmod-after-extract step is required: the +file is already a runnable executable with the appropriate permission bits +preserved by the build pipeline. + +The compiled executable MUST honour the SDK coordinator protocol — +``--comm=`` / ``--logs=`` socket-based IPC. + +See :class:`~airflow.sdk.coordinators.executable.ExecutableCoordinator` +for the consumer-side coordinator. + +Inspection +---------- + +Because the bundle is a single executable rather than an archive, +inspecting the embedded source and manifest requires a small CLI rather +than an off-the-shelf ``unzip``. The Go SDK's ``airflow-go-pack`` tool +provides an ``inspect`` subcommand that dumps both regions; equivalent +helpers are expected from each language's packer. + +Compatibility and Versioning +---------------------------- + +- The current bundle-spec format version is ``1.0``; the current trailer + format version is ``1`` (``footer_ver = 1``). +- Backward-incompatible bundle-spec changes increment the major component + of ``format_version`` and are gated behind an explicit opt-in on the + consumer side. +- New optional manifest fields MAY be added in minor versions and MUST be + ignored by older consumers. +- New trailer-format versions append fields after ``footer_ver`` (consuming + the reserved region) or extend the trailer with additional trailing + blobs ahead of the magic. Older readers MUST reject unknown + ``footer_ver`` rather than guessing. diff --git a/task-sdk/docs/index.rst b/task-sdk/docs/index.rst index d1b26544c8855..04e509369a648 100644 --- a/task-sdk/docs/index.rst +++ b/task-sdk/docs/index.rst @@ -175,3 +175,4 @@ For the full public API reference, see the :doc:`api` page. deferred-vs-async-operators api concepts + bundle-spec diff --git a/task-sdk/src/airflow/sdk/coordinators/executable/__init__.py b/task-sdk/src/airflow/sdk/coordinators/executable/__init__.py new file mode 100644 index 0000000000000..f3bdcc9bfd216 --- /dev/null +++ b/task-sdk/src/airflow/sdk/coordinators/executable/__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. +"""Native executable runtime coordinator for the Apache Airflow Task SDK.""" + +from __future__ import annotations + +from airflow.sdk.coordinators.executable.coordinator import ExecutableCoordinator + +__all__ = ["ExecutableCoordinator", "__version__"] + +__version__ = "0.1.0" diff --git a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py new file mode 100644 index 0000000000000..d63adfb7a93c5 --- /dev/null +++ b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py @@ -0,0 +1,324 @@ +# +# 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. +"""Native executable coordinator that launches a binary subprocess for task execution.""" + +from __future__ import annotations + +import os +import pathlib +import selectors +import socket +import struct +import subprocess +import time +from typing import TYPE_CHECKING, Any, NamedTuple, cast + +import attrs +import psutil +import structlog +import yaml + +from airflow.sdk.execution_time.coordinator import BaseCoordinator +from airflow.sdk.execution_time.supervisor import ActivitySubprocess + +if TYPE_CHECKING: + from collections.abc import Sequence + + from structlog.typing import FilteringBoundLogger + from typing_extensions import Self + + from airflow.sdk.api.client import Client + from airflow.sdk.api.datamodels._generated import BundleInfo + from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO + +log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators.executable") + + +FOOTER_MAGIC = b"AFBNDL01" +FOOTER_SIZE = 32 +FOOTER_VERSION = 1 + + +class _Footer(NamedTuple): + source_len: int + metadata_len: int + footer_ver: int + + +def _read_footer(path: pathlib.Path) -> _Footer | None: + try: + size = path.stat().st_size + except OSError: + return None + if size < FOOTER_SIZE: + return None + try: + with open(path, "rb") as f: + f.seek(size - FOOTER_SIZE) + trailer = f.read(FOOTER_SIZE) + except OSError: + return None + if len(trailer) != FOOTER_SIZE or trailer[24:32] != FOOTER_MAGIC: + return None + source_len, metadata_len, footer_ver = struct.unpack_from(" dict[str, Any] | None: + try: + footer = _read_footer(path) + except ValueError: + return None + if footer is None: + return None + metadata_start = path.stat().st_size - FOOTER_SIZE - footer.metadata_len + with open(path, "rb") as f: + f.seek(metadata_start) + metadata_bytes = f.read(footer.metadata_len) + try: + data = yaml.safe_load(metadata_bytes.decode("utf-8")) + except (UnicodeDecodeError, yaml.YAMLError): + return None + if not isinstance(data, dict): + return None + return data + + +def _dag_ids(metadata: dict[str, Any]) -> set[str]: + dags = metadata.get("dags") + if not isinstance(dags, dict): + return set() + return set(dags.keys()) + + +@attrs.define +class _Bundle: + path: pathlib.Path + + @classmethod + def find(cls, executables_root: Sequence[pathlib.Path], dag_id: str) -> Self: + for root in executables_root: + for p in root.iterdir(): + if not p.is_file() or not os.access(p, os.X_OK): + continue + if (metadata := _read_bundle_metadata(p)) is None: + continue + if dag_id in _dag_ids(metadata): + return cls(p.resolve()) + resolved_paths = os.pathsep.join(str(r.resolve()) for r in executables_root) + raise FileNotFoundError( + f"cannot find executable bundle containing dag_id={dag_id!r} in {resolved_paths}" + ) + + +def _start_server() -> socket.socket: + server = socket.socket() + server.bind(("127.0.0.1", 0)) + server.setblocking(True) + server.listen(1) # Just need to listen to the child process. + return server + + +def _accept_connections( + servers: dict[str, socket.socket], + proc: subprocess.Popen, + *, + max_wait: float = 10.0, +) -> dict[str, socket.socket]: + """Block until the executable process connects to servers.""" + accepted: dict[str, socket.socket] = {} + with selectors.DefaultSelector() as sel: + for key, soc in servers.items(): + sel.register(soc, selectors.EVENT_READ, data=key) + deadline = time.monotonic() + max_wait + while len(accepted) < len(servers): + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("process did not connect within timeout") + if proc.poll() is not None: + raise RuntimeError(f"process exited with {proc.returncode} before connecting") + for event, _ in sel.select(timeout=min(remaining, 1.0)): + log.debug("Accepting child process connection", key=(key := event.data)) + conn, _ = cast("socket.socket", event.fileobj).accept() + sel.unregister(servers[key]) + accepted[key] = conn + return accepted + + +@attrs.define(kw_only=True) +class _ExecutableActivitySubprocess(ActivitySubprocess): + """Native executable task runner process.""" + + _comm_server: socket.socket + _logs_server: socket.socket + _child_process: subprocess.Popen + + # Keep track of channels used to pipe subprocess stdout and stderr so we can + # close them on exit. The "read" side is handled by _register_pipe_readers + # callbacks so we don't need to worry about them. + _stdout_w: socket.socket + _stderr_w: socket.socket + + @classmethod + def start( # type: ignore[override] + cls, + *, + what: TaskInstanceDTO, + dag_rel_path: str | os.PathLike[str], + bundle_info, + logger: FilteringBoundLogger | None = None, + sentry_integration: str = "", + executable_path: str, + **kwargs, + ) -> Self: + comm_server = _start_server() + logs_server = _start_server() + + stdout_r, stdout_w = socket.socketpair() + stderr_r, stderr_w = socket.socketpair() + + comm_host, comm_port = comm_server.getsockname() + logs_host, logs_port = logs_server.getsockname() + + proc = subprocess.Popen( + [ + executable_path, + f"--comm={comm_host}:{comm_port}", + f"--logs={logs_host}:{logs_port}", + ], + stdout=stdout_w.makefile("wb", buffering=0).fileno(), + stderr=stderr_w.makefile("wb", buffering=0).fileno(), + ) + log.info("Starting subprocess", pid=proc.pid, executable=executable_path) + socks = _accept_connections({"comm": comm_server, "logs": logs_server}, proc) + + self = cls( + id=what.id, + pid=proc.pid, + process=psutil.Process(proc.pid), + process_log=logger or structlog.get_logger(logger_name="task").bind(), + start_time=time.monotonic(), + stdin=socks["comm"], + child_process=proc, + comm_server=comm_server, + logs_server=logs_server, + stdout_w=stdout_w, + stderr_w=stderr_w, + **kwargs, + ) + self._register_pipe_readers(stdout_r, stderr_r, socks["comm"], socks["logs"]) + self._on_child_started( + ti=what, + dag_rel_path=dag_rel_path, + bundle_info=bundle_info, + sentry_integration=sentry_integration, + ) + return self + + def wait(self) -> int: + code = super().wait() + self._close_unused_sockets(self._comm_server, self._logs_server, self._stdout_w, self._stderr_w) + return code + + +def _convert_executables_root( + value: None | os.PathLike[str] | pathlib.Path | list[os.PathLike[str] | pathlib.Path], +) -> list[pathlib.Path]: + if value is None: + return [] + if isinstance(value, (str, os.PathLike, pathlib.Path)): + return [pathlib.Path(value)] + return [pathlib.Path(v) for v in value] + + +@attrs.define(kw_only=True) +class ExecutableCoordinator(BaseCoordinator): + """ + Coordinator that launches a native executable subprocess for task execution. + + Configuration is taken from the ``[sdk] coordinators`` entry that constructs + this instance:: + + { + "name": "go", + "classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator", + "kwargs": { + "executables_root": ["~/airflow/executable-bundles"], + }, + } + + :param executables_root: A list of directories scanned for executable + bundles when a Python stub DAG delegates task execution to a native + runtime. + """ + + sdk: str = "executable" + file_extension: str = "" + executables_root: list[pathlib.Path] = attrs.field(converter=_convert_executables_root, factory=list) + + def _resolve_executable(self, *, what: TaskInstanceDTO) -> str: + """ + Resolve the executable path for *what*. + + Looks up the bundle whose embedded manifest declares ``what.dag_id`` + in the configured ``executables_root`` directories. + """ + if not self.executables_root: + raise ValueError( + "The executables_root kwarg must be set on the ExecutableCoordinator " + "to resolve the executable for task execution." + ) + return str(_Bundle.find(self.executables_root, what.dag_id).path) + + def execute_task( + self, + *, + what: TaskInstanceDTO, + dag_rel_path: str | os.PathLike[str], + bundle_info: BundleInfo, + client: Client, + logger: FilteringBoundLogger | None = None, + sentry_integration: str = "", + subprocess_logs_to_stdout: bool, + **kwargs, + ) -> BaseCoordinator.ExecutionResult: + executable_path = self._resolve_executable(what=what) + process = _ExecutableActivitySubprocess.start( + what=what, + dag_rel_path=dag_rel_path, + bundle_info=bundle_info, + client=client, + logger=logger, + subprocess_logs_to_stdout=subprocess_logs_to_stdout, + sentry_integration=sentry_integration, + executable_path=executable_path, + ) + exit_code = process.wait() + return self.ExecutionResult(exit_code, process.final_state) diff --git a/task-sdk/src/airflow/sdk/coordinators/java/__init__.py b/task-sdk/src/airflow/sdk/coordinators/java/__init__.py new file mode 100644 index 0000000000000..daf8fce338d23 --- /dev/null +++ b/task-sdk/src/airflow/sdk/coordinators/java/__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. +"""Java runtime coordinator for the Apache Airflow Task SDK.""" + +from __future__ import annotations + +from airflow.sdk.coordinators.java.coordinator import JavaCoordinator + +__all__ = ["JavaCoordinator", "__version__"] + +__version__ = "0.1.0" diff --git a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py new file mode 100644 index 0000000000000..9806533028ff8 --- /dev/null +++ b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py @@ -0,0 +1,261 @@ +# +# 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. +"""Java runtime coordinator that launches a JVM subprocess for Dag file processing and task execution.""" + +from __future__ import annotations + +import email +import os +import pathlib +import selectors +import socket +import subprocess +import time +import zipfile +from typing import TYPE_CHECKING, cast + +import attrs +import psutil +import structlog + +from airflow.sdk.execution_time.coordinator import BaseCoordinator +from airflow.sdk.execution_time.supervisor import ActivitySubprocess + +if TYPE_CHECKING: + from collections.abc import Sequence + + from structlog.typing import FilteringBoundLogger + from typing_extensions import Self + + from airflow.sdk.api.client import Client + from airflow.sdk.api.datamodels._generated import BundleInfo + from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO + +log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators.java") + + +def _start_server() -> socket.socket: + server = socket.socket() + server.bind(("127.0.0.1", 0)) + server.setblocking(True) + server.listen(1) # Just need to listen to the child process. + return server + + +def _calculate_classpath(jars_root: Sequence[pathlib.Path]) -> str: + jars = (p.as_posix() for root in jars_root for p in root.iterdir() if p.suffix == ".jar") + return os.pathsep.join(jars) + + +@attrs.define +class _MainJar: + path: pathlib.Path + main_class: str + schema_version: str | None + + @classmethod + def find(cls, jars_root: Sequence[pathlib.Path]) -> Self: + for root in jars_root: + for p in root.iterdir(): + if p.suffix != ".jar": + continue + with zipfile.ZipFile(p) as zf: + with zf.open("META-INF/MANIFEST.MF") as f: + manifest = email.message_from_binary_file(f) + if main_class := manifest["Main-Class"]: + return cls(p, main_class, manifest.get("Airflow-SDK-Supervisor-Schema-Version")) + resolved_paths = os.pathsep.join(str(p.resolve()) for p in jars_root) + raise FileNotFoundError(f"cannot fine main class in {resolved_paths}") + + +def _accept_connections( + servers: dict[str, socket.socket], + proc: subprocess.Popen, + *, + max_wait: float = 10.0, +) -> dict[str, socket.socket]: + """Block until the Java process connects to servers.""" + accepted: dict[str, socket.socket] = {} + with selectors.DefaultSelector() as sel: + for key, soc in servers.items(): + sel.register(soc, selectors.EVENT_READ, data=key) + deadline = time.monotonic() + max_wait + while len(accepted) < len(servers): + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("process did not connect within timeout") + if proc.poll() is not None: + raise RuntimeError(f"process exited with {proc.returncode} before connecting") + for event, _ in sel.select(timeout=min(remaining, 1.0)): + log.debug("Accepting child process connection", key=(key := event.data)) + conn, _ = cast("socket.socket", event.fileobj).accept() + sel.unregister(servers[key]) + accepted[key] = conn + return accepted + + +@attrs.define(kw_only=True) +class _JavaActivitySubprocess(ActivitySubprocess): + """Java task runner process.""" + + _comm_server: socket.socket + _logs_server: socket.socket + _subprocess: subprocess.Popen + + # Keep track of channels used to pipe subprocess stdout and stderr so we can + # close them on exit. The "read" side is handled by _register_pipe_readers + # callbacks so we don't need to worry about them. + _stdout_w: socket.socket + _stderr_w: socket.socket + + @classmethod + def start( # type: ignore[override] + cls, + *, + what: TaskInstanceDTO, + dag_rel_path: str | os.PathLike[str], + bundle_info, + logger: FilteringBoundLogger | None = None, + sentry_integration: str = "", + java_executable: str, + jvm_args: list[str], + jars_root: Sequence[pathlib.Path], + **kwargs, + ) -> Self: + jar = _MainJar.find(jars_root) + + comm_server = _start_server() + logs_server = _start_server() + + stdout_r, stdout_w = socket.socketpair() + stderr_r, stderr_w = socket.socketpair() + + comm_host, comm_port = comm_server.getsockname() + logs_host, logs_port = logs_server.getsockname() + + proc = subprocess.Popen( + [ + java_executable, + "-classpath", + _calculate_classpath(jars_root), + *jvm_args, + jar.main_class, + # Arguments to MainClass... + f"--comm={comm_host}:{comm_port}", + f"--logs={logs_host}:{logs_port}", + ], + stdout=stdout_w.makefile("wb", buffering=0).fileno(), + stderr=stderr_w.makefile("wb", buffering=0).fileno(), + ) + log.info("Starting subprocess", pid=proc.pid) + socks = _accept_connections({"comm": comm_server, "logs": logs_server}, proc) + + self = cls( + id=what.id, + pid=proc.pid, + process=psutil.Process(proc.pid), + process_log=logger or structlog.get_logger(logger_name="task").bind(), + start_time=time.monotonic(), + stdin=socks["comm"], + subprocess=proc, + subprocess_schema_version=jar.schema_version, + comm_server=comm_server, + logs_server=logs_server, + stdout_w=stdout_w, + stderr_w=stderr_w, + **kwargs, + ) + self._register_pipe_readers(stdout_r, stderr_r, socks["comm"], socks["logs"]) + self._on_child_started( + ti=what, + dag_rel_path=dag_rel_path, + bundle_info=bundle_info, + sentry_integration=sentry_integration, + ) + return self + + def wait(self) -> int: + code = super().wait() + self._close_unused_sockets(self._comm_server, self._logs_server, self._stdout_w, self._stderr_w) + return code + + +def _convert_jars_root( + value: None | os.PathLike[str] | pathlib.Path | list[os.PathLike[str] | pathlib.Path], +) -> list[pathlib.Path]: + if value is None: + return [] + if isinstance(value, (str, os.PathLike, pathlib.Path)): + return [pathlib.Path(value)] + return [pathlib.Path(v) for v in value] + + +@attrs.define(kw_only=True) +class JavaCoordinator(BaseCoordinator): + """ + Coordinator that launches a JVM subprocess for DAG parsing and task execution. + + Configuration is taken from the ``[sdk] coordinators`` entry that constructs + this instance:: + + { + "name": "jdk-17", + "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", + "kwargs": { + "java_executable": "/usr/lib/jvm/java-17-openjdk/bin/java", + "jvm_args": ["-Xmx1024m"], + "jars_root": ["~/airflow/jars"], + }, + } + + :param java_executable: Path to the ``java`` binary (defaults to ``"java"``, + which relies on ``$PATH``). + :param jvm_args: Extra arguments passed to the JVM (e.g. ``["-Xmx512m"]``). + :param jars_root: A list of directories scanned for JAR bundles. + """ + + java_executable: str = "java" + jvm_args: list[str] = attrs.field(factory=list) + jars_root: list[pathlib.Path] = attrs.field(converter=_convert_jars_root, factory=list) + + def execute_task( + self, + *, + what: TaskInstanceDTO, + dag_rel_path: str | os.PathLike[str], + bundle_info: BundleInfo, + client: Client, + logger: FilteringBoundLogger | None = None, + sentry_integration: str = "", + subprocess_logs_to_stdout: bool, + **kwargs, + ) -> BaseCoordinator.ExecutionResult: + process = _JavaActivitySubprocess.start( + what=what, + dag_rel_path=dag_rel_path, + bundle_info=bundle_info, + client=client, + logger=logger, + subprocess_logs_to_stdout=subprocess_logs_to_stdout, + sentry_integration=sentry_integration, + java_executable=self.java_executable, + jvm_args=self.jvm_args, + jars_root=self.jars_root, + ) + exit_code = process.wait() + return self.ExecutionResult(exit_code, process.final_state) diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py index 2364e942ed044..b98ea47ae4aea 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -80,7 +80,6 @@ PreviousTIResponse, PrevSuccessfulDagRunResponse, TaskBreadcrumbsResponse, - TaskInstance, TaskInstanceState, TaskStateResponse, TaskStatesResponse, @@ -98,6 +97,9 @@ XComSequenceSliceResponse, ) from airflow.sdk.exceptions import ErrorType +from airflow.sdk.execution_time.workloads.task import ( + TaskInstanceDTO, # noqa: TC001 -- Pydantic needs this at runtime +) try: from socket import recv_fds @@ -334,7 +336,7 @@ def _get_response(self) -> ReceiveMsgType | None: class StartupDetails(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) - ti: TaskInstance + ti: TaskInstanceDTO dag_rel_path: str bundle_info: BundleInfo start_date: datetime diff --git a/task-sdk/src/airflow/sdk/execution_time/coordinator.py b/task-sdk/src/airflow/sdk/execution_time/coordinator.py new file mode 100644 index 0000000000000..8d8e911bc0c68 --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/coordinator.py @@ -0,0 +1,227 @@ +# +# 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. +""" +Runtime coordinator for non-Python DAG file processing and task execution. + +Provides :class:`BaseCoordinator`, the base class for +SDK-specific coordinators that bridge subprocess I/O between the +Airflow supervisor and an external-SDK runtime (Java, Go, Rust, etc.), +and :class:`CoordinatorManager`, the registry that loads coordinator +instances from the ``[sdk] coordinators`` configuration. + +The coordinator's :meth:`~BaseCoordinator.run_task_execution` handles the full +lifecycle: + +1. Creates TCP servers for comm and logs channels, and a socketpair for stderr. +2. Calls :meth:`~BaseCoordinator.task_execution_cmd` (provided by the subclass) + to obtain the subprocess command. +3. Spawns the subprocess and accepts TCP connections from it. +4. Runs a selector-based bridge that transparently forwards bytes + between fd 0 (supervisor) and the subprocess comm socket, and + re-emits the subprocess's log and stderr output through structlog. +""" + +from __future__ import annotations + +import contextlib +import functools +from typing import TYPE_CHECKING, Any + +import attrs +import pydantic + +from airflow.sdk._shared.module_loading import import_string +from airflow.sdk.configuration import conf + +if TYPE_CHECKING: + from collections.abc import Mapping + from os import PathLike + + from structlog.typing import FilteringBoundLogger + from typing_extensions import Self + + from airflow.sdk.api.client import Client + from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO + +__all__ = [ + "BaseCoordinator", + "CoordinatorManager", + "get_coordinator_manager", + "reset_coordinator_manager", +] + + +class BaseCoordinator: + """ + Base coordinator for runtime-specific DAG file processing and task execution. + + Coordinators are instantiated from the ``[sdk] coordinators`` configuration + (see :class:`CoordinatorManager`) — each entry's ``classpath`` is resolved + via :func:`~airflow.sdk._shared.module_loading.import_string` and + constructed with the entry's ``kwargs``. + """ + + @attrs.define(slots=True) + class ExecutionResult: + """Return value for :meth:`BaseCoordinator.execute_task`.""" + + exit_code: Any + final_state: str + + def execute_task( + self, + *, + what: TaskInstanceDTO, + dag_rel_path: str | PathLike[str], + bundle_info, + client: Client, + logger: FilteringBoundLogger | None = None, + sentry_integration: str = "", + subprocess_logs_to_stdout: bool, + **kwargs, + ) -> ExecutionResult: + """ + Start task execution. + + This should execute the task and return a result. + """ + raise NotImplementedError + + +class _CoordinatorSpec(pydantic.BaseModel): + classpath: str + kwargs: dict[str, Any] + + +class _PythonCoordinator(BaseCoordinator): + """ + Coordinator implementation to execute Python tasks. + + This is not supposed to be specified by users directly, but the fallback + used by default when nothing is specified. + """ + + def execute_task( + self, + *, + what: TaskInstanceDTO, + dag_rel_path: str | PathLike[str], + bundle_info, + client: Client, + logger: FilteringBoundLogger | None = None, + sentry_integration: str = "", + subprocess_logs_to_stdout: bool, + **kwargs, + ) -> BaseCoordinator.ExecutionResult: + # TODO: Move this to somewhere that makes more sense. + from airflow.sdk.execution_time.supervisor import ActivitySubprocess + + process = ActivitySubprocess.start( + dag_rel_path=dag_rel_path, + what=what, + client=client, + logger=logger, + bundle_info=bundle_info, + subprocess_logs_to_stdout=subprocess_logs_to_stdout, + sentry_integration=sentry_integration, + ) + exit_code = process.wait() + return self.ExecutionResult(exit_code, process.final_state) + + +@functools.cache +def _build_python_coordinator() -> _PythonCoordinator: + return _PythonCoordinator() + + +@attrs.define(kw_only=True) +class CoordinatorManager: + """ + Registry of coordinator instances loaded from ``[sdk]`` configurations. + + The ``[sdk] coordinators`` value is a JSON object keyed by coordinator name:: + + { + "jdk-11": { + "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", + "kwargs": {"java_executable": "/usr/lib/jvm/jdk-11/bin/java", ...}, + } + } + + The ``classpath`` is resolved via + :func:`~airflow.sdk._shared.module_loading.import_string` (no + :class:`ProvidersManager` involvement) and constructed with ``kwargs`` on + first use. A coordinator entry that is never looked up incurs no startup + cost. At most one coordinator object can be created from each entry. + + The ``[sdk] queue_to_coordinator`` config maps queue names to a key in the + object, which lets users reuse existing queue assignments to route tasks to + a specific coordinator instance (for example, a ``"legacy-java"`` queue + routed to a JDK 11 coordinator, and a ``"modern-java"`` queue routed to a + JDK 17 coordinator). + + :meta private: + """ + + _coordinator_specs: Mapping[str, _CoordinatorSpec] + _queue_to_coordinator: Mapping[str, str] + + _created_coordinators: dict[str, BaseCoordinator] = attrs.field(init=False, factory=dict) + + @classmethod + def from_config(cls) -> Self: + """Load coordinator specs from configuration without initialization.""" + coordinator_specs = { + k: _CoordinatorSpec.model_validate(v) + for k, v in conf.getjson("sdk", "coordinators", fallback={}).items() + } + queue_to_coordinator = conf.getjson("sdk", "queue_to_coordinator", fallback={}) + for key in queue_to_coordinator.values(): + if key not in coordinator_specs: + raise ValueError(f"[sdk] queue_to_coordinator references invalid coordinator key: {key!r}") + return cls(coordinator_specs=coordinator_specs, queue_to_coordinator=queue_to_coordinator) + + def _for_queue_internal(self, queue: str) -> BaseCoordinator: + key = self._queue_to_coordinator[queue] + with contextlib.suppress(KeyError): + return self._created_coordinators[key] + spec = self._coordinator_specs[key] + coordinator = self._created_coordinators[key] = import_string(spec.classpath)(**spec.kwargs) + return coordinator + + def for_queue(self, queue: str) -> BaseCoordinator: + """ + Find the coordinator for *queue*. + + If an entry is not registered, a Python coordinator is returned. + """ + try: + return self._for_queue_internal(queue) + except KeyError: + return _build_python_coordinator() + + +@functools.cache +def get_coordinator_manager() -> CoordinatorManager: + """Return the process-wide :class:`CoordinatorManager`, loaded from config on first use.""" + return CoordinatorManager.from_config() + + +def reset_coordinator_manager() -> None: + """Clear the cached :class:`CoordinatorManager` (test helper).""" + get_coordinator_manager.cache_clear() 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..38981dfea76cf --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/schema/__init__.py @@ -0,0 +1,129 @@ +# 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 typing import TYPE_CHECKING, Annotated, Any, get_args, get_origin + +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 pydantic import BaseModel + + +def _members_of_discriminated_union(union_type: object) -> tuple[type, ...]: + """Return the BaseModel classes in an ``Annotated[A | B | ..., Field(...)]``.""" + # ``Annotated[X | Y, Field(...)]`` -> the first ``get_args`` arg is the union. + if get_origin(union_type) is Annotated: + union_type = get_args(union_type)[0] + members = get_args(union_type) + return tuple(m for m in members if isinstance(m, type)) + + +@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 pydantic import BaseModel + + from airflow.dag_processing.processor import ToDagProcessor, ToManager + from airflow.sdk.execution_time.comms import ToSupervisor, ToTask + + by_name: dict[str, type[BaseModel]] = {} + for union in (ToTask, ToSupervisor, ToManager, ToDagProcessor): + for member in _members_of_discriminated_union(union): + if not issubclass(member, BaseModel): + raise RuntimeError( + f"Invalid supervisor schema body {member!r}: " + f"union member {member!r} is not a Pydantic model class" + ) + existing = by_name.get(member.__name__) + if existing is None: + by_name[member.__name__] = member + elif existing is not member: + raise RuntimeError( + f"Duplicate supervisor schema body name {member.__name__!r}: " + f"both {existing!r} and {member!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..ff7314ebdbe01 --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/schema/migrator.py @@ -0,0 +1,214 @@ +# 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 + + +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. + """ + + __slots__ = ("body",) + + def __init__(self, body: dict[str, Any]) -> None: + # Copy so the caller's mapping survives intact when the + # instruction chain mutates ``info.body`` in place. + self.body = dict(body) + + +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, + trarget_schema_version: str, + **dump_opts: dict[str, Any], + ) -> BaseModel: + """ + Downgrade *msg* from server to *trarget_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 trarget_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 plain dict shaped for *trarget_schema_version*. + """ + model = type(msg) + trarget_schema_version = self._resolve_version(trarget_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 <= trarget_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(trarget_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 *trarget_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..e7e9cba5964eb --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -0,0 +1,6078 @@ +{ + "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", + "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" + }, + "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": "GetNumberXComs", + "default": "GetNumberXComs", + "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": { + "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" + ], + "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" + }, + "TaskInstanceDTO": { + "description": "Task SDK TaskInstanceDTO.", + "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" + }, + "executor_config": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Executor Config" + }, + "id": { + "format": "uuid", + "title": "Id", + "type": "string" + }, + "map_index": { + "default": -1, + "title": "Map Index", + "type": "integer" + }, + "parent_context_carrier": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Parent Context Carrier" + }, + "pool_slots": { + "title": "Pool Slots", + "type": "integer" + }, + "priority_weight": { + "title": "Priority Weight", + "type": "integer" + }, + "queue": { + "title": "Queue", + "type": "string" + }, + "run_id": { + "title": "Run Id", + "type": "string" + }, + "task_id": { + "title": "Task Id", + "type": "string" + }, + "try_number": { + "title": "Try Number", + "type": "integer" + } + }, + "required": [ + "id", + "dag_version_id", + "task_id", + "dag_id", + "run_id", + "try_number", + "pool_slots", + "queue", + "priority_weight" + ], + "title": "TaskInstanceDTO", + "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/TaskInstanceDTO" + }, + "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" + }, + "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": "XComLengthResponse", + "default": "XComLengthResponse", + "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 f7af97ffd8cb2..dad638ef3ba28 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 @@ -130,6 +130,7 @@ _RequestFrame, _ResponseFrame, ) +from airflow.sdk.execution_time.coordinator import get_coordinator_manager from airflow.sdk.execution_time.request_handlers import ( handle_delete_variable, handle_delete_xcom, @@ -151,6 +152,7 @@ handle_put_variable, handle_set_xcom, ) +from airflow.sdk.execution_time.schema import get_schema_version_migrator, resolve_body_class try: from socket import send_fds @@ -169,6 +171,7 @@ from airflow.executors.workloads import BundleInfo from airflow.sdk.bases.secrets_backend import BaseSecretsBackend from airflow.sdk.definitions.connection import Connection + from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO from airflow.sdk.types import RuntimeTaskInstanceProtocol as RuntimeTI __all__ = ["ActivitySubprocess", "WatchedSubprocess", "supervise", "supervise_task"] @@ -555,6 +558,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( @@ -737,35 +742,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_kwargs=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 @@ -1194,7 +1215,7 @@ class ActivitySubprocess(WatchedSubprocess): def start( # type: ignore[override] cls, *, - what: TaskInstance, + what: TaskInstanceDTO, dag_rel_path: str | os.PathLike[str], bundle_info, client: Client, @@ -1223,7 +1244,7 @@ def start( # type: ignore[override] def _on_child_started( self, *, - ti: TaskInstance, + ti: TaskInstanceDTO, dag_rel_path: str | os.PathLike[str], bundle_info, sentry_integration: str, @@ -2250,7 +2271,7 @@ def _configure_logging(log_path: str, client: Client) -> tuple[FilteringBoundLog def supervise_task( *, - ti: TaskInstance, + ti: TaskInstanceDTO, bundle_info: BundleInfo, dag_rel_path: str | os.PathLike[str], token: str, @@ -2336,27 +2357,26 @@ def supervise_task( reset_secrets_masker() try: - process = ActivitySubprocess.start( - dag_rel_path=dag_rel_path, + coordinator = get_coordinator_manager().for_queue(ti.queue) + result = coordinator.execute_task( what=ti, + dag_rel_path=dag_rel_path, + bundle_info=bundle_info, client=client, logger=logger, - bundle_info=bundle_info, - subprocess_logs_to_stdout=subprocess_logs_to_stdout, sentry_integration=sentry_integration, + subprocess_logs_to_stdout=subprocess_logs_to_stdout, ) - - exit_code = process.wait() end = time.monotonic() log.info( "Workload finished", workload_type="ExecuteTask", workload_id=str(ti.id), - exit_code=exit_code, + exit_code=result.exit_code, duration=end - start, - final_state=process.final_state, + final_state=result.final_state, ) - return exit_code + return result.exit_code finally: if log_path and log_file_descriptor: log_file_descriptor.close() diff --git a/task-sdk/tests/task_sdk/coordinators/__init__.py b/task-sdk/tests/task_sdk/coordinators/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/task-sdk/tests/task_sdk/coordinators/__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/coordinators/executable/__init__.py b/task-sdk/tests/task_sdk/coordinators/executable/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/task-sdk/tests/task_sdk/coordinators/executable/__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/coordinators/executable/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py new file mode 100644 index 0000000000000..2335189d6fb06 --- /dev/null +++ b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py @@ -0,0 +1,538 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import contextlib +import pathlib +import socket +import stat +import struct +import subprocess +import threading +import time +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +import yaml +from uuid6 import uuid7 + +from airflow.sdk.coordinators.executable.coordinator import ( + FOOTER_MAGIC, + FOOTER_SIZE, + ExecutableCoordinator, + _accept_connections, + _Bundle, + _ExecutableActivitySubprocess, + _start_server, +) +from airflow.sdk.execution_time.coordinator import BaseCoordinator +from airflow.sdk.execution_time.supervisor import ActivitySubprocess +from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO + +from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS + +if not AIRFLOW_V_3_3_PLUS: + pytest.skip("Coordinator is only compatible with Airflow >= 3.3.0", allow_module_level=True) + +_DEFAULT_BINARY_PAYLOAD = b"\x7fELF" + b"binary-stub-payload" + + +def _make_metadata(dag_ids, source_filename: str = "example.go") -> dict: + return { + "format_version": "1.0", + "sdk": {"language": "go", "version": "0.1.0"}, + "source": source_filename, + "dags": {dag_id: {"tasks": ["task1"]} for dag_id in dag_ids}, + } + + +def _build_bundle( + path: Path, + *, + dag_ids=("tutorial_dag",), + source: str | bytes = "package main\n\nfunc main() {}\n", + source_filename: str = "example.go", + metadata: dict | bytes | None = None, + binary_bytes: bytes = _DEFAULT_BINARY_PAYLOAD, + footer_ver: int = 1, + magic: bytes = FOOTER_MAGIC, + reserved: bytes = b"\x00" * 12, +) -> Path: + if isinstance(source, str): + source_bytes = source.encode("utf-8") + else: + source_bytes = source + + if metadata is None: + metadata_dict = _make_metadata(dag_ids, source_filename=source_filename) + metadata_bytes = yaml.safe_dump(metadata_dict, sort_keys=True).encode("utf-8") + elif isinstance(metadata, (bytes, bytearray)): + metadata_bytes = bytes(metadata) + else: + metadata_bytes = yaml.safe_dump(metadata, sort_keys=True).encode("utf-8") + + if len(reserved) != 12: + raise ValueError("reserved must be exactly 12 bytes") + trailer = struct.pack(" Path: + path.write_bytes(b"#!/bin/sh\nexit 0\n") + path.chmod(path.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH) + return path + + +def _make_ti(dag_id: str = "tutorial_dag", queue: str = "executable") -> TaskInstanceDTO: + return TaskInstanceDTO( + id=uuid7(), + dag_version_id=uuid7(), + task_id="task_1", + dag_id=dag_id, + run_id="run_1", + try_number=1, + map_index=-1, + pool_slots=1, + queue=queue, + priority_weight=1, + ) + + +class TestStartServer: + def test_returns_listening_socket(self): + server = _start_server() + try: + host, port = server.getsockname() + finally: + server.close() + assert host == "127.0.0.1" + assert port > 0 + + def test_two_calls_return_different_ports(self): + s1 = _start_server() + s2 = _start_server() + try: + _, port1 = s1.getsockname() + _, port2 = s2.getsockname() + finally: + s1.close() + s2.close() + assert port1 != port2 + + +class TestAcceptConnections: + def _connect_after_delay(self, addr: tuple[str, int], delay: float = 0.0) -> None: + def _connect(): + time.sleep(delay) + c = socket.socket() + with contextlib.suppress(OSError): # Server may already be closed in teardown. + c.connect(addr) + + threading.Thread(target=_connect, daemon=True).start() + + def test_accepts_multiple_servers(self): + comm_server = _start_server() + logs_server = _start_server() + _, comm_port = comm_server.getsockname() + _, logs_port = logs_server.getsockname() + + self._connect_after_delay(("127.0.0.1", comm_port)) + self._connect_after_delay(("127.0.0.1", logs_port)) + + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.poll.return_value = None + + try: + accepted = _accept_connections({"comm": comm_server, "logs": logs_server}, mock_proc) + assert set(accepted) == {"comm", "logs"} + for sock in accepted.values(): + sock.close() + finally: + comm_server.close() + logs_server.close() + + def test_raises_timeout_when_no_connection(self): + server = _start_server() + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.poll.return_value = None + try: + with pytest.raises(TimeoutError, match="did not connect within timeout"): + _accept_connections({"comm": server}, mock_proc, max_wait=0.05) + finally: + server.close() + + def test_raises_runtime_error_if_process_exits_before_connecting(self): + server = _start_server() + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.poll.return_value = 1 + mock_proc.returncode = 1 + try: + with pytest.raises(RuntimeError, match="process exited with 1"): + _accept_connections({"comm": server}, mock_proc) + finally: + server.close() + + +class TestBundleFind: + def test_finds_matching_dag_id(self, tmp_path): + binary = _build_bundle(tmp_path / "my_bundle", dag_ids=["tutorial_dag", "other_dag"]) + + bundle = _Bundle.find([tmp_path], "tutorial_dag") + assert bundle.path == binary.resolve() + + def test_picks_matching_bundle_among_many(self, tmp_path): + _build_bundle(tmp_path / "alpha", dag_ids=["alpha_dag"]) + beta = _build_bundle(tmp_path / "beta", dag_ids=["beta_dag"]) + _build_bundle(tmp_path / "gamma", dag_ids=["gamma_dag"]) + + bundle = _Bundle.find([tmp_path], "beta_dag") + assert bundle.path == beta.resolve() + + def test_searches_multiple_roots(self, tmp_path): + root_a = tmp_path / "a" + root_b = tmp_path / "b" + root_a.mkdir() + root_b.mkdir() + _build_bundle(root_a / "alpha", dag_ids=["alpha_dag"]) + target = _build_bundle(root_b / "beta", dag_ids=["beta_dag"]) + + bundle = _Bundle.find([root_a, root_b], "beta_dag") + assert bundle.path == target.resolve() + + def test_skips_non_bundle_files(self, tmp_path): + (tmp_path / "README.md").write_text("not a bundle") + _make_executable(tmp_path / "stray_executable") + binary = _build_bundle(tmp_path / "real_bundle", dag_ids=["tutorial_dag"]) + + bundle = _Bundle.find([tmp_path], "tutorial_dag") + assert bundle.path == binary.resolve() + + def test_skips_non_executable_files(self, tmp_path): + non_exec = _build_bundle(tmp_path / "non_exec", dag_ids=["tutorial_dag"]) + non_exec.chmod(non_exec.stat().st_mode & ~(stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)) + + with pytest.raises(FileNotFoundError, match="cannot find executable bundle"): + _Bundle.find([tmp_path], "tutorial_dag") + + def test_raises_when_not_found(self, tmp_path): + with pytest.raises(FileNotFoundError, match="cannot find executable bundle"): + _Bundle.find([tmp_path], "nonexistent_dag") + + def test_raises_when_directory_missing(self, tmp_path): + with pytest.raises(FileNotFoundError, match="cannot find executable bundle"): + _Bundle.find([tmp_path / "does_not_exist"], "tutorial_dag") + + +class TestExecutableCoordinatorAttributes: + def test_default_kwargs(self): + coordinator = ExecutableCoordinator() + assert coordinator.sdk == "executable" + assert coordinator.file_extension == "" + assert coordinator.executables_root == [] + + def test_executables_root_accepts_single_path(self, tmp_path): + coordinator = ExecutableCoordinator(executables_root=str(tmp_path)) + assert coordinator.executables_root == [tmp_path] + + def test_executables_root_accepts_list(self, tmp_path): + other = tmp_path / "other" + coordinator = ExecutableCoordinator(executables_root=[str(tmp_path), other]) + assert coordinator.executables_root == [tmp_path, other] + + def test_executables_root_none_becomes_empty_list(self): + coordinator = ExecutableCoordinator(executables_root=None) + assert coordinator.executables_root == [] + + +class TestResolveExecutable: + def test_resolves_via_executables_root(self, tmp_path): + binary = _build_bundle(tmp_path / "my_bundle", dag_ids=["tutorial_dag"]) + ti = _make_ti(dag_id="tutorial_dag") + + coordinator = ExecutableCoordinator(executables_root=[tmp_path]) + resolved = coordinator._resolve_executable(what=ti) + assert resolved == str(binary.resolve()) + + def test_raises_when_executables_root_missing(self): + ti = _make_ti(dag_id="tutorial_dag") + coordinator = ExecutableCoordinator() + with pytest.raises(ValueError, match="executables_root kwarg must be set"): + coordinator._resolve_executable(what=ti) + + def test_raises_when_dag_id_not_found(self, tmp_path): + _build_bundle(tmp_path / "my_bundle", dag_ids=["other_dag"]) + ti = _make_ti(dag_id="tutorial_dag") + + coordinator = ExecutableCoordinator(executables_root=[tmp_path]) + with pytest.raises(FileNotFoundError, match="cannot find executable bundle"): + coordinator._resolve_executable(what=ti) + + +@pytest.fixture +def bundles_dir(tmp_path): + _build_bundle(tmp_path / "my_bundle", dag_ids=["tutorial_dag"]) + return tmp_path + + +@pytest.fixture +def mock_client(make_ti_context): + client = MagicMock() + client.task_instances.start.return_value = make_ti_context() + return client + + +class TestExecutableCoordinatorExecuteTask: + def _captured_popen_cmd(self, bundles_dir: pathlib.Path, mock_client) -> list[str]: + """Run execute_task with mocked subprocess and return the command list.""" + ti = _make_ti(dag_id="tutorial_dag") + coordinator = ExecutableCoordinator(executables_root=[bundles_dir]) + + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.pid = 12345 + comm_sock = MagicMock(spec=socket.socket) + logs_sock = MagicMock(spec=socket.socket) + popen_calls: list = [] + + def capture_popen(cmd, **kwargs): + popen_calls.append(cmd) + return mock_proc + + with ( + patch( + "airflow.sdk.coordinators.executable.coordinator.subprocess.Popen", + side_effect=capture_popen, + ), + patch( + "airflow.sdk.coordinators.executable.coordinator._accept_connections", + return_value={"comm": comm_sock, "logs": logs_sock}, + ), + patch.object(ActivitySubprocess, "_register_pipe_readers"), + patch.object(ActivitySubprocess, "_on_child_started"), + patch.object(ActivitySubprocess, "wait", return_value=0), + patch("psutil.Process"), + ): + coordinator.execute_task( + what=ti, + dag_rel_path="my_bundle", + bundle_info=MagicMock(), + client=mock_client, + subprocess_logs_to_stdout=False, + ) + + assert popen_calls, "subprocess.Popen was not called" + return popen_calls[0] + + def test_executable_path_is_first_arg(self, bundles_dir, mock_client): + cmd = self._captured_popen_cmd(bundles_dir, mock_client) + expected = str((bundles_dir / "my_bundle").resolve()) + assert cmd[0] == expected + + def test_comm_and_logs_args_present(self, bundles_dir, mock_client): + cmd = self._captured_popen_cmd(bundles_dir, mock_client) + comm_args = [a for a in cmd if a.startswith("--comm=")] + logs_args = [a for a in cmd if a.startswith("--logs=")] + assert len(comm_args) == 1 + assert len(logs_args) == 1 + + def test_comm_and_logs_contain_port(self, bundles_dir, mock_client): + cmd = self._captured_popen_cmd(bundles_dir, mock_client) + comm_arg = next(a for a in cmd if a.startswith("--comm=")) + logs_arg = next(a for a in cmd if a.startswith("--logs=")) + assert ":" in comm_arg.split("=", 1)[1] + assert ":" in logs_arg.split("=", 1)[1] + + def test_returns_execution_result(self, bundles_dir, mock_client): + ti = _make_ti(dag_id="tutorial_dag") + coordinator = ExecutableCoordinator(executables_root=[bundles_dir]) + + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.pid = 99999 + comm_sock = MagicMock(spec=socket.socket) + logs_sock = MagicMock(spec=socket.socket) + + with ( + patch("subprocess.Popen", return_value=mock_proc), + patch( + "airflow.sdk.coordinators.executable.coordinator._accept_connections", + return_value={"comm": comm_sock, "logs": logs_sock}, + ), + patch.object(ActivitySubprocess, "_register_pipe_readers"), + patch.object(ActivitySubprocess, "_on_child_started"), + patch.object(ActivitySubprocess, "wait", return_value=0), + patch("psutil.Process"), + ): + result = coordinator.execute_task( + what=ti, + dag_rel_path="my_bundle", + bundle_info=MagicMock(), + client=mock_client, + subprocess_logs_to_stdout=False, + ) + + assert isinstance(result, BaseCoordinator.ExecutionResult) + assert result.exit_code == 0 + + +class TestExecutableActivitySubprocessStart: + """ + Unit tests for _ExecutableActivitySubprocess.start(). + + These tests mock subprocess.Popen and _accept_connections to verify that + start() wires up the right command and stores the right sockets, + without requiring a real native binary to launch. + """ + + def _start_with_mocks( + self, + executable_path: str, + mock_client, + *, + ti: TaskInstanceDTO | None = None, + ): + ti = ti or _make_ti() + + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.pid = 12345 + comm_sock = MagicMock(spec=socket.socket) + logs_sock = MagicMock(spec=socket.socket) + + with ( + patch( + "airflow.sdk.coordinators.executable.coordinator.subprocess.Popen", + return_value=mock_proc, + ) as popen_mock, + patch( + "airflow.sdk.coordinators.executable.coordinator._accept_connections", + return_value={"comm": comm_sock, "logs": logs_sock}, + ), + patch.object(ActivitySubprocess, "_register_pipe_readers"), + patch.object(ActivitySubprocess, "_on_child_started"), + patch("psutil.Process"), + ): + proc = _ExecutableActivitySubprocess.start( + what=ti, + dag_rel_path="my_bundle", + bundle_info=MagicMock(), + client=mock_client, + executable_path=executable_path, + subprocess_logs_to_stdout=False, + ) + return proc, popen_mock + + def test_stdout_write_socket_stored_for_cleanup(self, bundles_dir, mock_client): + proc, _ = self._start_with_mocks(str(bundles_dir / "my_bundle"), mock_client) + assert proc._stdout_w is not None + + def test_stderr_write_socket_stored_for_cleanup(self, bundles_dir, mock_client): + proc, _ = self._start_with_mocks(str(bundles_dir / "my_bundle"), mock_client) + assert proc._stderr_w is not None + + def test_stdout_and_stderr_write_sockets_are_distinct(self, bundles_dir, mock_client): + proc, _ = self._start_with_mocks(str(bundles_dir / "my_bundle"), mock_client) + assert proc._stdout_w is not proc._stderr_w + + def test_stdin_is_comm_socket(self, bundles_dir, mock_client): + """stdin (used by send_msg) must be the accepted comm socket.""" + ti = _make_ti() + comm_sock = MagicMock(spec=socket.socket) + logs_sock = MagicMock(spec=socket.socket) + + with ( + patch("airflow.sdk.coordinators.executable.coordinator.subprocess.Popen") as popen_mock, + patch( + "airflow.sdk.coordinators.executable.coordinator._accept_connections", + return_value={"comm": comm_sock, "logs": logs_sock}, + ), + patch.object(ActivitySubprocess, "_register_pipe_readers"), + patch.object(ActivitySubprocess, "_on_child_started"), + patch("psutil.Process"), + ): + popen_mock.return_value.pid = 12345 + proc = _ExecutableActivitySubprocess.start( + what=ti, + dag_rel_path="my_bundle", + bundle_info=MagicMock(), + client=MagicMock(), + executable_path=str(bundles_dir / "my_bundle"), + subprocess_logs_to_stdout=False, + ) + + assert proc.stdin is comm_sock + + def test_pid_taken_from_popen(self, bundles_dir, mock_client): + proc, _ = self._start_with_mocks(str(bundles_dir / "my_bundle"), mock_client) + assert proc.pid == 12345 + + def test_on_child_started_called(self, bundles_dir, mock_client): + ti = _make_ti() + with ( + patch("airflow.sdk.coordinators.executable.coordinator.subprocess.Popen") as popen_mock, + patch( + "airflow.sdk.coordinators.executable.coordinator._accept_connections", + return_value={"comm": MagicMock(), "logs": MagicMock()}, + ), + patch.object(ActivitySubprocess, "_register_pipe_readers"), + patch.object(ActivitySubprocess, "_on_child_started") as mock_on_started, + patch("psutil.Process"), + ): + popen_mock.return_value.pid = 12345 + _ExecutableActivitySubprocess.start( + what=ti, + dag_rel_path="my_bundle", + bundle_info=MagicMock(), + client=mock_client, + executable_path=str(bundles_dir / "my_bundle"), + subprocess_logs_to_stdout=False, + ) + + mock_on_started.assert_called_once() + kwargs = mock_on_started.call_args.kwargs + assert kwargs["ti"] is ti + assert kwargs["dag_rel_path"] == "my_bundle" + + def test_register_pipe_readers_called_with_four_sockets(self, bundles_dir, mock_client): + """Both socketpair read-ends and both TCP sockets must be registered.""" + with ( + patch("airflow.sdk.coordinators.executable.coordinator.subprocess.Popen") as popen_mock, + patch( + "airflow.sdk.coordinators.executable.coordinator._accept_connections", + return_value={"comm": MagicMock(), "logs": MagicMock()}, + ), + patch.object(ActivitySubprocess, "_register_pipe_readers") as mock_register, + patch.object(ActivitySubprocess, "_on_child_started"), + patch("psutil.Process"), + ): + popen_mock.return_value.pid = 12345 + _ExecutableActivitySubprocess.start( + what=_make_ti(), + dag_rel_path="my_bundle", + bundle_info=MagicMock(), + client=mock_client, + executable_path=str(bundles_dir / "my_bundle"), + subprocess_logs_to_stdout=False, + ) + + mock_register.assert_called_once() + args = mock_register.call_args.args + # positional: stdout, stderr, comm, logs — all four must be sockets + assert len(args) == 4 diff --git a/task-sdk/tests/task_sdk/coordinators/java/__init__.py b/task-sdk/tests/task_sdk/coordinators/java/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/task-sdk/tests/task_sdk/coordinators/java/__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/coordinators/java/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py new file mode 100644 index 0000000000000..382c907a5f451 --- /dev/null +++ b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py @@ -0,0 +1,601 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import contextlib +import os +import pathlib +import re +import socket +import subprocess +import threading +import time +import zipfile +from unittest.mock import MagicMock, patch + +import pytest +from uuid6 import uuid7 + +from airflow.sdk.coordinators.java.coordinator import ( + JavaCoordinator, + _accept_connections, + _calculate_classpath, + _JavaActivitySubprocess, + _MainJar, + _start_server, +) +from airflow.sdk.execution_time.coordinator import BaseCoordinator +from airflow.sdk.execution_time.supervisor import ActivitySubprocess +from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO + +from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS + +if not AIRFLOW_V_3_3_PLUS: + pytest.skip("Coordinator is only compatible with Airflow >= 3.3.0", allow_module_level=True) + +METADATA_YAML_PATH = "META-INF/airflow-metadata.yaml" +DAG_CODE_PATH = "dag_source.py" +TEST_MAIN_CLASS = "com.example.MyBundle" + + +def _make_ti(dag_id: str = "test_dag", queue: str = "java") -> TaskInstanceDTO: + return TaskInstanceDTO( + id=uuid7(), + dag_version_id=uuid7(), + task_id="task_1", + dag_id=dag_id, + run_id="run_1", + try_number=1, + map_index=-1, + pool_slots=1, + queue=queue, + priority_weight=1, + ) + + +def _make_jar( + path: pathlib.Path, + *, + main_class: str | None = "com.example.Main", + schema_version: str | None = None, +) -> pathlib.Path: + """Write a minimal JAR with (optionally) a Main-Class manifest entry.""" + lines = ["Manifest-Version: 1.0"] + if main_class: + lines.append(f"Main-Class: {main_class}") + if schema_version: + lines.append(f"Airflow-SDK-Supervisor-Schema-Version: {schema_version}") + manifest = "\n".join(lines) + "\n\n" + with zipfile.ZipFile(path, "w") as zf: + zf.writestr("META-INF/MANIFEST.MF", manifest) + return path + + +class TestStartServer: + def test_returns_listening_socket(self): + server = _start_server() + try: + _, port = server.getsockname() + finally: + server.close() + assert port > 0 + + def test_two_calls_return_different_ports(self): + s1 = _start_server() + s2 = _start_server() + try: + _, port1 = s1.getsockname() + _, port2 = s2.getsockname() + finally: + s1.close() + s2.close() + assert port1 != port2 + + def test_accepts_connection(self): + conn = client = None + server = _start_server() + try: + _, port = server.getsockname() + client = socket.socket() + client.connect(("127.0.0.1", port)) + conn, _ = server.accept() + conn.sendall(b"ping") + received = client.recv(4) + finally: + if conn: + conn.close() + if client: + client.close() + server.close() + assert received == b"ping" + + +class TestCalculateClasspath: + def test_single_jar(self, tmp_path): + jar = tmp_path.joinpath("app.jar") + jar.write_bytes(b"") + result = _calculate_classpath([tmp_path]) + assert result == jar.as_posix() + + def test_multiple_jars_all_included(self, tmp_path): + tmp_path.joinpath("a.jar").write_bytes(b"") + tmp_path.joinpath("b.jar").write_bytes(b"") + tmp_path.joinpath("c.jar").write_bytes(b"") + result = _calculate_classpath([tmp_path]) + entries = set(result.split(os.pathsep)) + assert entries == { + tmp_path.joinpath("a.jar").as_posix(), + tmp_path.joinpath("b.jar").as_posix(), + tmp_path.joinpath("c.jar").as_posix(), + } + + def test_non_jar_files_excluded(self, tmp_path): + jar = tmp_path.joinpath("app.jar") + jar.write_bytes(b"") + tmp_path.joinpath("readme.txt").write_bytes(b"") + tmp_path.joinpath("config.yaml").write_bytes(b"") + result = _calculate_classpath([tmp_path]) + assert result == jar.as_posix() + + def test_empty_directory_returns_empty_string(self, tmp_path): + result = _calculate_classpath([tmp_path]) + assert result == "" + + +class TestMainJar: + def test_returns_main_class_from_jar(self, tmp_path): + _make_jar(tmp_path.joinpath("app.jar"), main_class="com.example.Main") + assert _MainJar.find([tmp_path]) == _MainJar(tmp_path.joinpath("app.jar"), "com.example.Main", None) + + def test_no_jars_raises_file_not_found(self, tmp_path): + with pytest.raises(FileNotFoundError, match=re.escape(str(tmp_path.resolve()))): + _MainJar.find([tmp_path]) + + def test_jar_without_main_class_not_returned(self, tmp_path): + _make_jar(tmp_path.joinpath("app.jar"), main_class=None) + with pytest.raises(FileNotFoundError): + _MainJar.find([tmp_path]) + + def test_non_jar_files_skipped(self, tmp_path): + tmp_path.joinpath("readme.txt").write_bytes(b"not a jar") + _make_jar(tmp_path.joinpath("app.jar"), main_class="com.example.Main") + assert _MainJar.find([tmp_path]) == _MainJar(tmp_path.joinpath("app.jar"), "com.example.Main", None) + + def test_first_jar_missing_main_class_falls_through_to_second(self, tmp_path): + # Alphabetically: a.jar (no Main-Class), b.jar (has Main-Class). + _make_jar(tmp_path.joinpath("a.jar"), main_class=None) + _make_jar(tmp_path.joinpath("b.jar"), main_class="com.example.Fallback") + assert _MainJar.find([tmp_path]) == _MainJar(tmp_path.joinpath("b.jar"), "com.example.Fallback", None) + + def test_fully_qualified_class_name_preserved(self, tmp_path): + _make_jar(tmp_path.joinpath("app.jar"), main_class="org.apache.airflow.sdk.java.TaskRunner") + assert _MainJar.find([tmp_path]) == _MainJar( + path=tmp_path.joinpath("app.jar"), + main_class="org.apache.airflow.sdk.java.TaskRunner", + schema_version=None, + ) + + +class TestAcceptConnections: + def _connect_after_delay(self, addr: tuple[str, int], delay: float = 0.0) -> None: + def _connect(): + time.sleep(delay) + c = socket.socket() + with contextlib.suppress(OSError): # Server may already be closed in teardown. + c.connect(addr) + + threading.Thread(target=_connect, daemon=True).start() + + def test_accepts_single_server(self): + server = _start_server() + _, port = server.getsockname() + self._connect_after_delay(("127.0.0.1", port)) + + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.poll.return_value = None + + try: + accepted = _accept_connections({"comm": server}, mock_proc) + assert "comm" in accepted + accepted["comm"].close() + finally: + server.close() + + def test_accepts_multiple_servers(self): + comm_server = _start_server() + logs_server = _start_server() + _, comm_port = comm_server.getsockname() + _, logs_port = logs_server.getsockname() + + self._connect_after_delay(("127.0.0.1", comm_port)) + self._connect_after_delay(("127.0.0.1", logs_port)) + + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.poll.return_value = None + + try: + accepted = _accept_connections({"comm": comm_server, "logs": logs_server}, mock_proc) + assert set(accepted) == {"comm", "logs"} + for sock in accepted.values(): + sock.close() + finally: + comm_server.close() + logs_server.close() + + def test_raises_timeout_when_no_connection(self): + server = _start_server() + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.poll.return_value = None + try: + with pytest.raises(TimeoutError, match="did not connect within timeout"): + _accept_connections({"comm": server}, mock_proc, max_wait=0.05) + finally: + server.close() + + def test_raises_runtime_error_if_process_exits_before_connecting(self): + server = _start_server() + mock_proc = MagicMock(spec=subprocess.Popen) + # proc has already exited + mock_proc.poll.return_value = 1 + mock_proc.returncode = 1 + try: + with pytest.raises(RuntimeError, match="process exited with 1"): + _accept_connections({"comm": server}, mock_proc) + finally: + server.close() + + def test_returned_sockets_are_connected(self): + """Accepted sockets should be real, usable connections.""" + server = _start_server() + _, port = server.getsockname() + + client = socket.socket() + client.connect(("127.0.0.1", port)) + + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.poll.return_value = None + + try: + accepted = _accept_connections({"comm": server}, mock_proc) + accepted["comm"].sendall(b"hello") + assert client.recv(5) == b"hello" + accepted["comm"].close() + client.close() + finally: + server.close() + + +class TestJavaCoordinatorAttributes: + def test_default_kwargs(self): + coordinator = JavaCoordinator(jars_root="/airflow/java-bundles") + assert coordinator.java_executable == "java" + assert coordinator.jvm_args == [] + assert coordinator.jars_root == [pathlib.Path("/airflow/java-bundles")] + + def test_custom_kwargs(self): + coordinator = JavaCoordinator( + java_executable="/opt/java/bin/java", + jvm_args=["-Xmx512m", "-Xms256m"], + jars_root=["/airflow/java-bundles"], + ) + assert coordinator.java_executable == "/opt/java/bin/java" + assert coordinator.jvm_args == ["-Xmx512m", "-Xms256m"] + assert coordinator.jars_root == [pathlib.Path("/airflow/java-bundles")] + + +@pytest.fixture +def jars_root(tmp_path): + _make_jar(tmp_path.joinpath("app.jar"), main_class="com.example.TaskRunner") + return tmp_path + + +@pytest.fixture +def mock_client(make_ti_context): + client = MagicMock() + client.task_instances.start.return_value = make_ti_context() + return client + + +class TestJavaCoordinatorExecuteTask: + def _captured_popen_cmd( + self, + jars_root: pathlib.Path, + mock_client, + *, + java_executable: str = "java", + jvm_args: list[str] | None = None, + ) -> list[str]: + """Run execute_task with mocked subprocess and return the command list.""" + ti = _make_ti() + coordinator = JavaCoordinator( + java_executable=java_executable, + jvm_args=jvm_args or [], + jars_root=jars_root, + ) + + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.pid = 12345 + comm_sock = MagicMock(spec=socket.socket) + logs_sock = MagicMock(spec=socket.socket) + popen_calls: list = [] + + def capture_popen(cmd, **kwargs): + popen_calls.append(cmd) + return mock_proc + + with ( + patch( + "airflow.sdk.coordinators.java.coordinator.subprocess.Popen", + side_effect=capture_popen, + ), + patch( + "airflow.sdk.coordinators.java.coordinator._accept_connections", + return_value={"comm": comm_sock, "logs": logs_sock}, + ), + patch.object(ActivitySubprocess, "_register_pipe_readers"), + patch.object(ActivitySubprocess, "_on_child_started"), + patch.object(ActivitySubprocess, "wait", return_value=0), + patch("psutil.Process"), + ): + coordinator.execute_task( + what=ti, + dag_rel_path="dags/test.jar", + bundle_info=MagicMock(), + client=mock_client, + subprocess_logs_to_stdout=False, + ) + + assert popen_calls, "subprocess.Popen was not called" + return popen_calls[0] + + def test_java_executable_is_first_arg(self, jars_root, mock_client): + cmd = self._captured_popen_cmd( + jars_root, mock_client, java_executable="/usr/lib/jvm/java-17/bin/java" + ) + assert cmd[0] == "/usr/lib/jvm/java-17/bin/java" + + def test_classpath_flag_and_value_present(self, jars_root, mock_client): + cmd = self._captured_popen_cmd(jars_root, mock_client) + assert "-classpath" in cmd + cp_idx = cmd.index("-classpath") + classpath = cmd[cp_idx + 1] + assert jars_root.joinpath("app.jar").as_posix() in classpath + + def test_main_class_present(self, jars_root, mock_client): + cmd = self._captured_popen_cmd(jars_root, mock_client) + assert "com.example.TaskRunner" in cmd + + def test_comm_and_logs_args_present(self, jars_root, mock_client): + cmd = self._captured_popen_cmd(jars_root, mock_client) + comm_args = [a for a in cmd if a.startswith("--comm=")] + logs_args = [a for a in cmd if a.startswith("--logs=")] + assert len(comm_args) == 1 + assert len(logs_args) == 1 + + def test_comm_and_logs_contain_port(self, jars_root, mock_client): + cmd = self._captured_popen_cmd(jars_root, mock_client) + comm_arg = next(a for a in cmd if a.startswith("--comm=")) + logs_arg = next(a for a in cmd if a.startswith("--logs=")) + # format is host:port + assert ":" in comm_arg.split("=", 1)[1] + assert ":" in logs_arg.split("=", 1)[1] + + def test_jvm_args_inserted_before_main_class(self, jars_root, mock_client): + cmd = self._captured_popen_cmd(jars_root, mock_client, jvm_args=["-Xmx512m", "-Dsome.prop=value"]) + main_idx = cmd.index("com.example.TaskRunner") + for jvm_arg in ["-Xmx512m", "-Dsome.prop=value"]: + assert jvm_arg in cmd + assert cmd.index(jvm_arg) < main_idx + + def test_comm_and_logs_after_main_class(self, jars_root, mock_client): + cmd = self._captured_popen_cmd(jars_root, mock_client) + main_idx = cmd.index("com.example.TaskRunner") + comm_idx = next(i for i, a in enumerate(cmd) if a.startswith("--comm=")) + logs_idx = next(i for i, a in enumerate(cmd) if a.startswith("--logs=")) + assert comm_idx > main_idx + assert logs_idx > main_idx + + def test_returns_execution_result(self, jars_root, mock_client): + ti = _make_ti() + coordinator = JavaCoordinator(jars_root=jars_root) + + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.pid = 99999 + comm_sock = MagicMock(spec=socket.socket) + logs_sock = MagicMock(spec=socket.socket) + + with ( + patch("subprocess.Popen", return_value=mock_proc), + patch( + "airflow.sdk.coordinators.java.coordinator._accept_connections", + return_value={"comm": comm_sock, "logs": logs_sock}, + ), + patch.object(ActivitySubprocess, "_register_pipe_readers"), + patch.object(ActivitySubprocess, "_on_child_started"), + patch.object(ActivitySubprocess, "wait", return_value=0), + patch("psutil.Process"), + ): + result = coordinator.execute_task( + what=ti, + dag_rel_path="dags/test.jar", + bundle_info=MagicMock(), + client=mock_client, + subprocess_logs_to_stdout=False, + ) + + assert isinstance(result, BaseCoordinator.ExecutionResult) + assert result.exit_code == 0 + + +class TestJavaActivitySubprocessStart: + """ + Unit tests for _JavaActivitySubprocess.start(). + + These tests mock subprocess.Popen and _accept_connections to verify that + start() wires up the right command and stores the right sockets, + without requiring a real Java runtime. + """ + + def _start_with_mocks( + self, + jars_root: pathlib.Path, + mock_client, + *, + java_executable: str = "java", + jvm_args: list[str] | None = None, + ti: TaskInstanceDTO | None = None, + ): + """Call _JavaActivitySubprocess.start() with all subprocess machinery mocked out.""" + ti = ti or _make_ti() + + mock_proc = MagicMock(spec=subprocess.Popen) + mock_proc.pid = 12345 + comm_sock = MagicMock(spec=socket.socket) + logs_sock = MagicMock(spec=socket.socket) + + with ( + patch( + "airflow.sdk.coordinators.java.coordinator.subprocess.Popen", + return_value=mock_proc, + ) as popen_mock, + patch( + "airflow.sdk.coordinators.java.coordinator._accept_connections", + return_value={"comm": comm_sock, "logs": logs_sock}, + ), + patch.object(ActivitySubprocess, "_register_pipe_readers"), + patch.object(ActivitySubprocess, "_on_child_started"), + patch("psutil.Process"), + ): + proc = _JavaActivitySubprocess.start( + what=ti, + dag_rel_path="dags/test.jar", + bundle_info=MagicMock(), + client=mock_client, + java_executable=java_executable, + jvm_args=jvm_args or [], + jars_root=[jars_root], + subprocess_logs_to_stdout=False, + ) + + return proc, popen_mock + + def test_stdout_write_socket_stored_for_cleanup(self, jars_root, mock_client): + proc, _ = self._start_with_mocks(jars_root, mock_client) + # _stdout_w must be stored so wait() can close it + assert proc._stdout_w is not None + + def test_stderr_write_socket_stored_for_cleanup(self, jars_root, mock_client): + proc, _ = self._start_with_mocks(jars_root, mock_client) + assert proc._stderr_w is not None + + def test_stdout_and_stderr_write_sockets_are_distinct(self, jars_root, mock_client): + proc, _ = self._start_with_mocks(jars_root, mock_client) + assert proc._stdout_w is not proc._stderr_w + + def test_stdin_is_comm_socket(self, jars_root, mock_client): + """stdin (used by send_msg) must be the accepted comm socket.""" + ti = _make_ti() + comm_sock = MagicMock(spec=socket.socket) + logs_sock = MagicMock(spec=socket.socket) + + with ( + patch("airflow.sdk.coordinators.java.coordinator.subprocess.Popen") as popen_mock, + patch( + "airflow.sdk.coordinators.java.coordinator._accept_connections", + return_value={"comm": comm_sock, "logs": logs_sock}, + ), + patch.object(ActivitySubprocess, "_register_pipe_readers"), + patch.object(ActivitySubprocess, "_on_child_started"), + patch("psutil.Process"), + ): + popen_mock.return_value.pid = 12345 + proc = _JavaActivitySubprocess.start( + what=ti, + dag_rel_path="dags/test.jar", + bundle_info=MagicMock(), + client=MagicMock(), + java_executable="java", + jvm_args=[], + jars_root=[jars_root], + subprocess_logs_to_stdout=False, + ) + + assert proc.stdin is comm_sock + + def test_pid_taken_from_popen(self, jars_root, mock_client): + proc, _ = self._start_with_mocks(jars_root, mock_client) + assert proc.pid == 12345 + + def test_on_child_started_called(self, jars_root, mock_client): + ti = _make_ti() + with ( + patch("airflow.sdk.coordinators.java.coordinator.subprocess.Popen") as popen_mock, + patch( + "airflow.sdk.coordinators.java.coordinator._accept_connections", + return_value={"comm": MagicMock(), "logs": MagicMock()}, + ), + patch.object(ActivitySubprocess, "_register_pipe_readers"), + patch.object(ActivitySubprocess, "_on_child_started") as mock_on_started, + patch("psutil.Process"), + ): + popen_mock.return_value.pid = 12345 + _JavaActivitySubprocess.start( + what=ti, + dag_rel_path="dags/test.jar", + bundle_info=MagicMock(), + client=mock_client, + java_executable="java", + jvm_args=[], + jars_root=[jars_root], + subprocess_logs_to_stdout=False, + ) + + mock_on_started.assert_called_once() + kwargs = mock_on_started.call_args.kwargs + assert kwargs["ti"] is ti + assert kwargs["dag_rel_path"] == "dags/test.jar" + + def test_register_pipe_readers_called_with_four_sockets(self, jars_root, mock_client): + """Both socketpair read-ends and both TCP sockets must be registered.""" + with ( + patch("airflow.sdk.coordinators.java.coordinator.subprocess.Popen") as popen_mock, + patch( + "airflow.sdk.coordinators.java.coordinator._accept_connections", + return_value={"comm": MagicMock(), "logs": MagicMock()}, + ), + patch.object(ActivitySubprocess, "_register_pipe_readers") as mock_register, + patch.object(ActivitySubprocess, "_on_child_started"), + patch("psutil.Process"), + ): + popen_mock.return_value.pid = 12345 + _JavaActivitySubprocess.start( + what=_make_ti(), + dag_rel_path="dags/test.jar", + bundle_info=MagicMock(), + client=mock_client, + java_executable="java", + jvm_args=[], + jars_root=[jars_root], + subprocess_logs_to_stdout=False, + ) + + mock_register.assert_called_once() + args = mock_register.call_args.args + # positional: stdout, stderr, comm, logs — all four must be sockets + assert len(args) == 4 diff --git a/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py b/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py index 2b34fac6ea0f9..93c5cc19aed47 100644 --- a/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py +++ b/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py @@ -680,14 +680,14 @@ def mock_comms_response(msg): ("tg.t2", 0): ["a", "b"], ("tg.t2", 1): [4], ("tg.t2", 2): ["z"], - ("t3", None): [["a", "b"], [4], ["z"]], + ("t3", -1): [["a", "b"], [4], ["z"]], } # We hard-code the number of expansions here as the server is in charge of that. expansion_per_task_id = { "tg.t1": range(3), "tg.t2": range(3), - "t3": [None], + "t3": [-1], } for task in dag.tasks: for map_index in expansion_per_task_id[task.task_id]: diff --git a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py index af487851b07cb..6014d26f2208b 100644 --- a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py +++ b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py @@ -344,7 +344,7 @@ def xcom_get(msg): mock_supervisor_comms.send.side_effect = xcom_get # Run "pull_one" and "pull_all". - assert run_ti(dag, "pull_all", None) == TaskInstanceState.SUCCESS + assert run_ti(dag, "pull_all", -1) == TaskInstanceState.SUCCESS assert all_results == ["a", "b", "c", 1, 2] states = [run_ti(dag, "pull_one", map_index) for map_index in range(5)] diff --git a/task-sdk/tests/task_sdk/docs/test_public_api.py b/task-sdk/tests/task_sdk/docs/test_public_api.py index a21424ea101f6..0961458131735 100644 --- a/task-sdk/tests/task_sdk/docs/test_public_api.py +++ b/task-sdk/tests/task_sdk/docs/test_public_api.py @@ -42,30 +42,31 @@ def test_airflow_sdk_no_unexpected_exports(): ignore = { "__getattr__", "__lazy_imports", - "SecretCache", "TYPE_CHECKING", + "SecretCache", "annotations", "api", "bases", + "configuration", + "coordinators", + "crypto", "definitions", + "exceptions", "execution_time", "io", + "lineage", + "listener", "log", - "exceptions", - "timezone", - "secrets_masker", - "configuration", "module_loading", - "yaml", - "serde", "observability", "plugins_manager", - "listener", - "crypto", "providers_manager_runtime", - "lineage", + "secrets_masker", + "serde", + "timezone", "types", "state", + "yaml", } unexpected = actual - public - ignore assert not unexpected, f"Unexpected exports in airflow.sdk: {sorted(unexpected)}" 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..d89eb0f933adf --- /dev/null +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_integration.py @@ -0,0 +1,387 @@ +# 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. + +Notes on differences from PR #1577's original tests +---------------------------------------------------- +* The ``_subprocess_schema_version`` attribute (not ``lang_sdk_msg_schema_version``) + controls per-subprocess version pinning in the reimplementation. +* ``SchemaVersionMigrator`` is constructed with keyword arguments and an explicit + ``supervisor_version``. +* ``get_schema_version_migrator`` is imported directly into + ``airflow.sdk.execution_time.supervisor``; the monkeypatch therefore targets + ``airflow.sdk.execution_time.supervisor.get_schema_version_migrator`` (the + local binding in the supervisor module) in addition to the canonical location + in ``schema``. +* ``_send_startup_details`` from the coordinator module was removed from the + AIP-108 scope before this branch landed; the corresponding test class has + been omitted. +* ``WatchedSubprocess._serialize_response`` passes ``dump_kwargs=dump_opts`` to + ``downgrade`` instead of ``**dump_opts``, which would cause Pydantic's + ``model_dump`` to reject the unexpected keyword argument when any ``dump_opts`` + are present (including an empty dict). The ``mock_version_migrator`` fixture + patches ``_serialize_response`` with a corrected implementation so the + integration tests exercise the migration logic rather than this incidental bug. +""" + +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..74b8277f32045 --- /dev/null +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py @@ -0,0 +1,352 @@ +# 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. + +Notes on differences from PR #1577's original tests +----------------------------------------------------- +The reimplementation in ``supervisor-schema-migration`` differs from the +originally proposed design in a few ways that required test adaptations: + +* ``SchemaVersionMigrator`` uses ``attrs.define(kw_only=True)`` so it + must be constructed with keyword arguments (``bundle=``, + ``supervisor_version=``). There is no default for + ``supervisor_version``; callers must supply it explicitly or use + :func:`get_schema_version_migrator` which pins to the latest dated + entry. +* ``downgrade`` returns the Cadwyn-versioned Pydantic **model instance** + (not a plain dict). Tests that assert on output fields use + ``.model_dump()`` to normalise to a dict. +* ``_supervisor_version`` is a private attrs attribute; there is no + public ``supervisor_version`` property. Tests access it as + ``migrator._supervisor_version``. +* ``_resolve_version`` validates only that the version string is present + in the bundle; it does not enforce YYYY-MM-DD format. Version- + rejection tests therefore match the "not found in supervisor schema + bundle" error message rather than a format hint. +* ``_versioned_class`` uses Cadwyn's ``generate_versioned_models`` which + only generates classes for models mentioned in at least one + ``VersionChange`` instruction. Passing a model that is not registered + in the bundle raises ``KeyError``; the PR's "pass-through for + unregistered models" tests are therefore not applicable to this + implementation and have been omitted. +""" + +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 _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.""" + + 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 diff --git a/task-sdk/tests/task_sdk/execution_time/test_comms.py b/task-sdk/tests/task_sdk/execution_time/test_comms.py index 5c6d88439250c..37a91dd0ecc28 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_comms.py +++ b/task-sdk/tests/task_sdk/execution_time/test_comms.py @@ -86,6 +86,9 @@ def test_recv_StartupDetails(self): "run_id": "b", "dag_id": "c", "dag_version_id": uuid.UUID("4d828a62-a417-4936-a7a6-2b3fabacecab"), + "pool_slots": 1, + "queue": "default", + "priority_weight": 1, }, "ti_context": { "dag_run": { diff --git a/task-sdk/tests/task_sdk/execution_time/test_coordinator.py b/task-sdk/tests/task_sdk/execution_time/test_coordinator.py new file mode 100644 index 0000000000000..8e3bcfa102172 --- /dev/null +++ b/task-sdk/tests/task_sdk/execution_time/test_coordinator.py @@ -0,0 +1,122 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import json + +import pytest + +from airflow.sdk.configuration import conf +from airflow.sdk.execution_time.coordinator import ( + BaseCoordinator, + CoordinatorManager, + _PythonCoordinator, + get_coordinator_manager, + reset_coordinator_manager, +) + + +class _CoordinatorA(BaseCoordinator): + def __init__(self, *, label: str = "a"): + self.label = label + + +class _CoordinatorB(BaseCoordinator): + pass + + +@pytest.fixture +def sdk_config(monkeypatch): + """Set the ``[sdk]`` env vars consumed by :meth:`CoordinatorManager.from_config`. + + :return: Callable ``apply(*, coordinators=None, queue_to_coordinator=None)`` -- + each argument is the raw JSON string for the matching env var, or ``None`` + to unset it. The conf cache is invalidated after each call (and again on + teardown) so ``from_config()`` re-reads the values just set. + """ + from airflow.sdk.configuration import conf + + def _apply(*, coordinators: str | None = None, queue_to_coordinator: str | None = None) -> None: + if coordinators is None: + monkeypatch.delenv("AIRFLOW__SDK__COORDINATORS", raising=False) + else: + monkeypatch.setenv("AIRFLOW__SDK__COORDINATORS", coordinators) + if queue_to_coordinator is None: + monkeypatch.delenv("AIRFLOW__SDK__QUEUE_TO_COORDINATOR", raising=False) + else: + monkeypatch.setenv("AIRFLOW__SDK__QUEUE_TO_COORDINATOR", queue_to_coordinator) + conf.invalidate_cache() + + yield _apply + conf.invalidate_cache() + + +class TestCoordinatorManager: + @pytest.fixture(autouse=True) + def _reset_cache(self): + reset_coordinator_manager() + yield + reset_coordinator_manager() + + def test_from_config_loads_specs_and_resolves_instances(self, sdk_config): + sdk_config( + coordinators=json.dumps( + { + "alpha": { + "classpath": f"{_CoordinatorA.__module__}._CoordinatorA", + "kwargs": {"label": "alpha-label"}, + }, + "beta": {"classpath": f"{_CoordinatorB.__module__}._CoordinatorB", "kwargs": {}}, + } + ), + queue_to_coordinator=json.dumps({"queue-a": "alpha"}), + ) + manager = CoordinatorManager.from_config() + assert manager._queue_to_coordinator == {"queue-a": "alpha"} + assert manager._created_coordinators == {} + + coordinator_for_queue_a = manager.for_queue("queue-a") + assert isinstance(coordinator_for_queue_a, _CoordinatorA) + assert manager.for_queue("queue-a") is coordinator_for_queue_a, "instance should be cached" + assert manager._created_coordinators == {"alpha": coordinator_for_queue_a} + + coordinator_for_queue_missing = manager.for_queue("queue-1") + assert isinstance(coordinator_for_queue_missing, _PythonCoordinator) + assert manager.for_queue("queue-1") is coordinator_for_queue_missing + assert manager._created_coordinators == {"alpha": coordinator_for_queue_a} + + def test_from_config_empty(self, monkeypatch): + monkeypatch.delenv("AIRFLOW__SDK__COORDINATORS", raising=False) + monkeypatch.delenv("AIRFLOW__SDK__QUEUE_TO_COORDINATOR", raising=False) + conf.invalidate_cache() + + manager = CoordinatorManager.from_config() + assert manager._coordinator_specs == {} + assert manager._queue_to_coordinator == {} + + def test_get_coordinator_manager_is_cached(self, monkeypatch): + monkeypatch.delenv("AIRFLOW__SDK__COORDINATORS", raising=False) + + from airflow.sdk.configuration import conf + + conf.invalidate_cache() + + m1 = get_coordinator_manager() + m2 = get_coordinator_manager() + assert m1 is m2 diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index 811101dbf4c02..e578b6023f27a 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -64,7 +64,6 @@ DagRunState, DagRunType, PreviousTIResponse, - TaskInstance, TaskInstanceState, ) from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType, TaskAlreadyRunningError @@ -169,6 +168,7 @@ supervise_task, ) from airflow.sdk.execution_time.task_runner import run +from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO from tests_common.test_utils.config import conf_vars @@ -229,13 +229,16 @@ def test_supervise( """ Test that the supervisor validates server URL and dry_run parameter combinations correctly. """ - ti = TaskInstance( + ti = TaskInstanceDTO( id=uuid7(), task_id="async", dag_id="super_basic_deferred_run", run_id="d", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ) bundle_info = BundleInfo(name="my-bundle", version=None) @@ -322,13 +325,16 @@ def subprocess_main(): proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstance( + what=TaskInstanceDTO( id="4d828a62-a417-4936-a7a6-2b3fabacecab", task_id="b", dag_id="c", run_id="d", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), client=client_with_ti_start, target=subprocess_main, @@ -397,13 +403,16 @@ def subprocess_main(): proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstance( + what=TaskInstanceDTO( id="4d828a62-a417-4936-a7a6-2b3fabacecab", task_id="b", dag_id="c", run_id="d", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), client=client_with_ti_start, target=subprocess_main, @@ -494,13 +503,16 @@ def on_kill(self) -> None: proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstance( + what=TaskInstanceDTO( id=ti_id, task_id="b", dag_id="c", run_id="d", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), client=make_client(transport=httpx.MockTransport(handle_request)), target=subprocess_main, @@ -523,13 +535,16 @@ def subprocess_main(): proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstance( + what=TaskInstanceDTO( id="4d828a62-a417-4936-a7a6-2b3fabacecab", task_id="b", dag_id="c", run_id="d", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), client=client_with_ti_start, target=subprocess_main, @@ -558,8 +573,16 @@ def subprocess_main(): proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstance( - id=uuid7(), task_id="b", dag_id="c", run_id="d", try_number=1, dag_version_id=uuid7() + what=TaskInstanceDTO( + id=uuid7(), + task_id="b", + dag_id="c", + run_id="d", + try_number=1, + dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), client=mock_client, target=subprocess_main, @@ -597,13 +620,16 @@ def test_resume_start_date_from_context(self, mocker, make_ti_context, start_dat proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstance( + what=TaskInstanceDTO( id=uuid7(), task_id="b", dag_id="c", run_id="d", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), client=mock_client, target=lambda: None, @@ -640,8 +666,16 @@ def subprocess_main(): proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstance( - id=ti_id, task_id="b", dag_id="c", run_id="d", try_number=1, dag_version_id=uuid7() + what=TaskInstanceDTO( + id=ti_id, + task_id="b", + dag_id="c", + run_id="d", + try_number=1, + dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), client=sdk_client.Client(base_url="", dry_run=True, token=""), target=subprocess_main, @@ -677,8 +711,16 @@ def _on_child_started(self, *args, **kwargs): proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstance( - id=ti_id, task_id="b", dag_id="c", run_id="d", try_number=1, dag_version_id=uuid7() + what=TaskInstanceDTO( + id=ti_id, + task_id="b", + dag_id="c", + run_id="d", + try_number=1, + dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), client=sdk_client.Client(base_url="", dry_run=True, token=""), target=subprocess_main, @@ -693,13 +735,16 @@ def test_run_simple_dag(self, test_dags_dir, captured_logs, time_machine, mocker time_machine.move_to(instant, tick=False) dagfile_path = test_dags_dir - ti = TaskInstance( + ti = TaskInstanceDTO( id=uuid7(), task_id="hello", dag_id="super_basic_run", run_id="c", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ) bundle_info = BundleInfo(name="my-bundle", version=None) @@ -734,13 +779,16 @@ def test_supervise_handles_deferred_task( """ instant = timezone.datetime(2024, 11, 7, 12, 34, 56, 0) - ti = TaskInstance( + ti = TaskInstanceDTO( id=uuid7(), task_id="async", dag_id="super_basic_deferred_run", run_id="d", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ) # Create a mock client to assert calls to the client @@ -861,8 +909,16 @@ def handle_request(request: httpx.Request) -> httpx.Response: proc = ActivitySubprocess.start( dag_rel_path=os.devnull, - what=TaskInstance( - id=ti_id, task_id="b", dag_id="c", run_id="d", try_number=1, dag_version_id=uuid7() + what=TaskInstanceDTO( + id=ti_id, + task_id="b", + dag_id="c", + run_id="d", + try_number=1, + dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), client=make_client(transport=httpx.MockTransport(handle_request)), target=subprocess_main, @@ -939,8 +995,16 @@ def subprocess_main(): ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstance( - id=ti_id, task_id="b", dag_id="c", run_id="d", try_number=1, dag_version_id=uuid7() + what=TaskInstanceDTO( + id=ti_id, + task_id="b", + dag_id="c", + run_id="d", + try_number=1, + dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), client=make_client(transport=httpx.MockTransport(handle_request)), target=subprocess_main, @@ -1144,13 +1208,16 @@ def subprocess_main(): proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstance( + what=TaskInstanceDTO( id="4d828a62-a417-4936-a7a6-2b3fabacecab", task_id="b", dag_id="c", run_id="d", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), client=client_with_ti_start, target=subprocess_main, @@ -1301,8 +1368,16 @@ def _handler(sig, frame): proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstance( - id=ti_id, task_id="b", dag_id="c", run_id="d", try_number=1, dag_version_id=uuid7() + what=TaskInstanceDTO( + id=ti_id, + task_id="b", + dag_id="c", + run_id="d", + try_number=1, + dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), client=client_with_ti_start, target=subprocess_main, @@ -3688,13 +3763,16 @@ def subprocess_main(): proc = ActivitySubprocess.start( dag_rel_path=os.devnull, bundle_info=FAKE_BUNDLE, - what=TaskInstance( + what=TaskInstanceDTO( id="4d828a62-a417-4936-a7a6-2b3fabacecab", task_id="b", dag_id="c", run_id="d", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), client=client_with_ti_start, target=subprocess_main, diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index b37d2569ea4f1..2f0759133ba11 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -165,6 +165,7 @@ run, startup, ) +from airflow.sdk.execution_time.workloads.task import TaskInstanceDTO from airflow.sdk.execution_time.xcom import XCom from airflow.sdk.serde import deserialize from airflow.triggers.base import BaseEventTrigger, BaseTrigger, TriggerEvent @@ -196,13 +197,16 @@ def execute(self, context): def test_parse(test_dags_dir: Path, make_ti_context): """Test that checks parsing of a basic dag with an un-mocked parse.""" what = StartupDetails( - ti=TaskInstance( + ti=TaskInstanceDTO( id=uuid7(), task_id="a", dag_id="super_basic", run_id="c", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), dag_rel_path="super_basic.py", bundle_info=BundleInfo(name="my-bundle", version=None), @@ -243,13 +247,16 @@ def test_parse_dag_bag(mock_dagbag, test_dags_dir: Path, make_ti_context): mock_dag.task_dict = {"a": mock_task} what = StartupDetails( - ti=TaskInstance( + ti=TaskInstanceDTO( id=uuid7(), task_id="a", dag_id="super_basic", run_id="c", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), dag_rel_path="super_basic.py", bundle_info=BundleInfo(name="my-bundle", version=None), @@ -303,13 +310,16 @@ def test_parse_dag_bag(mock_dagbag, test_dags_dir: Path, make_ti_context): def test_parse_not_found(test_dags_dir: Path, make_ti_context, dag_id, task_id, expected_error): """Check for nice error messages on dag not found.""" what = StartupDetails( - ti=TaskInstance( + ti=TaskInstanceDTO( id=uuid7(), task_id=task_id, dag_id=dag_id, run_id="c", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), dag_rel_path="super_basic.py", bundle_info=BundleInfo(name="my-bundle", version=None), @@ -349,13 +359,16 @@ def test_parse_not_found_does_not_reschedule_when_max_attempts_reached(test_dags and should surface as a hard failure (SystemExit in the task runner process). """ what = StartupDetails( - ti=TaskInstance( + ti=TaskInstanceDTO( id=uuid7(), task_id="a", dag_id="madeup_dag_id", run_id="c", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), dag_rel_path="super_basic.py", bundle_info=BundleInfo(name="my-bundle", version=None), @@ -410,13 +423,16 @@ def test_main_sends_reschedule_task_when_startup_reschedules( mock_comms_instance.socket = None mock_comms_decoder_cls.__getitem__.return_value.return_value = mock_comms_instance what = StartupDetails( - ti=TaskInstance( + ti=TaskInstanceDTO( id=uuid7(), task_id="my_task", dag_id="test_dag", run_id="test_run", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, context_carrier={}, ), dag_rel_path="", @@ -583,13 +599,16 @@ def test_task_span_is_child_of_dag_run_span(make_ti_context): # Step 3: build StartupDetails with ti.context_carrier = ti_carrier. what = StartupDetails( - ti=TaskInstance( + ti=TaskInstanceDTO( id=uuid7(), task_id="my_task", dag_id="test_dag", run_id="test_run", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, context_carrier=ti_carrier, ), dag_rel_path="", @@ -651,13 +670,16 @@ def test_task_span_no_parent_when_no_context_carrier(make_ti_context): provider.add_span_processor(SimpleSpanProcessor(in_mem_exporter)) what = StartupDetails( - ti=TaskInstance( + ti=TaskInstanceDTO( id=uuid7(), task_id="standalone_task", dag_id="test_dag", run_id="test_run", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, context_carrier=None, ), dag_rel_path="", @@ -692,13 +714,16 @@ def test_parse_module_in_bundle_root(tmp_path: Path, make_ti_context): dag1_path.write_text(textwrap.dedent(dag1_code)) what = StartupDetails( - ti=TaskInstance( + ti=TaskInstanceDTO( id=uuid7(), task_id="a", dag_id="dag_name", run_id="c", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), dag_rel_path="path_test.py", bundle_info=BundleInfo(name="my-bundle", version=None), @@ -1139,13 +1164,16 @@ def test_basic_templated_dag(mocked_parse, make_ti_context, mock_supervisor_comm ) what = StartupDetails( - ti=TaskInstance( + ti=TaskInstanceDTO( id=uuid7(), task_id="templated_task", dag_id="basic_templated_dag", run_id="c", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), bundle_info=FAKE_BUNDLE, dag_rel_path="", @@ -1255,13 +1283,16 @@ def execute(self, context): instant = timezone.datetime(2024, 12, 3, 10, 0) what = StartupDetails( - ti=TaskInstance( + ti=TaskInstanceDTO( id=uuid7(), task_id="templated_task", dag_id="basic_dag", run_id="c", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), dag_rel_path="", bundle_info=FAKE_BUNDLE, @@ -1303,13 +1334,16 @@ def execute(self, context): instant = timezone.datetime(2024, 12, 3, 10, 0) what = StartupDetails( - ti=TaskInstance( + ti=TaskInstanceDTO( id=uuid7(), task_id="impersonation_task", dag_id="basic_dag", run_id="c", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), dag_rel_path="", bundle_info=FAKE_BUNDLE, @@ -1351,13 +1385,16 @@ def execute(self, context): instant = timezone.datetime(2024, 12, 3, 10, 0) what = StartupDetails( - ti=TaskInstance( + ti=TaskInstanceDTO( id=uuid7(), task_id="impersonation_task", dag_id="basic_dag", run_id="c", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), dag_rel_path="", bundle_info=FAKE_BUNDLE, @@ -1391,13 +1428,16 @@ def execute(self, context): instant = timezone.datetime(2024, 12, 3, 10, 0) what = StartupDetails( - ti=TaskInstance( + ti=TaskInstanceDTO( id=uuid7(), task_id="impersonation_task", dag_id="basic_dag", run_id="c", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), dag_rel_path="", bundle_info=FAKE_BUNDLE, @@ -1564,8 +1604,16 @@ def test_dag_parsing_context(make_ti_context, mock_supervisor_comms, monkeypatch task_id = "conditional_task" what = StartupDetails( - ti=TaskInstance( - id=uuid7(), task_id=task_id, dag_id=dag_id, run_id="c", try_number=1, dag_version_id=uuid7() + ti=TaskInstanceDTO( + id=uuid7(), + task_id=task_id, + dag_id=dag_id, + run_id="c", + try_number=1, + dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), dag_rel_path="dag_parsing_context.py", bundle_info=BundleInfo(name="my-bundle", version=None), @@ -2204,8 +2252,10 @@ def execute(self, context): test_task_id = "pull_task" task = CustomOperator(task_id=test_task_id) - # In case of the specific map_index or None we should check it is passed to TI - extra_for_ti = {"map_index": map_indexes} if map_indexes in (1, None) else {} + # In case of the specific map_index we should check it is passed to TI. + # ``None`` is not a valid TaskInstanceDTO.map_index value, but xcom_pull's + # behaviour with ``map_indexes=None`` is independent of the TI's own map_index. + extra_for_ti = {"map_index": map_indexes} if isinstance(map_indexes, int) else {} runtime_ti = create_runtime_ti(task=task, **extra_for_ti) ser_value = BaseXCom.serialize_value(xcom_values) @@ -4100,13 +4150,16 @@ def execute(self, context): task_id="test_task_runner_calls_listeners", do_xcom_push=True, multiple_outputs=True ) what = StartupDetails( - ti=TaskInstance( + ti=TaskInstanceDTO( id=uuid7(), task_id="templated_task", dag_id="basic_dag", run_id="c", try_number=1, dag_version_id=uuid7(), + pool_slots=1, + queue="default", + priority_weight=1, ), dag_rel_path="", bundle_info=FAKE_BUNDLE, @@ -4755,7 +4808,8 @@ class CustomOperator(BaseOperator): class TestTriggerDagRunOperator: """Tests to verify various aspects of TriggerDagRunOperator""" - @time_machine.travel("2025-01-01 00:00:00", tick=False) + # make timetravel timezone-aware + @time_machine.travel(datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc), tick=False) def test_handle_trigger_dag_run(self, create_runtime_ti, mock_supervisor_comms): """Test that TriggerDagRunOperator (with default args) sends the correct message to the Supervisor""" from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator @@ -4803,7 +4857,7 @@ def test_handle_trigger_dag_run(self, create_runtime_ti, mock_supervisor_comms): (False, TaskInstanceState.FAILED), ], ) - @time_machine.travel("2025-01-01 00:00:00", tick=False) + @time_machine.travel(datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc), tick=False) def test_handle_trigger_dag_run_conflict( self, skip_when_already_exists, expected_state, create_runtime_ti, mock_supervisor_comms ): @@ -4847,7 +4901,7 @@ def test_handle_trigger_dag_run_conflict( ([DagRunState.SUCCESS], None, DagRunState.FAILED, DagRunState.FAILED), ], ) - @time_machine.travel("2025-01-01 00:00:00", tick=False) + @time_machine.travel(datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc), tick=False) def test_handle_trigger_dag_run_wait_for_completion( self, allowed_states, @@ -4968,7 +5022,7 @@ def test_handle_trigger_dag_run_deferred( assert state == intermediate_state - @time_machine.travel("2025-01-01 00:00:00", tick=False) + @time_machine.travel(datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc), tick=False) def test_handle_trigger_dag_run_deferred_with_reset_uses_run_id_only( self, create_runtime_ti, mock_supervisor_comms ): diff --git a/uv.lock b/uv.lock index 7766a2d3cf466..891e7b8c7e159 100644 --- a/uv.lock +++ b/uv.lock @@ -457,7 +457,7 @@ wheels = [ [[package]] name = "aiofile" -version = "3.11.0" +version = "3.11.1" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.15' and sys_platform == 'win32'", @@ -474,9 +474,9 @@ resolution-markers = [ dependencies = [ { name = "caio", marker = "python_full_version >= '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ec/c6/01a1e5c556d806a865dd891daaffee93aee8c32840f1858f0c8312413f85/aiofile-3.11.0.tar.gz", hash = "sha256:7f164b85fb767ae7889707d964373ca1277c2aa75f544d564f8577b2c71b5130", size = 19579, upload-time = "2026-05-16T07:09:30.632Z" } +sdist = { url = "https://files.pythonhosted.org/packages/48/41/2fea7e193e061ce54eacc3b7bc0e6a99e4fcff43c78cf0a76dd781ed8334/aiofile-3.11.1.tar.gz", hash = "sha256:1f91912c6643d2a4e49ca4ae3514f0bf3867ce948a36d99a6411b8f4755f4cf9", size = 19342, upload-time = "2026-05-16T08:18:33.538Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/dd/ab/c6ab3a18a7b4d16c5d50ac88c31819a14a7c1866a653f6f7285d7e8ed69e/aiofile-3.11.0-py3-none-any.whl", hash = "sha256:7560061eee6773e767fd9b5a274fbe4ec93a840eae6de00bb61edb38fab3a57e", size = 20663, upload-time = "2026-05-16T07:09:29.281Z" }, + { url = "https://files.pythonhosted.org/packages/67/cd/0d76dfc5de72bde52f55f53e925c7d152d9c7906634ec1e0cbc7e8d4ad93/aiofile-3.11.1-py3-none-any.whl", hash = "sha256:ce77d14ac07f77bc2b757834a5c129321f3f705c474593deed5ab209079a52c9", size = 20446, upload-time = "2026-05-16T08:18:32.051Z" }, ] [[package]] @@ -8922,40 +8922,42 @@ wheels = [ [[package]] name = "ast-serialize" -version = "0.4.0" +version = "0.5.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e2/1f/50f241d4e01fe75f4bba6a209edd4047c4b26acf70992ff885fd161f79cb/ast_serialize-0.4.0.tar.gz", hash = "sha256:74e4e634ab82d1466acf0be27043178570b98ebeaa3165f9240a6fad4c286471", size = 60687, upload-time = "2026-05-14T22:44:38.251Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b0/85/232631c59b5ca7152c08f026e9a46f47d852298acff74edd04a1fc1d0005/ast_serialize-0.4.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:a6f26937ce0293aafbece0e39019e020369a5a70486ff4088227f0cc888844a9", size = 1182685, upload-time = "2026-05-14T22:43:40.205Z" }, - { url = "https://files.pythonhosted.org/packages/5d/5e/4838d4d3ddc4425555601467d4e2a565e4340899e45feee4e32c80fbc911/ast_serialize-0.4.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:074032142777e3e6091977dc3c5146a8ca58ae6825b7f64e9a0b604153ddabd8", size = 1173113, upload-time = "2026-05-14T22:43:41.937Z" }, - { url = "https://files.pythonhosted.org/packages/22/fc/d622b19fc1c79a62028ec17f4ad4323177af25b174d32b07c84d61ef9d47/ast_serialize-0.4.0-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:404f3462b4532e13a70b8849bba241dbd82e30043ff58d98c7e762fd925b116a", size = 1234117, upload-time = "2026-05-14T22:43:43.977Z" }, - { url = "https://files.pythonhosted.org/packages/d5/b5/72f8c8659da0b64562e6d97f852d5c2022c74577df27c922e1e7065039ce/ast_serialize-0.4.0-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:97c55336e16f5c4ca2bde7be94cca4b8f7d665d64f7008925a82e02707ba14ac", size = 1231703, upload-time = "2026-05-14T22:43:46.064Z" }, - { url = "https://files.pythonhosted.org/packages/7b/98/ccc51ee4f90f97a1ed0a0848bd4c9d77a80969849db8a262b7d2970a6a15/ast_serialize-0.4.0-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:732b4ef76adcb0f298a7d18c4558336d83b1384f9ae0c7eaa1dc8d031b0a4390", size = 1441574, upload-time = "2026-05-14T22:43:47.784Z" }, - { url = "https://files.pythonhosted.org/packages/6a/ce/668c4efe79e09c9cc97a4d0a1c29e61fe6f78857fe1e57c086772af55f89/ast_serialize-0.4.0-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b3db87c4772097c0782250bcd550d66b1189a8c889793c7bcf153f4fee70005c", size = 1254040, upload-time = "2026-05-14T22:43:49.879Z" }, - { url = "https://files.pythonhosted.org/packages/3d/be/38b27bc2909b7236939801ca9f0d97cdc6198da4f435a81658e0db506fdb/ast_serialize-0.4.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:43729a5e369ebbe7750635c0c206bc616fcd36e703cb9c4497d6b4df0291ee64", size = 1257847, upload-time = "2026-05-14T22:43:51.607Z" }, - { url = "https://files.pythonhosted.org/packages/68/df/360ebccc361235c167a8be2a0476870cb9ef44c42413bf1289b885684052/ast_serialize-0.4.0-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:91d3786f3929786cdc4eeedfd110abb4603e7f6c1390c5af398f333a947b742d", size = 1298683, upload-time = "2026-05-14T22:43:53.606Z" }, - { url = "https://files.pythonhosted.org/packages/51/5c/7d5e0b4d47aafa1600c19e3670f962f81a9bf3da1bc25a1382529a447cf3/ast_serialize-0.4.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:7fba7315fd4bd87cb5560792709f6e66e0606402d362c0a38dd32dfb66ba6066", size = 1409438, upload-time = "2026-05-14T22:43:55.316Z" }, - { url = "https://files.pythonhosted.org/packages/b5/3d/8875b2f1af3ec1539b88ff193dfbfa5573084ef7fcab27ea4cd09b6dc829/ast_serialize-0.4.0-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:4db9769d57deb5545ce56ebbbbe3436dcc0ae2688ce14c295cd14e106624ece7", size = 1507922, upload-time = "2026-05-14T22:43:56.959Z" }, - { url = "https://files.pythonhosted.org/packages/55/c3/5ec6927eb493ece7ba64263cdc556be889e0c62a013b1851bbe674a0dcda/ast_serialize-0.4.0-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:dcd04f85a29deb80400e8987cfaceb9907140f763453cbffdbd6ff36f1b32c12", size = 1502817, upload-time = "2026-05-14T22:43:59.081Z" }, - { url = "https://files.pythonhosted.org/packages/9d/c8/40cb818a08396b1f34d6189c0c42aec917dd331e11fb7c3b870cc61b795a/ast_serialize-0.4.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:905fc11940831454d93589bd7ce2acb6a5eb01c2936156f751d2a21087c98cd3", size = 1454318, upload-time = "2026-05-14T22:44:01.377Z" }, - { url = "https://files.pythonhosted.org/packages/74/d5/d51494b60cc52f4792be5ddc951631cddb17a2990154634549abdbdbb5bf/ast_serialize-0.4.0-cp314-cp314t-win32.whl", hash = "sha256:3bdde2c4570143791f636aed4e3ef868f5b46eb90a18f8d5c41dd045aab08bef", size = 1060098, upload-time = "2026-05-14T22:44:03.265Z" }, - { url = "https://files.pythonhosted.org/packages/7a/c9/b0086257c79ff95743a3621448a01fc71b234ae359d3d54cda383aa43939/ast_serialize-0.4.0-cp314-cp314t-win_amd64.whl", hash = "sha256:6551d55b8607b97a7755683d743200b398c61a0b71a11b7f00c89c335a11d0f4", size = 1101015, upload-time = "2026-05-14T22:44:05.055Z" }, - { url = "https://files.pythonhosted.org/packages/3d/6d/3dfddef4990fda47745af6615a3e51c4de711eda56c3a8072a0d8b6181c7/ast_serialize-0.4.0-cp314-cp314t-win_arm64.whl", hash = "sha256:7234ff086cb152ea2a3b7ef895b5ebeb6d80779df049d5c6431c8e3536d5b03c", size = 1074495, upload-time = "2026-05-14T22:44:07.186Z" }, - { url = "https://files.pythonhosted.org/packages/be/d5/044c5f995ef75807a0effb56fc288cfdedeeb571222450fb6f7d94fd52f1/ast_serialize-0.4.0-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:dcded5056d9f3d201df7833082c07ebcbc566ffc3d4105c9fc9fe278fa086ecb", size = 1189800, upload-time = "2026-05-14T22:44:09.333Z" }, - { url = "https://files.pythonhosted.org/packages/a9/5a/52163557789d59a8197c10912ab4a1791c9143731ba0e3d9283ac0791db6/ast_serialize-0.4.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:bd50d201098aae0d202805fe9606c0545492f69a3ec4403337e32c54ad29fc41", size = 1181713, upload-time = "2026-05-14T22:44:11.286Z" }, - { url = "https://files.pythonhosted.org/packages/2c/c3/678ce3b6cb594b01c361da87f6c5679d26c1dae1583a082a8cd190e7232e/ast_serialize-0.4.0-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6615b39cd747967c3aabe68bf3f5f26748e823cc6b474ddc1510ed188a824149", size = 1243258, upload-time = "2026-05-14T22:44:13.345Z" }, - { url = "https://files.pythonhosted.org/packages/3d/dd/4810fbeb81c47b7e4e65db15ca65c71330efc59b460bd10c12338dc6012e/ast_serialize-0.4.0-cp39-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:91362c0a9fdf1c344b7f50a5b0508b11a0732102998fbd754a191f7187e77031", size = 1239226, upload-time = "2026-05-14T22:44:15.811Z" }, - { url = "https://files.pythonhosted.org/packages/28/38/13a88d90b664c009ed208346ec2ed248b0ab2cb0b582ae467acaa7f44fa4/ast_serialize-0.4.0-cp39-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:70d9c5d527bbfa69bd3c7d17dac11fb6781e36186a434a06d7d5892e0b2f88f9", size = 1448867, upload-time = "2026-05-14T22:44:17.99Z" }, - { url = "https://files.pythonhosted.org/packages/4c/19/a069dba1a634b703bf07fb49df8f7e3c04e9ba8ef3f0d9f4495f72630f92/ast_serialize-0.4.0-cp39-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:4738790cf54d8b416de992b87ee567056980bc82134d52458bd4985f389d1658", size = 1264135, upload-time = "2026-05-14T22:44:19.8Z" }, - { url = "https://files.pythonhosted.org/packages/2d/4c/76ec4279fecd7e78b60c3c99321f944c43cd11e5ff09c952746f5f9c0f4c/ast_serialize-0.4.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:faa008dccfcb793ae9101325e4d6d026caaa5d845c2182f03749c759834b0a3a", size = 1269060, upload-time = "2026-05-14T22:44:21.894Z" }, - { url = "https://files.pythonhosted.org/packages/33/c5/9230ef7481e5cb63b93a1f7738e959586202b081caf32b8bc5d9f673ef56/ast_serialize-0.4.0-cp39-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1c5245228e65d38cb48e1251f0ca71b0fa417e527141491e8c92f740e8e2d121", size = 1309654, upload-time = "2026-05-14T22:44:23.725Z" }, - { url = "https://files.pythonhosted.org/packages/b9/54/7d7397528d181ad68e476e0c81aa3ceff7d1f1b5c7fa958d6be28628ef16/ast_serialize-0.4.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:8f5153e9c44a02e61f4042c5f9249d2e8a759773d621a0b2f445a899e536e181", size = 1418855, upload-time = "2026-05-14T22:44:25.415Z" }, - { url = "https://files.pythonhosted.org/packages/b8/8f/87d6428adaa0986b817404f09329b64f8d2614cfe061ebf4951b4a7e0d19/ast_serialize-0.4.0-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:1e1fb90def261f6a0db885876f7e1a49ad2dbac38ad9f2f62dba2f9543af16e7", size = 1516040, upload-time = "2026-05-14T22:44:27.535Z" }, - { url = "https://files.pythonhosted.org/packages/b5/bb/5aaa41a21314c8b0d6dee54867b16535682c6660dd28cac64dba1380062d/ast_serialize-0.4.0-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:cf2ff7b654c8e95143e20f5d75878cbb78b65b928b26c4d58ef71cdba9d6d981", size = 1511450, upload-time = "2026-05-14T22:44:29.522Z" }, - { url = "https://files.pythonhosted.org/packages/87/16/cc729b5bb4b21da99db1379266cc367512e82ba10f9b3300a6f3e9941325/ast_serialize-0.4.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:90fc5c0d35a22f1a92dd33635508626d50f8fc64deb897c23e78e666a60804c9", size = 1463654, upload-time = "2026-05-14T22:44:31.265Z" }, - { url = "https://files.pythonhosted.org/packages/43/97/7198321b0244d011093387b41affea934d58bda08d59a2adfde72976b6c4/ast_serialize-0.4.0-cp39-abi3-win32.whl", hash = "sha256:9ecd6a1fc1b86f1f4e8ae206759b6319c10019706b3496b01b54d02b9b2cd918", size = 1068636, upload-time = "2026-05-14T22:44:33.189Z" }, - { url = "https://files.pythonhosted.org/packages/10/09/3b868f6d8df4bbe452903a5e0e039ebcec9ea0045f1a77951546205097e8/ast_serialize-0.4.0-cp39-abi3-win_amd64.whl", hash = "sha256:79c8d015c771c8bfdb1208003b227b27c40034790a2c29c09f2317a041825ce2", size = 1107137, upload-time = "2026-05-14T22:44:35.304Z" }, - { url = "https://files.pythonhosted.org/packages/fd/78/9387dffccdc55a12734f83aaccc4a987404a217a2a12a1920d8d4585950b/ast_serialize-0.4.0-cp39-abi3-win_arm64.whl", hash = "sha256:1026f565a7ab846337c630909089b3346a2fe417bf1552b1581ab01852137407", size = 1079199, upload-time = "2026-05-14T22:44:36.816Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/81/9d/09e27731bd5864a9ce04e3244074e674bb8936bf62b45e0357248717adac/ast_serialize-0.5.0.tar.gz", hash = "sha256:5880091bfe6f4f986f22866375c2e884843e7a0b6343ae41aeea659613d879b6", size = 61157, upload-time = "2026-05-17T17:48:29.429Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c0/9a/13dde51ba9e15f8b97957ab7cb0120d0e381524d651c6bd630b9c359227f/ast_serialize-0.5.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:8f5c14f169eb0972c0c21bada5358b23d6047c76583b005234f865b11f1fa00a", size = 1183520, upload-time = "2026-05-17T17:47:30.831Z" }, + { url = "https://files.pythonhosted.org/packages/37/de/5a7f0a9fe68944f536632a5af84676739c7d2582be42deb082634bf3a754/ast_serialize-0.5.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:7d1a2de9de5be04652f0ed60738356ef94f66db37924a9499fffe98dc491aa0b", size = 1175779, upload-time = "2026-05-17T17:47:32.551Z" }, + { url = "https://files.pythonhosted.org/packages/9c/81/0bb853e76e4f6e9a1855d569003c59e19ffac45f7079d91505d1bb212f92/ast_serialize-0.5.0-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:be5173fb66f9b49026d9d5a2ff0fc7c7009077107c0eb285b2d60fdf1fe10bd1", size = 1233750, upload-time = "2026-05-17T17:47:34.731Z" }, + { url = "https://files.pythonhosted.org/packages/e5/d3/4cf705beeccc08754d0bbda99aefff26110e209b9a07ac8a6b60eec48531/ast_serialize-0.5.0-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:f8015cd071ac1339924ee2b8098c93e00e155f30a16f40ec9816fcf84f4753f6", size = 1235942, upload-time = "2026-05-17T17:47:36.287Z" }, + { url = "https://files.pythonhosted.org/packages/26/c8/ee097e437ea27dd2b8b227865c875492b585650a5802a22d82b304c8201b/ast_serialize-0.5.0-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5499e8797edff2a9186aa313ed382c6b422e798e9332d9953badcee6e69a88f2", size = 1442517, upload-time = "2026-05-17T17:47:38.17Z" }, + { url = "https://files.pythonhosted.org/packages/ff/bd/68063442838f1ba68ec72b5436430bc75b3bb17a1a3c3063f09b0c05ae2b/ast_serialize-0.5.0-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:6848f2a093fb5548751a9a09bff8fcd229e2bbeb0e3331f391b6ae6d26cd9903", size = 1254081, upload-time = "2026-05-17T17:47:39.826Z" }, + { url = "https://files.pythonhosted.org/packages/50/e2/1e520793bc6a4e4524a6ab022391e827825eaa0c3811828bfdc6852eca26/ast_serialize-0.5.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:832d4c998e0b091fd60a6d6bceee535483c4d490de9ba85003af835225719261", size = 1259910, upload-time = "2026-05-17T17:47:41.369Z" }, + { url = "https://files.pythonhosted.org/packages/4e/e1/49b60f467979979cfe6913b43948ff25bca971ad0591d181812f163a988e/ast_serialize-0.5.0-cp314-cp314t-manylinux_2_31_riscv64.whl", hash = "sha256:16db7c62ec0b8efe1d7afd283a388d8f74f2605d56032e5a37747d2de8dba027", size = 1250678, upload-time = "2026-05-17T17:47:43.702Z" }, + { url = "https://files.pythonhosted.org/packages/74/ba/66ab9555de6275677566f6574e5ef6c29cb185ea866f643bc06f8280a8ee/ast_serialize-0.5.0-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:baf5eb061eb5bccade4128ad42da33787d72f6013809cd1b590376ece8b3c937", size = 1301603, upload-time = "2026-05-17T17:47:46.256Z" }, + { url = "https://files.pythonhosted.org/packages/66/42/6aca9b9abc710014b2be9059689e5dd1679339e78f567ffb4d255a9e2050/ast_serialize-0.5.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:104e4a35bd7c124173c41760ef9aaea17ddb3f86c65cb643671d59afbe3ee94c", size = 1410332, upload-time = "2026-05-17T17:47:47.899Z" }, + { url = "https://files.pythonhosted.org/packages/47/68/2f76594432a22581ecf878b5e75a9b8601c24b2241cf0bbeb1e21fcf370c/ast_serialize-0.5.0-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:36be371028fc1675acb38a331bde160dbab7ff907fdf00b67eb6911aa106951b", size = 1509979, upload-time = "2026-05-17T17:47:50.942Z" }, + { url = "https://files.pythonhosted.org/packages/40/ac/a93c9b58292653f6c595752f677a08e608f903b710594909e9231a389b3b/ast_serialize-0.5.0-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:061ee58bdb52341c8201a6df41182a977736bae3b7ded87ca7176ca25a8a47ab", size = 1505002, upload-time = "2026-05-17T17:47:54.093Z" }, + { url = "https://files.pythonhosted.org/packages/14/2e/b278f68c497ee2f1d1576cbbef8db5281cd4a5f2db040537592ac9c8862e/ast_serialize-0.5.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:b15219e9cdc9f53f6f4cb51c009203507228226148c05c5e8fe451c28b435eb3", size = 1456231, upload-time = "2026-05-17T17:47:56.311Z" }, + { url = "https://files.pythonhosted.org/packages/0b/43/419be1c566a4c504cd8fd60ce2f84e790f295495c0f327cfaeadf3d51012/ast_serialize-0.5.0-cp314-cp314t-win32.whl", hash = "sha256:842d1c004bb466c7df036f95fabef789570541922b10976b12f5592a69cf0b38", size = 1058668, upload-time = "2026-05-17T17:47:58.305Z" }, + { url = "https://files.pythonhosted.org/packages/03/6f/c9d4d549295ed05111aeb8853232d1afd9d0a179fddb01eeffbb3a4a6842/ast_serialize-0.5.0-cp314-cp314t-win_amd64.whl", hash = "sha256:b0c06d760909b095cc466356dfccd05a1c7233a6ca191c020dca2c6a6f16c24c", size = 1101075, upload-time = "2026-05-17T17:48:00.35Z" }, + { url = "https://files.pythonhosted.org/packages/d0/8e/d00c5ab30c58222e07d62956fca86c59d91b9ad32997e633c38b526623a3/ast_serialize-0.5.0-cp314-cp314t-win_arm64.whl", hash = "sha256:787baedb0262cc49e8ce37cc15c00ae818e46a165a3b36f5e21ed174998104cb", size = 1075347, upload-time = "2026-05-17T17:48:01.753Z" }, + { url = "https://files.pythonhosted.org/packages/e0/9e/dc2530acb3a60dc6e46d65abf27d1d9f86721694757906a148d90a6860de/ast_serialize-0.5.0-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:0668aa9459cfa8c9c49ddd2163ebcf43088ba045ef7492af6fe22e0098303101", size = 1191380, upload-time = "2026-05-17T17:48:03.738Z" }, + { url = "https://files.pythonhosted.org/packages/26/0a/bd3d18a582f273d6c843d16bb9e22e9e16365ff7991e92f18f798e9f1224/ast_serialize-0.5.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:bf683d6363edf2b39eed6b6d4fe22d34b6203867a67e27134d9e2a2680c4bc4a", size = 1183879, upload-time = "2026-05-17T17:48:05.463Z" }, + { url = "https://files.pythonhosted.org/packages/40/ae/1f919100f8620887af58fcc381c61a1f218cdf89c6e155f87b213e61010a/ast_serialize-0.5.0-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9cc22cf0c9be65e71cf88fda130af60d61eb4a79370ad4cfe7900d48a4aa2211", size = 1244529, upload-time = "2026-05-17T17:48:07.008Z" }, + { url = "https://files.pythonhosted.org/packages/c6/ca/6376559dcce707cdbc1d0d9a13c8d3baaaa501e949ce0ebdc4230cd881aa/ast_serialize-0.5.0-cp39-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:f66173891548c9f2726bf27957b41cabce12fa679dc6da505ddbde4d4b3b31cf", size = 1240560, upload-time = "2026-05-17T17:48:08.46Z" }, + { url = "https://files.pythonhosted.org/packages/35/b2/a620e206b5aeb7efbf2710336df57d457cffbb3991076bbcc1147ef9abd4/ast_serialize-0.5.0-cp39-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:e42d729ef2be96a14efbad355093284739e3670ece3e534f82cc8832790911d9", size = 1451172, upload-time = "2026-05-17T17:48:09.922Z" }, + { url = "https://files.pythonhosted.org/packages/fa/e0/4ad5c04c24a40481b2935ce9a0ccdb6023dc8b667167d06ae530cc3512f2/ast_serialize-0.5.0-cp39-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b725026bafa801dbd7310eb13a75f0a2e370e7e51b2cb225f9d21fcfadf919ee", size = 1265072, upload-time = "2026-05-17T17:48:11.469Z" }, + { url = "https://files.pythonhosted.org/packages/b2/71/4d1d479aa56d0101c40e17720c3d6ac2af7269ea0487a80b18e7bfd1a5b7/ast_serialize-0.5.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b54f60c1d78767a53b67eaa663f0dfac3afe606aa07f1301572f588b73d64809", size = 1270488, upload-time = "2026-05-17T17:48:13.575Z" }, + { url = "https://files.pythonhosted.org/packages/6d/4f/0de1bbe06f6edef9fde4ed12ca8e7b3ec7e6e2bd4e672c5af487f7957665/ast_serialize-0.5.0-cp39-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:27d51654fc240a1e87e742d353d98eb45b75f62f129086b3596ab53df2ac2a43", size = 1260702, upload-time = "2026-05-17T17:48:15.141Z" }, + { url = "https://files.pythonhosted.org/packages/75/61/e00872439cfdddcc3c1b6cdaa6e5d904ba8e26a18807c67c4e14409d0ca8/ast_serialize-0.5.0-cp39-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:2782c36237c46dd1674542f2109740ea5ea485a169bf1431939ada0434e17934", size = 1311182, upload-time = "2026-05-17T17:48:16.779Z" }, + { url = "https://files.pythonhosted.org/packages/76/8e/699a5b955f7926956c95e9e1d74132acad73c2fe7a426f94da89123c20aa/ast_serialize-0.5.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:1943db345233cc7194a470f13afa9c59772c0b123dea0c9414c4d4ca54369759", size = 1421410, upload-time = "2026-05-17T17:48:18.527Z" }, + { url = "https://files.pythonhosted.org/packages/a9/ae/d5b7626874478997adc7a29ab28accf21e596fb590c944290401dfd0b29e/ast_serialize-0.5.0-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:df1c00022cbbcb064bfaa505aa9c9295362443ce5dacb459d1331d3da353f887", size = 1516587, upload-time = "2026-05-17T17:48:20.133Z" }, + { url = "https://files.pythonhosted.org/packages/0c/ce/b59e02a82d9c4244d64cde502e0b00e83e38816abe19155ceb5437402c7f/ast_serialize-0.5.0-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:cae65289fc456fde04af979a2be09302ef5d8ab92ef23e596d6746dc267ada27", size = 1515171, upload-time = "2026-05-17T17:48:21.921Z" }, + { url = "https://files.pythonhosted.org/packages/8b/38/d8d90042747d05aa08d4efcf1c99035a5f670a6bf4c214d31644392afbca/ast_serialize-0.5.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:239a4c354e8d676e9d94631d1d4a64edc6b266f86ff3a5a80aedd344f342c01d", size = 1464668, upload-time = "2026-05-17T17:48:23.544Z" }, + { url = "https://files.pythonhosted.org/packages/dd/51/5b840c4df7334104cecffa28f23904fe81ca89ca223d2450e288de39fd3c/ast_serialize-0.5.0-cp39-abi3-win32.whl", hash = "sha256:143a4ef63285a075871908fda3672dc21864b83a8ec3ee12304aa3e4c5387b9a", size = 1068311, upload-time = "2026-05-17T17:48:25.027Z" }, + { url = "https://files.pythonhosted.org/packages/41/11/ca5672c7d491825bc4cd6702dea106a6b60d928707712ec257c7833ae476/ast_serialize-0.5.0-cp39-abi3-win_amd64.whl", hash = "sha256:cf25572c526add400f26a4750dc6ce0c3bb93fc1f75e7ae0cad4ce4f2cd5c590", size = 1108931, upload-time = "2026-05-17T17:48:26.591Z" }, + { url = "https://files.pythonhosted.org/packages/45/19/cc8bd127d28a43da249aa955cfd164cf8fd534e79e42cea96c4854d72fd0/ast_serialize-0.5.0-cp39-abi3-win_arm64.whl", hash = "sha256:92a31c9c20d25a076edaeec76b128a3535d74a24f340b9a8a7e96c9b86dc9642", size = 1081181, upload-time = "2026-05-17T17:48:28.122Z" }, ] [[package]] @@ -9856,7 +9858,7 @@ wheels = [ [[package]] name = "black" -version = "26.3.1" +version = "26.5.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, @@ -9868,34 +9870,34 @@ dependencies = [ { name = "tomli", marker = "python_full_version < '3.11'" }, { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e1/c5/61175d618685d42b005847464b8fb4743a67b1b8fdb75e50e5a96c31a27a/black-26.3.1.tar.gz", hash = "sha256:2c50f5063a9641c7eed7795014ba37b0f5fa227f3d408b968936e24bc0566b07", size = 666155, upload-time = "2026-03-12T03:36:03.593Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/32/a8/11170031095655d36ebc6664fe0897866f6023892396900eec0e8fdc4299/black-26.3.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:86a8b5035fce64f5dcd1b794cf8ec4d31fe458cf6ce3986a30deb434df82a1d2", size = 1866562, upload-time = "2026-03-12T03:39:58.639Z" }, - { url = "https://files.pythonhosted.org/packages/69/ce/9e7548d719c3248c6c2abfd555d11169457cbd584d98d179111338423790/black-26.3.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:5602bdb96d52d2d0672f24f6ffe5218795736dd34807fd0fd55ccd6bf206168b", size = 1703623, upload-time = "2026-03-12T03:40:00.347Z" }, - { url = "https://files.pythonhosted.org/packages/7f/0a/8d17d1a9c06f88d3d030d0b1d4373c1551146e252afe4547ed601c0e697f/black-26.3.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6c54a4a82e291a1fee5137371ab488866b7c86a3305af4026bdd4dc78642e1ac", size = 1768388, upload-time = "2026-03-12T03:40:01.765Z" }, - { url = "https://files.pythonhosted.org/packages/52/79/c1ee726e221c863cde5164f925bacf183dfdf0397d4e3f94889439b947b4/black-26.3.1-cp310-cp310-win_amd64.whl", hash = "sha256:6e131579c243c98f35bce64a7e08e87fb2d610544754675d4a0e73a070a5aa3a", size = 1412969, upload-time = "2026-03-12T03:40:03.252Z" }, - { url = "https://files.pythonhosted.org/packages/73/a5/15c01d613f5756f68ed8f6d4ec0a1e24b82b18889fa71affd3d1f7fad058/black-26.3.1-cp310-cp310-win_arm64.whl", hash = "sha256:5ed0ca58586c8d9a487352a96b15272b7fa55d139fc8496b519e78023a8dab0a", size = 1220345, upload-time = "2026-03-12T03:40:04.892Z" }, - { url = "https://files.pythonhosted.org/packages/17/57/5f11c92861f9c92eb9dddf515530bc2d06db843e44bdcf1c83c1427824bc/black-26.3.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:28ef38aee69e4b12fda8dba75e21f9b4f979b490c8ac0baa7cb505369ac9e1ff", size = 1851987, upload-time = "2026-03-12T03:40:06.248Z" }, - { url = "https://files.pythonhosted.org/packages/54/aa/340a1463660bf6831f9e39646bf774086dbd8ca7fc3cded9d59bbdf4ad0a/black-26.3.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:bf9bf162ed91a26f1adba8efda0b573bc6924ec1408a52cc6f82cb73ec2b142c", size = 1689499, upload-time = "2026-03-12T03:40:07.642Z" }, - { url = "https://files.pythonhosted.org/packages/f3/01/b726c93d717d72733da031d2de10b92c9fa4c8d0c67e8a8a372076579279/black-26.3.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:474c27574d6d7037c1bc875a81d9be0a9a4f9ee95e62800dab3cfaadbf75acd5", size = 1754369, upload-time = "2026-03-12T03:40:09.279Z" }, - { url = "https://files.pythonhosted.org/packages/e3/09/61e91881ca291f150cfc9eb7ba19473c2e59df28859a11a88248b5cbbc4d/black-26.3.1-cp311-cp311-win_amd64.whl", hash = "sha256:5e9d0d86df21f2e1677cc4bd090cd0e446278bcbbe49bf3659c308c3e402843e", size = 1413613, upload-time = "2026-03-12T03:40:10.943Z" }, - { url = "https://files.pythonhosted.org/packages/16/73/544f23891b22e7efe4d8f812371ab85b57f6a01b2fc45e3ba2e52ba985b8/black-26.3.1-cp311-cp311-win_arm64.whl", hash = "sha256:9a5e9f45e5d5e1c5b5c29b3bd4265dcc90e8b92cf4534520896ed77f791f4da5", size = 1219719, upload-time = "2026-03-12T03:40:12.597Z" }, - { url = "https://files.pythonhosted.org/packages/dc/f8/da5eae4fc75e78e6dceb60624e1b9662ab00d6b452996046dfa9b8a6025b/black-26.3.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b5e6f89631eb88a7302d416594a32faeee9fb8fb848290da9d0a5f2903519fc1", size = 1895920, upload-time = "2026-03-12T03:40:13.921Z" }, - { url = "https://files.pythonhosted.org/packages/2c/9f/04e6f26534da2e1629b2b48255c264cabf5eedc5141d04516d9d68a24111/black-26.3.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:41cd2012d35b47d589cb8a16faf8a32ef7a336f56356babd9fcf70939ad1897f", size = 1718499, upload-time = "2026-03-12T03:40:15.239Z" }, - { url = "https://files.pythonhosted.org/packages/04/91/a5935b2a63e31b331060c4a9fdb5a6c725840858c599032a6f3aac94055f/black-26.3.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0f76ff19ec5297dd8e66eb64deda23631e642c9393ab592826fd4bdc97a4bce7", size = 1794994, upload-time = "2026-03-12T03:40:17.124Z" }, - { url = "https://files.pythonhosted.org/packages/e7/0a/86e462cdd311a3c2a8ece708d22aba17d0b2a0d5348ca34b40cdcbea512e/black-26.3.1-cp312-cp312-win_amd64.whl", hash = "sha256:ddb113db38838eb9f043623ba274cfaf7d51d5b0c22ecb30afe58b1bb8322983", size = 1420867, upload-time = "2026-03-12T03:40:18.83Z" }, - { url = "https://files.pythonhosted.org/packages/5b/e5/22515a19cb7eaee3440325a6b0d95d2c0e88dd180cb011b12ae488e031d1/black-26.3.1-cp312-cp312-win_arm64.whl", hash = "sha256:dfdd51fc3e64ea4f35873d1b3fb25326773d55d2329ff8449139ebaad7357efb", size = 1230124, upload-time = "2026-03-12T03:40:20.425Z" }, - { url = "https://files.pythonhosted.org/packages/f5/77/5728052a3c0450c53d9bb3945c4c46b91baa62b2cafab6801411b6271e45/black-26.3.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:855822d90f884905362f602880ed8b5df1b7e3ee7d0db2502d4388a954cc8c54", size = 1895034, upload-time = "2026-03-12T03:40:21.813Z" }, - { url = "https://files.pythonhosted.org/packages/52/73/7cae55fdfdfbe9d19e9a8d25d145018965fe2079fa908101c3733b0c55a0/black-26.3.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:8a33d657f3276328ce00e4d37fe70361e1ec7614da5d7b6e78de5426cb56332f", size = 1718503, upload-time = "2026-03-12T03:40:23.666Z" }, - { url = "https://files.pythonhosted.org/packages/e1/87/af89ad449e8254fdbc74654e6467e3c9381b61472cc532ee350d28cfdafb/black-26.3.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f1cd08e99d2f9317292a311dfe578fd2a24b15dbce97792f9c4d752275c1fa56", size = 1793557, upload-time = "2026-03-12T03:40:25.497Z" }, - { url = "https://files.pythonhosted.org/packages/43/10/d6c06a791d8124b843bf325ab4ac7d2f5b98731dff84d6064eafd687ded1/black-26.3.1-cp313-cp313-win_amd64.whl", hash = "sha256:c7e72339f841b5a237ff14f7d3880ddd0fc7f98a1199e8c4327f9a4f478c1839", size = 1422766, upload-time = "2026-03-12T03:40:27.14Z" }, - { url = "https://files.pythonhosted.org/packages/59/4f/40a582c015f2d841ac24fed6390bd68f0fc896069ff3a886317959c9daf8/black-26.3.1-cp313-cp313-win_arm64.whl", hash = "sha256:afc622538b430aa4c8c853f7f63bc582b3b8030fd8c80b70fb5fa5b834e575c2", size = 1232140, upload-time = "2026-03-12T03:40:28.882Z" }, - { url = "https://files.pythonhosted.org/packages/d5/da/e36e27c9cebc1311b7579210df6f1c86e50f2d7143ae4fcf8a5017dc8809/black-26.3.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:2d6bfaf7fd0993b420bed691f20f9492d53ce9a2bcccea4b797d34e947318a78", size = 1889234, upload-time = "2026-03-12T03:40:30.964Z" }, - { url = "https://files.pythonhosted.org/packages/0e/7b/9871acf393f64a5fa33668c19350ca87177b181f44bb3d0c33b2d534f22c/black-26.3.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:f89f2ab047c76a9c03f78d0d66ca519e389519902fa27e7a91117ef7611c0568", size = 1720522, upload-time = "2026-03-12T03:40:32.346Z" }, - { url = "https://files.pythonhosted.org/packages/03/87/e766c7f2e90c07fb7586cc787c9ae6462b1eedab390191f2b7fc7f6170a9/black-26.3.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b07fc0dab849d24a80a29cfab8d8a19187d1c4685d8a5e6385a5ce323c1f015f", size = 1787824, upload-time = "2026-03-12T03:40:33.636Z" }, - { url = "https://files.pythonhosted.org/packages/ac/94/2424338fb2d1875e9e83eed4c8e9c67f6905ec25afd826a911aea2b02535/black-26.3.1-cp314-cp314-win_amd64.whl", hash = "sha256:0126ae5b7c09957da2bdbd91a9ba1207453feada9e9fe51992848658c6c8e01c", size = 1445855, upload-time = "2026-03-12T03:40:35.442Z" }, - { url = "https://files.pythonhosted.org/packages/86/43/0c3338bd928afb8ee7471f1a4eec3bdbe2245ccb4a646092a222e8669840/black-26.3.1-cp314-cp314-win_arm64.whl", hash = "sha256:92c0ec1f2cc149551a2b7b47efc32c866406b6891b0ee4625e95967c8f4acfb1", size = 1258109, upload-time = "2026-03-12T03:40:36.832Z" }, - { url = "https://files.pythonhosted.org/packages/8e/0d/52d98722666d6fc6c3dd4c76df339501d6efd40e0ff95e6186a7b7f0befd/black-26.3.1-py3-none-any.whl", hash = "sha256:2bd5aa94fc267d38bb21a70d7410a89f1a1d318841855f698746f8e7f51acd1b", size = 207542, upload-time = "2026-03-12T03:36:01.668Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/22/58/0a9d9b1195c159d206000c541c3e05897e339be754f0e4d8b29445ab536e/black-26.5.0.tar.gz", hash = "sha256:5cbe4cc4037ffca34cdb0a6a9a046f104b262d0bd63c30fd4a88c7adc2049b1d", size = 677762, upload-time = "2026-05-16T17:57:12.54Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8b/2a/f0bdee0d043b9e860fc1ae35596aa6d663d334b195d87019532afe97f29f/black-26.5.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:402454bfdd7a940be00455e87309438a24b328b7ba7d80b7207e8a87b32ffc29", size = 1983871, upload-time = "2026-05-16T18:00:43.834Z" }, + { url = "https://files.pythonhosted.org/packages/fe/a9/3dbf82806bc3b884ccb116a0f3b34f94ee2e0e6d5477d7abd215b1704907/black-26.5.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:4863b2a2c382661a018bf2213f2b957fa34511df131259ffaa8d54859620ac31", size = 1806039, upload-time = "2026-05-16T18:00:45.681Z" }, + { url = "https://files.pythonhosted.org/packages/57/10/150f2b66a08f840b89824dc5750363ee834e73e6b1b31050cfe4e76e13f3/black-26.5.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:490b623006a75c0ea59c1ecf91cc76ecb9d66df1482c3a53f4f7de95a7c85e10", size = 1856443, upload-time = "2026-05-16T18:00:47.89Z" }, + { url = "https://files.pythonhosted.org/packages/c4/71/d1f562c52c7a55060783e82b07b47c7eb09384f3f2759f868028a8a8aba7/black-26.5.0-cp310-cp310-win_amd64.whl", hash = "sha256:6f53deb3d1108a523212da5c79e5c0cd76abcc548948f2d8415e62929c81a569", size = 1474602, upload-time = "2026-05-16T18:00:49.678Z" }, + { url = "https://files.pythonhosted.org/packages/1e/35/a0e0a1e57bd72099fc72b52e96fbfdc52af273254526e6783bcf136ae207/black-26.5.0-cp310-cp310-win_arm64.whl", hash = "sha256:a62f9d069ac27de20c6fa3dbf60d7c951141c4025bb9755274802d05b1aa418b", size = 1273042, upload-time = "2026-05-16T18:00:51.949Z" }, + { url = "https://files.pythonhosted.org/packages/6b/71/17d04d49a406640f531f6d12e0f15858e0d337b7dbd4a5a05476cd04b229/black-26.5.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:862945b2a08193cdff9f632f51bdadbb11e6852da1d31c306a3508449dc81b84", size = 1965325, upload-time = "2026-05-16T18:00:53.755Z" }, + { url = "https://files.pythonhosted.org/packages/1a/a6/0739015dbd9df669529657bf6bef1185679a0eb8ba93bb6e160561f57652/black-26.5.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:03102aa97c279e5f62e1e1ab828cfe8aa72c3af4cf86f9448e5537b2519cbfea", size = 1786840, upload-time = "2026-05-16T18:00:55.55Z" }, + { url = "https://files.pythonhosted.org/packages/09/23/6cd101b4bc2234708120450d8ac54f6580d6ae52f6dce1098e040e6f259c/black-26.5.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:990ee0e1d96dd8ca623f19dd3f339c138bdc02f74e4fea01cc64aee38944ea2b", size = 1840560, upload-time = "2026-05-16T18:00:57.103Z" }, + { url = "https://files.pythonhosted.org/packages/73/8a/ded16f0183e370d44a4042a731f61669ad5e171f6d3ae98f8bb52182f917/black-26.5.0-cp311-cp311-win_amd64.whl", hash = "sha256:209fabb250681900502b3b6a03e31d8cac606c9ef9629fd0fbd5d33235647c00", size = 1475629, upload-time = "2026-05-16T18:00:59.209Z" }, + { url = "https://files.pythonhosted.org/packages/95/80/9191f47b6a7e7e752e55b6b01122594135f12ccad60aad27d4c206a38ad6/black-26.5.0-cp311-cp311-win_arm64.whl", hash = "sha256:dbb6fc70f8bd9821981fd47efb68a5be0eee9055f400eb3bf2dbebf49f9ec4fe", size = 1274370, upload-time = "2026-05-16T18:01:01.711Z" }, + { url = "https://files.pythonhosted.org/packages/22/89/feb65d2b11f8ccf60307b589e091e928011bde37751a451012e246a2e3dd/black-26.5.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b92983a6674c133ca61d6b4fea17f76cbbaac582ea583002792ee1094dbece49", size = 2007091, upload-time = "2026-05-16T18:01:03.624Z" }, + { url = "https://files.pythonhosted.org/packages/07/13/3684a1ba34c06ba9d5cf63ecdc3cd3635cdf347b7a9fbc67e0c31724f047/black-26.5.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1f80998e73fcfc67fc1d222060cf34ab213f1ae7e131b5c8199d93405890c13a", size = 1811228, upload-time = "2026-05-16T18:01:05.458Z" }, + { url = "https://files.pythonhosted.org/packages/7e/ea/6aa8f74867d1f7bc5d182ccd51ceaff9f48eb121d0b91c11030e554cca91/black-26.5.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:081df4dc908702e2becd66d714f125a954cbf1c6dbe2ad83a6be313368c7c2db", size = 1880889, upload-time = "2026-05-16T18:01:07.327Z" }, + { url = "https://files.pythonhosted.org/packages/d6/91/22e1222946dc566a05e62d2d0880ac3228ca07272eb3d4c490a48c788a56/black-26.5.0-cp312-cp312-win_amd64.whl", hash = "sha256:cf015b38829ca32a699312fdcfb8c15bd0b156192f5400bd0b559c6bfef25236", size = 1483664, upload-time = "2026-05-16T18:01:08.875Z" }, + { url = "https://files.pythonhosted.org/packages/a4/56/b238209a41209e1c9c7e05dfbc63e656516a5db31acb3248890e538a3e79/black-26.5.0-cp312-cp312-win_arm64.whl", hash = "sha256:828db2292848cf427592fcd162f02d770849d20ea4bdda2806e9494b3a15d481", size = 1285804, upload-time = "2026-05-16T18:01:10.812Z" }, + { url = "https://files.pythonhosted.org/packages/b2/0e/328992a8ce73c93605e7fe7325bcf38d3f1bc9b0118b514873699a5ed379/black-26.5.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c2b64ce9841e8b8254c3d702ebccdaf5c520607df8aa4176f5732b7f9af1e6f6", size = 2003830, upload-time = "2026-05-16T18:01:12.853Z" }, + { url = "https://files.pythonhosted.org/packages/15/07/0ded3f1c10306c0d4c5b112ec7c75bd323a199b96d9a0c61f4116ab985e8/black-26.5.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:0a789a41b386f0f83711785f182f2977138ba9cc1f41ad0f6fbc8faac4d2639e", size = 1810249, upload-time = "2026-05-16T18:01:14.367Z" }, + { url = "https://files.pythonhosted.org/packages/65/71/b5cf00e7d8e5b168bfc389e3b937b8d1250cfdda0c6c607f91dba0d5c2a7/black-26.5.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f69837f7e26d67b1d1e9d0ed49231a14a0469f266e44cd142873e0552f325395", size = 1879117, upload-time = "2026-05-16T18:01:16.095Z" }, + { url = "https://files.pythonhosted.org/packages/b5/0e/01baec29dd65ecca6be69d721b90dfff473b0e49fb49bb1b5b3fa470ab9d/black-26.5.0-cp313-cp313-win_amd64.whl", hash = "sha256:c5b08371561dae9c90391fe7f2138fe7fa495437d3bb134eb865839036e65784", size = 1486102, upload-time = "2026-05-16T18:01:17.78Z" }, + { url = "https://files.pythonhosted.org/packages/36/4b/6f9623c8cd5a3c6883318800e2073761fd9db1e859f594ee42e95c18fcd6/black-26.5.0-cp313-cp313-win_arm64.whl", hash = "sha256:3968ce82ca0bd4914769518490d91a9b0ef2ff2fc68e2122d22b5915a0342eaa", size = 1286888, upload-time = "2026-05-16T18:01:19.275Z" }, + { url = "https://files.pythonhosted.org/packages/75/d1/40d151b65b659848001ec8b8226323a6f25ee535a2f9d441392e1d86933b/black-26.5.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:ea8a0c4505486c132c6640e4e108d25f41360a06d844db5a76477c3dbae1b616", size = 1998941, upload-time = "2026-05-16T18:01:20.788Z" }, + { url = "https://files.pythonhosted.org/packages/bd/d1/991d741faf172502f17966ad8abb7e5b6ce06560855938000564dcf8e1f1/black-26.5.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:2178a70e7c45fb85999b687d8326abceef1e7227463d5d7e07ef125c9fbb9c5c", size = 1810853, upload-time = "2026-05-16T18:01:22.369Z" }, + { url = "https://files.pythonhosted.org/packages/aa/6c/6bb8ab3fa60074d5295162493482b4ed01c33dd19acf1754497fd506caed/black-26.5.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f3ad14d7c24c40eafecf4fb212d9c01e7c7b2ab05c8646b351c93728f499c555", size = 1874114, upload-time = "2026-05-16T18:01:23.973Z" }, + { url = "https://files.pythonhosted.org/packages/b9/3b/d9dc4206bbd9313d5c3761bd88e9bece5c85e909e5870c46bb7f835ecbcb/black-26.5.0-cp314-cp314-win_amd64.whl", hash = "sha256:8ea767bae9c4f331ea9ad2e08895c951e600dffd550a42624d5210a908720b39", size = 1508463, upload-time = "2026-05-16T18:01:25.878Z" }, + { url = "https://files.pythonhosted.org/packages/3c/65/2c5fc4152fc3bf79aa498bce429581b87aca340da2fde92423c0b6ce74bd/black-26.5.0-cp314-cp314-win_arm64.whl", hash = "sha256:d658f4ee6167797b08be07ee4bbf6045753ddabfc676c3cb0eec23752ca83eff", size = 1312669, upload-time = "2026-05-16T18:01:27.503Z" }, + { url = "https://files.pythonhosted.org/packages/14/c8/13da5c6a37b46a690199e0895c33a758ba4f2ec3cd81d1d72ebb373509a8/black-26.5.0-py3-none-any.whl", hash = "sha256:241f25bf59f5ca17f5121031e310e089b84cd22bb4eca47360099ea825544f17", size = 212907, upload-time = "2026-05-16T17:57:10.792Z" }, ] [[package]] @@ -10381,14 +10383,14 @@ wheels = [ [[package]] name = "click" -version = "8.3.3" +version = "8.4.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "colorama", marker = "sys_platform == 'win32'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/bb/63/f9e1ea081ce35720d8b92acde70daaedace594dc93b693c869e0d5910718/click-8.3.3.tar.gz", hash = "sha256:398329ad4837b2ff7cbe1dd166a4c0f8900c3ca3a218de04466f38f6497f18a2", size = 328061, upload-time = "2026-04-22T15:11:27.506Z" } +sdist = { url = "https://files.pythonhosted.org/packages/23/e4/796662cd90cf80e3a363c99db2b88e0e394b988a575f60a17e16440cd011/click-8.4.0.tar.gz", hash = "sha256:638f1338fe1235c8f4e008e4a8a254fb5c5fbdcbb40ece3c9142ebb78e792973", size = 350843, upload-time = "2026-05-17T00:47:58.425Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ae/44/c1221527f6a71a01ec6fbad7fa78f1d50dfa02217385cf0fa3eec7087d59/click-8.3.3-py3-none-any.whl", hash = "sha256:a2bf429bb3033c89fa4936ffb35d5cb471e3719e1f3c8a7c3fff0b8314305613", size = 110502, upload-time = "2026-04-22T15:11:25.044Z" }, + { url = "https://files.pythonhosted.org/packages/ee/ae/8e92f8058baf87f6c7d86ee7e457668690195cc77efedb8d3797a06e3940/click-8.4.0-py3-none-any.whl", hash = "sha256:40c50b7c6c6adac2823d411041ec84f3f103f1b280d5e9ce0d7f998995832f81", size = 116147, upload-time = "2026-05-17T00:47:56.842Z" }, ] [[package]] @@ -10979,11 +10981,11 @@ wheels = [ [[package]] name = "decorator" -version = "5.2.1" +version = "5.3.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/43/fa/6d96a0978d19e17b68d634497769987b16c8f4cd0a7a05048bec693caa6b/decorator-5.2.1.tar.gz", hash = "sha256:65f266143752f734b0a7cc83c46f4618af75b8c5911b00ccb61d0ac9b6da0360", size = 56711, upload-time = "2025-02-24T04:41:34.073Z" } +sdist = { url = "https://files.pythonhosted.org/packages/5c/50/a39dd7ab407e93978dfa07d109b7d633e37958c89f30cbcec061b77b3ebc/decorator-5.3.0.tar.gz", hash = "sha256:95fda3122972c847cf0ff7e0ce2829bf25136f2526b627b3da85b60ca5f485c0", size = 58431, upload-time = "2026-05-17T06:59:57.258Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/4e/8c/f3147f5c4b73e7550fe5f9352eaa956ae838d5c51eb58e7a25b9f3e2643b/decorator-5.2.1-py3-none-any.whl", hash = "sha256:d316bb415a2d9e2d2b3abcc4084c6502fc09240e292cd76a76afc106a1c8e04a", size = 9190, upload-time = "2025-02-24T04:41:32.565Z" }, + { url = "https://files.pythonhosted.org/packages/d5/6f/f8d0bba4dc2a69817d74f640d504650241ebf2f9f7263426f1b953b344d4/decorator-5.3.0-py3-none-any.whl", hash = "sha256:f8c2d71ede92f073144ddd7f3e9fbbc3bd0f2f29522c9d75ee648d66553834f4", size = 11104, upload-time = "2026-05-17T06:59:54.676Z" }, ] [[package]] @@ -11476,11 +11478,11 @@ wheels = [ [[package]] name = "fastcore" -version = "1.13.0" +version = "1.13.2" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/68/71/f96a5070cdc4b85ce3f9673f8e3121c081dee7517ef67986abbeefefb3f7/fastcore-1.13.0.tar.gz", hash = "sha256:933960662fec19ada98e704a8feb2b6d5e9a7828f27712dd57014ad55b4b3160", size = 100301, upload-time = "2026-05-13T02:11:46.662Z" } +sdist = { url = "https://files.pythonhosted.org/packages/70/65/99f599a285033febf95f9c608d91d629ac5d9995f57e5b3ac3397097f440/fastcore-1.13.2.tar.gz", hash = "sha256:f660b3448de48ba31973b2866c994ea3cd5e0a654847f57d6911a1a4bffda777", size = 100337, upload-time = "2026-05-17T06:02:24.383Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/96/9c/c86a018d1cbb3c117f7a83fcc86e74b648edaeacc849b8dbc71dd84171bf/fastcore-1.13.0-py3-none-any.whl", hash = "sha256:d56d5727a65835c1326299afc33a166cfa9977a6720b5e9580314c0eeb0ebbfe", size = 105027, upload-time = "2026-05-13T02:11:44.995Z" }, + { url = "https://files.pythonhosted.org/packages/c7/41/2c368f804bb9bd918da3b61324207fc4b410d0f32352c372c0680fc1f670/fastcore-1.13.2-py3-none-any.whl", hash = "sha256:2103c9e9e613311c0b36eab17299a221e778fd214ec526e8df1d32908928277c", size = 105060, upload-time = "2026-05-17T06:02:22.28Z" }, ] [[package]] @@ -17849,14 +17851,14 @@ wheels = [ [[package]] name = "pipdeptree" -version = "2.35.2" +version = "2.35.3" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "packaging" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/fd/10/1c7c74fd7e6072a5dc0f981d69dfcf7acf6bc7f12593b3a7c3d232f569d5/pipdeptree-2.35.2.tar.gz", hash = "sha256:5f338ca966f0596c089245324dd6b27031073746d345a6b2b7594450bea82c4a", size = 74605, upload-time = "2026-05-01T03:35:49.716Z" } +sdist = { url = "https://files.pythonhosted.org/packages/22/6e/17cf2810ff928751bb678cbff3b44cf02b6a4229c45130e48116f04625cd/pipdeptree-2.35.3.tar.gz", hash = "sha256:73238b3336698032abdabaa5508c404ce8c293ec7dcaa41e96c3d14734ce9f72", size = 78327, upload-time = "2026-05-17T15:54:44.311Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/03/be/d33b6a259d8e41af04d19300b4a289d6d08ad696cf5f5e9e64c01d0e18bf/pipdeptree-2.35.2-py3-none-any.whl", hash = "sha256:c8e67055c055cc0966751dc1275c93b5ae05eedee4207cdef543ff4c907061dc", size = 48974, upload-time = "2026-05-01T03:35:48.296Z" }, + { url = "https://files.pythonhosted.org/packages/56/1c/9857773ae300313349a1a44fcdbdd28d661a4feff68c7d61738cacc64a20/pipdeptree-2.35.3-py3-none-any.whl", hash = "sha256:51fe2bfc9fec359d97c47fc6ff6144d82668044d578bf1939989c608e50a4c69", size = 50879, upload-time = "2026-05-17T15:54:42.947Z" }, ] [[package]] @@ -18390,7 +18392,7 @@ wheels = [ [package.optional-dependencies] filetree = [ { name = "aiofile", version = "3.9.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "aiofile", version = "3.11.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "aiofile", version = "3.11.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "anyio" }, ] keyring = [ @@ -19713,11 +19715,11 @@ sdist = { url = "https://files.pythonhosted.org/packages/0c/88/8d2797decc42e1c1c [[package]] name = "python-multipart" -version = "0.0.28" +version = "0.0.29" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/82/54/a85eb421fbdd5007bc5af39d0f4ed9fa609e0fedbfdc2adcf0b34526870e/python_multipart-0.0.28.tar.gz", hash = "sha256:8550da197eac0f7ab748961fc9509b999fa2662ea25cef857f05249f6893c0f8", size = 45314, upload-time = "2026-05-10T11:05:16.596Z" } +sdist = { url = "https://files.pythonhosted.org/packages/4e/fe/70bd71a6738b09a0bdf6480ca6436b167469ca4578b2a0efbe390b4b0e70/python_multipart-0.0.29.tar.gz", hash = "sha256:643e93849196645e2dbdd81a0f8829a23123ad7f797a84a364c6fb3563f18904", size = 45678, upload-time = "2026-05-17T17:29:47.654Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f3/a2/43bbc5860b5034e2af4ef99a0e04d726ff329c43e192ef3abaa8d7ecfce5/python_multipart-0.0.28-py3-none-any.whl", hash = "sha256:10faac07eb966c3f48dc415f9dee46c04cb10d58d30a35677db8027c825ed9b6", size = 29438, upload-time = "2026-05-10T11:05:15.052Z" }, + { url = "https://files.pythonhosted.org/packages/8f/cb/769cfc37177252872a45a71f3fbdde9d51b471a3f3c14bfe95dde3407386/python_multipart-0.0.29-py3-none-any.whl", hash = "sha256:2ddcc971cef266225f54f552d8fa10bcfbb1f14446caec199060daac59ff2d69", size = 29640, upload-time = "2026-05-17T17:29:45.69Z" }, ] [[package]]