diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py new file mode 100644 index 0000000000000..69b6453bb1c8e --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py @@ -0,0 +1,95 @@ +# 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. +""" +Positional-argument binding spec for stub (foreign-runtime) tasks. + +Captured at parse time from the ``@task.stub`` TaskFlow call, stored in the serialized +Dag, and delivered to the lang-SDK runtime via ``TIRunContext.arg_bindings``. +""" + +from __future__ import annotations + +from functools import cache +from typing import Annotated, Literal + +from pydantic import Field, JsonValue, TypeAdapter +from typing_extensions import TypeAliasType + +from airflow.api_fastapi.core_api.base import BaseModel + +# A named, titled alias (like TaskArgBinding below) kept as free-form JSON rather than a +# typed model, so unknown JSON-schema keywords survive re-serialization along the way. +ArgValueSchema = TypeAliasType( + "ArgValueSchema", Annotated[dict[str, JsonValue], Field(title="ArgValueSchema")] +) +"""JSON-schema fragment constraining the value a stub-task argument binds to; generated +by pydantic from the stub annotation, carried verbatim, unknown keywords ignored.""" + + +class XComArgBinding(BaseModel): + """One positional stub-task argument pulled from an upstream task's XCom.""" + + # No default: it would drop ``kind`` from ``required``, and the generated task-sdk + # client then types it ``Literal | None``, invalid as a tagged-union discriminator. + kind: Literal["xcom"] + + name: str + """The stub function's parameter name this binding fills, in declaration order.""" + + value_schema: ArgValueSchema | None = None + """Schema fragment from the stub function's annotation; omitted when unconstrained.""" + + task_id: str + """Upstream task id whose ``return_value`` XCom is pulled.""" + + +class LiteralArgBinding(BaseModel): + """One positional stub-task argument carrying an inline literal from the Dag file.""" + + kind: Literal["literal"] + """No default, for the same generated-client reason as ``XComArgBinding.kind``.""" + + name: str + """The stub function's parameter name this binding fills, in declaration order.""" + + value_schema: ArgValueSchema | None = None + """Schema fragment from the stub function's annotation; omitted when unconstrained.""" + + value: JsonValue | None = None + """The literal value from the Dag file.""" + + from_default: bool = False + """True when the value was filled from the stub signature's default rather than passed in the call.""" + + +# A named alias with an explicit title so the union lands in every schema as its own +# named definition, which the supervisor-schema dump dedups with its task-sdk twin by title. +TaskArgBinding = TypeAliasType( + "TaskArgBinding", + Annotated[XComArgBinding | LiteralArgBinding, Field(discriminator="kind", title="TaskArgBinding")], +) +"""One positional argument of a stub (foreign-runtime) task, in declaration order.""" + + +@cache +def get_arg_bindings_adapter() -> TypeAdapter[list[TaskArgBinding]]: + """ + Build (lazily, then cache) the adapter validating serialized dicts into ``TaskArgBinding``. + + Only the stub-task path in the execution API needs it, so regular runs never pay for it. + """ + return TypeAdapter(list[TaskArgBinding]) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py index ad051b3e6d340..5e09e0ac06619 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py @@ -36,6 +36,7 @@ from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel from airflow.api_fastapi.execution_api.datamodels.asset import AssetProfile from airflow.api_fastapi.execution_api.datamodels.connection import ConnectionResponse +from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import TaskArgBinding from airflow.api_fastapi.execution_api.datamodels.variable import VariableResponse from airflow.utils.state import ( DagRunState, @@ -435,6 +436,13 @@ class TIRunContext(BaseModel): always reflects when the task *first* started, not when it was rescheduled/resumed. """ + arg_bindings: list[TaskArgBinding] | None = None + """ + Ordered positional-argument binding spec for stub (foreign-runtime) tasks. + + ``None`` for regular tasks and for stub tasks that declare no parameters. + """ + class PrevSuccessfulDagRunResponse(BaseModel): """Schema for response with previous successful DagRun information for Task Template Context.""" diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index 34c3dc35406f8..694a4f4729cf5 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -32,7 +32,7 @@ from opentelemetry import trace from opentelemetry.trace import StatusCode from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator -from pydantic import JsonValue +from pydantic import JsonValue, ValidationError from sqlalchemy import and_, func, or_, tuple_, update from sqlalchemy.engine import CursorResult from sqlalchemy.exc import DataError, NoResultFound, SQLAlchemyError @@ -49,6 +49,7 @@ from airflow.api_fastapi.common.types import UtcDateTime from airflow.api_fastapi.compat import HTTP_422_UNPROCESSABLE_CONTENT from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import get_arg_bindings_adapter from airflow.api_fastapi.execution_api.datamodels.taskinstance import ( InactiveAssetsResponse, PreviousTIResponse, @@ -75,6 +76,8 @@ get_team_name_for_ti, require_auth, ) +from airflow.api_fastapi.execution_api.services.task_instances import STUB_TASK_TYPE, get_arg_bindings +from airflow.api_fastapi.execution_api.versions import bundle from airflow.configuration import conf from airflow.exceptions import InvalidPartitionKeyError, TaskNotFound from airflow.models.asset import AssetActive @@ -110,6 +113,22 @@ log = structlog.get_logger(__name__) tracer = trace.get_tracer(__name__) +# The first execution API version whose TIRunContext carries ``arg_bindings``. +ARG_BINDINGS_API_VERSION = "2026-10-30" + + +def _client_supports_arg_bindings() -> bool: + """ + Whether the request's negotiated API version can receive ``arg_bindings``. + + Clients on older versions never see the field (the version migration strips it from + the response), so the derivation -- and the structured failures it raises for + undeliverable mapped-stub specs -- must not run for them: a stub Dag that ran before + arg bindings existed keeps running against those clients. + """ + version = bundle.api_version_var.get(None) + return version is None or str(version) >= ARG_BINDINGS_API_VERSION + @ti_id_router.patch( "/{task_instance_id}/run", @@ -163,6 +182,8 @@ def ti_run( TI.hostname, TI.unixname, TI.pid, + TI.operator, + TI.dag_version_id, # This selects the raw JSON value, bypassing the deserialization -- we want that to happen on the # client column("next_kwargs", JSON), @@ -310,6 +331,30 @@ def ti_run( should_retry=_is_eligible_to_retry(previous_state, ti.try_number, ti.max_tries), ) + # Only set for stub (foreign-runtime) tasks with a captured TaskFlow arg + # spec; the route excludes unset fields, keeping regular responses lean. + if ( + ti.operator == STUB_TASK_TYPE + and _client_supports_arg_bindings() + and (arg_bindings := get_arg_bindings(dag_bag, ti, session=session)) + ): + try: + context.arg_bindings = get_arg_bindings_adapter().validate_python(arg_bindings) + except ValidationError: + log.exception( + "Serialized arg_bindings spec failed validation", + dag_id=ti.dag_id, + task_id=ti.task_id, + dag_version_id=ti.dag_version_id, + ) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={ + "reason": "invalid_arg_bindings", + "message": "The serialized TaskFlow arg spec for this stub task is not valid.", + }, + ) + # Only set if they are non-null if ti.next_method: context.next_method = ti.next_method diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/services/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/services/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/services/__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/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py new file mode 100644 index 0000000000000..e06db8ffa70be --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py @@ -0,0 +1,49 @@ +# 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. +"""Business logic backing the task-instance execution routes.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from sqlalchemy.orm import Session + + from airflow.models.dagbag import DBDagBag + +# Task type recorded on the TI row (``TaskInstance.operator``) for +# ``airflow.providers.standard.decorators.stub._StubOperator``. Used to gate the +# serialized-Dag lookup for ``arg_bindings`` so regular tasks never pay for it. +# The gate matches the exact class name; a subclass would need its own entry here. +STUB_TASK_TYPE = "_StubOperator" + + +def get_arg_bindings(dag_bag: DBDagBag, ti: Any, *, session: Session) -> list | None: + """ + Extract the stub task's TaskFlow arg spec from its Dag version. + + Mapped (``.expand()``) stubs never capture a parse-time spec, so they resolve to + ``None`` here and keep the legacy ignored-args behavior; per-map-index delivery + lands in a follow-up. + """ + if ti.dag_version_id is None: + return None + if (dag := dag_bag.get_dag(ti.dag_version_id, session=session)) is None: + return None + if (task := dag.task_dict.get(ti.task_id)) is None: + return None + return getattr(task, "_arg_bindings", None) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py index dc7035d31e3c9..d56ec735c8f13 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py @@ -51,9 +51,11 @@ AddTeamNameField, AddVariableKeysEndpoint, ) +from airflow.api_fastapi.execution_api.versions.v2026_10_30 import AddArgBindingsToTIRunContext bundle = VersionBundle( HeadVersion(), + Version("2026-10-30", AddArgBindingsToTIRunContext), Version( "2026-06-30", AddVariableKeysEndpoint, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py new file mode 100644 index 0000000000000..2b456eae4da1f --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py @@ -0,0 +1,40 @@ +# 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 ( + ResponseInfo, + VersionChange, + convert_response_to_previous_version_for, + schema, +) + +from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIRunContext + + +class AddArgBindingsToTIRunContext(VersionChange): + """Add the ``arg_bindings`` argument-binding spec for stub (foreign-runtime) tasks.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = (schema(TIRunContext).field("arg_bindings").didnt_exist,) + + @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type] + def remove_arg_bindings_field(response: ResponseInfo) -> None: # type: ignore[misc] + """Strip ``arg_bindings`` from the run context for older clients.""" + response.body.pop("arg_bindings", None) diff --git a/airflow-core/src/airflow/serialization/schema.json b/airflow-core/src/airflow/serialization/schema.json index 872c3a1331ee3..fb954f0615b9d 100644 --- a/airflow-core/src/airflow/serialization/schema.json +++ b/airflow-core/src/airflow/serialization/schema.json @@ -142,6 +142,48 @@ "description": "A python dictionary containing values of any type", "type": "object" }, + "typed_dict": { + "type": "object", + "properties": { + "__type": { + "type": "string", + "const": "dict" + }, + "__var": { "$ref": "#/definitions/dict" } + }, + "required": [ + "__type", + "__var" + ], + "additionalProperties": false + }, + "arg_binding": { + "$comment": "One captured TaskFlow call argument of a @task.stub task, in dict-encoded form. The inner object stays open so future binding fields keep validating on older cores", + "type": "object", + "properties": { + "__type": { + "type": "string", + "const": "dict" + }, + "__var": { + "type": "object", + "properties": { + "name": { "type": "string" }, + "kind": { "type": "string" }, + "value_schema": { "$ref": "#/definitions/typed_dict" }, + "task_id": { "type": "string" }, + "value": {}, + "from_default": { "type": "boolean" } + }, + "required": [ "name", "kind" ] + } + }, + "required": [ + "__type", + "__var" + ], + "additionalProperties": false + }, "color": { "type": "string", "pattern": "^#[a-fA-F0-9]{3,6}$" @@ -345,7 +387,12 @@ "is_teardown": {"type": "boolean", "default": false}, "on_failure_fail_dagrun": {"type": "boolean", "default": false}, "max_active_tis_per_dag": {"type": "integer"}, - "max_active_tis_per_dagrun": {"type": "integer"} + "max_active_tis_per_dagrun": {"type": "integer"}, + "_arg_bindings": { + "$comment": "Only present on @task.stub tasks called with TaskFlow arguments", + "type": "array", + "items": { "$ref": "#/definitions/arg_binding" } + } }, "dependencies": { "expand_input": ["partial_kwargs", "_is_mapped"], diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index 8a152bebe0d3f..e5dfea5360521 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -32,6 +32,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import StatusCode from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator +from pydantic import ValidationError from sqlalchemy import select, update from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import Session @@ -371,6 +372,135 @@ async def workload_token(request: Request) -> TIToken: assert extras["scope"] == "execution" assert extras["sub"] == str(ti.id) + def test_ti_run_returns_arg_bindings_for_stub_task(self, client, dag_maker): + """A stub task's TaskFlow arg spec is extracted from the serialized Dag and returned.""" + with dag_maker("test_arg_bindings_dag", serialized=True): + + @task.stub + def extract(): ... + + @task.stub + def transform(country: str, extracted: dict, limit: int = 10): ... + + transform("uk", extract()) + + dr = dag_maker.create_dagrun() + tis = {ti.task_id: ti for ti in dr.get_task_instances()} + for ti in tis.values(): + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + payload = { + "state": "running", + "hostname": "random-hostname", + "unixname": "random-unixname", + "pid": 100, + "start_date": "2024-09-30T12:00:00Z", + } + + response = client.patch(f"/execution/task-instances/{tis['transform'].id}/run", json=payload) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "extract", + }, + { + "name": "limit", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 10, + "from_default": True, + }, + ] + + # An argless stub has no captured spec, so the field stays unset. + response = client.patch(f"/execution/task-instances/{tis['extract'].id}/run", json=payload) + assert response.status_code == 200 + assert "arg_bindings" not in response.json() + + @mock.patch( + "airflow.api_fastapi.execution_api.routes.task_instances.get_arg_bindings", + autospec=True, + return_value=[{"name": "country", "kind": "hologram", "value": "uk"}], + ) + def test_ti_run_reports_invalid_arg_bindings_spec(self, _, client, dag_maker): + """A serialized spec this core version cannot validate fails with a structured error, not a bare 500.""" + with dag_maker("test_invalid_arg_bindings_dag", serialized=True): + + @task.stub + def transform(country: str): ... + + transform("uk") + + dr = dag_maker.create_dagrun() + (ti,) = dr.get_task_instances() + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch( + f"/execution/task-instances/{ti.id}/run", + json={ + "state": "running", + "hostname": "random-hostname", + "unixname": "random-unixname", + "pid": 100, + "start_date": "2024-09-30T12:00:00Z", + }, + ) + + assert response.status_code == 500 + assert response.json()["detail"]["reason"] == "invalid_arg_bindings" + + RUN_PAYLOAD = { + "state": "running", + "hostname": "random-hostname", + "unixname": "random-unixname", + "pid": 100, + "start_date": "2024-09-30T12:00:00Z", + } + + def test_ti_run_returns_no_arg_bindings_for_mapped_stub(self, client, dag_maker): + """Mapped stubs keep the legacy ignored-args behavior until per-map-index delivery lands.""" + with dag_maker("test_mapped_stub_ignored_args", serialized=True): + + @task.stub + def transform(country: str): ... + + transform.expand(country=["uk", "fr"]) + + dr = dag_maker.create_dagrun() + ti = next(t for t in dr.get_task_instances() if t.map_index == 0) + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 200 + assert "arg_bindings" not in response.json() + + def test_arg_bindings_adapter_rejects_unknown_kind(self): + """The discriminated union refuses serialized specs with an unrecognised kind.""" + from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import get_arg_bindings_adapter + + with pytest.raises(ValidationError, match="does not match any of the expected tags"): + get_arg_bindings_adapter().validate_python( + [{"name": "country", "kind": "template", "value": "x"}] + ) + + def test_arg_bindings_adapter_carries_value_schema_fragments_verbatim(self): + """The fragment is free-form JSON schema: every keyword the provider generated must + survive validation untouched -- a typed model would silently strip what it doesn't know.""" + from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import get_arg_bindings_adapter + + fragment = {"anyOf": [{"type": "array", "items": {"type": "string"}}, {"type": "null"}]} + (binding,) = get_arg_bindings_adapter().validate_python( + [{"name": "tags", "kind": "literal", "value_schema": fragment, "value": ["a"]}] + ) + assert binding.value_schema == fragment + def test_dynamic_task_mapping_with_parse_time_value(self, client, dag_maker): """Test that dynamic task mapping works correctly with parse-time values.""" with dag_maker("test_dynamic_task_mapping_with_parse_time_value", serialized=True): diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/__init__.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/__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/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py new file mode 100644 index 0000000000000..a4b98bd10206e --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py @@ -0,0 +1,100 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import pytest + +from airflow.sdk import task +from airflow.utils.state import State + +from tests_common.test_utils.db import clear_db_runs + +pytestmark = pytest.mark.db_test + +TIMESTAMP_STR = "2024-09-30T12:00:00Z" + +RUN_PATCH_BODY = { + "state": "running", + "hostname": "h", + "unixname": "u", + "pid": 1, + "start_date": TIMESTAMP_STR, +} + + +@pytest.fixture +def old_ver_client(client): + """Execution API version immediately before ``arg_bindings`` was added.""" + client.headers["Airflow-API-Version"] = "2026-06-30" + return client + + +class TestArgBindingsFieldBackwardCompat: + @pytest.fixture(autouse=True) + def _freeze_time(self, time_machine): + time_machine.move_to(TIMESTAMP_STR, tick=False) + + def setup_method(self): + clear_db_runs() + + def teardown_method(self): + clear_db_runs() + + @pytest.fixture + def stub_ti(self, dag_maker): + with dag_maker("test_arg_bindings_compat_dag", serialized=True): + + @task.stub + def extract(): ... + + @task.stub + def transform(country: str, extracted: dict, limit: int = 10): ... + + transform("uk", extract()) + + dr = dag_maker.create_dagrun() + tis = {ti.task_id: ti for ti in dr.get_task_instances()} + for ti in tis.values(): + ti.set_state(State.QUEUED) + dag_maker.session.flush() + return tis["transform"] + + def test_old_version_strips_arg_bindings_even_when_set(self, old_ver_client, stub_ti): + response = old_ver_client.patch(f"/execution/task-instances/{stub_ti.id}/run", json=RUN_PATCH_BODY) + assert response.status_code == 200 + assert "arg_bindings" not in response.json() + + def test_head_version_includes_arg_bindings(self, client, stub_ti): + response = client.patch(f"/execution/task-instances/{stub_ti.id}/run", json=RUN_PATCH_BODY) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "extract", + }, + { + "name": "limit", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 10, + "from_default": True, + }, + ] diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index 4708365e846e2..889131a152032 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -77,7 +77,7 @@ from airflow.serialization.definitions.param import SerializedParam from airflow.serialization.definitions.xcom_arg import SchedulerPlainXComArg from airflow.serialization.encoders import ensure_serialized_asset -from airflow.serialization.enums import Encoding +from airflow.serialization.enums import DagAttributeTypes as DAT, Encoding from airflow.serialization.json_schema import load_dag_schema_dict from airflow.serialization.serialized_objects import ( BaseSerialization, @@ -3405,6 +3405,60 @@ def inner(): assert serialized3["python_callable_name"] == "empty_function" +def test_stub_task_args_round_trip(): + """The stub task's TaskFlow arg spec (``_arg_bindings``) survives Dag serialization.""" + from airflow.sdk import task + + with DAG(dag_id="arg_bindings_dag", schedule=None) as dag: + + @task.stub + def extract(): ... + + @task.stub + def transform(country: str, extracted: dict): ... + + transform("uk", extract()) + + ser_dag = DagSerialization.to_dict(dag) + encoded_tasks = {t[Encoding.VAR]["task_id"]: t[Encoding.VAR] for t in ser_dag["dag"]["tasks"]} + assert "_arg_bindings" not in encoded_tasks["extract"], "argless stubs must not serialize a spec" + assert encoded_tasks["transform"]["_arg_bindings"] == [ + { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: { + "name": "country", + "kind": "literal", + "value_schema": {Encoding.TYPE: DAT.DICT, Encoding.VAR: {"type": "string"}}, + "value": "uk", + }, + }, + { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: { + "name": "extracted", + "kind": "xcom", + "value_schema": { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: {"type": "object", "additionalProperties": True}, + }, + "task_id": "extract", + }, + }, + ] + + round_tripped = DagSerialization.from_dict(ser_dag) + assert round_tripped.task_dict["transform"]._arg_bindings == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "extract", + }, + ] + assert not hasattr(round_tripped.task_dict["extract"], "_arg_bindings") + + def test_handle_v1_serdag(): v1 = { "__version": 1, diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_taskflow_binding.py b/airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_taskflow_binding.py new file mode 100644 index 0000000000000..b4229a10dee3c --- /dev/null +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_taskflow_binding.py @@ -0,0 +1,162 @@ +# 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. +"""E2E test for the Go SDK ``taskflow_binding_dag`` example. + +The stub Dag's single mixed positional/keyword TaskFlow call carries literals +of every scalar type, an array literal, a defaulted ``None``, and XComs from +two upstream Go tasks (an object bound onto a strict Go struct and an array +bound onto ``[]int``). The Go ``via_flat_args`` task verifies every bound +value and errors on any mismatch, so a green run *is* the binding assertion; +the tests here check the run outcome and the summary XCom it pushes. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime, timezone + +import pytest + +from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient + +# Three short Go tasks; allow room for coordinator startup. +_GO_TASK_TIMEOUT = 300 + +_DAG_ID = "taskflow_binding_dag" + + +@dataclass +class _CompletedRun: + """The single ``taskflow_binding_dag`` run shared across this module's tests.""" + + client: AirflowClient + run_id: str + state: str + ti_states: dict[str, str] + + def xcom(self, task_id: str, key: str = "return_value"): + return self.client.get_xcom_value(dag_id=_DAG_ID, task_id=task_id, run_id=self.run_id, key=key).get( + "value" + ) + + +@pytest.fixture(scope="module") +def completed_run() -> _CompletedRun: + """Trigger ``taskflow_binding_dag`` once and wait for it to finish.""" + client = AirflowClient() + resp = client.trigger_dag(_DAG_ID, json={"logical_date": datetime.now(timezone.utc).isoformat()}) + run_id = resp["dag_run_id"] + state = client.wait_for_dag_run(dag_id=_DAG_ID, run_id=run_id, timeout=_GO_TASK_TIMEOUT) + ti_resp = client.get_task_instances(dag_id=_DAG_ID, run_id=run_id) + ti_states = {ti["task_id"]: ti.get("state") for ti in ti_resp.get("task_instances", [])} + return _CompletedRun(client=client, run_id=run_id, state=state, ti_states=ti_states) + + +def test_all_tasks_succeeded(completed_run: _CompletedRun): + """The Go ``via_flat_args`` task errors on any mis-bound argument, so success here + proves every literal, XCom, keyword, and defaulted-None binding was correct.""" + assert completed_run.state == "success", ( + f"expected the run to succeed; got {completed_run.state!r}. task states: {completed_run.ti_states}" + ) + for task_id in ( + "make_config", + "make_numbers", + "make_region", + "via_flat_args", + "via_struct_no_tags", + "via_struct_arg_tag", + "via_struct_unmatched_arg", + "via_flat_map", + "via_struct_map", + ): + assert completed_run.ti_states.get(task_id) == "success", completed_run.ti_states + + +def test_upstream_xcoms_keep_their_shapes(completed_run: _CompletedRun): + """The Go struct arrives as an object XCom, the ``[]int`` as an array, the region as a string.""" + assert completed_run.xcom("make_config") == { + "environment": "production", + "region": "eu-west-1", + "debug": True, + } + assert completed_run.xcom("make_numbers") == [1, 1, 2, 3, 5, 8] + assert completed_run.xcom("make_region") == "eu-west-1" + + +def test_via_flat_args_summary_reflects_bound_arguments(completed_run: _CompletedRun): + """``via_flat_args`` re-emits every bound value, confirming types survived the + Python literal / XCom -> Go parameter -> XCom round trip.""" + assert completed_run.xcom("via_flat_args") == { + "name": "summary", + "count": 3, + "ratio": 2.5, + "enabled": True, + "tags": ["metrics", "hourly"], + "environment": "production", + "debug": True, + "sum": 20, + "note_was_null": True, + } + + +def test_via_struct_no_tags_reflects_bound_arguments(completed_run: _CompletedRun): + """``via_struct_no_tags`` demonstrates the Go SDK's name-based struct binding + with no field tags at all: each field binds the TaskFlow argument spelled + exactly like its Go field name (``RegionCode``, ``Threshold``). The region is + ``make_region``'s XCom, so a struct field binds an XCom-sourced value here.""" + assert completed_run.xcom("via_struct_no_tags") == { + "region_code": "eu-west-1", + "threshold": 0.75, + } + + +def test_via_struct_arg_tag_reflects_bound_arguments(completed_run: _CompletedRun): + """``via_struct_arg_tag`` demonstrates explicit ``arg:`` tags: ``Region`` is + genuinely renamed to ``region_code`` (bound from ``make_region``'s XCom), and + ``Threshold`` is tagged ``threshold`` to pull the snake_case literal its + verbatim field name would miss.""" + assert completed_run.xcom("via_struct_arg_tag") == { + "region": "eu-west-1", + "threshold": 0.75, + } + + +def test_via_struct_unmatched_arg_reflects_zero_valued_field(completed_run: _CompletedRun): + """``via_struct_unmatched_arg`` demonstrates mismatch tolerance in both directions: + a struct field whose name has no corresponding TaskFlow call argument stays at its + Go zero value instead of failing the task (kwarg-style, an unpassed name simply + isn't bound), and the stub's defaulted ``sample_rate`` -- captured into the spec as + ``from_default`` -- needs no matching struct field. The task succeeding at all + proves the second half.""" + assert completed_run.xcom("via_struct_unmatched_arg") == { + "region": "eu-west-1", + "missing_was_empty": True, + } + + +def test_via_flat_map_decodes_single_dict_whole(completed_run: _CompletedRun): + """``via_flat_map`` passes one dict literal whose argument name matches no Go + struct field, so the whole map is decoded into the struct (flat binding).""" + assert completed_run.xcom("via_flat_map") == {"region": "eu-west-1", "count": 3} + + +def test_via_struct_map_binds_single_dict_onto_map_field(completed_run: _CompletedRun): + """``via_struct_map`` passes one dict literal whose argument name binds by name + onto a Go struct's ``map`` field (struct-based binding).""" + assert completed_run.xcom("via_struct_map") == { + "payload": {"region": "eu-west-1", "count": 3}, + } diff --git a/go-sdk/README.md b/go-sdk/README.md index a84373bd29183..fb4ba0ef15c23 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -105,6 +105,17 @@ A task is an ordinary Go function. The runtime inspects its signature and inject `sdk.VariableClient`). An optional `(any, error)` return becomes the task's XCom; an `error` return marks the task failed. +Any other parameter is a **data parameter**: in declaration order, data parameters receive the +positional arguments of the Python stub Dag's TaskFlow call. A JSON-serializable literal in the Dag +file (`transform("uk", ...)`) decodes straight into the parameter; an upstream task output +(`transform(..., extract())`) is pulled from that task's XCom in the current Dag run and decoded into +the parameter's type (independent XCom pulls run concurrently). The runtime fails the task loudly when +the argument count doesn't match the number of data parameters or a declared type can't bind to the Go +type. Data parameters must be JSON-decodable (no func/chan/unsafe-pointer, no non-empty interfaces) — +checked once at registration. TaskFlow argument binding arrives over the coordinator protocol, so it is +coordinator-mode only today; on the Edge Worker path a task with data parameters fails with the arity +error (and a name-bound struct with bindable fields fails the same way, since nothing can fill them). + ```go func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (any, error) { conn, err := client.GetConnection(ctx, "test_http") @@ -112,12 +123,17 @@ func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (any, er return map[string]any{"go_version": runtime.Version()}, nil } -func transform(ctx sdk.TIRunContext, client sdk.VariableClient, log *slog.Logger) error { +// The stub Dag calls transform("uk", extract()): "uk" binds onto country and +// extract's return-value XCom is pulled into extracted. +func transform( + ctx sdk.TIRunContext, client sdk.VariableClient, log *slog.Logger, + country string, extracted map[string]any, +) error { val, err := client.GetVariable(ctx, "my_variable") if err != nil { return err } - log.Info("Obtained variable", "my_variable", val) + log.Info("Obtained variable", "my_variable", val, "country", country) return nil } ``` @@ -126,6 +142,59 @@ Asking for the narrowest interface a task needs (e.g. `sdk.VariableClient` inste unit testing easier and documents which Airflow features the task touches. `RegisterDags` is the single source of truth for which `dag_id`s and `task_id`s a bundle can run. +### Name-based struct binding + +When a task function's **sole data parameter** is a struct, its fields bind **by name** +(keyword-argument style) instead of positionally — an ergonomic alternative to a long flat +parameter list. There is no marker to add: being the only data parameter is the opt-in. + +Conceptually, a plain flat parameter list is **positional-argument** binding: order matters, and +every parameter must be filled or the task fails before its body runs. A sole struct parameter is +closer to **keyword-argument** binding: fields match by name instead of position, and (see the +`arg:` bullet below) a field whose name has no corresponding TaskFlow call argument is simply left +at its Go zero value rather than failing the task — the same way an unpassed keyword argument falls +back to a caller-side default in a kwargs-style call. + +```go +type CombineInput struct { + Region string `arg:"region_code"` // named lookup against the TaskFlow call argument "region_code" + Threshold float64 `arg:"threshold"` // tags also bridge Go's UpperCamelCase to a snake_case argument +} + +func Combine(ctx sdk.TIRunContext, log *slog.Logger, input CombineInput) (any, error) { + // input.Region and input.Threshold are both populated. + return nil, nil +} +``` + +Each exported field binds from the TaskFlow call argument named by its optional `arg:""` tag +(matched against the stub function's Python parameter name, independent of declaration order on +either side). With no tag, the field's own Go name is matched verbatim — an untagged `Threshold` +only binds an argument literally spelled `Threshold`, so a snake_case Python parameter needs an +explicit tag. If no TaskFlow call argument carries that name, the field is simply left at its Go +zero value — it does not fail the task, kwarg-style (see `ViaStructUnmatchedArg` below). + +The matching is checked in the other direction too: every argument the Dag author **explicitly +passed** in the TaskFlow call must be claimed by some field, so a typo'd field name fails the task +instead of silently dropping the value. Stub parameters the author left at their Python defaults +are the exception — the Python side captures them into the spec (marked `from_default` on the +wire), and the struct is free not to mirror them, the same way a Python callee never sees which +defaulted kwargs went unpassed. And when no argument spec arrives at all (an argless stub call, or +the Edge Worker path) a struct with bindable fields fails loudly rather than running fully +zero-valued. + +Two related cases keep whole-value decoding available. A struct that is **not** the sole data +parameter — one among several flat parameters — is decoded whole from its one positional argument, +as `Config` is in +[`example/bundle/taskflowbinding/taskflowbinding.go`](./example/bundle/taskflowbinding/taskflowbinding.go). +And a sole struct parameter that instead receives exactly one explicitly passed argument no field +claims falls back to decoding that argument whole into the struct — so a task can still take an +upstream object as a single argument. Because `arg:` tags only take effect in the name-based path, a +struct carrying `arg:` tags must be the only data parameter; pairing it with other data parameters +is a registration error rather than a silent whole-value decode. See +[`ViaStructNoTags`, `ViaStructArgTag`, and `ViaStructUnmatchedArg`](./example/bundle/taskflowbinding/taskflowbinding.go) +for a full worked example of each field-binding mode — and the unmatched-field case — in isolation. + ### Reading the task runtime context Declare an `sdk.TIRunContext` parameter on a task to read the identifiers and scheduling timestamps of the diff --git a/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md b/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md index a78c7bb92e4e1..bb0c24fd29827 100644 --- a/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md +++ b/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md @@ -217,7 +217,11 @@ Supervisor Bundle binary (Go) │ │ ├── StartupDetails ────────────────────►│ │ (ti, dag_rel_path, bundle_info, │ - │ start_date, ti_context) │ + │ start_date, ti_context; the │ + │ ti_context carries arg_bindings, │ + │ the positional-argument spec │ + │ captured from the stub Dag's │ + │ TaskFlow call) │ │ │ │ ├── lookup task: │ │ bundle.dags[ti.dag_id] @@ -225,6 +229,12 @@ Supervisor Bundle binary (Go) │ │ (returns TaskState{state:"removed"} │ │ if not found, mirroring Java) │ │ + │ ├── bind arg_bindings onto the task + │ │ fn's data parameters (literals + │ │ decode directly; xcom refs pull + │ │ below); arity/type mismatch + │ │ fails the task + │ │ │ ├── construct sdk.Client whose │ │ GetConnection / GetVariable / │ │ GetXCom / SetXCom calls block on diff --git a/go-sdk/bundle/bundlev1/task.go b/go-sdk/bundle/bundlev1/task.go index d31fea84b73f3..cb18a4e951519 100644 --- a/go-sdk/bundle/bundlev1/task.go +++ b/go-sdk/bundle/bundlev1/task.go @@ -25,30 +25,58 @@ import ( "runtime" "github.com/apache/airflow/go-sdk/pkg/api" + "github.com/apache/airflow/go-sdk/pkg/binding" "github.com/apache/airflow/go-sdk/pkg/sdkcontext" "github.com/apache/airflow/go-sdk/sdk" ) +// TaskWithArgs is implemented by tasks that can bind positional arguments +// captured from the Dag's TaskFlow call (delivered per execution in +// coordinator mode). Execute(ctx, logger) is equivalent to +// ExecuteArgs(ctx, logger, nil). +type TaskWithArgs interface { + Task + ExecuteArgs(ctx context.Context, logger *slog.Logger, args []binding.Arg) error +} + type taskFunction struct { fn reflect.Value fullName string + plan *binding.Plan } -var _ Task = (*taskFunction)(nil) +var _ TaskWithArgs = (*taskFunction)(nil) // NewTaskFunction wraps a plain Go function as a Task, validating its signature -// (injectable parameters, and a return of error or (result, error)). Bundle -// authors normally use Dag.AddTask, which calls this for them; use it directly -// only when building a Task outside the registry. +// (injectable parameters, data parameters that task arguments can decode into, +// and a return of error or (result, error)). Bundle authors normally use +// Dag.AddTask, which calls this for them; use it directly only when building a +// Task outside the registry. func NewTaskFunction(fn any) (Task, error) { v := reflect.ValueOf(fn) fullName := runtime.FuncForPC(v.Pointer()).Name() - f := &taskFunction{v, fullName} - return f, f.validateFn(v.Type()) + f := &taskFunction{fn: v, fullName: fullName} + if err := f.validateFn(v.Type()); err != nil { + // A half-built task (nil binding plan) would panic in Execute; a caller + // that mishandles the error must not be able to run it. + return nil, err + } + return f, nil } func (f *taskFunction) Execute(ctx context.Context, logger *slog.Logger) error { - fnType := f.fn.Type() + return f.ExecuteArgs(ctx, logger, nil) +} + +// ExecuteArgs resolves the function's parameters — injectables from the +// context and args onto the data parameters — and invokes it. A resolution +// error (arity or type mismatch, xcom pull/decode failure) fails the task +// before its body runs. +func (f *taskFunction) ExecuteArgs( + ctx context.Context, + logger *slog.Logger, + args []binding.Arg, +) error { var sdkClient sdk.Client if injected, ok := ctx.Value(sdkcontext.SdkClientContextKey).(sdk.Client); ok { sdkClient = injected @@ -56,41 +84,14 @@ func (f *taskFunction) Execute(ctx context.Context, logger *slog.Logger) error { sdkClient = sdk.NewClient() } - reflectArgs := make([]reflect.Value, fnType.NumIn()) - for i := range reflectArgs { - in := fnType.In(i) - - switch { - case isTIRunContext(in): - // sdk.TIRunContext embeds context.Context, so it also satisfies - // isContext - this case must come first. The runtime stores the - // identifiers/timestamps under RuntimeContextKey; rebuild the - // value around the live task context here. - var ti sdk.TaskInstance - var dagRun sdk.DagRun - if stored, ok := ctx.Value(sdkcontext.RuntimeContextKey).(sdk.TIRunContext); ok { - ti, dagRun = stored.TaskInstance(), stored.DagRun() - } - reflectArgs[i] = reflect.ValueOf(sdk.NewTIRunContext(ctx, ti, dagRun)) - case isContext(in): - // Plain context.Context injection is retained for the Edge Worker - // runtime path, which does not populate the task runtime context - // (TI/DagRun) that sdk.TIRunContext carries. New tasks should - // declare sdk.TIRunContext instead. - reflectArgs[i] = reflect.ValueOf(ctx) - case isLogger(in): - reflectArgs[i] = reflect.ValueOf(logger) - case isClient(in): - reflectArgs[i] = reflect.ValueOf(sdkClient) - default: - // TODO: deal with other value types. For now they will all be Zero values unless it's a context - reflectArgs[i] = reflect.Zero(in) - } + reflectArgs, err := f.plan.Resolve(ctx, logger, sdkClient, args) + if err != nil { + return err } slog.Debug("Attempting to call fn", "fn", f.fn, "args", reflectArgs) retValues := f.fn.Call(reflectArgs) - var err error + err = nil if errResult := retValues[len(retValues)-1].Interface(); errResult != nil { var ok bool if err, ok = errResult.(error); !ok { @@ -150,11 +151,11 @@ func (f *taskFunction) validateFn(fnType reflect.Type) error { ) } - for i := range fnType.NumIn() { - if err := validateParam(fnType.In(i)); err != nil { - return fmt.Errorf("task function %s parameter %d: %w", f.fullName, i, err) - } + plan, err := binding.Analyze(fnType, f.fullName) + if err != nil { + return err } + f.plan = plan return nil } @@ -168,75 +169,8 @@ func isValidResultType(inType reflect.Type) bool { return true } -var ( - errorType = reflect.TypeFor[error]() - contextType = reflect.TypeFor[context.Context]() - tiRunContextType = reflect.TypeFor[sdk.TIRunContext]() - slogLoggerType = reflect.TypeFor[*slog.Logger]() - - clientType = reflect.TypeFor[sdk.Client]() -) +var errorType = reflect.TypeFor[error]() func isError(inType reflect.Type) bool { return inType != nil && inType.Implements(errorType) } - -func isContext(inType reflect.Type) bool { - return inType != nil && inType.Implements(contextType) -} - -func isTIRunContext(inType reflect.Type) bool { - return inType == tiRunContextType -} - -func isLogger(inType reflect.Type) bool { - return inType != nil && inType.AssignableTo(slogLoggerType) -} - -// isClient reports whether inType's method set is a subset of sdk.Client's, -// keeping new client capabilities injectable without a hand-kept list. -func isClient(inType reflect.Type) bool { - return inType != nil && inType.Kind() == reflect.Interface && - inType.NumMethod() > 0 && clientType.Implements(inType) -} - -// validateParam rejects interface parameters Execute cannot inject; they -// would be bound to nil and panic on first use. -func validateParam(in reflect.Type) error { - if in.Kind() != reflect.Interface || isTIRunContext(in) || isClient(in) { - return nil - } - if isContext(in) { - // The plain task context injected here cannot satisfy extra methods. - if contextType.Implements(in) { - return nil - } - return fmt.Errorf( - "interface %s adds methods on top of context.Context; declare sdk.TIRunContext or a separate parameter instead", - in, - ) - } - return fmt.Errorf( - "interface %s is not injectable (want context.Context, sdk.TIRunContext, or a subset of sdk.Client): %s", - in, - explainClientMismatch(in), - ) -} - -// explainClientMismatch returns why in is not a subset of sdk.Client. -func explainClientMismatch(in reflect.Type) string { - if in.NumMethod() == 0 { - return "empty interfaces cannot be injected" - } - for i := range in.NumMethod() { - m := in.Method(i) - cm, ok := clientType.MethodByName(m.Name) - if !ok { - return fmt.Sprintf("sdk.Client has no method %s", m.Name) - } - if cm.Type != m.Type { - return fmt.Sprintf("method %s is %s on sdk.Client, not %s", m.Name, cm.Type, m.Type) - } - } - return "its method set is not a subset of sdk.Client" -} diff --git a/go-sdk/bundle/bundlev1/task_test.go b/go-sdk/bundle/bundlev1/task_test.go index 3f58e930b8721..3daeb449ddef7 100644 --- a/go-sdk/bundle/bundlev1/task_test.go +++ b/go-sdk/bundle/bundlev1/task_test.go @@ -20,11 +20,11 @@ package bundlev1 import ( "context" "log/slog" - "reflect" "testing" "github.com/stretchr/testify/suite" + "github.com/apache/airflow/go-sdk/pkg/binding" "github.com/apache/airflow/go-sdk/pkg/logging" "github.com/apache/airflow/go-sdk/pkg/sdkcontext" "github.com/apache/airflow/go-sdk/sdk" @@ -141,21 +141,9 @@ func (s *TaskSuite) TestClientSubsetInjection() { s.Require().NoError(task.Execute(context.Background(), slog.New(logging.NewTeeLogger()))) } -// TestNamedClientInterfacesAreInjectable guards against sdk.Client dropping an -// embedded interface, which would break tasks declaring it. -func (s *TaskSuite) TestNamedClientInterfacesAreInjectable() { - for name, typ := range map[string]reflect.Type{ - "Client": reflect.TypeFor[sdk.Client](), - "VariableClient": reflect.TypeFor[sdk.VariableClient](), - "ConnectionClient": reflect.TypeFor[sdk.ConnectionClient](), - "XComClient": reflect.TypeFor[sdk.XComClient](), - } { - s.True(isClient(typ), "sdk.%s must stay injectable", name) - } -} - // TestNonInjectableParamsAreRejected checks registration fails fast on -// interface parameters Execute cannot inject. +// parameters Execute can neither inject nor bind a task argument to. This +// replaces the historical silent zero-fill of unrecognized parameters. func (s *TaskSuite) TestNonInjectableParamsAreRejected() { cases := map[string]struct { fn any @@ -174,9 +162,9 @@ func (s *TaskSuite) TestNonInjectableParamsAreRejected() { }, "method GetVariable is func(context.Context, string) (string, error) on sdk.Client", }, - "empty-interface": { - func(x any) error { return nil }, - "empty interfaces cannot be injected", + "func-param": { + func(cb func()) error { return nil }, + "cannot receive a task argument", }, "context-with-extra-methods": { func(x interface { @@ -201,6 +189,65 @@ func (s *TaskSuite) TestNonInjectableParamsAreRejected() { } } +// TestExecuteArgsBindsDataParameters covers the TaskFlow path end to end at the +// task level: literals decode onto data parameters interleaved with +// injectables, and Execute (nil args) keeps working for argless functions. +func (s *TaskSuite) TestExecuteArgsBindsDataParameters() { + var gotCountry string + var gotMeta map[string]any + task, err := NewTaskFunction(func(log *slog.Logger, country string, meta map[string]any) error { + gotCountry = country + gotMeta = meta + return nil + }) + s.Require().NoError(err) + + tw, ok := task.(TaskWithArgs) + s.Require().True(ok, "taskFunction must implement TaskWithArgs") + + err = tw.ExecuteArgs(context.Background(), slog.New(logging.NewTeeLogger()), []binding.Arg{ + binding.LiteralArg{Value: "uk"}, + binding.LiteralArg{Value: map[string]any{"k": "v"}}, + }) + s.Require().NoError(err) + s.Equal("uk", gotCountry) + s.Equal(map[string]any{"k": "v"}, gotMeta) +} + +// TestExecuteWithoutArgsFailsForDataParameters: a function with data +// parameters run through the argless Execute path (e.g. the Edge Worker, or a +// stub Dag that passes no arguments) fails loudly on the arity check instead +// of silently zero-filling. +func (s *TaskSuite) TestExecuteWithoutArgsFailsForDataParameters() { + task, err := NewTaskFunction(func(country string) error { return nil }) + s.Require().NoError(err) + + err = task.Execute(context.Background(), slog.New(logging.NewTeeLogger())) + if s.Assert().Error(err) { + s.Contains(err.Error(), "argument count mismatch") + } +} + +// TestExecuteArgsArityMismatch fails loudly when the Dag passes more arguments +// than the function declares data parameters. +func (s *TaskSuite) TestExecuteArgsArityMismatch() { + task, err := NewTaskFunction(func(country string) error { return nil }) + s.Require().NoError(err) + + err = task.(TaskWithArgs).ExecuteArgs( + context.Background(), + slog.New(logging.NewTeeLogger()), + []binding.Arg{ + binding.LiteralArg{Value: "uk"}, + binding.LiteralArg{Value: "de"}, + }, + ) + if s.Assert().Error(err) { + s.Contains(err.Error(), "argument count mismatch") + s.Contains(err.Error(), "passes 2 positional argument(s)") + } +} + // probeKey is an unexported context key used to confirm the live task context // (not a freshly built one) backs the injected sdk.TIRunContext. type probeKeyType struct{} diff --git a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go index 77725b0ac4efc..7eaccb8195a10 100644 --- a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go +++ b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go @@ -34,6 +34,7 @@ import ( "github.com/apache/airflow/go-sdk/internal/airflowmetadata" "github.com/apache/airflow/go-sdk/internal/bundlefooter" + "github.com/apache/airflow/go-sdk/pkg/execution" ) // crossArchFor returns an architecture different from the host that the Go @@ -142,7 +143,7 @@ func TestPack_CrossArchExecutableWithMetadataFile(t *testing.T) { sdk: language: "go" version: "` + sdkVersion + `" - supervisor_schema_version: "2026-06-16" + supervisor_schema_version: "` + execution.SupervisorSchemaVersion + `" source: "main.go" dags: concurrent_xcom_dag: @@ -153,6 +154,17 @@ dags: - "extract" - "transform" - "load" + taskflow_binding_dag: + tasks: + - "make_config" + - "make_numbers" + - "make_region" + - "via_flat_args" + - "via_struct_no_tags" + - "via_struct_arg_tag" + - "via_struct_unmatched_arg" + - "via_flat_map" + - "via_struct_map" ` assert.Equal(t, expectedManifest, string(metadata)) diff --git a/go-sdk/dags/go_examples.py b/go-sdk/dags/go_examples.py index 23e02dd5e49dc..dc47605782446 100644 --- a/go-sdk/dags/go_examples.py +++ b/go-sdk/dags/go_examples.py @@ -17,9 +17,14 @@ """ Python stub Dags mirroring the Go SDK example bundle (``go-sdk/example/bundle``). -Two Dags, both backed by the same Go bundle: ``simple_dag`` (extract/transform/ -load, below) and ``concurrent_xcom_dag`` (one ``pull_xcoms_concurrently`` task -timing sequential vs goroutine XCom pulls). +Three Dags, all backed by the same Go bundle: ``simple_dag`` (extract/transform/ +load, below), ``concurrent_xcom_dag`` (one ``pull_xcoms_concurrently`` task +timing sequential vs goroutine XCom pulls), and ``taskflow_binding_dag`` +(stressing the TaskFlow argument-binding surface -- the flat, positional +parameter list ``via_flat_args`` binds onto, plus three name-based +(keyword-style) struct examples, ``via_struct_no_tags``/``via_struct_arg_tag``/ +``via_struct_unmatched_arg``, each isolating one field-binding mode; see its +Dag function below). ``simple_dag`` sandwiches the Go tasks between two native Python tasks so the run exercises XCom across the language boundary, the same way @@ -33,6 +38,10 @@ routed to the ``ExecutableCoordinator``, which locates the bundle by dag_id and runs the binary in coordinator mode. ``extract`` returns a map (pushed as its ``return_value`` XCom); ``transform`` reads the ``my_variable`` variable. +* ``transform`` is called TaskFlow-style -- ``transform("uk", extract())`` -- so + the stub captures a positional-argument spec (a literal plus an XCom + reference) that the Go runtime binds onto the Go function's ``country`` and + ``extracted`` parameters, pulling ``extract``'s XCom on demand. * ``load`` (``retries=1``) returns an error on its first attempt and succeeds on the retry, exercising the UP_FOR_RETRY path through the Go coordinator. It is a leaf (not upstream of ``python_task_2``) so its retry is observable @@ -65,7 +74,7 @@ def extract(): ... @task.stub(queue="golang") -def transform(): ... +def transform(country: str, extracted: dict): ... # ``load`` fails on its first attempt and succeeds on the retry, exercising the @@ -86,7 +95,10 @@ def python_task_2(extracted): @dag(dag_id="simple_dag") def simple_dag(): extracted = extract() - transformed = transform() + # TaskFlow-style call: "uk" is captured as a literal argument and + # ``extracted`` as an XCom reference; both bind onto the Go function's + # data parameters at execution time (this also wires extract >> transform). + transformed = transform("uk", extracted) python_task_1() >> extracted >> transformed # ``load`` fails once then succeeds on retry; keep it a leaf (not upstream # of python_task_2) so its retry is observable without affecting the Python @@ -107,3 +119,118 @@ def concurrent_xcom_dag(): concurrent_xcom_dag() + + +@task.stub(queue="golang") +def make_config(): ... + + +@task.stub(queue="golang") +def make_numbers(): ... + + +@task.stub(queue="golang") +def make_region(): ... + + +@task.stub(queue="golang") +def via_flat_args( + name: str, + count: int, + ratio: float, + enabled: bool, + tags: list, + config: dict, + numbers: list, + note: str | None = None, +): ... + + +# Capitalized parameters on purpose: with no ``arg:`` tags on the Go side, each +# struct field binds the argument spelled exactly like its Go field name. +@task.stub(queue="golang") +def via_struct_no_tags(RegionCode: str, Threshold: float): ... + + +@task.stub(queue="golang") +def via_struct_arg_tag(region_code: str, threshold: float): ... + + +@task.stub(queue="golang") +def via_struct_unmatched_arg(region_code: str, sample_rate: float = 0.1): ... + + +@task.stub(queue="golang") +def via_flat_map(config: dict): ... + + +@task.stub(queue="golang") +def via_struct_map(payload: dict): ... + + +@dag(dag_id="taskflow_binding_dag") +def taskflow_binding_dag(): + """ + Stress the TaskFlow argument-binding surface beyond ``simple_dag``'s transform. + + Conceptually, the flat parameter list is *positional-argument* binding: order + matters, and every parameter must be filled or the task fails before it runs. + A task whose sole data parameter is a struct is closer to *keyword-argument* + binding: fields match by name, and (see ``via_struct_unmatched_arg`` below) a + field whose name has no corresponding TaskFlow call argument simply stays at + its zero value instead of failing the task -- the same way an unpassed keyword + argument falls back to a default in kwargs-style calls. + + ``via_flat_args``'s one mixed positional/keyword call carries literals of every + scalar type plus an array literal, and fans in XComs from *two* upstream Go + tasks: ``make_config`` returns an object that binds onto a strictly-decoded Go + struct, ``make_numbers`` an array that binds onto ``[]int``. ``note`` is not + passed, so its ``None`` default is captured and arrives in Go as a nil + ``*string``. The Go ``via_flat_args`` (``go-sdk/example/bundle/taskflowbinding``) + verifies every bound value and fails the task on any mismatch. + + Three further tasks demonstrate the Go SDK's name-based struct binding, used + when a struct is the task's sole data parameter. Each call mixes a literal + (``threshold``) with an XCom reference: the + ``region_code`` argument is ``make_region``'s output, so every struct example + also proves an XCom-sourced value binds onto a struct field. One field-binding + mode at a time: + + * ``via_struct_no_tags``: no ``arg:`` tags at all -- each struct field binds + the argument spelled exactly like its Go field name, hence this stub's + capitalized ``RegionCode``/``Threshold`` parameters. + * ``via_struct_arg_tag``: every field names its argument via an explicit + ``arg:`` tag -- ``Region`` is genuinely renamed to ``region_code``, and + ``Threshold`` is tagged ``threshold`` to pull the snake_case argument its + verbatim field name would miss. + * ``via_struct_unmatched_arg``: the mismatch tolerance in both directions. + The Go struct declares a field with no corresponding argument in this + TaskFlow call at all -- it stays at its Go zero value rather than failing + the task. And the stub's defaulted ``sample_rate`` is never passed, so its + captured-from-default entry needs no matching struct field (an explicitly + passed argument no field claims would fail the task instead). + + ``via_flat_map`` and ``via_struct_map`` each pass a single dict literal to + show the two ways one map binds on the Go side: ``via_flat_map``'s ``config`` + argument matches no struct field, so the whole dict is decoded into a Go + struct (flat), while ``via_struct_map``'s ``payload`` argument binds by name + onto a Go struct's ``map`` field (struct-based). + """ + via_flat_args( + "summary", + 3, + 2.5, + True, + ["metrics", "hourly"], + config=make_config(), + numbers=make_numbers(), + ) + region = make_region() + via_struct_no_tags(RegionCode=region, Threshold=0.75) + via_struct_arg_tag(region_code=region, threshold=0.75) + via_struct_unmatched_arg(region_code=region) + via_flat_map(config={"region": "eu-west-1", "count": 3}) + via_struct_map(payload={"region": "eu-west-1", "count": 3}) + + +taskflow_binding_dag() diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 23e60bd1dd46e..7ef64aa77405f 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -27,6 +27,7 @@ import ( v1 "github.com/apache/airflow/go-sdk/bundle/bundlev1" "github.com/apache/airflow/go-sdk/bundle/bundlev1/bundlev1server" "github.com/apache/airflow/go-sdk/example/bundle/concurrentxcom" + "github.com/apache/airflow/go-sdk/example/bundle/taskflowbinding" "github.com/apache/airflow/go-sdk/sdk" ) @@ -55,6 +56,17 @@ func (m *myBundle) RegisterDags(dagbag v1.Registry) error { concurrentDag := dagbag.AddDag("concurrent_xcom_dag") concurrentDag.AddTaskWithName("pull_xcoms_concurrently", concurrentxcom.PullXComsConcurrently) + bindingDag := dagbag.AddDag("taskflow_binding_dag") + bindingDag.AddTaskWithName("make_config", taskflowbinding.MakeConfig) + bindingDag.AddTaskWithName("make_numbers", taskflowbinding.MakeNumbers) + bindingDag.AddTaskWithName("make_region", taskflowbinding.MakeRegion) + bindingDag.AddTaskWithName("via_flat_args", taskflowbinding.ViaFlatArgs) + bindingDag.AddTaskWithName("via_struct_no_tags", taskflowbinding.ViaStructNoTags) + bindingDag.AddTaskWithName("via_struct_arg_tag", taskflowbinding.ViaStructArgTag) + bindingDag.AddTaskWithName("via_struct_unmatched_arg", taskflowbinding.ViaStructUnmatchedArg) + bindingDag.AddTaskWithName("via_flat_map", taskflowbinding.ViaFlatMap) + bindingDag.AddTaskWithName("via_struct_map", taskflowbinding.ViaStructMap) + return nil } @@ -127,12 +139,28 @@ func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (any, er return ret, nil } -func transform(ctx sdk.TIRunContext, client sdk.VariableClient, log *slog.Logger) error { +// transform demonstrates TaskFlow-style argument binding: the Python stub Dag +// calls “transform("uk", extract())“, so the runtime binds the "uk" literal +// onto country and pulls extract's return-value XCom into extracted -- the +// injectable parameters (runtime context, client, logger) are filled by type +// as before, in any position. +func transform( + ctx sdk.TIRunContext, + client sdk.VariableClient, + log *slog.Logger, + country string, + extracted map[string]any, +) error { // This function takes a VariableClient and not a Client to make unit testing it easier. See // `./main_test.go` for an example unit of this task fn. Functionally taking a `sdk.Client` is the same (as // Client includes VariableClient) but by using the dedicated type it can be easier to write unit tests. // // It also gives a better indication of what features the tasks use + log.InfoContext(ctx, "Bound TaskFlow arguments", + "country", country, + "extracted_go_version", extracted["go_version"], + "extracted_timestamp", extracted["timestamp"], + ) key := "my_variable" val, err := client.GetVariable(ctx, key) if err != nil { diff --git a/go-sdk/example/bundle/main_test.go b/go-sdk/example/bundle/main_test.go index 474a8406d6224..84bb174533812 100644 --- a/go-sdk/example/bundle/main_test.go +++ b/go-sdk/example/bundle/main_test.go @@ -52,8 +52,10 @@ var _ sdk.VariableClient = (*mockVars)(nil) func Test_transform(t *testing.T) { log := slog.Default() // This is not the best test, but it is a good proof of concept -- you can just call the function. - // sdk.NewTIRunContext wraps any context to build a TIRunContext in a test. + // sdk.NewTIRunContext wraps any context to build a TIRunContext in a test. The data parameters + // (country, extracted) are passed directly, exactly as the runtime would bind them from the + // stub Dag's TaskFlow call. ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) - err := transform(ctx, &mockVars{}, log) + err := transform(ctx, &mockVars{}, log, "uk", map[string]any{"go_version": "go1.24"}) assert.NoError(t, err) } diff --git a/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go b/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go new file mode 100644 index 0000000000000..59a3549aa82d2 --- /dev/null +++ b/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go @@ -0,0 +1,328 @@ +// 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. + +// Package taskflowbinding holds the taskflow_binding_dag tasks. ViaFlatArgs is +// positional-argument binding pushed to its limit: literals of every scalar +// type, an array literal, keyword arguments, a defaulted null, and XCom fan-in +// from two upstream Go tasks decoded into a strict struct and a typed slice -- +// where simple_dag's transform shows the minimal case (one literal, one +// XCom), this shows the full argument surface. The ViaStruct* functions +// instead show name-based (keyword-argument) binding, used when a struct is the +// task's sole data parameter: fields match by name and an unmatched name is +// left at its zero value rather than failing the task -- one field-binding +// mode at a time: ViaStructNoTags (verbatim field-name fallback), +// ViaStructArgTag (explicit `arg:` naming), and ViaStructUnmatchedArg (a +// field whose name has no corresponding TaskFlow call argument at all). Each +// ViaStruct* call binds MakeRegion's XCom onto its region field alongside a +// literal, so struct fields are exercised with both argument sources. ViaFlatMap +// and ViaStructMap round out the two ways a single dict argument binds: whole +// into a struct (flat) versus onto a struct's map field (by name). +package taskflowbinding + +import ( + "fmt" + "log/slog" + "reflect" + + "github.com/apache/airflow/go-sdk/sdk" +) + +// Config is the object make_config returns as its XCom; via_flat_args declares +// the same struct as a parameter, so the round trip exercises strict struct +// decoding (an unknown or renamed key fails the task rather than silently +// zeroing a field). +type Config struct { + Environment string `json:"environment"` + Region string `json:"region"` + Debug bool `json:"debug"` +} + +// MakeConfig pushes an object XCom that via_flat_args binds onto its Config parameter. +func MakeConfig(log *slog.Logger) (any, error) { + cfg := Config{Environment: "production", Region: "eu-west-1", Debug: true} + log.Info( + "Pushing config", + "environment", + cfg.Environment, + "region", + cfg.Region, + "debug", + cfg.Debug, + ) + return cfg, nil +} + +// MakeNumbers pushes an array XCom that via_flat_args binds onto its []int parameter. +func MakeNumbers(log *slog.Logger) (any, error) { + numbers := []int{1, 1, 2, 3, 5, 8} + log.Info("Pushing numbers", "numbers", fmt.Sprint(numbers)) + return numbers, nil +} + +// MakeRegion pushes a string XCom that every ViaStruct* task binds onto a +// struct field, so each field-binding mode is exercised with an XCom-sourced +// argument and not just literals. +func MakeRegion(log *slog.Logger) (any, error) { + region := "eu-west-1" + log.Info("Pushing region", "region", region) + return region, nil +} + +// ViaFlatArgs receives every argument shape the stub Dag can express as plain, +// positional data parameters. The Python side calls it as +// +// via_flat_args("summary", 3, 2.5, True, ["metrics", "hourly"], +// config=make_config(), numbers=make_numbers()) +// +// so the bound values are fixed; any mismatch below is a binding regression +// and fails the task loudly. note is never passed and falls back to the stub's +// None default, arriving as a nil *string. +func ViaFlatArgs( + ctx sdk.TIRunContext, + log *slog.Logger, + name string, + count int, + ratio float64, + enabled bool, + tags []string, + config Config, + numbers []int, + note *string, +) (any, error) { + if name != "summary" || count != 3 || ratio != 2.5 || !enabled { + return nil, fmt.Errorf( + "scalar literals bound incorrectly: name=%q count=%d ratio=%v enabled=%v", + name, count, ratio, enabled, + ) + } + if want := []string{"metrics", "hourly"}; !reflect.DeepEqual(tags, want) { + return nil, fmt.Errorf("array literal bound incorrectly: tags=%v, want %v", tags, want) + } + if want := (Config{Environment: "production", Region: "eu-west-1", Debug: true}); config != want { + return nil, fmt.Errorf("object XCom bound incorrectly: config=%+v, want %+v", config, want) + } + if want := []int{1, 1, 2, 3, 5, 8}; !reflect.DeepEqual(numbers, want) { + return nil, fmt.Errorf("array XCom bound incorrectly: numbers=%v, want %v", numbers, want) + } + if note != nil { + return nil, fmt.Errorf("defaulted None bound incorrectly: note=%q, want nil", *note) + } + + sum := 0 + for _, n := range numbers { + sum += n + } + log.InfoContext(ctx, "Bound TaskFlow arguments", + "name", name, + "count", count, + "ratio", ratio, + "enabled", enabled, + "tags", fmt.Sprint(tags), + "environment", config.Environment, + "sum", sum, + ) + return map[string]any{ + "name": name, + "count": count, + "ratio": ratio, + "enabled": enabled, + "tags": tags, + "environment": config.Environment, + "debug": config.Debug, + "sum": sum, + "note_was_null": note == nil, + }, nil +} + +// ViaStructNoTagsInput demonstrates name-based binding with no field tags at +// all: each field binds the TaskFlow call argument spelled exactly like its Go +// field name ("RegionCode", "Threshold"), which is why the stub declares +// capitalized parameters. +type ViaStructNoTagsInput struct { + RegionCode string + Threshold float64 +} + +// ViaStructNoTags is called as +// +// via_struct_no_tags(RegionCode=make_region(), Threshold=0.75) +// +// so RegionCode arrives via make_region's XCom and Threshold as a literal. +func ViaStructNoTags( + ctx sdk.TIRunContext, + log *slog.Logger, + input ViaStructNoTagsInput, +) (any, error) { + if input.RegionCode != "eu-west-1" || input.Threshold != 0.75 { + return nil, fmt.Errorf( + "struct fields bound incorrectly: region_code=%q threshold=%v", + input.RegionCode, + input.Threshold, + ) + } + + log.InfoContext(ctx, "Bound struct (no tags)", + "region_code", input.RegionCode, + "threshold", input.Threshold, + ) + return map[string]any{ + "region_code": input.RegionCode, + "threshold": input.Threshold, + }, nil +} + +// ViaStructArgTagInput demonstrates name-based binding with explicit arg: tags: +// Region binds to the "region_code" TaskFlow argument under a renamed Go field, +// proving the tag remaps the name rather than coincidentally matching it; +// Threshold is intentionally tagged "threshold" because an untagged field would +// only match an argument spelled exactly "Threshold". +type ViaStructArgTagInput struct { + Region string `arg:"region_code"` + Threshold float64 `arg:"threshold"` +} + +// ViaStructArgTag is called as +// +// via_struct_arg_tag(region_code=make_region(), threshold=0.75) +// +// so Region arrives via make_region's XCom and Threshold as a literal. +func ViaStructArgTag( + ctx sdk.TIRunContext, + log *slog.Logger, + input ViaStructArgTagInput, +) (any, error) { + if input.Region != "eu-west-1" || input.Threshold != 0.75 { + return nil, fmt.Errorf( + "struct fields bound incorrectly: region=%q threshold=%v", + input.Region, + input.Threshold, + ) + } + + log.InfoContext(ctx, "Bound struct (arg: tag)", + "region", input.Region, + "threshold", input.Threshold, + ) + return map[string]any{ + "region": input.Region, + "threshold": input.Threshold, + }, nil +} + +// ViaStructUnmatchedArgInput demonstrates the mismatch tolerance in both +// directions. A field whose name has no corresponding TaskFlow call argument +// at all is left at its Go zero value rather than failing the task: Region +// binds normally, but Missing's arg name is never among this call's +// arguments -- conceptually, an unpassed keyword argument falling back to +// its default in a kwargs-style call. The reverse also holds: the stub's +// defaulted sample_rate parameter arrives marked from_default, so this +// struct is free not to mirror it (an explicitly passed argument no field +// claims would fail the task instead). +type ViaStructUnmatchedArgInput struct { + Region string `arg:"region_code"` + Missing string `arg:"does_not_exist"` +} + +// ViaStructUnmatchedArg is called as +// +// via_struct_unmatched_arg(region_code=make_region()) +// +// -- the stub declares region_code (bound from make_region's XCom) plus a +// defaulted sample_rate this struct deliberately omits, so Missing's arg +// name never appears among the call's arguments and stays at its Go zero +// value (""), while sample_rate's from_default entry goes unclaimed without +// failing the task. +func ViaStructUnmatchedArg( + ctx sdk.TIRunContext, log *slog.Logger, input ViaStructUnmatchedArgInput, +) (any, error) { + if input.Region != "eu-west-1" { + return nil, fmt.Errorf("struct field bound incorrectly: region=%q", input.Region) + } + if input.Missing != "" { + return nil, fmt.Errorf( + "expected the unmatched field to stay at its Go zero value, got missing=%q", + input.Missing, + ) + } + + log.InfoContext(ctx, "Bound struct (unmatched arg)", + "region", input.Region, + "missing_was_empty", input.Missing == "", + ) + return map[string]any{ + "region": input.Region, + "missing_was_empty": input.Missing == "", + }, nil +} + +// FlatMapConfig is decoded whole from a single dict argument -- the flat +// (whole-value) side of single-map injection. The TaskFlow call passes one +// object whose keys land on these fields; because the argument's name matches +// no field, the runtime decodes it whole rather than binding field-by-field. +type FlatMapConfig struct { + Region string `json:"region"` + Count int `json:"count"` +} + +// ViaFlatMap is called as +// +// via_flat_map(config={"region": "eu-west-1", "count": 3}) +// +// The single "config" argument matches no FlatMapConfig field, so the whole +// dict is decoded into the struct -- one map bound flat into a struct. +func ViaFlatMap( + ctx sdk.TIRunContext, log *slog.Logger, config FlatMapConfig, +) (any, error) { + if config.Region != "eu-west-1" || config.Count != 3 { + return nil, fmt.Errorf( + "whole-value map bound incorrectly: region=%q count=%d", + config.Region, + config.Count, + ) + } + + log.InfoContext(ctx, "Bound whole map into struct", + "region", config.Region, + "count", config.Count, + ) + return map[string]any{"region": config.Region, "count": config.Count}, nil +} + +// StructMapInput is a name-bound struct whose single field is itself a map, so +// one dict argument binds onto that field by name -- the struct-based side of +// single-map injection. +type StructMapInput struct { + Payload map[string]any `arg:"payload"` +} + +// ViaStructMap is called as +// +// via_struct_map(payload={"region": "eu-west-1", "count": 3}) +// +// The "payload" argument matches the Payload field by name, so the dict binds +// onto the map field -- one map bound by name onto a struct field. +func ViaStructMap( + ctx sdk.TIRunContext, log *slog.Logger, input StructMapInput, +) (any, error) { + region, _ := input.Payload["region"].(string) + if region != "eu-west-1" { + return nil, fmt.Errorf("map field bound incorrectly: payload=%v", input.Payload) + } + + log.InfoContext(ctx, "Bound map onto struct field", "payload", fmt.Sprint(input.Payload)) + return map[string]any{"payload": input.Payload}, nil +} diff --git a/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go b/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go new file mode 100644 index 0000000000000..55c84762aaf88 --- /dev/null +++ b/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go @@ -0,0 +1,166 @@ +// 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. + +package taskflowbinding + +import ( + "context" + "log/slog" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/apache/airflow/go-sdk/sdk" +) + +// Like example/bundle/main_test.go, this shows a task fn is unit-testable by +// passing the data parameters directly, exactly as the runtime binds them. +func TestViaFlatArgs(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + got, err := ViaFlatArgs(ctx, slog.Default(), + "summary", 3, 2.5, true, + []string{"metrics", "hourly"}, + Config{Environment: "production", Region: "eu-west-1", Debug: true}, + []int{1, 1, 2, 3, 5, 8}, + nil, + ) + require.NoError(t, err) + + summary, ok := got.(map[string]any) + require.True(t, ok, "ViaFlatArgs should return a map summary, got %T", got) + assert.Equal(t, 20, summary["sum"]) + assert.Equal(t, true, summary["note_was_null"]) +} + +func TestViaFlatArgsRejectsWrongBinding(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + _, err := ViaFlatArgs(ctx, slog.Default(), + "summary", 3, 2.5, true, + []string{"metrics", "hourly"}, + Config{}, + []int{1, 1, 2, 3, 5, 8}, + nil, + ) + assert.ErrorContains(t, err, "object XCom bound incorrectly") +} + +func TestViaStructNoTags(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + got, err := ViaStructNoTags(ctx, slog.Default(), ViaStructNoTagsInput{ + RegionCode: "eu-west-1", + Threshold: 0.75, + }) + require.NoError(t, err) + + summary, ok := got.(map[string]any) + require.True(t, ok, "ViaStructNoTags should return a map summary, got %T", got) + assert.Equal(t, "eu-west-1", summary["region_code"]) +} + +func TestViaStructNoTagsRejectsWrongBinding(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + _, err := ViaStructNoTags(ctx, slog.Default(), ViaStructNoTagsInput{ + RegionCode: "wrong-region", + Threshold: 0.75, + }) + assert.ErrorContains(t, err, "struct fields bound incorrectly") +} + +func TestViaStructArgTag(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + got, err := ViaStructArgTag(ctx, slog.Default(), ViaStructArgTagInput{ + Region: "eu-west-1", + Threshold: 0.75, + }) + require.NoError(t, err) + + summary, ok := got.(map[string]any) + require.True(t, ok, "ViaStructArgTag should return a map summary, got %T", got) + assert.Equal(t, "eu-west-1", summary["region"]) +} + +func TestViaStructArgTagRejectsWrongBinding(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + _, err := ViaStructArgTag(ctx, slog.Default(), ViaStructArgTagInput{ + Region: "wrong-region", + Threshold: 0.75, + }) + assert.ErrorContains(t, err, "struct fields bound incorrectly") +} + +func TestViaStructUnmatchedArg(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + // Missing is left at its Go zero value, exactly as binding.Resolve leaves an + // unmatched struct field -- this task fn is unit-testable independent of + // the binding package precisely because it declares that expectation itself. + got, err := ViaStructUnmatchedArg(ctx, slog.Default(), ViaStructUnmatchedArgInput{ + Region: "eu-west-1", + Missing: "", + }) + require.NoError(t, err) + + summary, ok := got.(map[string]any) + require.True(t, ok, "ViaStructUnmatchedArg should return a map summary, got %T", got) + assert.Equal(t, true, summary["missing_was_empty"]) +} + +func TestViaStructUnmatchedArgRejectsNonZeroMissingField(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + _, err := ViaStructUnmatchedArg(ctx, slog.Default(), ViaStructUnmatchedArgInput{ + Region: "eu-west-1", + Missing: "unexpected", + }) + assert.ErrorContains(t, err, "expected the unmatched field to stay at its Go zero value") +} + +func TestViaFlatMap(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + got, err := ViaFlatMap(ctx, slog.Default(), FlatMapConfig{Region: "eu-west-1", Count: 3}) + require.NoError(t, err) + + summary, ok := got.(map[string]any) + require.True(t, ok, "ViaFlatMap should return a map summary, got %T", got) + assert.Equal(t, "eu-west-1", summary["region"]) + assert.Equal(t, 3, summary["count"]) +} + +func TestViaFlatMapRejectsWrongBinding(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + _, err := ViaFlatMap(ctx, slog.Default(), FlatMapConfig{Region: "wrong-region", Count: 3}) + assert.ErrorContains(t, err, "whole-value map bound incorrectly") +} + +func TestViaStructMap(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + got, err := ViaStructMap(ctx, slog.Default(), StructMapInput{ + Payload: map[string]any{"region": "eu-west-1", "count": 3}, + }) + require.NoError(t, err) + + summary, ok := got.(map[string]any) + require.True(t, ok, "ViaStructMap should return a map summary, got %T", got) + assert.Equal(t, map[string]any{"region": "eu-west-1", "count": 3}, summary["payload"]) +} + +func TestViaStructMapRejectsWrongBinding(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + _, err := ViaStructMap(ctx, slog.Default(), StructMapInput{ + Payload: map[string]any{"region": "wrong-region"}, + }) + assert.ErrorContains(t, err, "map field bound incorrectly") +} diff --git a/go-sdk/pkg/binding/binding.go b/go-sdk/pkg/binding/binding.go new file mode 100644 index 0000000000000..7a6c511dccf98 --- /dev/null +++ b/go-sdk/pkg/binding/binding.go @@ -0,0 +1,836 @@ +// 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. + +// Package binding turns a task function's parameter list into the concrete +// argument values it is called with at execution time. +// +// Parameters fall into two groups: +// +// - Injectable runtime values: context.Context, sdk.TIRunContext, +// *slog.Logger, and any interface whose method set is a subset of +// sdk.Client. These are filled by type, in any position. +// - Data parameters: everything else, in declaration order. They receive the +// positional arguments the Python stub Dag captured at parse time from the +// TaskFlow call (“transform("uk", extract())“) and delivered in +// StartupDetails. A literal argument decodes directly; an XCom argument is +// pulled from the named upstream task in the current Dag run first +// (independent pulls run concurrently). A struct data parameter is decoded +// whole from its one positional argument. +// +// Data parameters are positional-argument binding: order matters, every +// parameter must be filled, and Resolve fails the task before its body runs on +// an arity or type mismatch. A declared type of "any" (or a Go parameter typed +// any) opts that argument out of the type check; the decode step still fails +// loudly on unusable values. +// +// One shape is treated specially: a function whose sole data parameter is a +// (pointer-to-)struct binds by field name (kwarg-style) instead, using the +// per-execution argument spec. Each exported field claims the argument named by +// its `arg:""` tag, or its verbatim Go field name when untagged -- so an +// untagged Name binds the argument "Name", and a snake_case argument like +// "count" needs an explicit tag. A field whose name has no corresponding +// argument is left at its Go zero value rather than failing the task, the same +// way an unpassed keyword argument falls back to its default; conversely every +// explicitly passed argument must be claimed by some field (catching typo'd +// field names), while entries the Python side captured from the stub +// signature's defaults (from_default on the wire) may go unclaimed. When the +// sole struct instead receives exactly one explicitly passed argument that no +// field claims, it falls back to a whole-value decode of that argument -- so a +// struct can still receive an upstream object as one argument. A bindable-field +// struct fails loudly when no argument spec arrives at all (e.g. the Edge +// Worker path, where nothing could ever fill it), matching the flat-parameter +// arity check. +// +// Because `arg:` tags only take effect in that sole-struct name-binding path, a +// struct data parameter that carries an `arg:` tag must be the only data +// parameter; Analyze rejects a signature that pairs a tagged struct with other +// data parameters rather than silently ignoring the tags. +// +// Analyze inspects a function once at registration and returns a Plan; Resolve +// builds the call arguments for each execution from that Plan and the +// per-execution argument spec. The flat-vs-name-bound choice for a sole struct +// parameter is made per execution, so the same function can bind either way in +// different Dags. +// +// Mapped (.expand()) stubs are out of scope: their specs carry no bindings on +// this wire version, so an XCom pull always targets the unmapped upstream +// instance. Per-map-index delivery arrives with the mapped follow-up. +package binding + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "reflect" + "strings" + "sync" + + "github.com/apache/airflow/go-sdk/pkg/api" + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" + "github.com/apache/airflow/go-sdk/pkg/sdkcontext" + "github.com/apache/airflow/go-sdk/sdk" +) + +// Arg is one positional argument for a task function's data parameters, in +// declaration order: an XComArg or a LiteralArg. A sealed sum type over the +// wire model's XComArgBinding/LiteralArgBinding split; each variant is +// defined in terms of its generated schema struct so the fields cannot drift +// from the coordinator protocol. +type Arg interface { + // ArgName is the stub function's parameter name this binding fills; used + // to match a struct field's `arg:` tag (or its verbatim field-name + // fallback) in name-based binding. + ArgName() string + // Schema is the JSON-schema fragment the Dag declared for the argument + // (generated by pydantic from the stub annotation), or nil when the + // argument is unconstrained; the type check reads its "type" keyword and + // falls back to a decode-only check when nil or absent. + Schema() *genmodels.ArgValueSchema + // sealedArg restricts implementations to this package, keeping + // resolveOne's type switch the single exhaustive consumer. + sealedArg() +} + +// XComArg sources the argument from an upstream task's return-value XCom. +// Kind carries the wire discriminant ("xcom") from the generated shape; the +// resolve path dispatches on the Go type itself and never reads it. +type XComArg genmodels.XComArgBinding + +// LiteralArg carries an inline value from the Dag file. Kind carries the wire +// discriminant ("literal") from the generated shape; the resolve path +// dispatches on the Go type itself and never reads it. +type LiteralArg genmodels.LiteralArgBinding + +func (a XComArg) ArgName() string { return a.Name } +func (a LiteralArg) ArgName() string { return a.Name } + +func (a XComArg) Schema() *genmodels.ArgValueSchema { return a.ValueSchema } +func (a LiteralArg) Schema() *genmodels.ArgValueSchema { return a.ValueSchema } + +func (XComArg) sealedArg() {} +func (LiteralArg) sealedArg() {} + +// paramKind classifies how a task-function parameter is filled at execution. +type paramKind int + +const ( + paramTIRunContext paramKind = iota + paramContext + paramLogger + paramClient + paramData + // paramLoneStruct is the function's sole data parameter and is a + // (pointer-to-)struct. Resolve decides per execution whether to bind its + // fields by name (kwarg-style) or decode one argument whole into it. + paramLoneStruct +) + +// structField describes how Resolve fills one exported field of a name-bound +// struct parameter. Precomputed once by Analyze. +type structField struct { + // structIndex is the field's index within the struct, for + // reflect.Value.Field. + structIndex int + // goName is the Go field name, for error messages. + goName string + fieldType reflect.Type + // argName is the name to claim from the argument spec: the field's `arg:` + // tag, or its Go field name verbatim when the tag is omitted. + argName string +} + +// paramPlan describes how Resolve fills a single task-function parameter. +type paramPlan struct { + kind paramKind + // typ is the declared Go type of a data parameter (kind == paramData) or + // the struct type (kind == paramLoneStruct). + typ reflect.Type + // index is the parameter's position in the function signature, for error + // messages. + index int + // fields describes each exported field's binding. Set only for + // kind == paramLoneStruct. + fields []structField +} + +// Plan is the precomputed recipe for filling a task function's parameters. It +// is built once by Analyze and reused for every execution of that function. +type Plan struct { + fnName string + params []paramPlan + numData int + // loneStruct is true when the sole data parameter is a struct resolved + // per execution (kind == paramLoneStruct). + loneStruct bool +} + +// Analyze inspects the parameters of a task function type and builds a Plan. +// fnName appears in error messages only. Every parameter must be an injectable +// runtime type or a type that can receive a task argument (JSON-decodable); +// anything else is a registration error. +func Analyze(fnType reflect.Type, fnName string) (*Plan, error) { + p := &Plan{fnName: fnName, params: make([]paramPlan, fnType.NumIn())} + var dataIdxs []int + for i := range fnType.NumIn() { + plan, err := classifyParam(fnName, fnType.In(i), i) + if err != nil { + return nil, err + } + if plan.kind == paramData { + p.numData++ + dataIdxs = append(dataIdxs, i) + } + p.params[i] = plan + } + + // A function whose sole data parameter is a struct binds its fields by name + // at execution (see resolveLoneStructParam); any other struct data + // parameter is decoded whole from its positional slot. + if p.numData == 1 { + i := dataIdxs[0] + if st := structParamType(p.params[i].typ); st != nil { + fields, err := buildStructFields(fnName, st, i) + if err != nil { + return nil, err + } + p.params[i].kind = paramLoneStruct + p.params[i].fields = fields + p.numData = 0 + p.loneStruct = true + return p, nil + } + } + + // `arg:` tags only take effect in the sole-struct name-binding path above; + // a tagged struct alongside other data parameters would silently ignore its + // tags (it is decoded whole), so reject it as a likely mistake. + for _, i := range dataIdxs { + if st := structParamType(p.params[i].typ); st != nil && hasArgTag(st) { + return nil, fmt.Errorf( + "task function %s: parameter %d: a struct with `arg:` tags must be the function's "+ + "only data parameter (its fields bind TaskFlow arguments by name); it cannot be "+ + "combined with other data parameters", + fnName, i, + ) + } + } + return p, nil +} + +// Resolve builds the ordered argument values for one call. Injectable +// parameters receive values derived from ctx, logger, or client. Plain flat +// data parameters consume args in declaration order; a sole struct data +// parameter is instead resolved by field name or decoded whole (see Analyze +// and resolveLoneStructParam). An error fails the task before its body runs. +// +// client must be the full sdk.Client -- not just the sdk.XComClient the +// resolve helpers narrow to -- because a paramClient parameter receives the +// client itself, verbatim. +func (p *Plan) Resolve( + ctx context.Context, + logger *slog.Logger, + client sdk.Client, + args []Arg, +) ([]reflect.Value, error) { + out := make([]reflect.Value, len(p.params)) + for i, plan := range p.params { + switch plan.kind { + case paramTIRunContext: + // The runtime stores the identifiers/timestamps under + // RuntimeContextKey; rebuild the value around the live task context. + var ti sdk.TaskInstance + var dagRun sdk.DagRun + if stored, ok := ctx.Value(sdkcontext.RuntimeContextKey).(sdk.TIRunContext); ok { + ti, dagRun = stored.TaskInstance(), stored.DagRun() + } + out[i] = reflect.ValueOf(sdk.NewTIRunContext(ctx, ti, dagRun)) + case paramContext: + out[i] = reflect.ValueOf(ctx) + case paramLogger: + out[i] = reflect.ValueOf(logger) + case paramClient: + out[i] = reflect.ValueOf(client) + case paramData, paramLoneStruct: + // Filled below, once the spec is matched and XComs are pulled. + } + } + if p.loneStruct { + return p.resolveLoneStructParam(ctx, client, args, out) + } + return p.resolveFlatParams(ctx, client, args, out) +} + +// resolveFlatParams fills the plain data parameters positionally -- args[0] +// onto the first data parameter, and so on -- with strict arity (see the +// package doc comment). +func (p *Plan) resolveFlatParams( + ctx context.Context, + c sdk.XComClient, + args []Arg, + out []reflect.Value, +) ([]reflect.Value, error) { + if len(args) != p.numData { + return nil, fmt.Errorf( + "task function %s: argument count mismatch: the Dag passes %d positional argument(s) "+ + "but the Go function declares %d data parameter(s)", + p.fnName, len(args), p.numData, + ) + } + raws, err := p.fetchArgValues(ctx, c, args, nil) + if err != nil { + return nil, err + } + flatIdx := 0 + for i, plan := range p.params { + if plan.kind != paramData { + continue + } + v, err := p.decodeArg( + args[flatIdx], raws[flatIdx], plan.typ, + fmt.Sprintf("argument %d (parameter %d)", flatIdx, plan.index), + ) + if err != nil { + return nil, err + } + out[i] = v + flatIdx++ + } + return out, nil +} + +// resolveLoneStructParam fills the function's sole struct data parameter. It +// binds by field name (kwarg-style) by default: every explicitly passed spec +// entry must be claimed by some field; entries the Python side captured from +// the stub signature's defaults (FromDefault) may go unclaimed, the same way an +// unpassed keyword argument never reaches the callee. When instead exactly one +// explicitly passed argument arrives that no field claims, the struct is +// decoded whole from it (see resolveWholeStructParam). +func (p *Plan) resolveLoneStructParam( + ctx context.Context, + c sdk.XComClient, + args []Arg, + out []reflect.Value, +) ([]reflect.Value, error) { + var paramIdx int + var plan paramPlan + for i, pl := range p.params { + if pl.kind == paramLoneStruct { + paramIdx, plan = i, pl + break + } + } + + byName := make(map[string]int, len(args)) + for i, a := range args { + if a != nil { + byName[a.ArgName()] = i + } + } + + claimed := make([]bool, len(args)) + type fieldBind struct { + field structField + argIdx int + } + binds := make([]fieldBind, 0, len(plan.fields)) + for _, sf := range plan.fields { + idx, ok := byName[sf.argName] + if !ok { + // No TaskFlow call argument carries this name -- kwarg-style, an + // unpassed name leaves the field at its Go zero value rather than + // failing the task (see the package doc comment). + continue + } + claimed[idx] = true + binds = append(binds, fieldBind{field: sf, argIdx: idx}) + } + + // Whole-value fallback: a single explicitly passed argument that no field + // claims decodes whole into the struct, so a sole struct parameter can still + // receive an upstream object as one positional argument. + if len(binds) == 0 && isLoneWholeValueArg(args) { + return p.resolveWholeStructParam(ctx, c, args, plan, paramIdx, out) + } + + if len(args) == 0 && len(plan.fields) > 0 { + return nil, fmt.Errorf( + "task function %s: no TaskFlow arg bindings arrived but the struct declares "+ + "%d bindable field(s); nothing can fill them on this execution path", + p.fnName, len(plan.fields), + ) + } + + var unclaimed []string + for i, c := range claimed { + if c { + continue + } + if lit, ok := args[i].(LiteralArg); ok && lit.FromDefault { + // The Dag author never passed this argument; the Python side filled + // it from the stub signature's default. A struct that does not + // mirror the defaulted parameter is fine. + continue + } + name := "" + if args[i] != nil { + name = fmt.Sprintf("%q", args[i].ArgName()) + } + unclaimed = append(unclaimed, name) + } + if len(unclaimed) > 0 { + return nil, fmt.Errorf( + "task function %s: %d TaskFlow call argument(s) not claimed by any struct "+ + "field: %s", + p.fnName, len(unclaimed), strings.Join(unclaimed, ", "), + ) + } + + raws, err := p.fetchArgValues(ctx, c, args, claimed) + if err != nil { + return nil, err + } + + structType := plan.typ + isPtr := structType.Kind() == reflect.Pointer + if isPtr { + structType = structType.Elem() + } + structVal := reflect.New(structType).Elem() + for _, b := range binds { + v, err := p.decodeArg( + args[b.argIdx], raws[b.argIdx], b.field.fieldType, + fmt.Sprintf("struct field %s (parameter %d)", b.field.goName, plan.index), + ) + if err != nil { + return nil, err + } + structVal.Field(b.field.structIndex).Set(v) + } + + if isPtr { + out[paramIdx] = structVal.Addr() + } else { + out[paramIdx] = structVal + } + return out, nil +} + +// isLoneWholeValueArg reports whether args is a single explicitly passed +// argument (not a captured default) -- the case a sole struct parameter decodes +// whole rather than binding field-by-field. +func isLoneWholeValueArg(args []Arg) bool { + if len(args) != 1 || args[0] == nil { + return false + } + lit, ok := args[0].(LiteralArg) + return !ok || !lit.FromDefault +} + +// resolveWholeStructParam decodes the sole struct parameter's one positional +// argument whole into it, honouring pointer-ness (decodeArg allocates through +// plan.typ, the same as a flat data parameter). +func (p *Plan) resolveWholeStructParam( + ctx context.Context, + c sdk.XComClient, + args []Arg, + plan paramPlan, + paramIdx int, + out []reflect.Value, +) ([]reflect.Value, error) { + raws, err := p.fetchArgValues(ctx, c, args, nil) + if err != nil { + return nil, err + } + v, err := p.decodeArg( + args[0], raws[0], plan.typ, + fmt.Sprintf("argument %q (parameter %d)", args[0].ArgName(), plan.index), + ) + if err != nil { + return nil, err + } + out[paramIdx] = v + return out, nil +} + +// fetchArgValues produces the raw (pre-decode) value for each argument the +// caller will consume: a literal's inline value, or the upstream task's +// return-value XCom pulled over the API -- independent pulls run +// concurrently. needed selects which entries to fetch; nil means all. +func (p *Plan) fetchArgValues( + ctx context.Context, + c sdk.XComClient, + args []Arg, + needed []bool, +) ([]any, error) { + raws := make([]any, len(args)) + var xcomIdxs []int + for i, a := range args { + if needed != nil && !needed[i] { + continue + } + switch a := a.(type) { + case LiteralArg: + raws[i] = a.Value + case XComArg: + xcomIdxs = append(xcomIdxs, i) + } + // A nil or foreign Arg implementation fails in decodeArg, which names + // the destination parameter/field in the error. + } + if len(xcomIdxs) == 0 { + return raws, nil + } + + workload, ok := ctx.Value(sdkcontext.WorkloadContextKey).(api.ExecuteTaskWorkload) + if !ok { + return nil, fmt.Errorf( + "task function %s: no workload in context, cannot resolve xcom arguments", p.fnName, + ) + } + // Always the return-value XCom -- a stub Dag cannot reference any other + // key. Pull from the upstream's unmapped instance (map_index nil); mapped + // upstream fan-in is out of scope for now. + pull := func(i int) error { + a := args[i].(XComArg) + raw, err := c.GetXCom( + ctx, workload.TI.DagId, workload.TI.RunId, a.TaskID, nil, api.XComReturnValueKey, nil, + ) + if err != nil { + return fmt.Errorf( + "task function %s: argument %q: pulling xcom from task %q: %w", + p.fnName, a.Name, a.TaskID, err, + ) + } + raws[i] = raw + return nil + } + if len(xcomIdxs) == 1 { + if err := pull(xcomIdxs[0]); err != nil { + return nil, err + } + return raws, nil + } + var wg sync.WaitGroup + errs := make([]error, len(xcomIdxs)) + for j, i := range xcomIdxs { + wg.Add(1) + go func() { + defer wg.Done() + errs[j] = pull(i) + }() + } + wg.Wait() + if err := errors.Join(errs...); err != nil { + return nil, err + } + return raws, nil +} + +// decodeArg decodes one argument-spec entry's raw value into targetType: +// type-check against the declared Dag type, then a strict decode. errCtx +// names the destination parameter or struct field for error messages. +func (p *Plan) decodeArg( + arg Arg, + raw any, + targetType reflect.Type, + errCtx string, +) (reflect.Value, error) { + if arg == nil { + return reflect.Value{}, fmt.Errorf( + "task function %s: %s: nil argument binding", p.fnName, errCtx, + ) + } + if err := checkValueType(arg.Schema(), targetType); err != nil { + return reflect.Value{}, fmt.Errorf("task function %s: %s: %w", p.fnName, errCtx, err) + } + var source string + switch a := arg.(type) { + case LiteralArg: + source = "literal value" + case XComArg: + source = fmt.Sprintf("xcom from task %q", a.TaskID) + default: + return reflect.Value{}, fmt.Errorf( + "task function %s: %s: unsupported argument binding %T", p.fnName, errCtx, arg, + ) + } + v, err := decodeValue(raw, targetType) + if err != nil { + return reflect.Value{}, fmt.Errorf( + "task function %s: %s: decoding %s into %s: %w", + p.fnName, errCtx, source, targetType, err, + ) + } + return v, nil +} + +// classifyParam decides how a single parameter is filled. Injectable runtime +// types map to their paramKind; anything else is a data parameter and must be +// a type a task argument can decode into. +func classifyParam(fnName string, in reflect.Type, index int) (paramPlan, error) { + switch { + case isTIRunContext(in): + // sdk.TIRunContext embeds context.Context, so it also satisfies + // isContext - this case must come first. + return paramPlan{kind: paramTIRunContext, index: index}, nil + case isContext(in): + // The plain task context injected here cannot satisfy extra methods. + if !contextType.Implements(in) { + return paramPlan{}, fmt.Errorf( + "task function %s: parameter %d: interface %s adds methods on top of "+ + "context.Context; declare sdk.TIRunContext or a separate parameter instead", + fnName, index, in, + ) + } + return paramPlan{kind: paramContext, index: index}, nil + case isLogger(in): + return paramPlan{kind: paramLogger, index: index}, nil + case isClient(in): + return paramPlan{kind: paramClient, index: index}, nil + } + if in.Kind() == reflect.Interface && in.NumMethod() > 0 { + return paramPlan{}, fmt.Errorf( + "task function %s: parameter %d: interface %s is not injectable "+ + "(want context.Context, sdk.TIRunContext, or a subset of sdk.Client): %s", + fnName, index, in, explainClientMismatch(in), + ) + } + if !isDecodableType(in) { + return paramPlan{}, fmt.Errorf( + "task function %s: parameter %d: type %s cannot receive a task argument "+ + "(func/chan/unsafe-pointer values cannot be decoded)", + fnName, index, in, + ) + } + return paramPlan{kind: paramData, typ: in, index: index}, nil +} + +// structParamType reports whether in (after dereferencing one pointer level, +// matching checkValueType's convention) is a struct, returning that struct type +// or nil. +func structParamType(in reflect.Type) reflect.Type { + t := in + if t.Kind() == reflect.Pointer { + t = t.Elem() + } + if t.Kind() != reflect.Struct { + return nil + } + return t +} + +// hasArgTag reports whether structType has an exported field carrying a +// non-empty `arg:` tag -- the signal that its author means for it to bind +// TaskFlow arguments by name. +func hasArgTag(structType reflect.Type) bool { + for i := range structType.NumField() { + f := structType.Field(i) + if f.IsExported() && f.Tag.Get("arg") != "" { + return true + } + } + return false +} + +// buildStructFields validates and precomputes the by-name field-binding plan +// for a sole struct parameter. It runs once at registration time so a +// misconfigured struct fails loudly before any task ever executes. +func buildStructFields( + fnName string, + structType reflect.Type, + paramIndex int, +) ([]structField, error) { + var fields []structField + seenArgNames := make(map[string]string) // resolved arg name -> Go field name that claims it + + for i := range structType.NumField() { + f := structType.Field(i) + if !f.IsExported() { + continue + } + + if !isDecodableType(f.Type) { + return nil, fmt.Errorf( + "task function %s: parameter %d: struct field %s: type %s cannot receive a task "+ + "argument (func/chan/unsafe-pointer values cannot be decoded)", + fnName, + paramIndex, + f.Name, + f.Type, + ) + } + + sf := structField{structIndex: i, goName: f.Name, fieldType: f.Type} + sf.argName = f.Tag.Get("arg") + if sf.argName == "" { + sf.argName = f.Name + } + if existing, ok := seenArgNames[sf.argName]; ok { + return nil, fmt.Errorf( + "task function %s: parameter %d: struct fields %s and %s both bind arg name %q", + fnName, paramIndex, existing, f.Name, sf.argName, + ) + } + seenArgNames[sf.argName] = f.Name + fields = append(fields, sf) + } + return fields, nil +} + +// checkValueType verifies the Dag-declared value schema can bind to the Go +// parameter type. The schema is an open-vocabulary JSON-schema fragment; only +// its "type" keyword is inspected, and one pointer level is dereferenced first. +// A nil schema, a schema whose "type" is absent or not a plain string (e.g. a +// union's ["string","null"]), an unrecognized type, or an `any` parameter all +// skip the check -- the strict decode still fails loudly on unusable values. +func checkValueType(schema *genmodels.ArgValueSchema, target reflect.Type) error { + if schema == nil { + return nil + } + jsonType, ok := (*schema)["type"].(string) + if !ok { + return nil + } + t := target + if t.Kind() == reflect.Pointer { + t = t.Elem() + } + if t.Kind() == reflect.Interface { + return nil + } + bindable := false + switch jsonType { + case "string": + bindable = t.Kind() == reflect.String + case "integer": + switch t.Kind() { + case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, + reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64: + bindable = true + } + case "number": + bindable = t.Kind() == reflect.Float32 || t.Kind() == reflect.Float64 + case "boolean": + bindable = t.Kind() == reflect.Bool + case "object": + bindable = t.Kind() == reflect.Struct || t.Kind() == reflect.Map + case "array": + bindable = t.Kind() == reflect.Slice || t.Kind() == reflect.Array + default: + // A type keyword this runtime does not recognize; JSON schema is + // open-vocabulary, so leave it to the strict decode. + return nil + } + if !bindable { + return fmt.Errorf( + "the Dag declares JSON-schema type %q which cannot bind to Go parameter type %s", + jsonType, target, + ) + } + return nil +} + +// decodeValue decodes a raw (generically deserialised) value into target. +// Decoding into a struct is strict: unknown/renamed keys fail rather than +// silently leaving fields zero. Decoding into a map / interface accepts any +// shape, so authors opt into loose decoding by typing the parameter +// map[string]any or any. A null value is allowed only for a nilable target. +func decodeValue(raw any, target reflect.Type) (reflect.Value, error) { + out := reflect.New(target) + + if raw == nil { + switch target.Kind() { + case reflect.Pointer, reflect.Slice, reflect.Map, reflect.Interface: + return out.Elem(), nil + default: + return reflect.Value{}, fmt.Errorf( + "value is null but the parameter type %s is not nilable", target, + ) + } + } + + blob, err := json.Marshal(raw) + if err != nil { + return reflect.Value{}, err + } + dec := json.NewDecoder(bytes.NewReader(blob)) + dec.DisallowUnknownFields() + if err := dec.Decode(out.Interface()); err != nil { + return reflect.Value{}, err + } + return out.Elem(), nil +} + +// isDecodableType reports whether a value can be JSON-decoded into inType. It +// rejects kinds json cannot target (func, chan, unsafe pointer) and non-empty +// interfaces (only the empty interface `any` is a valid decode target). +func isDecodableType(inType reflect.Type) bool { + switch inType.Kind() { + case reflect.Func, reflect.Chan, reflect.UnsafePointer: + return false + case reflect.Interface: + return inType.NumMethod() == 0 + } + return true +} + +var ( + contextType = reflect.TypeFor[context.Context]() + tiRunContextType = reflect.TypeFor[sdk.TIRunContext]() + slogLoggerType = reflect.TypeFor[*slog.Logger]() + clientType = reflect.TypeFor[sdk.Client]() +) + +func isContext(inType reflect.Type) bool { + return inType != nil && inType.Implements(contextType) +} + +func isTIRunContext(inType reflect.Type) bool { + return inType == tiRunContextType +} + +func isLogger(inType reflect.Type) bool { + return inType != nil && inType.AssignableTo(slogLoggerType) +} + +// isClient reports whether inType's method set is a subset of sdk.Client's, +// keeping new client capabilities injectable without a hand-kept list. +func isClient(inType reflect.Type) bool { + return inType != nil && inType.Kind() == reflect.Interface && + inType.NumMethod() > 0 && clientType.Implements(inType) +} + +// explainClientMismatch returns why in is not a subset of sdk.Client. +func explainClientMismatch(in reflect.Type) string { + if in.NumMethod() == 0 { + return "empty interfaces cannot be injected" + } + for i := range in.NumMethod() { + m := in.Method(i) + cm, ok := clientType.MethodByName(m.Name) + if !ok { + return fmt.Sprintf("sdk.Client has no method %s", m.Name) + } + if cm.Type != m.Type { + return fmt.Sprintf("method %s is %s on sdk.Client, not %s", m.Name, cm.Type, m.Type) + } + } + return "its method set is not a subset of sdk.Client" +} diff --git a/go-sdk/pkg/binding/binding_test.go b/go-sdk/pkg/binding/binding_test.go new file mode 100644 index 0000000000000..b33857cb48ff0 --- /dev/null +++ b/go-sdk/pkg/binding/binding_test.go @@ -0,0 +1,689 @@ +// 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. + +package binding + +import ( + "context" + "log/slog" + "reflect" + "sync" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/suite" + + "github.com/apache/airflow/go-sdk/pkg/api" + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" + "github.com/apache/airflow/go-sdk/pkg/sdkcontext" + "github.com/apache/airflow/go-sdk/sdk" +) + +type BindingSuite struct { + suite.Suite +} + +func TestBindingSuite(t *testing.T) { + suite.Run(t, &BindingSuite{}) +} + +// argSchema builds a minimal JSON-schema fragment declaring only a "type", +// matching what the Python stub decorator emits for a scalar/container +// annotation. +func argSchema(jsonType string) *genmodels.ArgValueSchema { + s := genmodels.ArgValueSchema{"type": jsonType} + return &s +} + +// fakeXComClient records GetXCom calls and returns preconfigured values. +// Resolve pulls XComs concurrently, so recording is mutex-guarded. +type fakeXComClient struct { + sdk.Client + + values map[string]any // "/" -> raw value + mu sync.Mutex + calls []fakeXComCall + err error +} + +type fakeXComCall struct { + dagID, runID, taskID, key string + mapIndex *int +} + +func (f *fakeXComClient) GetXCom( + ctx context.Context, + dagID, runID, taskID string, + mapIndex *int, + key string, + _ any, +) (any, error) { + f.mu.Lock() + f.calls = append(f.calls, fakeXComCall{dagID, runID, taskID, key, mapIndex}) + f.mu.Unlock() + if f.err != nil { + return nil, f.err + } + return f.values[taskID+"/"+key], nil +} + +// workloadCtx returns a context carrying an ExecuteTaskWorkload the resolver +// reads the Dag/run identifiers from. +func workloadCtx() context.Context { + return context.WithValue( + context.Background(), + sdkcontext.WorkloadContextKey, + api.ExecuteTaskWorkload{ + TI: api.TaskInstance{ + Id: uuid.New(), + DagId: "dag1", + RunId: "run1", + TaskId: "transform", + }, + }, + ) +} + +func analyze(s *BindingSuite, fn any) *Plan { + plan, err := Analyze(reflect.TypeOf(fn), "testFn") + s.Require().NoError(err) + return plan +} + +func (s *BindingSuite) resolve(fn any, args []Arg, client sdk.Client) ([]reflect.Value, error) { + plan := analyze(s, fn) + return plan.Resolve(workloadCtx(), slog.Default(), client, args) +} + +func (s *BindingSuite) TestAnalyzeClassification() { + plan := analyze( + s, + func(ctx sdk.TIRunContext, log *slog.Logger, c sdk.VariableClient, country string, extracted map[string]any) error { + return nil + }, + ) + s.Equal(2, plan.numData) + + s.Zero(analyze(s, func() error { return nil }).numData) + s.Equal( + 1, + analyze(s, func(x any) error { return nil }).numData, + "an `any` parameter is a data parameter", + ) +} + +func (s *BindingSuite) TestAnalyzeRejections() { + cases := map[string]struct { + fn any + errContains string + }{ + "func-param": { + func(cb func()) error { return nil }, + "cannot receive a task argument", + }, + "chan-param": { + func(ch chan int) error { return nil }, + "cannot receive a task argument", + }, + "non-client-interface": { + func(x interface{ NotAClientMethod() }) error { return nil }, + "sdk.Client has no method NotAClientMethod", + }, + "context-with-extra-methods": { + func(x interface { + context.Context + TaskInstance() sdk.TaskInstance + }, + ) error { + return nil + }, + "adds methods on top of context.Context", + }, + } + for name, tt := range cases { + s.Run(name, func() { + _, err := Analyze(reflect.TypeOf(tt.fn), "testFn") + if s.Assert().Error(err) { + s.Assert().Contains(err.Error(), tt.errContains) + } + }) + } +} + +// TestNamedClientInterfacesAreInjectable guards against sdk.Client dropping an +// embedded interface, which would break tasks declaring it. +func (s *BindingSuite) TestNamedClientInterfacesAreInjectable() { + for name, typ := range map[string]reflect.Type{ + "Client": reflect.TypeFor[sdk.Client](), + "VariableClient": reflect.TypeFor[sdk.VariableClient](), + "ConnectionClient": reflect.TypeFor[sdk.ConnectionClient](), + "XComClient": reflect.TypeFor[sdk.XComClient](), + } { + s.True(isClient(typ), "sdk.%s must stay injectable", name) + } +} + +func (s *BindingSuite) TestResolveArityMismatch() { + fn := func(country string) error { return nil } + _, err := s.resolve(fn, nil, &fakeXComClient{}) + if s.Assert().Error(err) { + s.Contains(err.Error(), "argument count mismatch") + s.Contains(err.Error(), "passes 0 positional argument(s)") + s.Contains(err.Error(), "declares 1 data parameter(s)") + } + + _, err = s.resolve( + func() error { return nil }, + []Arg{LiteralArg{Value: "uk"}}, + &fakeXComClient{}, + ) + if s.Assert().Error(err) { + s.Contains(err.Error(), "argument count mismatch") + } +} + +func (s *BindingSuite) TestResolveLiterals() { + fn := func(country string, count int, ratio float64, on bool, tags []string, meta map[string]any) error { + return nil + } + got, err := s.resolve(fn, []Arg{ + LiteralArg{Value: "uk", ValueSchema: argSchema("string")}, + LiteralArg{Value: 3, ValueSchema: argSchema("integer")}, + LiteralArg{Value: 1.5, ValueSchema: argSchema("number")}, + LiteralArg{Value: true, ValueSchema: argSchema("boolean")}, + LiteralArg{Value: []any{"a", "b"}, ValueSchema: argSchema("array")}, + LiteralArg{Value: map[string]any{"k": "v"}, ValueSchema: argSchema("object")}, + }, &fakeXComClient{}) + s.Require().NoError(err) + s.Equal("uk", got[0].Interface()) + s.Equal(3, got[1].Interface()) + s.Equal(1.5, got[2].Interface()) + s.Equal(true, got[3].Interface()) + s.Equal([]string{"a", "b"}, got[4].Interface()) + s.Equal(map[string]any{"k": "v"}, got[5].Interface()) +} + +func (s *BindingSuite) TestResolveInterleavedInjectables() { + fn := func(log *slog.Logger, country string, ctx context.Context, meta map[string]any) error { + return nil + } + got, err := s.resolve(fn, []Arg{ + LiteralArg{Value: "uk", ValueSchema: argSchema("string")}, + LiteralArg{Value: map[string]any{"k": "v"}, ValueSchema: argSchema("object")}, + }, &fakeXComClient{}) + s.Require().NoError(err) + s.NotNil(got[0].Interface().(*slog.Logger)) + s.Equal("uk", got[1].Interface()) + s.NotNil(got[2].Interface().(context.Context)) + s.Equal(map[string]any{"k": "v"}, got[3].Interface()) +} + +func (s *BindingSuite) TestCheckValueTypeMatrix() { + unionType := &genmodels.ArgValueSchema{"type": []any{"string", "null"}} + cases := map[string]struct { + schema *genmodels.ArgValueSchema + target reflect.Type + errContains string + }{ + "string-ok": {argSchema("string"), reflect.TypeFor[string](), ""}, + "string-ptr-ok": {argSchema("string"), reflect.TypeFor[*string](), ""}, + "string-vs-int": {argSchema("string"), reflect.TypeFor[int](), "cannot bind"}, + "integer-ok": {argSchema("integer"), reflect.TypeFor[int64](), ""}, + "integer-uint-ok": {argSchema("integer"), reflect.TypeFor[uint32](), ""}, + "integer-vs-float": {argSchema("integer"), reflect.TypeFor[float64](), "cannot bind"}, + "number-ok": {argSchema("number"), reflect.TypeFor[float32](), ""}, + "number-vs-int": {argSchema("number"), reflect.TypeFor[int](), "cannot bind"}, + "boolean-ok": {argSchema("boolean"), reflect.TypeFor[bool](), ""}, + "boolean-vs-string": {argSchema("boolean"), reflect.TypeFor[string](), "cannot bind"}, + "object-map-ok": {argSchema("object"), reflect.TypeFor[map[string]int](), ""}, + "object-struct-ok": {argSchema("object"), reflect.TypeFor[struct{ A int }](), ""}, + "object-vs-slice": {argSchema("object"), reflect.TypeFor[[]int](), "cannot bind"}, + "array-slice-ok": {argSchema("array"), reflect.TypeFor[[]string](), ""}, + "array-array-ok": {argSchema("array"), reflect.TypeFor[[2]int](), ""}, + "array-vs-map": {argSchema("array"), reflect.TypeFor[map[string]any](), "cannot bind"}, + // A nil schema, a fragment without a plain-string "type" (a union or a + // keyword-less fragment), an unrecognized type, or an `any` target all + // skip the check -- the strict decode still guards the value. + "nil-schema-skips": {nil, reflect.TypeFor[chan int](), ""}, + "no-type-skips": {&genmodels.ArgValueSchema{}, reflect.TypeFor[string](), ""}, + "union-type-skips": {unionType, reflect.TypeFor[int](), ""}, + "unknown-type-skips": {argSchema("uuid"), reflect.TypeFor[string](), ""}, + "any-target-skips": {argSchema("string"), reflect.TypeFor[any](), ""}, + } + for name, tt := range cases { + s.Run(name, func() { + err := checkValueType(tt.schema, tt.target) + if tt.errContains == "" { + s.NoError(err) + } else if s.Assert().Error(err) { + s.Contains(err.Error(), tt.errContains) + } + }) + } +} + +func (s *BindingSuite) TestResolveTypeMismatchFailsLoudly() { + fn := func(count int) error { return nil } + _, err := s.resolve( + fn, + []Arg{LiteralArg{Value: "uk", ValueSchema: argSchema("string")}}, + &fakeXComClient{}, + ) + if s.Assert().Error(err) { + s.Contains( + err.Error(), + `the Dag declares JSON-schema type "string" which cannot bind to Go parameter type int`, + ) + } +} + +func (s *BindingSuite) TestResolveLiteralDecodeFailure() { + // The Dag declared "any", so the type check passes but the JSON decode of a + // string into an int must still fail loudly. + fn := func(count int) error { return nil } + _, err := s.resolve(fn, []Arg{LiteralArg{Value: "uk"}}, &fakeXComClient{}) + if s.Assert().Error(err) { + s.Contains(err.Error(), "decoding literal value into int") + } +} + +type extractResult struct { + GoVersion string `json:"go_version"` + Timestamp int64 `json:"timestamp"` +} + +// simpleInput is the minimal name-bound struct: one field, no tags, so it +// falls back to matching its own field name, verbatim. +type simpleInput struct { + Name string +} + +// twoFieldInput has one field the args always match (Name) and one whose +// arg name is never present in the tests that use it (Missing), to prove an +// unmatched field is left at its Go zero value instead of failing the task. +type twoFieldInput struct { + Name string + Missing string `arg:"missing"` +} + +// wholeConfig carries only `json:` tags, so it names no TaskFlow argument by +// field. As a flat data parameter it is decoded whole; as a sole struct +// parameter it is the whole-value fallback target. +type wholeConfig struct { + Environment string `json:"environment"` + Region string `json:"region"` +} + +// combineInput exercises both field-binding modes side by side: Name falls +// back to its verbatim field name, Count is explicitly named via its `arg:` +// tag. +type combineInput struct { + Name string + Count int `arg:"count"` +} + +// reportInput deliberately declares Ratio before Region, the reverse of the +// wire order those names appear in, to prove field declaration order is +// irrelevant to by-name claiming. +type reportInput struct { + Ratio float64 + Region string `arg:"region"` +} + +func (s *BindingSuite) TestResolveXComArgs() { + client := &fakeXComClient{values: map[string]any{ + "extract/return_value": map[string]any{"go_version": "go1.24", "timestamp": int64(42)}, + "probe/return_value": "probe-value", + }} + + fn := func(res extractResult, probe string) error { return nil } + got, err := s.resolve(fn, []Arg{ + XComArg{TaskID: "extract", ValueSchema: argSchema("object")}, + XComArg{TaskID: "probe", ValueSchema: argSchema("string")}, + }, client) + s.Require().NoError(err) + s.Equal(extractResult{GoVersion: "go1.24", Timestamp: 42}, got[0].Interface()) + s.Equal("probe-value", got[1].Interface()) + + s.Require().Len(client.calls, 2) + taskIDs := make([]string, 0, 2) + for _, call := range client.calls { + // Pulls run concurrently, so assert per-call properties order-independently. + taskIDs = append(taskIDs, call.taskID) + s.Equal("dag1", call.dagID) + s.Equal("run1", call.runID) + s.Equal( + api.XComReturnValueKey, + call.key, + "an XCom argument always pulls the return-value key", + ) + s.Nil(call.mapIndex, "v1 always pulls the unmapped upstream instance") + } + s.ElementsMatch([]string{"extract", "probe"}, taskIDs) +} + +func (s *BindingSuite) TestResolveXComStrictStructDecode() { + client := &fakeXComClient{values: map[string]any{ + "extract/return_value": map[string]any{"go_version": "go1.24", "renamed_field": 1}, + }} + fn := func(res extractResult) error { return nil } + _, err := s.resolve(fn, []Arg{XComArg{TaskID: "extract"}}, client) + if s.Assert().Error(err) { + s.Contains(err.Error(), `decoding xcom from task "extract"`) + s.Contains(err.Error(), "unknown field") + } +} + +func (s *BindingSuite) TestResolveXComPullFailure() { + client := &fakeXComClient{err: sdk.XComNotFound} + fn := func(res map[string]any) error { return nil } + _, err := s.resolve(fn, []Arg{XComArg{TaskID: "extract"}}, client) + if s.Assert().Error(err) { + s.Contains(err.Error(), `pulling xcom from task "extract"`) + } +} + +func (s *BindingSuite) TestResolveXComWithoutWorkload() { + plan := analyze(s, func(res map[string]any) error { return nil }) + _, err := plan.Resolve( + context.Background(), slog.Default(), &fakeXComClient{}, + []Arg{XComArg{TaskID: "extract"}}, + ) + if s.Assert().Error(err) { + s.Contains(err.Error(), "no workload in context") + } +} + +func (s *BindingSuite) TestResolveNullHandling() { + fn := func(meta map[string]any) error { return nil } + got, err := s.resolve( + fn, + []Arg{LiteralArg{Value: nil, ValueSchema: argSchema("object")}}, + &fakeXComClient{}, + ) + s.Require().NoError(err) + s.Nil(got[0].Interface()) + + fnStr := func(country string) error { return nil } + _, err = s.resolve(fnStr, []Arg{LiteralArg{Value: nil}}, &fakeXComClient{}) + if s.Assert().Error(err) { + s.Contains(err.Error(), "not nilable") + } +} + +// fakeArg is an out-of-catalogue Arg variant: the compiler seals the sum type +// to this package, so the defensive default branch can only be reached from +// inside it. +type fakeArg struct{} + +func (fakeArg) ArgName() string { return "fake" } +func (fakeArg) Schema() *genmodels.ArgValueSchema { return nil } +func (fakeArg) sealedArg() {} + +func (s *BindingSuite) TestResolveUnsupportedVariant() { + fn := func(country string) error { return nil } + _, err := s.resolve(fn, []Arg{fakeArg{}}, &fakeXComClient{}) + if s.Assert().Error(err) { + s.Contains(err.Error(), "unsupported argument binding binding.fakeArg") + } +} + +func (s *BindingSuite) TestResolveNilArg() { + fn := func(country string) error { return nil } + _, err := s.resolve(fn, []Arg{nil}, &fakeXComClient{}) + if s.Assert().Error(err) { + s.Contains(err.Error(), "nil argument binding") + } +} + +func (s *BindingSuite) TestResolveTIRunContextRebuild() { + ti := sdk.TaskInstance{DagID: "dag1", RunID: "run1", TaskID: "transform"} + dagRun := sdk.DagRun{DagID: "dag1", RunID: "run1"} + ctx := context.WithValue( + workloadCtx(), + sdkcontext.RuntimeContextKey, + sdk.NewTIRunContext(context.Background(), ti, dagRun), + ) + + plan := analyze(s, func(rc sdk.TIRunContext, country string) error { return nil }) + got, err := plan.Resolve(ctx, slog.Default(), &fakeXComClient{}, []Arg{ + LiteralArg{Value: "uk", ValueSchema: argSchema("string")}, + }) + s.Require().NoError(err) + rc := got[0].Interface().(sdk.TIRunContext) + s.Equal(ti, rc.TaskInstance()) + s.Equal(dagRun, rc.DagRun()) + s.Equal("uk", got[1].Interface()) +} + +func (s *BindingSuite) TestAnalyzeLoneStructClassification() { + plan := analyze(s, func(input simpleInput) error { return nil }) + s.True(plan.loneStruct, "a sole struct data parameter is resolved by name at execution") + s.Zero(plan.numData) + + ptrPlan := analyze(s, func(input *simpleInput) error { return nil }) + s.True(ptrPlan.loneStruct, "a pointer to a sole struct is detected the same way") + s.Zero(ptrPlan.numData) + + flatPlan := analyze(s, func(prefix string, cfg wholeConfig) error { return nil }) + s.False(flatPlan.loneStruct, "a struct alongside another data parameter is a flat slot") + s.Equal(2, flatPlan.numData) + + scalarPlan := analyze(s, func(name string) error { return nil }) + s.False(scalarPlan.loneStruct, "a sole non-struct data parameter is plain positional") + s.Equal(1, scalarPlan.numData) +} + +func (s *BindingSuite) TestAnalyzeMultipleStructsAreFlat() { + // With no marker, the old "only one struct" / "cannot mix" restrictions are + // gone: any struct that is not the sole data parameter is a flat whole-value + // slot, so these signatures are now accepted rather than rejected. + plan := analyze(s, func(a wholeConfig, b wholeConfig) error { return nil }) + s.False(plan.loneStruct) + s.Equal(2, plan.numData) +} + +func (s *BindingSuite) TestAnalyzeStructValidation() { + type duplicateArgNames struct { + A string + B string `arg:"A"` + } + type nonDecodableField struct { + Bad chan int + } + + cases := map[string]struct { + fn any + errContains string + }{ + "duplicate-arg-names": { + func(input duplicateArgNames) error { return nil }, + `fields A and B both bind arg name "A"`, + }, + "non-decodable-field": { + func(input nonDecodableField) error { return nil }, + "cannot receive a task argument", + }, + // A struct carrying `arg:` tags must be the sole data parameter, else its + // tags would be silently ignored (the struct would be decoded whole). + "tagged-struct-not-sole": { + func(prefix string, input combineInput) error { return nil }, + "must be the function's only data parameter", + }, + "tagged-struct-trailing": { + func(input combineInput, suffix string) error { return nil }, + "must be the function's only data parameter", + }, + } + for name, tt := range cases { + s.Run(name, func() { + _, err := Analyze(reflect.TypeOf(tt.fn), "testFn") + if s.Assert().Error(err) { + s.Assert().Contains(err.Error(), tt.errContains) + } + }) + } +} + +func (s *BindingSuite) TestResolveStructAllFields() { + fn := func(input combineInput) error { return nil } + got, err := s.resolve(fn, []Arg{ + LiteralArg{Name: "Name", Value: "widget", ValueSchema: argSchema("string")}, + LiteralArg{Name: "count", Value: 7, ValueSchema: argSchema("integer")}, + }, &fakeXComClient{}) + s.Require().NoError(err) + + input := got[0].Interface().(combineInput) + s.Equal("widget", input.Name, "the untagged field claims its verbatim field name") + s.Equal(7, input.Count, "the `arg:` tag claims its named entry") +} + +func (s *BindingSuite) TestResolveStructXComArg() { + fn := func(log *slog.Logger, input reportInput) error { return nil } + got, err := s.resolve(fn, []Arg{ + XComArg{Name: "region", TaskID: "make_region", ValueSchema: argSchema("string")}, + LiteralArg{Name: "Ratio", Value: 0.5, ValueSchema: argSchema("number")}, + }, &fakeXComClient{values: map[string]any{"make_region/return_value": "east"}}) + s.Require().NoError(err) + + input := got[1].Interface().(reportInput) + s.Equal("east", input.Region, "Region resolves by name despite being declared after Ratio") + s.Equal(0.5, input.Ratio) +} + +func (s *BindingSuite) TestResolveStructSingleClaimedArgBindsByName() { + // A single argument whose name a field claims binds by name -- the tie-break + // that keeps a one-field struct name-bound rather than whole-value decoded. + fn := func(input simpleInput) error { return nil } + got, err := s.resolve(fn, []Arg{ + LiteralArg{Name: "Name", Value: "widget", ValueSchema: argSchema("string")}, + }, &fakeXComClient{}) + s.Require().NoError(err) + s.Equal("widget", got[0].Interface().(simpleInput).Name) +} + +func (s *BindingSuite) TestResolveStructPointer() { + fn := func(input *simpleInput) error { return nil } + got, err := s.resolve(fn, []Arg{ + LiteralArg{Name: "Name", Value: "widget", ValueSchema: argSchema("string")}, + }, &fakeXComClient{}) + s.Require().NoError(err) + input := got[0].Interface().(*simpleInput) + s.Require().NotNil(input) + s.Equal("widget", input.Name) +} + +func (s *BindingSuite) TestResolveLoneStructBothModes() { + // The same one-struct signature resolves either way, chosen per execution + // from the argument spec: field-by-field when the argument names match the + // struct's fields, or whole from a single argument no field claims. + fn := func(cfg wholeConfig) error { return nil } + want := wholeConfig{Environment: "production", Region: "eu-west-1"} + + // Struct-based: two arguments named like the fields bind by name. + named, err := s.resolve(fn, []Arg{ + LiteralArg{Name: "Environment", Value: "production", ValueSchema: argSchema("string")}, + LiteralArg{Name: "Region", Value: "eu-west-1", ValueSchema: argSchema("string")}, + }, &fakeXComClient{}) + s.Require().NoError(err) + s.Equal(want, named[0].Interface(), "argument names matching the fields bind field-by-field") + + // Flat-based: one argument no field claims decodes whole into the struct. + whole, err := s.resolve(fn, []Arg{ + LiteralArg{ + Name: "cfg", + Value: map[string]any{"environment": "production", "region": "eu-west-1"}, + ValueSchema: argSchema("object"), + }, + }, &fakeXComClient{}) + s.Require().NoError(err) + s.Equal(want, whole[0].Interface(), "a single unclaimed argument decodes whole into the struct") +} + +func (s *BindingSuite) TestResolveStructUnclaimedArgFailsLoudly() { + fn := func(input combineInput) error { return nil } + _, err := s.resolve(fn, []Arg{ + LiteralArg{Name: "Name", Value: "widget", ValueSchema: argSchema("string")}, + LiteralArg{Name: "typo", Value: "x", ValueSchema: argSchema("string")}, + }, &fakeXComClient{}) + // Name is claimed, so this is name-binding (not the whole-value fallback); + // the leftover "typo" argument fails the task rather than being dropped. + if s.Assert().Error(err) { + s.Contains(err.Error(), `not claimed by any struct field: "typo"`) + } +} + +func (s *BindingSuite) TestResolveStructUnclaimedFromDefaultAllowed() { + fn := func(input combineInput) error { return nil } + got, err := s.resolve(fn, []Arg{ + LiteralArg{Name: "Name", Value: "widget", ValueSchema: argSchema("string")}, + // The Dag author never passed "threshold"; Python captured it from the + // stub signature's default. The struct need not mirror it. + LiteralArg{ + Name: "threshold", + Value: 0.75, + ValueSchema: argSchema("number"), + FromDefault: true, + }, + }, &fakeXComClient{}) + s.Require().NoError(err) + s.Equal("widget", got[0].Interface().(combineInput).Name) +} + +func (s *BindingSuite) TestResolveStructEmptySpecFailsLoudly() { + fn := func(input simpleInput) error { return nil } + for name, args := range map[string][]Arg{"nil-spec": nil, "empty-spec": {}} { + s.Run(name, func() { + _, err := s.resolve(fn, args, &fakeXComClient{}) + // The Edge Worker path delivers no arg bindings; a struct with + // bindable fields must fail rather than run fully zero-valued. + if s.Assert().Error(err) { + s.Contains(err.Error(), "no TaskFlow arg bindings arrived") + } + }) + } +} + +func (s *BindingSuite) TestResolveStructOnlyDefaultsZeroValues() { + // A lone from_default argument that no field claims is neither whole-value + // decoded (defaults were never explicitly passed) nor an error; the fields + // keep their kwarg-style zero values. + fn := func(input twoFieldInput) error { return nil } + got, err := s.resolve(fn, []Arg{ + LiteralArg{ + Name: "threshold", + Value: 0.75, + ValueSchema: argSchema("number"), + FromDefault: true, + }, + }, &fakeXComClient{}) + s.Require().NoError(err) + input := got[0].Interface().(twoFieldInput) + s.Equal("", input.Name, "no explicit entry arrived; fields keep kwarg-style zero values") + s.Equal("", input.Missing) +} + +func (s *BindingSuite) TestResolveStructUnmatchedFieldZeroValued() { + fn := func(input twoFieldInput) error { return nil } + got, err := s.resolve(fn, []Arg{ + LiteralArg{Name: "Name", Value: "widget", ValueSchema: argSchema("string")}, + }, &fakeXComClient{}) + s.Require().NoError(err) + input := got[0].Interface().(twoFieldInput) + s.Equal("widget", input.Name, "the matched field binds normally") + s.Equal("", input.Missing, "the unmatched field is left at its Go zero value, not an error") +} diff --git a/go-sdk/pkg/execution/frames.go b/go-sdk/pkg/execution/frames.go index f9a246286efce..947346316c57c 100644 --- a/go-sdk/pkg/execution/frames.go +++ b/go-sdk/pkg/execution/frames.go @@ -62,6 +62,13 @@ func encodeRequest(id int64, body any) ([]byte, error) { var buf bytes.Buffer enc := msgpack.NewEncoder(&buf) enc.UseCompactInts(true) + // Honour `json` struct tags when encoding user-provided values (XCom and + // Variable payloads). Without this, msgpack uses Go field names, so a typed + // XCom pushed as a struct would cross the wire as e.g. "GoVersion" and fail + // to decode into the json-tagged "go_version" the value is read back with + // (and that the HTTP-backed client uses). `msgpack` tags still win where + // present, so the genmodels protocol frames are unaffected. + enc.SetCustomStructTag("json") if err := enc.EncodeArrayLen(2); err != nil { return nil, err diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index e6861d8c8add5..7697829481779 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -20,6 +20,10 @@ package genmodels import "time" +type ArgBindings []TaskArgBinding + +type ArgValueSchema map[string]JsonValue + // Schema for AssetAliasModel used in AssetEventDagRunReference. type AssetAliasReferenceAssetEventDagRun struct { // Name corresponds to the JSON schema field "name". @@ -370,6 +374,9 @@ type DagCallbackRequest struct { // Type corresponds to the JSON schema field "type". Type string `msgpack:"type,omitempty"` + + // VersionData corresponds to the JSON schema field "version_data". + VersionData *VersionData `msgpack:"version_data,omitempty"` } // Request for DAG File Parsing. @@ -749,6 +756,9 @@ type EmailRequest struct { // Type corresponds to the JSON schema field "type". Type string `msgpack:"type,omitempty"` + + // VersionData corresponds to the JSON schema field "version_data". + VersionData *VersionData `msgpack:"version_data,omitempty"` } type EmailRequestEmailType string @@ -808,12 +818,22 @@ type GetAssetEventByAsset struct { // Before corresponds to the JSON schema field "before". Before interface{} `msgpack:"before,omitempty"` + // Extra corresponds to the JSON schema field "extra". + Extra *Extra `msgpack:"extra,omitempty"` + // Limit corresponds to the JSON schema field "limit". Limit interface{} `msgpack:"limit,omitempty"` // Name corresponds to the JSON schema field "name". Name interface{} `msgpack:"name"` + // PartitionKey corresponds to the JSON schema field "partition_key". + PartitionKey interface{} `msgpack:"partition_key,omitempty"` + + // PartitionKeyRegexpPattern corresponds to the JSON schema field + // "partition_key_regexp_pattern". + PartitionKeyRegexpPattern interface{} `msgpack:"partition_key_regexp_pattern,omitempty"` + // Type corresponds to the JSON schema field "type". Type string `msgpack:"type,omitempty"` @@ -834,9 +854,19 @@ type GetAssetEventByAssetAlias struct { // Before corresponds to the JSON schema field "before". Before interface{} `msgpack:"before,omitempty"` + // Extra corresponds to the JSON schema field "extra". + Extra *Extra `msgpack:"extra,omitempty"` + // Limit corresponds to the JSON schema field "limit". Limit interface{} `msgpack:"limit,omitempty"` + // PartitionKey corresponds to the JSON schema field "partition_key". + PartitionKey interface{} `msgpack:"partition_key,omitempty"` + + // PartitionKeyRegexpPattern corresponds to the JSON schema field + // "partition_key_regexp_pattern". + PartitionKeyRegexpPattern interface{} `msgpack:"partition_key_regexp_pattern,omitempty"` + // Type corresponds to the JSON schema field "type". Type string `msgpack:"type,omitempty"` } @@ -1239,6 +1269,24 @@ type LazyDeserializedDAG struct { LastLoaded interface{} `msgpack:"last_loaded,omitempty"` } +// One positional stub-task argument carrying an inline literal from the Dag file. +type LiteralArgBinding struct { + // FromDefault corresponds to the JSON schema field "from_default". + FromDefault bool `msgpack:"from_default,omitempty"` + + // Kind corresponds to the JSON schema field "kind". + Kind string `msgpack:"kind"` + + // Name corresponds to the JSON schema field "name". + Name string `msgpack:"name"` + + // Value corresponds to the JSON schema field "value". + Value interface{} `msgpack:"value,omitempty"` + + // ValueSchema corresponds to the JSON schema field "value_schema". + ValueSchema *ArgValueSchema `msgpack:"value_schema,omitempty"` +} + type LogicalDates []time.Time // Add a new value to be redacted in task logs. @@ -1564,6 +1612,9 @@ type TICount struct { // Response schema for TaskInstance run context. type TIRunContext struct { + // ArgBindings corresponds to the JSON schema field "arg_bindings". + ArgBindings *ArgBindings `msgpack:"arg_bindings,omitempty"` + // Connections corresponds to the JSON schema field "connections". Connections []ConnectionResponse `msgpack:"connections,omitempty"` @@ -1596,6 +1647,8 @@ type TIRunContext struct { XcomKeysToClear []string `msgpack:"xcom_keys_to_clear,omitempty"` } +type TaskArgBinding interface{} + type TaskBreadcrumbsResult struct { // Breadcrumbs corresponds to the JSON schema field "breadcrumbs". Breadcrumbs []TaskBreadcrumbsResultBreadcrumbsElem `msgpack:"breadcrumbs"` @@ -1636,6 +1689,9 @@ type TaskCallbackRequest struct { // Type corresponds to the JSON schema field "type". Type string `msgpack:"type,omitempty"` + + // VersionData corresponds to the JSON schema field "version_data". + VersionData *VersionData `msgpack:"version_data,omitempty"` } type TaskIds []string @@ -1777,19 +1833,6 @@ type TriggerDagRun struct { type TriggerKwargs map[string]JsonValue -type Warnings []interface{} - -// Variable schema for responses with fields that are needed for Runtime. -type VariableResponse struct { - // Key corresponds to the JSON schema field "key". - Key string `msgpack:"key"` - - // Value corresponds to the JSON schema field "value". - Value interface{} `msgpack:"value"` -} - -type VersionData map[string]interface{} - // Update the response content part of an existing Human-in-the-loop response. type UpdateHITLDetail struct { // ChosenOptions corresponds to the JSON schema field "chosen_options". @@ -1805,6 +1848,17 @@ type UpdateHITLDetail struct { Type string `msgpack:"type,omitempty"` } +// Variable schema for responses with fields that are needed for Runtime. +type VariableResponse struct { + // Key corresponds to the JSON schema field "key". + Key string `msgpack:"key"` + + // Value corresponds to the JSON schema field "value". + Value interface{} `msgpack:"value"` +} + +type VersionData map[string]interface{} + type ValidateInletsAndOutlets struct { // TIID corresponds to the JSON schema field "ti_id". TIID string `msgpack:"ti_id"` @@ -1824,6 +1878,8 @@ type VariableKeysResult struct { Type string `msgpack:"type,omitempty"` } +type Warnings []interface{} + type VariableResult struct { // Key corresponds to the JSON schema field "key". Key string `msgpack:"key"` @@ -1835,6 +1891,21 @@ type VariableResult struct { Value interface{} `msgpack:"value,omitempty"` } +// One positional stub-task argument pulled from an upstream task's XCom. +type XComArgBinding struct { + // Kind corresponds to the JSON schema field "kind". + Kind string `msgpack:"kind"` + + // Name corresponds to the JSON schema field "name". + Name string `msgpack:"name"` + + // TaskID corresponds to the JSON schema field "task_id". + TaskID string `msgpack:"task_id"` + + // ValueSchema corresponds to the JSON schema field "value_schema". + ValueSchema *ArgValueSchema `msgpack:"value_schema,omitempty"` +} + type XComCountResponse struct { // Len corresponds to the JSON schema field "len". Len int `msgpack:"len"` diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 2559359453866..105d692f0876c 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -230,6 +230,387 @@ func TestTaskRunnerPanicRetry(t *testing.T) { assertRetryTask(t, result, "panic: something went wrong") } +// TestTaskRunnerBindsArgs covers the TaskFlow path through RunTask: the +// positional-argument spec in ti_context.arg_bindings binds literals onto the +// task function's data parameters. +func TestTaskRunnerBindsArgs(t *testing.T) { + var gotCountry string + var gotMeta map[string]any + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTaskWithName("transform", + func(log *slog.Logger, country string, meta map[string]any) error { + gotCountry = country + gotMeta = meta + return nil + }) + }) + + details := &genmodels.StartupDetails{ + TI: genmodels.TaskInstance{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "transform", + RunID: "run1", + MapIndex: ptr(-1), + }, + BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, + TIContext: genmodels.TIRunContext{ + ArgBindings: &genmodels.ArgBindings{ + map[string]any{ + "name": "country", + "kind": "literal", + "value_schema": map[string]any{"type": "string"}, + "value": "uk", + }, + map[string]any{ + "name": "meta", + "kind": "literal", + "value_schema": map[string]any{"type": "object"}, + "value": map[string]any{"k": "v"}, + }, + }, + }, + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(context.Background(), bundle, details, comm, logger) + assertSucceedTask(t, result) + assert.Equal(t, "uk", gotCountry) + assert.Equal(t, map[string]any{"k": "v"}, gotMeta) +} + +// TestTaskRunnerArgBindingsArityMismatch: an argument spec that does not match +// the function's data parameters fails the task loudly instead of running it +// with zero values. +func TestTaskRunnerArgBindingsArityMismatch(t *testing.T) { + ran := false + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTaskWithName("transform", + func(country string, meta map[string]any) error { + ran = true + return nil + }) + }) + + details := &genmodels.StartupDetails{ + TI: genmodels.TaskInstance{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "transform", + RunID: "run1", + MapIndex: ptr(-1), + }, + BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, + TIContext: genmodels.TIRunContext{ + ArgBindings: &genmodels.ArgBindings{ + map[string]any{ + "name": "country", + "kind": "literal", + "value_schema": map[string]any{"type": "string"}, + "value": "uk", + }, + }, + }, + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(context.Background(), bundle, details, comm, logger) + assertTaskState(t, result, genmodels.TaskStateStateFailed) + assert.False(t, ran, "the task body must not run on an arity mismatch") +} + +// combineInput is a sole struct parameter whose field claims a named entry +// out of ti_context.arg_bindings. +type combineInput struct { + Region string `arg:"region"` +} + +// TestTaskRunnerBindsStructArgs covers the TaskFlow path through RunTask for a +// name-bound struct parameter: convertArgBindings must propagate each spec's +// Name through to binding.Arg so the struct's `arg:"region"` field can claim it +// by name. +func TestTaskRunnerBindsStructArgs(t *testing.T) { + var got combineInput + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTaskWithName("transform", + func(input combineInput) error { + got = input + return nil + }) + }) + + details := &genmodels.StartupDetails{ + TI: genmodels.TaskInstance{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "transform", + RunID: "run1", + MapIndex: ptr(-1), + }, + BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, + TIContext: genmodels.TIRunContext{ + ArgBindings: &genmodels.ArgBindings{ + map[string]any{ + "name": "region", + "kind": "literal", + "value_schema": map[string]any{"type": "string"}, + "value": "eu-west-1", + }, + }, + }, + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(context.Background(), bundle, details, comm, logger) + assertSucceedTask(t, result) + assert.Equal(t, "eu-west-1", got.Region) +} + +// TestTaskRunnerStructIgnoresUnclaimedDefault: convertArgBindings must +// propagate from_default so a spec entry the Python side filled from the stub +// signature's default may go unclaimed by the name-bound struct. +func TestTaskRunnerStructIgnoresUnclaimedDefault(t *testing.T) { + var got combineInput + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTaskWithName("transform", + func(input combineInput) error { + got = input + return nil + }) + }) + + details := &genmodels.StartupDetails{ + TI: genmodels.TaskInstance{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "transform", + RunID: "run1", + MapIndex: ptr(-1), + }, + BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, + TIContext: genmodels.TIRunContext{ + ArgBindings: &genmodels.ArgBindings{ + map[string]any{ + "name": "region", + "kind": "literal", + "value_schema": map[string]any{"type": "string"}, + "value": "eu-west-1", + }, + map[string]any{ + "name": "threshold", + "kind": "literal", + "value_schema": map[string]any{"type": "number"}, + "value": 0.75, + "from_default": true, + }, + }, + }, + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(context.Background(), bundle, details, comm, logger) + assertSucceedTask(t, result) + assert.Equal(t, "eu-west-1", got.Region) +} + +// TestTaskRunnerArgBindingsTypeMismatch: a declared Dag type that cannot bind to +// the Go parameter type fails the task loudly before the body runs. +func TestTaskRunnerArgBindingsTypeMismatch(t *testing.T) { + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTaskWithName("transform", + func(count int) error { return nil }) + }) + + details := &genmodels.StartupDetails{ + TI: genmodels.TaskInstance{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "transform", + RunID: "run1", + MapIndex: ptr(-1), + }, + BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, + TIContext: genmodels.TIRunContext{ + ArgBindings: &genmodels.ArgBindings{ + map[string]any{ + "name": "count", + "kind": "literal", + "value_schema": map[string]any{"type": "string"}, + "value": "uk", + }, + }, + }, + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(context.Background(), bundle, details, comm, logger) + assertTaskState(t, result, genmodels.TaskStateStateFailed) +} + +// TestTaskRunnerArgBindingsUnknownKind: a wire spec whose kind is neither xcom +// nor literal fails the task before the body runs. +func TestTaskRunnerArgBindingsUnknownKind(t *testing.T) { + ran := false + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTaskWithName("transform", + func(country string) error { + ran = true + return nil + }) + }) + + details := &genmodels.StartupDetails{ + TI: genmodels.TaskInstance{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "transform", + RunID: "run1", + MapIndex: ptr(-1), + }, + BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, + TIContext: genmodels.TIRunContext{ + ArgBindings: &genmodels.ArgBindings{ + map[string]any{"name": "country", "kind": "template", "value": "x"}, + }, + }, + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(context.Background(), bundle, details, comm, logger) + assertTaskState(t, result, genmodels.TaskStateStateFailed) + assert.False(t, ran, "the task body must not run on an unknown binding kind") +} + +// TestTaskRunnerArgBindingsMalformedElement: a wire spec element that is not a +// map at all fails the task before the body runs. +func TestTaskRunnerArgBindingsMalformedElement(t *testing.T) { + ran := false + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTaskWithName("transform", + func(country string) error { + ran = true + return nil + }) + }) + + details := &genmodels.StartupDetails{ + TI: genmodels.TaskInstance{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "transform", + RunID: "run1", + MapIndex: ptr(-1), + }, + BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, + TIContext: genmodels.TIRunContext{ + ArgBindings: &genmodels.ArgBindings{"bogus"}, + }, + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(context.Background(), bundle, details, comm, logger) + assertTaskState(t, result, genmodels.TaskStateStateFailed) + assert.False(t, ran, "the task body must not run on a malformed binding element") +} + +// TestTaskRunnerArgBindingsMissingRequiredFields: a wire spec entry without a +// usable name, or an xcom entry without a task_id, fails the task before the +// body runs instead of silently binding empty strings. +func TestTaskRunnerArgBindingsMissingRequiredFields(t *testing.T) { + cases := []struct { + name string + spec map[string]any + }{ + {name: "missing name", spec: map[string]any{"kind": "literal", "value": "x"}}, + {name: "empty name", spec: map[string]any{"name": "", "kind": "literal", "value": "x"}}, + {name: "xcom missing task_id", spec: map[string]any{"name": "country", "kind": "xcom"}}, + { + name: "xcom empty task_id", + spec: map[string]any{"name": "country", "kind": "xcom", "task_id": ""}, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + ran := false + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTaskWithName("transform", + func(country string) error { + ran = true + return nil + }) + }) + + details := &genmodels.StartupDetails{ + TI: genmodels.TaskInstance{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "transform", + RunID: "run1", + MapIndex: ptr(-1), + }, + BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, + TIContext: genmodels.TIRunContext{ + ArgBindings: &genmodels.ArgBindings{tc.spec}, + }, + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(context.Background(), bundle, details, comm, logger) + assertTaskState(t, result, genmodels.TaskStateStateFailed) + assert.False(t, ran, "the task body must not run on an incomplete binding spec") + }) + } +} + +// TestTaskRunnerMalformedSpecHonorsShouldRetry: a spec that fails +// convertArgBindings terminates with the same retry semantics as a binding +// failure inside executeTask, not an unconditional FAILED. +func TestTaskRunnerMalformedSpecHonorsShouldRetry(t *testing.T) { + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTaskWithName("transform", + func(country string) error { return nil }) + }) + + details := &genmodels.StartupDetails{ + TI: genmodels.TaskInstance{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "transform", + RunID: "run1", + MapIndex: ptr(-1), + }, + BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, + TIContext: genmodels.TIRunContext{ + ShouldRetry: true, + ArgBindings: &genmodels.ArgBindings{ + map[string]any{"name": "country", "kind": "template", "value": "x"}, + }, + }, + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(context.Background(), bundle, details, comm, logger) + assertRetryTask(t, result, `unknown kind "template"`) +} + func TestRunTaskHonorsContextCancellation(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { r.AddDag("test_dag").AddTaskWithName("ctxcheck", diff --git a/go-sdk/pkg/execution/messages.go b/go-sdk/pkg/execution/messages.go index 72d451a866632..bb81d60c0a4ff 100644 --- a/go-sdk/pkg/execution/messages.go +++ b/go-sdk/pkg/execution/messages.go @@ -32,7 +32,7 @@ import ( // reported in a bundle's airflow-metadata manifest as // sdk.supervisor_schema_version so the supervisor can down/upgrade messages to // a shape the bundle understands. -const SupervisorSchemaVersion = "2026-06-16" +const SupervisorSchemaVersion = "2026-10-30" // The message-type discriminator strings (genmodels.Type*) are generated from the // schema's "type" consts in discriminators.gen.go; outbound messages stamp the diff --git a/go-sdk/pkg/execution/task_runner.go b/go-sdk/pkg/execution/task_runner.go index c656b4cfd71f2..6cf48c8c63b97 100644 --- a/go-sdk/pkg/execution/task_runner.go +++ b/go-sdk/pkg/execution/task_runner.go @@ -28,6 +28,7 @@ import ( "github.com/apache/airflow/go-sdk/bundle/bundlev1" "github.com/apache/airflow/go-sdk/pkg/api" + "github.com/apache/airflow/go-sdk/pkg/binding" "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" "github.com/apache/airflow/go-sdk/pkg/sdkcontext" "github.com/apache/airflow/go-sdk/sdk" @@ -124,7 +125,96 @@ func RunTask( ctx = context.WithValue(ctx, sdkcontext.SdkClientContextKey, sdk.Client(client)) ctx = context.WithValue(ctx, sdkcontext.RuntimeContextKey, runtimeContext) - return executeTask(ctx, task, details.TIContext.ShouldRetry, logger) + args, err := convertArgBindings(details.TIContext.ArgBindings) + if err != nil { + logger.Error("Invalid arg_bindings spec from supervisor", + "dag_id", details.TI.DagID, + "task_id", details.TI.TaskID, + "error", err, + ) + // Same retry semantics as a binding failure inside executeTask: an + // equally permanent spec error must not terminate differently. + if details.TIContext.ShouldRetry { + return genmodels.RetryTask{ + EndDate: time.Now().UTC(), + RetryReason: err.Error(), + } + } + return genmodels.TaskState{ + State: genmodels.TaskStateStateFailed, + EndDate: time.Now().UTC(), + } + } + + return executeTask(ctx, task, args, details.TIContext.ShouldRetry, logger) +} + +// convertArgBindings maps the wire-model positional-argument spec (captured from +// the Python stub Dag's TaskFlow call) onto the runtime binding sum type. The +// wire union generates untyped items (msgpack delivers each XComArgBinding / +// LiteralArgBinding as a plain map), so the kind dispatch and the optional +// value_schema fragment are unpacked here. +func convertArgBindings(specsPtr *genmodels.ArgBindings) ([]binding.Arg, error) { + if specsPtr == nil || len(*specsPtr) == 0 { + return nil, nil + } + specs := *specsPtr + args := make([]binding.Arg, len(specs)) + for i, raw := range specs { + m, ok := raw.(map[string]any) + if !ok { + return nil, fmt.Errorf("arg_bindings[%d]: unexpected wire shape %T", i, raw) + } + name, ok := m["name"].(string) + if !ok || name == "" { + return nil, fmt.Errorf("arg_bindings[%d]: missing or empty name", i) + } + valueSchema := argValueSchema(m["value_schema"]) + switch kind, _ := m["kind"].(string); kind { + case "xcom": + taskID, ok := m["task_id"].(string) + if !ok || taskID == "" { + return nil, fmt.Errorf( + "arg_bindings[%d] (%q): missing or empty task_id for xcom kind", + i, + name, + ) + } + args[i] = binding.XComArg{ + Kind: kind, + Name: name, + TaskID: taskID, + ValueSchema: valueSchema, + } + case "literal": + fromDefault, _ := m["from_default"].(bool) + args[i] = binding.LiteralArg{ + Kind: kind, + Name: name, + Value: m["value"], + ValueSchema: valueSchema, + FromDefault: fromDefault, + } + default: + return nil, fmt.Errorf("arg_bindings[%d]: unknown kind %q", i, kind) + } + } + return args, nil +} + +// argValueSchema unpacks the optional value_schema fragment msgpack delivers as +// a plain map into the generated ArgValueSchema type. A missing or empty +// fragment yields nil, which the binding type check treats as unconstrained. +func argValueSchema(raw any) *genmodels.ArgValueSchema { + m, ok := raw.(map[string]any) + if !ok || len(m) == 0 { + return nil + } + schema := make(genmodels.ArgValueSchema, len(m)) + for k, v := range m { + schema[k] = v + } + return &schema } // mapIndexPtr normalizes the supervisor's map_index into the optional form @@ -142,9 +232,15 @@ func mapIndexPtr(mapIndex *int) *int { // executeTask runs the task, handling success, failure, and panics, and returns // the terminal body: genmodels.SucceedTask, TaskState, or RetryTask. +// +// args carries the positional-argument spec from the stub Dag's TaskFlow call; +// tasks that implement bundlev1.TaskWithArgs bind it (an empty spec still runs +// the arity check), while a custom Task implementation that receives a +// non-empty spec fails loudly rather than silently dropping the arguments. func executeTask( ctx context.Context, task bundlev1.Task, + args []binding.Arg, shouldRetry bool, logger *slog.Logger, ) (result any) { @@ -168,7 +264,19 @@ func executeTask( } }() - if err := task.Execute(ctx, logger); err != nil { + var err error + if tw, ok := task.(bundlev1.TaskWithArgs); ok { + err = tw.ExecuteArgs(ctx, logger, args) + } else if len(args) > 0 { + err = fmt.Errorf( + "task received %d positional argument(s) from the Dag but its implementation "+ + "does not support argument binding (does not implement TaskWithArgs)", + len(args), + ) + } else { + err = task.Execute(ctx, logger) + } + if err != nil { logger.ErrorContext(ctx, "Task failed", "error", err) // A task that fails when ti_context.should_retry is set is reported as // UP_FOR_RETRY via RetryTask; otherwise it terminates as FAILED. diff --git a/providers/common/compat/src/airflow/providers/common/compat/sdk.py b/providers/common/compat/src/airflow/providers/common/compat/sdk.py index 93174df7b2a28..772650f5499e5 100644 --- a/providers/common/compat/src/airflow/providers/common/compat/sdk.py +++ b/providers/common/compat/src/airflow/providers/common/compat/sdk.py @@ -83,9 +83,13 @@ from airflow.sdk.bases.sensor import poke_mode_only as poke_mode_only from airflow.sdk.bases.skipmixin import SkipMixin as SkipMixin from airflow.sdk.configuration import conf as conf - from airflow.sdk.definitions.context import context_merge as context_merge + from airflow.sdk.definitions.context import ( + KNOWN_CONTEXT_KEYS as KNOWN_CONTEXT_KEYS, + context_merge as context_merge, + ) from airflow.sdk.definitions.mappedoperator import MappedOperator as MappedOperator from airflow.sdk.definitions.template import literal as literal + from airflow.sdk.definitions.xcom_arg import PlainXComArg as PlainXComArg from airflow.sdk.exceptions import ( AirflowConfigException as AirflowConfigException, AirflowException as AirflowException, @@ -192,6 +196,7 @@ "DAG": ("airflow.sdk", "airflow.models.dag"), "Param": ("airflow.sdk", "airflow.models.param"), "XComArg": ("airflow.sdk", "airflow.models.xcom_arg"), + "PlainXComArg": ("airflow.sdk.definitions.xcom_arg", "airflow.models.xcom_arg"), "DecoratedOperator": ("airflow.sdk.bases.decorator", "airflow.decorators.base"), "DecoratedMappedOperator": ("airflow.sdk.bases.decorator", "airflow.decorators.base"), "MappedOperator": ("airflow.sdk.definitions.mappedoperator", "airflow.models.mappedoperator"), @@ -246,6 +251,7 @@ # ============================================================================ "Context": ("airflow.sdk", "airflow.utils.context"), "context_merge": ("airflow.sdk.definitions.context", "airflow.utils.context"), + "KNOWN_CONTEXT_KEYS": ("airflow.sdk.definitions.context", "airflow.utils.context"), "context_to_airflow_vars": ("airflow.sdk.execution_time.context", "airflow.utils.operator_helpers"), "AIRFLOW_VAR_NAME_FORMAT_MAPPING": ( "airflow.sdk.execution_time.context", diff --git a/providers/standard/pyproject.toml b/providers/standard/pyproject.toml index b385a9e9eb7f6..373d16bba6c4e 100644 --- a/providers/standard/pyproject.toml +++ b/providers/standard/pyproject.toml @@ -60,7 +60,7 @@ requires-python = ">=3.10" # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build`` dependencies = [ "apache-airflow>=2.11.0", - "apache-airflow-providers-common-compat>=1.14.1", + "apache-airflow-providers-common-compat>=1.14.1", # use next version ] # The optional dependencies should be modified in place in the generated file diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index 08bcf163a56ad..c7f07810db27c 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -18,12 +18,33 @@ from __future__ import annotations import ast -from collections.abc import Callable +import copy +import datetime +import inspect +import json +import types +import typing +from collections.abc import Callable, Collection, Mapping +from functools import cache from typing import TYPE_CHECKING, Any +try: + from pydantic import PydanticInvalidForJsonSchema, PydanticSchemaGenerationError, TypeAdapter + from pydantic.json_schema import GenerateJsonSchema +except ImportError: + # Airflow 3 always ships pydantic but Airflow 2.x base installs do not; without it, + # stub args carry no value schemas and runtimes keep their decode-only fallback. + GenerateJsonSchema = object # type: ignore[assignment,misc] + TypeAdapter = None # type: ignore[assignment,misc] + PydanticInvalidForJsonSchema = PydanticSchemaGenerationError = None # type: ignore[assignment,misc] + from airflow.providers.common.compat.sdk import ( + KNOWN_CONTEXT_KEYS, + XCOM_RETURN_KEY, DecoratedOperator, + PlainXComArg, TaskDecorator, + XComArg, task_decorator_factory, ) @@ -31,6 +52,220 @@ from airflow.providers.common.compat.sdk import Context +class _ValueSchemaGenerator(GenerateJsonSchema): + """ + Pydantic's stock JSON-schema generation plus OpenAPI's fixed-width numeric formats. + + A foreign runtime decodes numbers into machine types, which the bare + ``integer``/``number`` type names cannot convey; ``format`` is an annotation per + JSON schema, so runtimes that don't know these names simply skip them. + """ + + def int_schema(self, schema): + return {**super().int_schema(schema), "format": "int64"} + + def float_schema(self, schema): + return {**super().float_schema(schema), "format": "double"} + + +# Most-derived first: datetime subclasses date, so it must be matched before date. +_TEMPORAL_BASES = (datetime.datetime, datetime.date, datetime.time, datetime.timedelta) + + +def _normalize_temporal_annotation(annotation: Any) -> Any: + """ + Map temporal subclasses (e.g. ``pendulum.DateTime``) to their stdlib base. + + Applied recursively through unions and containers, and only as a retry when direct + schema generation fails, so temporal types carrying their own pydantic schema keep it. + """ + # Parametrized generics must be detected before the plain-class branch: on Python + # 3.10, isinstance(list[X], type) is True and issubclass silently consults the + # origin, so the class branch would return list[X] unnormalized. + origin = typing.get_origin(annotation) + args = typing.get_args(annotation) + if origin is not None and args: + normalized = tuple(_normalize_temporal_annotation(arg) for arg in args) + if normalized == args: + return annotation + if origin in (typing.Union, types.UnionType): + return typing.Union[normalized] # noqa: UP007 -- runtime construction from a tuple + return origin[normalized] + if isinstance(annotation, type): + return next((base for base in _TEMPORAL_BASES if issubclass(annotation, base)), annotation) + return annotation + + +def _infer_value_schema(annotation: Any) -> dict[str, Any] | None: + """ + Build the JSON-schema fragment for one stub parameter annotation, via pydantic. + + The pydantic-generated schema ships verbatim, so runtimes must treat it as + open-vocabulary JSON schema. Returns ``None`` when the annotation constrains nothing + (missing, ``Any``, bare ``None``) or pydantic cannot generate a schema for it; the + binding then omits ``value_schema`` and the foreign runtime falls back to a + decode-only check. + """ + if TypeAdapter is None: + return None + if annotation is inspect.Parameter.empty or annotation is None or annotation is Any: + return None + if annotation is type(None): + # get_type_hints normalizes a bare ``None`` annotation to NoneType; a parameter + # that can only ever be None constrains nothing worth shipping. + return None + try: + schema = _generate_value_schema(annotation) + except TypeError: + # Unhashable annotations cannot key the cache; generate directly. + schema = _generate_value_schema.__wrapped__(annotation) + # Deep-copy so callers embedding the fragment never alias the cached dict. + return copy.deepcopy(schema) if schema else None + + +@cache +def _generate_value_schema(annotation: Any) -> dict[str, Any] | None: + """ + Generate the schema for one annotation, cached for the process lifetime. + + TypeAdapter construction is one of pydantic's most expensive operations and + annotations are static, so re-parses of the same Dag file must not re-pay it. + """ + try: + return TypeAdapter(annotation).json_schema(schema_generator=_ValueSchemaGenerator) + except (PydanticSchemaGenerationError, PydanticInvalidForJsonSchema): + normalized = _normalize_temporal_annotation(annotation) + if normalized is annotation: + return None + try: + return TypeAdapter(normalized).json_schema(schema_generator=_ValueSchemaGenerator) + except (PydanticSchemaGenerationError, PydanticInvalidForJsonSchema): + return None + + +def _validate_stub_signature(signature: inspect.Signature, task_id: str) -> None: + for param in signature.parameters.values(): + if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD): + raise ValueError( + f"@task.stub task {task_id!r} must declare a fixed number of parameters for the " + f"foreign runtime to bind against; *{param.name} is not supported" + ) + if param.name in KNOWN_CONTEXT_KEYS: + raise ValueError( + f"@task.stub task {task_id!r} parameter {param.name!r} is an Airflow context key; " + "stub signatures declare only data parameters -- the lang-SDK runtime injects its " + "own task context natively (e.g. the Go SDK's sdk.TIRunContext parameter)" + ) + + +def _resolve_param_annotations(python_callable: Callable, signature: inspect.Signature) -> dict[str, Any]: + """Map each parameter to its parse-time-resolvable annotation (``Parameter.empty`` when not).""" + try: + hints = typing.get_type_hints(python_callable) + except (NameError, TypeError): + # Annotations that cannot be resolved at parse time (e.g. names behind + # TYPE_CHECKING with ``from __future__ import annotations``) degrade to "any". + hints = {} + + def resolve(name: str, param: inspect.Parameter) -> Any: + if name in hints: + return hints[name] + if isinstance(param.annotation, str): + return inspect.Parameter.empty + return param.annotation + + return {name: resolve(name, param) for name, param in signature.parameters.items()} + + +def _ensure_json_literal(value: Any, task_id: str, name: str) -> None: + try: + json.dumps(value, allow_nan=False) + except (TypeError, ValueError): + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} received a literal of type " + f"{type(value).__name__} that is not JSON-serializable, so it cannot be passed " + "to the foreign runtime" + ) + + +def _validate_xcom_value(value: Any, task_id: str, name: str) -> bool: + """Validate an XComArg argument, returning True when it is a bindable direct upstream output.""" + if isinstance(value, PlainXComArg): + if value.key != XCOM_RETURN_KEY: + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} references the XCom key " + f"{value.key!r}; only an upstream task's return value can cross the language " + "boundary -- indexing an output by a custom key is not supported" + ) + if value.operator.is_mapped: + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} references the aggregated " + f"output of the mapped task {value.operator.task_id!r}; a foreign runtime " + "pulls single XCom rows, so a mapped upstream's combined output is not " + "supported -- use .expand() on the stub to consume it per element" + ) + return True + if isinstance(value, XComArg): + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} received a " + f"{type(value).__name__}; only direct upstream task outputs can cross the " + "language boundary -- .map()/.zip()/.concat() results are not supported" + ) + return False + + +def _build_arg_bindings( + python_callable: Callable, + op_args: Collection[Any], + op_kwargs: Mapping[str, Any], + task_id: str, +) -> list[dict[str, Any]] | None: + """ + Bind the TaskFlow call arguments to the stub signature and build the ordered arg spec. + + Each spec entry is a plain dict matching one variant of the execution API's + ``TaskArgBinding`` union: an ``XComArgBinding`` (``kind="xcom"``) for upstream TaskFlow + outputs, or a ``LiteralArgBinding`` (``kind="literal"``) for everything else. ``name`` is + always the stub function's parameter name, so a foreign runtime can bind by name (e.g. the + Go SDK's name-based struct fields) in addition to the existing positional order. + Returns ``None`` for argless calls: the binding contract (including the signature checks + below) applies only once a TaskFlow call actually passes arguments, so pre-TaskFlow stub + Dags whose call arguments were always ignored keep parsing. + """ + if not op_args and not op_kwargs: + return None + + signature = inspect.signature(python_callable) + _validate_stub_signature(signature, task_id) + + bound = signature.bind(*op_args, **op_kwargs) + explicitly_bound = set(bound.arguments) + bound.apply_defaults() + + annotations = _resolve_param_annotations(python_callable, signature) + + spec: list[dict[str, Any]] = [] + for name in signature.parameters: + value = bound.arguments[name] + value_schema = _infer_value_schema(annotations[name]) + if _validate_xcom_value(value, task_id, name): + xcom_entry: dict[str, Any] = {"name": name, "kind": "xcom", "task_id": value.operator.task_id} + if value_schema is not None: + xcom_entry["value_schema"] = value_schema + spec.append(xcom_entry) + continue + _ensure_json_literal(value, task_id, name) + entry: dict[str, Any] = {"name": name, "kind": "literal", "value": value} + if value_schema is not None: + # Key omission (never ``None``) is the wire contract for "unconstrained": + # ti_run responds with ``exclude_unset``, so an absent key stays absent. + entry["value_schema"] = value_schema + if name not in explicitly_bound: + entry["from_default"] = True + spec.append(entry) + return spec + + class _StubOperator(DecoratedOperator): custom_operator_name: str = "@task.stub" @@ -60,10 +295,10 @@ def __init__( module = ast.parse(self.get_python_source()) if len(module.body) != 1: - raise RuntimeError("Expected a single statement") + raise ValueError("Expected a single statement") fn = module.body[0] if not isinstance(fn, ast.FunctionDef): - raise RuntimeError("Expected a single sync function") + raise ValueError("Expected a single sync function") for stmt in fn.body: if isinstance(stmt, ast.Pass): continue @@ -75,7 +310,27 @@ def __init__( f"Functions passed to @task.stub must be an empty function (`pass`, or `...` only) (got {stmt})" ) - ... + # Bind the TaskFlow call to the *original* signature (DecoratedOperator mangles context + # key defaults, which stubs reject anyway) and persist the ordered arg spec so the + # execution API can hand it to the foreign runtime via StartupDetails. + self._arg_bindings = _build_arg_bindings(python_callable, self.op_args, self.op_kwargs, self.task_id) + + # Direct .expand() on the stub needs no parse-time spec (ti_run derives per-map-index + # bindings from the serialized expand input), but a mapped task group creates + # per-map-index instances of the tasks inside it with no expand input of their own, + # so their arg values are unresolvable both here and server-side. + in_mapped_group = getattr(self, "get_closest_mapped_task_group", lambda: None)() is not None + if self._arg_bindings is not None and in_mapped_group: + raise ValueError( + f"@task.stub task {self.task_id!r} passes TaskFlow call arguments inside a mapped " + "task group; the captured spec cannot carry values that resolve per map index at " + "runtime, so stub tasks with arguments are not supported under a task group's " + ".expand()" + ) + + @classmethod + def get_serialized_fields(cls): + return super().get_serialized_fields() | {"_arg_bindings"} def execute(self, context: Context) -> Any: raise RuntimeError( @@ -96,6 +351,14 @@ def stub( Stub tasks exist in the Dag graph only, but the execution must happen in an external environment via the Task Execution Interface. + Stub functions may declare parameters and be called TaskFlow-style with upstream task + outputs or JSON-serializable literals; the resulting argument-binding spec (parameter + names, value schemas, and values, in declaration order) is delivered to the foreign + runtime, which binds the values onto the native task function. + + Mapped (``.expand()``) stubs do not receive TaskFlow arguments yet -- their call args + keep the legacy ignored behavior; per-map-index delivery is part of + https://github.com/apache/airflow/issues/66937 and lands in a follow-up. """ return task_decorator_factory( decorated_operator_class=_StubOperator, diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index 2a17c3fdd82c1..68d804de4fc17 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -17,10 +17,16 @@ from __future__ import annotations import contextlib +import datetime +import typing +from typing import Any +from unittest import mock +import pendulum import pytest -from airflow.providers.standard.decorators.stub import stub +from airflow.providers.common.compat.sdk import DAG, task_group +from airflow.providers.standard.decorators.stub import _infer_value_schema, stub from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS @@ -69,3 +75,346 @@ def test_stub_rejects_retry_policy(): def test_stub_allows_retries(): stub(fn_pass, retries=5)() + + +def fn_extract(): ... + + +def fn_transform(country: str, extracted: dict, retries_num: int = 3): ... + + +def fn_untyped(a, b): ... + + +def fn_varargs(*args): ... + + +def fn_kwonly_varkw(**kwargs): ... + + +def fn_context_key(ti): ... + + +class TestStubTaskflowArgs: + """The TaskFlow call on a stub captures the ordered positional-arg spec (``_arg_bindings``).""" + + def test_literal_and_xcom_spec(self): + with DAG(dag_id="d"): + extracted = stub(fn_extract)() + result = stub(fn_transform)("uk", extracted) + + op = result.operator + assert op._arg_bindings == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "fn_extract", + }, + { + "name": "retries_num", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 3, + "from_default": True, + }, + ] + assert op.upstream_task_ids == {"fn_extract"} + + def test_kwargs_normalize_to_declaration_order(self): + with DAG(dag_id="d"): + extracted = stub(fn_extract)() + result = stub(fn_transform)(extracted=extracted, country="fr", retries_num=7) + + assert result.operator._arg_bindings == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "fr"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "fn_extract", + }, + { + "name": "retries_num", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 7, + }, + ] + + def test_explicitly_passing_the_default_value_is_not_from_default(self): + """The flag tracks provenance, not value equality: an author-passed argument is explicit + even when it equals the signature default, so keyword-style consumers must still claim it.""" + with DAG(dag_id="d"): + extracted = stub(fn_extract)() + result = stub(fn_transform)("uk", extracted, retries_num=3) + + assert result.operator._arg_bindings[2] == { + "name": "retries_num", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 3, + } + + def test_custom_xcom_key_rejected(self): + with DAG(dag_id="d"): + extracted = stub(fn_extract)() + with pytest.raises(ValueError, match="indexing an output by a custom key"): + stub(fn_transform)("uk", extracted["part"]) + + def test_zero_param_stub_has_no_spec(self): + assert stub(fn_pass)().operator._arg_bindings is None + + def test_untyped_params_omit_value_schema(self): + """Key absence (never ``None``) is the wire contract for an unconstrained argument.""" + with DAG(dag_id="d"): + result = stub(fn_untyped)(1, "x") + + assert result.operator._arg_bindings == [ + {"name": "a", "kind": "literal", "value": 1}, + {"name": "b", "kind": "literal", "value": "x"}, + ] + + def test_unresolvable_annotation_omits_value_schema(self): + def fn(x): ... + + fn.__annotations__ = {"x": "NotARealType"} + with DAG(dag_id="d"): + result = stub(fn)("v") + + assert result.operator._arg_bindings == [{"name": "x", "kind": "literal", "value": "v"}] + + def test_varargs_rejected(self): + with pytest.raises(ValueError, match="fixed number of parameters"): + stub(fn_varargs)(1, 2) + + def test_varkw_rejected(self): + with pytest.raises(ValueError, match="fixed number of parameters"): + stub(fn_kwonly_varkw)(x=1) + + def test_context_key_param_rejected(self): + with pytest.raises(ValueError, match="is an Airflow context key"): + stub(fn_context_key)(1) + + @pytest.mark.parametrize("fn", [fn_varargs, fn_kwonly_varkw, fn_context_key], ids=lambda f: f.__name__) + def test_argless_call_skips_signature_checks(self, fn): + """Pre-TaskFlow stub Dags never passed arguments; their signatures must keep parsing.""" + assert stub(fn)().operator._arg_bindings is None + + def test_argless_call_captures_no_spec_for_defaulted_params(self): + def fn(limit: int = 10): ... + + assert stub(fn)().operator._arg_bindings is None + + def test_non_json_literal_rejected(self): + with DAG(dag_id="d"), pytest.raises(ValueError, match="not JSON-serializable"): + stub(fn_transform)("uk", object()) + + def test_nan_literal_rejected(self): + with DAG(dag_id="d"), pytest.raises(ValueError, match="not JSON-serializable"): + stub(fn_transform)("uk", {"ratio": float("nan")}) + + def test_mapped_xcom_arg_rejected(self): + with DAG(dag_id="d"): + extracted = stub(fn_extract)() + with pytest.raises(ValueError, match="only direct upstream task outputs"): + stub(fn_transform)("uk", extracted.map(lambda v: v)) + + def test_mapped_upstream_aggregated_output_rejected(self): + def fn_produce(n: int): ... + + with DAG(dag_id="d"): + vals = stub(fn_produce).expand(n=[1, 2]) + with pytest.raises(ValueError, match="aggregated output of the mapped task"): + stub(fn_transform)("uk", vals) + + def test_arg_bindings_survive_dag_serialization_round_trip(self): + """The captured spec must survive whichever core serializer the provider runs against.""" + try: + from airflow.serialization.serialized_objects import DagSerialization + except ImportError: # Airflow 2 exposes the round-trip API on SerializedDAG + from airflow.serialization.serialized_objects import SerializedDAG as DagSerialization + + with DAG(dag_id="d") as dag: + extracted = stub(fn_extract)() + stub(fn_transform)("uk", extracted) + + round_tripped = DagSerialization.from_dict(DagSerialization.to_dict(dag)) + assert round_tripped.task_dict["fn_transform"]._arg_bindings == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "fn_extract", + }, + { + "name": "retries_num", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 3, + "from_default": True, + }, + ] + + def test_expand_builds_mapped_stub_without_parse_time_bindings(self): + """Mapped stubs capture no spec: their call args keep the legacy ignored behavior for now.""" + with DAG(dag_id="d"): + result = stub(fn_transform).expand(country=["uk", "fr"], extracted=[{}, {}]) + # op_kwargs_expand_input/partial_kwargs (not is_mapped) so the assertions also + # hold on the Airflow 2.x MappedOperator, which the provider still supports. + assert result.operator.op_kwargs_expand_input.value == { + "country": ["uk", "fr"], + "extracted": [{}, {}], + } + assert "_arg_bindings" not in result.operator.partial_kwargs + + def test_stub_with_args_inside_mapped_task_group_rejected(self): + @task_group + def group(n): + stub(fn_transform)("uk", {}) + + with DAG(dag_id="d"): + with pytest.raises(ValueError, match="mapped task group"): + group.expand(n=[1, 2]) + + def test_argless_stub_inside_mapped_task_group_allowed(self): + @task_group + def group(n): + stub(fn_extract)() + + with DAG(dag_id="d"): + group.expand(n=[1, 2]) + + +@pytest.mark.parametrize( + ("annotation", "expected"), + [ + pytest.param(str, {"type": "string"}, id="str"), + pytest.param(bool, {"type": "boolean"}, id="bool"), + pytest.param(int, {"type": "integer", "format": "int64"}, id="int"), + pytest.param(float, {"type": "number", "format": "double"}, id="float"), + pytest.param(dict, {"type": "object", "additionalProperties": True}, id="dict"), + pytest.param( + dict[str, int], + {"type": "object", "additionalProperties": {"type": "integer", "format": "int64"}}, + id="dict-parameterized", + ), + pytest.param( + typing.Mapping[str, int], + {"type": "object", "additionalProperties": {"type": "integer", "format": "int64"}}, + id="mapping", + ), + pytest.param(list, {"type": "array", "items": {}}, id="list"), + pytest.param( + list[int], + {"type": "array", "items": {"type": "integer", "format": "int64"}}, + id="list-parameterized", + ), + pytest.param(tuple, {"type": "array", "items": {}}, id="tuple"), + pytest.param(set, {"type": "array", "items": {}, "uniqueItems": True}, id="set"), + pytest.param( + typing.Sequence[int], + {"type": "array", "items": {"type": "integer", "format": "int64"}}, + id="sequence", + ), + pytest.param(datetime.datetime, {"type": "string", "format": "date-time"}, id="datetime"), + pytest.param(datetime.date, {"type": "string", "format": "date"}, id="date"), + pytest.param(datetime.time, {"type": "string", "format": "time"}, id="time"), + pytest.param(datetime.timedelta, {"type": "string", "format": "duration"}, id="timedelta"), + pytest.param(bytes, {"type": "string", "format": "binary"}, id="bytes"), + pytest.param( + typing.Literal["a", "b"], + {"type": "string", "enum": ["a", "b"]}, + id="literal", + ), + pytest.param(Any, None, id="any"), + pytest.param(None, None, id="none"), + pytest.param(type(None), None, id="nonetype"), + pytest.param( + pendulum.DateTime, + {"type": "string", "format": "date-time"}, + id="pendulum-datetime", + ), + pytest.param( + pendulum.DateTime | None, + {"anyOf": [{"type": "string", "format": "date-time"}, {"type": "null"}]}, + id="optional-pendulum-datetime", + ), + pytest.param( + list[pendulum.DateTime], + {"type": "array", "items": {"type": "string", "format": "date-time"}}, + id="list-pendulum-datetime", + ), + pytest.param(pendulum.Duration, {"type": "string", "format": "duration"}, id="pendulum-duration"), + pytest.param( + typing.Optional[str], # noqa: UP045 -- legacy form on purpose + {"anyOf": [{"type": "string"}, {"type": "null"}]}, + id="optional-str", + ), + pytest.param( + typing.Union[int, str], # noqa: UP007 -- legacy form on purpose + {"anyOf": [{"type": "integer", "format": "int64"}, {"type": "string"}]}, + id="union", + ), + pytest.param(str | None, {"anyOf": [{"type": "string"}, {"type": "null"}]}, id="pep604-optional"), + pytest.param( + int | None, + {"anyOf": [{"type": "integer", "format": "int64"}, {"type": "null"}]}, + id="optional-int", + ), + pytest.param( + datetime.datetime | None, + {"anyOf": [{"type": "string", "format": "date-time"}, {"type": "null"}]}, + id="optional-datetime", + ), + pytest.param( + dict | bool, + {"anyOf": [{"type": "object", "additionalProperties": True}, {"type": "boolean"}]}, + id="union-dict-bool", + ), + pytest.param( + str | int | None, + {"anyOf": [{"type": "string"}, {"type": "integer", "format": "int64"}, {"type": "null"}]}, + id="union-with-null", + ), + pytest.param(list | tuple, {"type": "array", "items": {}}, id="union-dedupes-equal-members"), + pytest.param( + datetime.datetime | str, + {"anyOf": [{"type": "string", "format": "date-time"}, {"type": "string"}]}, + id="mixed-format-union-keeps-both", + ), + pytest.param( + str | contextlib.AbstractContextManager, + None, + id="union-unclassifiable-member", + ), + pytest.param(contextlib.AbstractContextManager, None, id="custom-class"), + pytest.param(typing.Callable[[int], str], None, id="callable-invalid-for-json-schema"), + pytest.param( + pendulum.DateTime | contextlib.AbstractContextManager, + None, + id="union-temporal-and-unclassifiable", + ), + ], +) +def test_infer_value_schema(annotation, expected): + assert _infer_value_schema(annotation) == expected + + +@mock.patch("airflow.providers.standard.decorators.stub.TypeAdapter", None) +def test_infer_value_schema_without_pydantic(): + assert _infer_value_schema(str) is None + + +def test_infer_value_schema_cache_returns_isolated_copies(): + first = _infer_value_schema(dict) + second = _infer_value_schema(dict) + assert first == second + assert first is not second, "callers embed and serialize the fragment, so it must not alias the cache" + + +def test_infer_value_schema_unhashable_annotation_generates_uncached(): + annotation = typing.Annotated[int, {"unhashable": True}] + assert _infer_value_schema(annotation) == {"type": "integer", "format": "int64"} diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 8e5bfc1d076e5..d97d38613834c 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -27,7 +27,7 @@ from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, RootModel -API_VERSION: Final[str] = "2026-06-30" +API_VERSION: Final[str] = "2026-10-30" class AssetAliasReferenceAssetEventDagRun(BaseModel): @@ -608,6 +608,10 @@ class DagAttributeTypes(str, Enum): TASK_GROUP = "taskgroup" +class ArgValueSchema(RootModel[dict[str, JsonValue]]): + root: dict[str, JsonValue] + + class AssetReferenceAssetEventDagRun(BaseModel): """ Schema for AssetModel used in AssetEventDagRunReference. @@ -697,6 +701,18 @@ class HTTPValidationError(BaseModel): detail: Annotated[list[ValidationError] | None, Field(title="Detail")] = None +class LiteralArgBinding(BaseModel): + """ + One positional stub-task argument carrying an inline literal from the Dag file. + """ + + kind: Annotated[Literal["literal"], Field(title="Kind")] + name: Annotated[str, Field(title="Name")] + value_schema: ArgValueSchema | None = None + value: JsonValue | None = None + from_default: Annotated[bool | None, Field(title="From Default")] = False + + class TITerminalStatePayload(BaseModel): """ Schema for updating TaskInstance to a terminal state except SUCCESS state. @@ -710,6 +726,17 @@ class TITerminalStatePayload(BaseModel): rendered_map_index: Annotated[str | None, Field(title="Rendered Map Index")] = None +class XComArgBinding(BaseModel): + """ + One positional stub-task argument pulled from an upstream task's XCom. + """ + + kind: Annotated[Literal["xcom"], Field(title="Kind")] + name: Annotated[str, Field(title="Name")] + value_schema: ArgValueSchema | None = None + task_id: Annotated[str, Field(title="Task Id")] + + class AssetEventDagRunReference(BaseModel): """ Schema for AssetEvent model used in DagRun. @@ -782,6 +809,10 @@ class DagRun(BaseModel): team_name: Annotated[str | None, Field(title="Team Name")] = None +class TaskArgBinding(RootModel[XComArgBinding | LiteralArgBinding]): + root: Annotated[XComArgBinding | LiteralArgBinding, Field(discriminator="kind", title="TaskArgBinding")] + + class TIRunContext(BaseModel): """ Response schema for TaskInstance run context. @@ -797,3 +828,4 @@ class TIRunContext(BaseModel): xcom_keys_to_clear: Annotated[list[str] | None, Field(title="Xcom Keys To Clear")] = None should_retry: Annotated[bool | None, Field(title="Should Retry")] = False start_date: Annotated[AwareDatetime | None, Field(title="Start Date")] = None + arg_bindings: Annotated[list[TaskArgBinding] | None, Field(title="Arg Bindings")] = None diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 0ec8fe4e49a1c..8e47af32cf5d4 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -1,8 +1,15 @@ { "$schema": "https://json-schema.org/draft/2020-12/schema", - "api_version": "2026-06-16", + "api_version": "2026-10-30", "description": "Apache Airflow SDK Supervisor Schema", "$defs": { + "ArgValueSchema": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "ArgValueSchema", + "type": "object" + }, "AssetAliasReferenceAssetEventDagRun": { "additionalProperties": false, "description": "Schema for AssetAliasModel used in AssetEventDagRunReference.", @@ -4389,6 +4396,42 @@ "title": "VariableResult", "type": "object" }, + "XComArgBinding": { + "description": "One positional stub-task argument pulled from an upstream task's XCom.", + "properties": { + "kind": { + "const": "xcom", + "title": "Kind", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "value_schema": { + "anyOf": [ + { + "$ref": "#/$defs/ArgValueSchema" + }, + { + "type": "null" + } + ], + "default": null + }, + "task_id": { + "title": "Task Id", + "type": "string" + } + }, + "required": [ + "kind", + "name", + "task_id" + ], + "title": "XComArgBinding", + "type": "object" + }, "XComCountResponse": { "properties": { "len": { @@ -4563,6 +4606,71 @@ "title": "ConnectionResponse", "type": "object" }, + "LiteralArgBinding": { + "description": "One positional stub-task argument carrying an inline literal from the Dag file.", + "properties": { + "kind": { + "const": "literal", + "title": "Kind", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "value_schema": { + "anyOf": [ + { + "$ref": "#/$defs/ArgValueSchema" + }, + { + "type": "null" + } + ], + "default": null + }, + "value": { + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ], + "default": null + }, + "from_default": { + "default": false, + "title": "From Default", + "type": "boolean" + } + }, + "required": [ + "kind", + "name" + ], + "title": "LiteralArgBinding", + "type": "object" + }, + "TaskArgBinding": { + "discriminator": { + "mapping": { + "literal": "#/$defs/LiteralArgBinding", + "xcom": "#/$defs/XComArgBinding" + }, + "propertyName": "kind" + }, + "oneOf": [ + { + "$ref": "#/$defs/XComArgBinding" + }, + { + "$ref": "#/$defs/LiteralArgBinding" + } + ], + "title": "TaskArgBinding" + }, "AssetEventDagRunReference": { "additionalProperties": false, "description": "Schema for AssetEvent model used in DagRun.", @@ -4926,6 +5034,21 @@ ], "default": null, "title": "Start Date" + }, + "arg_bindings": { + "anyOf": [ + { + "items": { + "$ref": "#/$defs/TaskArgBinding" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Arg Bindings" } }, "required": [ 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 index 9491a8993fdc3..7e5ce93f86bdc 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py @@ -37,8 +37,13 @@ def get_bundle() -> VersionBundle: """ from cadwyn import HeadVersion, Version, VersionBundle + from airflow.sdk.execution_time.schema.versions.v2026_10_30 import ( + AddArgBindingsToSupervisorTIRunContext, + ) + return VersionBundle( HeadVersion(), + Version("2026-10-30", AddArgBindingsToSupervisorTIRunContext), Version("2026-06-16"), ) diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py new file mode 100644 index 0000000000000..e6b93f5dea805 --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py @@ -0,0 +1,36 @@ +# 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 VersionChange, schema + +from airflow.sdk.api.datamodels._generated import TIRunContext + + +class AddArgBindingsToSupervisorTIRunContext(VersionChange): + """ + Add the ``arg_bindings`` argument-binding spec for stub (foreign-runtime) tasks. + + Each entry is a discriminated union of ``XComArgBinding`` and ``LiteralArgBinding`` + keyed on ``kind``. The supervisor-schema mirror of the execution API's + ``AddArgBindingsToTIRunContext``, named apart so the two migrations are not confused. + """ + + description = __doc__ + + instructions_to_migrate_to_previous_version = (schema(TIRunContext).field("arg_bindings").didnt_exist,) 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 index 05218aded3d3b..8e2379e2fdbba 100644 --- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py @@ -105,12 +105,9 @@ def _backfill_sentry_trace(request): 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. + *field-level* migration behaviour independent of the real bundle's + contents. The real bundle's ``arg_bindings`` migration is covered by + :class:`TestRealBundleArgBindingsDowngrade` below. """ @pytest.fixture @@ -369,3 +366,103 @@ def test_accessing_bundle_loads_cadwyn(self): "assert 'cadwyn' in sys.modules, 'cadwyn should load when the bundle is accessed'" ) subprocess.run([sys.executable, "-c", code], check=True, capture_output=True, text=True) + + +class TestRealBundleArgBindingsDowngrade: + """ + Drive the *real* supervisor bundle through the ``arg_bindings`` migration. + + ``AddArgBindingsToSupervisorTIRunContext`` is the bundle's first ``schema(...)`` + instruction on a model *nested* inside a registered body + (``StartupDetails.ti_context``); this pins that the downgrade + re-validation strips the nested field on the wire for a runtime + pinned to the previous version, and keeps it at head. + """ + + @pytest.fixture + def startup_details(self): + import datetime + import uuid + + from airflow.sdk.api.datamodels._generated import ( + BundleInfo, + DagRun, + DagRunState, + DagRunType, + TaskInstance, + TIRunContext, + ) + from airflow.sdk.execution_time.comms import StartupDetails + + now = datetime.datetime.now(datetime.timezone.utc) + return StartupDetails( + ti=TaskInstance( + id=uuid.uuid4(), + task_id="transform", + dag_id="d", + run_id="r", + try_number=1, + dag_version_id=uuid.uuid4(), + ), + dag_rel_path="d.py", + bundle_info=BundleInfo(name="b", version=None), + start_date=now, + ti_context=TIRunContext( + dag_run=DagRun( + dag_id="d", + run_id="r", + logical_date=now, + start_date=now, + run_type=DagRunType.MANUAL, + state=DagRunState.RUNNING, + run_after=now, + consumed_asset_events=[], + ), + max_tries=1, + arg_bindings=[ + # No value_schema: the unconstrained ("any") case rides through the migrator too. + {"name": "country", "kind": "literal", "value": "uk"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object"}, + "task_id": "extract", + }, + { + "name": "limit", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 10, + "from_default": True, + }, + ], + ), + sentry_integration="", + ) + + @pytest.fixture + def real_migrator(self) -> SchemaVersionMigrator: + return get_schema_version_migrator() + + def test_downgrade_strips_arg_bindings_for_previous_version(self, real_migrator, startup_details): + out = real_migrator.downgrade(startup_details, "2026-06-16").model_dump() + assert "arg_bindings" not in out["ti_context"] + + def test_head_version_keeps_arg_bindings(self, real_migrator, startup_details): + from airflow.sdk.api.datamodels._generated import LiteralArgBinding, XComArgBinding + + out = real_migrator.downgrade(startup_details, "2026-10-30") + assert out.ti_context.arg_bindings is not None + literal, xcom, defaulted = (a.root for a in out.ti_context.arg_bindings) + assert isinstance(literal, LiteralArgBinding) + assert literal.value == "uk" + assert literal.name == "country" + assert literal.from_default is False + assert literal.value_schema is None + assert isinstance(xcom, XComArgBinding) + assert xcom.task_id == "extract" + assert xcom.name == "extracted" + assert xcom.value_schema.root == {"type": "object"} + assert isinstance(defaulted, LiteralArgBinding) + assert defaulted.from_default is True + assert defaulted.value_schema.root == {"type": "integer", "format": "int64"} diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index ab2632831ab95..af944e4e565d7 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -22,17 +22,17 @@ // // Re-run with: pnpm run generate:supervisor +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "JsonValue". + */ +export type JsonValue = unknown; export type Name = string; export type Id = number; export type Timestamp = string; export type Extra = { [k: string]: JsonValue; } | null; -/** - * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema - * via the `definition` "JsonValue". - */ -export type JsonValue = unknown; export type Name1 = string; export type Uri = string; export type Group = string; @@ -245,6 +245,18 @@ export type NextKwargs1 = export type XcomKeysToClear = string[]; export type ShouldRetry = boolean; export type StartDate2 = string | null; +export type ArgBindings = TaskArgBinding[] | null; +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskArgBinding". + */ +export type TaskArgBinding = XComArgBinding | LiteralArgBinding; +export type Kind = "xcom"; +export type Name8 = string; +export type TaskId1 = string; +export type Kind1 = "literal"; +export type Name9 = string; +export type FromDefault = boolean; export type Type13 = "TaskCallbackRequest"; export type Filepath2 = string; export type BundleName3 = string; @@ -310,7 +322,7 @@ export type NextKwargs2 = { } | null; export type RenderedMapIndex1 = string | null; export type Type20 = "DeferTask"; -export type Name8 = string; +export type Name10 = string; export type Key1 = string; export type Type21 = "DeleteAssetStateStoreByName"; export type Uri5 = string; @@ -324,7 +336,7 @@ export type Type24 = "DeleteVariable"; export type Key5 = string; export type DagId6 = string; export type RunId5 = string; -export type TaskId1 = string; +export type TaskId2 = string; export type MapIndex1 = number | null; export type Type25 = "DeleteXCom"; /** @@ -362,11 +374,11 @@ export type ErrorType1 = | "PERMISSION_DENIED" | "GENERIC_ERROR" | "API_SERVER_ERROR"; -export type Name9 = string; +export type Name11 = string; export type Type27 = "GetAssetByName"; export type Uri6 = string; export type Type28 = "GetAssetByUri"; -export type Name10 = string | null; +export type Name12 = string | null; export type Uri7 = string | null; export type After = string | null; export type Before = string | null; @@ -389,7 +401,7 @@ export type Extra8 = { [k: string]: string; } | null; export type Type30 = "GetAssetEventByAssetAlias"; -export type Name11 = string; +export type Name13 = string; export type Key6 = string; export type Type31 = "GetAssetStateStoreByName"; export type Uri8 = string; @@ -421,7 +433,7 @@ export type LogicalDate3 = string; export type State3 = string | null; export type Type41 = "GetPreviousDagRun"; export type DagId12 = string; -export type TaskId2 = string; +export type TaskId3 = string; export type LogicalDate4 = string | null; export type MapIndex2 = number; export type Type42 = "GetPreviousTI"; @@ -458,25 +470,25 @@ export type Type49 = "GetVariableKeys"; export type Key10 = string; export type DagId16 = string; export type RunId9 = string; -export type TaskId3 = string; +export type TaskId4 = string; export type MapIndex5 = number | null; export type IncludePriorDates = boolean; export type Type50 = "GetXCom"; export type Key11 = string; export type DagId17 = string; export type RunId10 = string; -export type TaskId4 = string; +export type TaskId5 = string; export type Type51 = "GetXComCount"; export type Key12 = string; export type DagId18 = string; export type RunId11 = string; -export type TaskId5 = string; +export type TaskId6 = string; export type Offset1 = number; export type Type52 = "GetXComSequenceItem"; export type Key13 = string; export type DagId19 = string; export type RunId12 = string; -export type TaskId6 = string; +export type TaskId7 = string; export type Start = number | null; export type Stop = number | null; export type Step = number | null; @@ -498,7 +510,7 @@ export type AssignedUsers1 = HITLUser[] | null; export type Type54 = "HITLDetailRequestResult"; export type InactiveAssets = AssetProfile[] | null; export type Type55 = "InactiveAssetsResult"; -export type Name12 = string | null; +export type Name14 = string | null; export type Type56 = "MaskSecret"; export type Ok = boolean; export type Type57 = "OKResponse"; @@ -508,7 +520,7 @@ export type StartDate4 = string | null; export type EndDate3 = string | null; export type Type58 = "PrevSuccessfulDagRunResult"; export type Type59 = "PreviousDagRunResult"; -export type TaskId7 = string; +export type TaskId8 = string; export type DagId20 = string; export type RunId13 = string; export type LogicalDate5 = string | null; @@ -536,7 +548,7 @@ export type RetryReason = string | null; export type Type64 = "RetryTask"; export type Type65 = "SentFDs"; export type Fds = number[]; -export type Name13 = string; +export type Name15 = string; export type Key15 = string; export type Type66 = "SetAssetStateStoreByName"; export type Uri9 = string; @@ -552,7 +564,7 @@ export type Type70 = "SetTaskStateStore"; export type Key18 = string; export type DagId21 = string; export type RunId14 = string; -export type TaskId8 = string; +export type TaskId9 = string; export type MapIndex7 = number | null; export type DagResult1 = boolean; export type MappedLength = number | null; @@ -624,6 +636,13 @@ export type Root = JsonValue[]; export type Type89 = "XComSequenceSliceResult"; export interface SupervisorWireSchema {} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "ArgValueSchema". + */ +export interface ArgValueSchema { + [k: string]: JsonValue; +} /** * Schema for AssetAliasModel used in AssetEventDagRunReference. * @@ -1003,6 +1022,7 @@ export interface TIRunContext { xcom_keys_to_clear?: XcomKeysToClear; should_retry?: ShouldRetry; start_date?: StartDate2; + arg_bindings?: ArgBindings; } /** * Variable schema for responses with fields that are needed for Runtime. @@ -1030,6 +1050,31 @@ export interface ConnectionResponse { port: Port1; extra: Extra6; } +/** + * One positional stub-task argument pulled from an upstream task's XCom. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "XComArgBinding". + */ +export interface XComArgBinding { + kind: Kind; + name: Name8; + value_schema?: ArgValueSchema | null; + task_id: TaskId1; +} +/** + * One positional stub-task argument carrying an inline literal from the Dag file. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "LiteralArgBinding". + */ +export interface LiteralArgBinding { + kind: Kind1; + name: Name9; + value_schema?: ArgValueSchema | null; + value?: unknown; + from_default?: FromDefault; +} /** * Email notification request for task failures/retries. * @@ -1149,7 +1194,7 @@ export interface DeferTask { * via the `definition` "DeleteAssetStateStoreByName". */ export interface DeleteAssetStateStoreByName { - name: Name8; + name: Name10; key: Key1; type?: Type21; } @@ -1187,7 +1232,7 @@ export interface DeleteXCom { key: Key5; dag_id: DagId6; run_id: RunId5; - task_id: TaskId1; + task_id: TaskId2; map_index?: MapIndex1; type?: Type25; } @@ -1205,7 +1250,7 @@ export interface ErrorResponse { * via the `definition` "GetAssetByName". */ export interface GetAssetByName { - name: Name9; + name: Name11; type?: Type27; } /** @@ -1221,7 +1266,7 @@ export interface GetAssetByUri { * via the `definition` "GetAssetEventByAsset". */ export interface GetAssetEventByAsset { - name: Name10; + name: Name12; uri: Uri7; after?: After; before?: Before; @@ -1252,7 +1297,7 @@ export interface GetAssetEventByAssetAlias { * via the `definition` "GetAssetStateStoreByName". */ export interface GetAssetStateStoreByName { - name: Name11; + name: Name13; key: Key6; type?: Type31; } @@ -1354,7 +1399,7 @@ export interface GetPreviousDagRun { */ export interface GetPreviousTI { dag_id: DagId12; - task_id: TaskId2; + task_id: TaskId3; logical_date?: LogicalDate4; map_index?: MapIndex2; state?: TaskInstanceState | null; @@ -1440,7 +1485,7 @@ export interface GetXCom { key: Key10; dag_id: DagId16; run_id: RunId9; - task_id: TaskId3; + task_id: TaskId4; map_index?: MapIndex5; include_prior_dates?: IncludePriorDates; type?: Type50; @@ -1455,7 +1500,7 @@ export interface GetXComCount { key: Key11; dag_id: DagId17; run_id: RunId10; - task_id: TaskId4; + task_id: TaskId5; type?: Type51; } /** @@ -1466,7 +1511,7 @@ export interface GetXComSequenceItem { key: Key12; dag_id: DagId18; run_id: RunId11; - task_id: TaskId5; + task_id: TaskId6; offset: Offset1; type?: Type52; } @@ -1478,7 +1523,7 @@ export interface GetXComSequenceSlice { key: Key13; dag_id: DagId19; run_id: RunId12; - task_id: TaskId6; + task_id: TaskId7; start: Start; stop: Stop; step: Step; @@ -1520,7 +1565,7 @@ export interface InactiveAssetsResult { */ export interface MaskSecret { value: JsonValue; - name?: Name12; + name?: Name14; type?: Type56; } /** @@ -1559,7 +1604,7 @@ export interface PreviousDagRunResult { * via the `definition` "PreviousTIResponse". */ export interface PreviousTIResponse { - task_id: TaskId7; + task_id: TaskId8; dag_id: DagId20; run_id: RunId13; logical_date?: LogicalDate5; @@ -1636,7 +1681,7 @@ export interface SentFDs { * via the `definition` "SetAssetStateStoreByName". */ export interface SetAssetStateStoreByName { - name: Name13; + name: Name15; key: Key15; value: JsonValue; type?: Type66; @@ -1694,7 +1739,7 @@ export interface SetXCom { value: JsonValue; dag_id: DagId21; run_id: RunId14; - task_id: TaskId8; + task_id: TaskId9; map_index?: MapIndex7; dag_result?: DagResult1; mapped_length?: MappedLength; @@ -1896,4 +1941,4 @@ export interface XComSequenceSliceResult { * (e.g. bundle metadata) and runs the migrator accordingly. * Exposed so the SDK author / operator can confirm which schema * version their build is pinned to. */ -export const SUPERVISOR_API_VERSION = "2026-06-16" as const; +export const SUPERVISOR_API_VERSION = "2026-10-30" as const;