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..432c7ce3a4e20 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py @@ -0,0 +1,58 @@ +# 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 a stub task's TaskFlow call (``@task.stub``), stored in the +serialized Dag, and delivered to the lang-SDK runtime through ``TIRunContext.arg_bindings`` +so it can bind the values onto the native task function's parameters. +""" + +from __future__ import annotations + +from typing import Literal + +from pydantic import JsonValue + +from airflow.api_fastapi.core_api.base import BaseModel + +ArgBindingDataType = Literal["string", "integer", "number", "boolean", "object", "array", "any"] +"""Language-neutral value type a stub-task argument binds to in the foreign runtime.""" + + +class TaskArgBinding(BaseModel): + """ + One positional argument of a stub (foreign-runtime) task, in declaration order. + + A deliberately flat shape (``kind`` discriminates instead of a union) so the JSON schema + generates a plain struct in the foreign-language SDKs consuming the supervisor schema. + """ + + kind: Literal["xcom", "literal"] + """Whether the value comes from an upstream task's XCom or is a literal from the Dag file.""" + + data_type: ArgBindingDataType = "any" + """Declared type from the stub function's annotation; runtimes type-check against it.""" + + task_id: str | None = None + """Upstream task id to pull the XCom from. Only set when ``kind`` is ``xcom``.""" + + key: str = "return_value" + """XCom key to pull. Only meaningful when ``kind`` is ``xcom``.""" + + value: JsonValue | None = None + """The literal value from the Dag file. Only set when ``kind`` is ``literal``.""" 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 c1bac7960234d..e3061e68b28c4 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 @@ -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 TaskArgBinding from airflow.api_fastapi.execution_api.datamodels.taskinstance import ( InactiveAssetsResponse, PreviousTIResponse, @@ -110,6 +111,36 @@ log = structlog.get_logger(__name__) tracer = trace.get_tracer(__name__) +# 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. +_STUB_TASK_TYPE = "_StubOperator" + + +def _get_arg_bindings(dag_version_id: UUID | None, task_id: str, *, session) -> list[dict] | None: + """Extract the stub task's serialized positional-arg spec from the serialized Dag blob.""" + # Imported here on purpose: only the Multi-Lang stub-task path touches the + # serialized-dag machinery, so keep it off the module's top-level imports. + from airflow.models.dag_version import DagVersion + from airflow.serialization.enums import Encoding + from airflow.serialization.serialized_objects import BaseSerialization + + if dag_version_id is None: + return None + dag_version = session.get(DagVersion, dag_version_id) + if dag_version is None or dag_version.serialized_dag is None: + return None + data = dag_version.serialized_dag.data + if not data: + return None + for task in data.get("dag", {}).get("tasks", []): + var = task.get(Encoding.VAR) or {} + if var.get("task_id") == task_id: + if encoded := var.get("_arg_bindings"): + return BaseSerialization.deserialize(encoded) + return None + return None + @ti_id_router.patch( "/{task_instance_id}/run", @@ -163,6 +194,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 +343,13 @@ 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 ( + arg_bindings := _get_arg_bindings(ti.dag_version_id, ti.task_id, session=session) + ): + context.arg_bindings = [TaskArgBinding.model_validate(arg) for arg in arg_bindings] + # 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/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py index dc7035d31e3c9..f4a08db4dc7fa 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 @@ -41,6 +41,7 @@ RemoveUpstreamMapIndexesField, ) from airflow.api_fastapi.execution_api.versions.v2026_06_30 import ( + AddArgBindingsToTIRunContext, AddAssetsByAliasEndpoint, AddAwaitingInputStatePayload, AddConnectionTestEndpoint, @@ -65,6 +66,7 @@ AddTaskAndAssetStateStoreEndpoints, AddAssetsByAliasEndpoint, AddPartitionDateField, + AddArgBindingsToTIRunContext, ), Version( "2026-04-06", diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py index e89e2ed04cc5d..15b04fc1db8d2 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py @@ -141,3 +141,16 @@ def remove_partition_date_from_dag_run(response: ResponseInfo) -> None: # type: """Strip ``partition_date`` from the nested ``dag_run`` payload for older clients.""" if "dag_run" in response.body and isinstance(response.body["dag_run"], dict): response.body["dag_run"].pop("partition_date", None) + + +class AddArgBindingsToTIRunContext(VersionChange): + """Add the ``arg_bindings`` positional-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/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 542ce7eaaf15b..ff6678a08ff4b 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 @@ -370,6 +370,43 @@ 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.""" + from airflow.providers.standard.decorators.stub import stub + + def extract(): ... + + def transform(country: str, extracted: dict): ... + + with dag_maker("test_arg_bindings_dag", serialized=True): + stub(transform)("uk", stub(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"] == [ + {"kind": "literal", "data_type": "string", "value": "uk"}, + {"kind": "xcom", "data_type": "object", "task_id": "extract", "key": "return_value"}, + ] + + # 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() + 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_04_17/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_04_17/test_task_instances.py index 71775cb50590a..e51444d239b86 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_04_17/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_04_17/test_task_instances.py @@ -84,3 +84,46 @@ def test_head_version_includes_team_name_field(self, client, session, create_tas response = client.patch(f"/execution/task-instances/{ti.id}/run", json=RUN_PATCH_BODY) assert response.status_code == 200 assert response.json()["dag_run"]["team_name"] is None + + +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): + from airflow.providers.standard.decorators.stub import stub + + def extract(): ... + + def transform(country: str, extracted: dict): ... + + with dag_maker("test_arg_bindings_compat_dag", serialized=True): + stub(transform)("uk", stub(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"] == [ + {"kind": "literal", "data_type": "string", "value": "uk"}, + {"kind": "xcom", "data_type": "object", "task_id": "extract", "key": "return_value"}, + ] diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index 4708365e846e2..f03ac2d64fc96 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,46 @@ 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.providers.standard.decorators.stub import stub + + def extract(): ... + + def transform(country: str, extracted: dict): ... + + with DAG(dag_id="arg_bindings_dag", schedule=None) as dag: + stub(transform)("uk", stub(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: {"kind": "literal", "data_type": "string", "value": "uk"}, + }, + { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: { + "kind": "xcom", + "data_type": "object", + "task_id": "extract", + "key": "return_value", + }, + }, + ] + + round_tripped = DagSerialization.from_dict(ser_dag) + assert round_tripped.task_dict["transform"]._arg_bindings == [ + {"kind": "literal", "data_type": "string", "value": "uk"}, + {"kind": "xcom", "data_type": "object", "task_id": "extract", "key": "return_value"}, + ] + assert not hasattr(round_tripped.task_dict["extract"], "_arg_bindings") or ( + round_tripped.task_dict["extract"]._arg_bindings is None + ) + + 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..2a6065801f105 --- /dev/null +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_taskflow_binding.py @@ -0,0 +1,102 @@ +# 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 ``combine`` 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 ``combine`` 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", "combine"): + 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 and the ``[]int`` as an array.""" + 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] + + +def test_combine_summary_reflects_bound_arguments(completed_run: _CompletedRun): + """``combine`` re-emits every bound value, confirming types survived the + Python literal / XCom -> Go parameter -> XCom round trip.""" + assert completed_run.xcom("combine") == { + "name": "summary", + "count": 3, + "ratio": 2.5, + "enabled": True, + "tags": ["metrics", "hourly"], + "environment": "production", + "debug": True, + "sum": 20, + "note_was_null": True, + } diff --git a/devel-common/src/tests_common/test_utils/version_compat.py b/devel-common/src/tests_common/test_utils/version_compat.py index 7eb25dec2b3cb..d96b9dce07b4d 100644 --- a/devel-common/src/tests_common/test_utils/version_compat.py +++ b/devel-common/src/tests_common/test_utils/version_compat.py @@ -42,6 +42,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_2_PLUS = get_base_airflow_version_tuple() >= (3, 2, 0) AIRFLOW_V_3_2_2_PLUS = get_base_airflow_version_tuple() >= (3, 2, 2) AIRFLOW_V_3_3_PLUS = get_base_airflow_version_tuple() >= (3, 3, 0) +AIRFLOW_V_3_4_PLUS = get_base_airflow_version_tuple() >= (3, 4, 0) if AIRFLOW_V_3_1_PLUS: from airflow.sdk import PokeReturnValue, timezone diff --git a/go-sdk/README.md b/go-sdk/README.md index 7bd6b3d13c81a..3f3b14ad44c8f 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -105,6 +105,16 @@ 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. 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. + ```go func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (any, error) { conn, err := client.GetConnection(ctx, "test_http") @@ -112,12 +122,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 } ``` diff --git a/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md b/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md index 82798bdeb10f4..e7535ab6bea75 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..c4d21083a1c0f 100644 --- a/go-sdk/bundle/bundlev1/task.go +++ b/go-sdk/bundle/bundlev1/task.go @@ -25,30 +25,53 @@ 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} + f := &taskFunction{fn: v, fullName: fullName} return f, f.validateFn(v.Type()) } 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 +79,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 +146,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 +164,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..1535fe5eb097c 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,69 @@ 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{ + {Kind: binding.ArgKindLiteral, Value: "uk", DataType: binding.DataTypeString}, + { + Kind: binding.ArgKindLiteral, + Value: map[string]any{"k": "v"}, + DataType: binding.DataTypeObject, + }, + }) + 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{ + {Kind: binding.ArgKindLiteral, Value: "uk"}, + {Kind: binding.ArgKindLiteral, 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..f6432595a0458 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,11 @@ dags: - "extract" - "transform" - "load" + taskflow_binding_dag: + tasks: + - "make_config" + - "make_numbers" + - "combine" ` assert.Equal(t, expectedManifest, string(metadata)) diff --git a/go-sdk/dags/go_examples.py b/go-sdk/dags/go_examples.py index 23e02dd5e49dc..39328747283c3 100644 --- a/go-sdk/dags/go_examples.py +++ b/go-sdk/dags/go_examples.py @@ -17,9 +17,10 @@ """ 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, 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 +34,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 +70,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 +91,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 +115,51 @@ 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 combine( + name: str, + count: int, + ratio: float, + enabled: bool, + tags: list, + config: dict, + numbers: list, + note: str | None = None, +): ... + + +@dag(dag_id="taskflow_binding_dag") +def taskflow_binding_dag(): + """ + Stress the TaskFlow argument-binding surface beyond ``simple_dag``'s transform. + + 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 ``combine`` + (``go-sdk/example/bundle/taskflowbinding``) verifies every bound value and fails + the task on any mismatch. + """ + combine( + "summary", + 3, + 2.5, + True, + ["metrics", "hourly"], + config=make_config(), + numbers=make_numbers(), + ) + + +taskflow_binding_dag() diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 23e60bd1dd46e..1eaf213666913 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,11 @@ 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("combine", taskflowbinding.Combine) + return nil } @@ -127,12 +133,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..a23b470ed93c6 --- /dev/null +++ b/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go @@ -0,0 +1,129 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// Package taskflowbinding holds the taskflow_binding_dag tasks. Where +// simple_dag's transform shows the minimal TaskFlow binding (one literal, one +// XCom), this Dag stresses the full argument surface: 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. +package taskflowbinding + +import ( + "fmt" + "log/slog" + "reflect" + + "github.com/apache/airflow/go-sdk/sdk" +) + +// Config is the object make_config returns as its XCom; combine 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 combine 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 combine 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 +} + +// Combine receives every argument shape the stub Dag can express. The Python +// side calls it as +// +// combine("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 Combine( + 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 +} 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..51168778ce9b8 --- /dev/null +++ b/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go @@ -0,0 +1,60 @@ +// 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 TestCombine(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + got, err := Combine(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, "Combine should return a map summary, got %T", got) + assert.Equal(t, 20, summary["sum"]) + assert.Equal(t, true, summary["note_was_null"]) +} + +func TestCombineRejectsWrongBinding(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + _, err := Combine(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") +} diff --git a/go-sdk/pkg/binding/binding.go b/go-sdk/pkg/binding/binding.go new file mode 100644 index 0000000000000..c8a621835e017 --- /dev/null +++ b/go-sdk/pkg/binding/binding.go @@ -0,0 +1,425 @@ +// 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. +// +// Two kinds of parameter are supported: +// +// - 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. +// +// 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, failing loudly on arity or type mismatches. 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. +// +// Mapped upstream fan-in is out of scope: XCom arguments always pull the +// unmapped upstream instance (map_index is never sent). +package binding + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "log/slog" + "reflect" + + "github.com/apache/airflow/go-sdk/pkg/api" + "github.com/apache/airflow/go-sdk/pkg/sdkcontext" + "github.com/apache/airflow/go-sdk/sdk" +) + +// ArgKind discriminates how one positional argument is sourced. +type ArgKind string + +const ( + ArgKindXCom ArgKind = "xcom" + ArgKindLiteral ArgKind = "literal" +) + +// DataType is the language-neutral value type the Dag declared for an +// argument (from the stub function's annotation on the Python side). +type DataType string + +const ( + DataTypeString DataType = "string" + DataTypeInteger DataType = "integer" + DataTypeNumber DataType = "number" + DataTypeBoolean DataType = "boolean" + DataTypeObject DataType = "object" + DataTypeArray DataType = "array" + DataTypeAny DataType = "any" +) + +// Arg is one positional argument for a task function's data parameters, in +// declaration order. It is a runtime-neutral mirror of the wire model so this +// package stays decoupled from the generated coordinator schema types. +type Arg struct { + Kind ArgKind + // TaskID is the upstream task to pull from. Set only for ArgKindXCom. + TaskID string + // Key is the XCom key to pull; empty means the return-value key. Set only + // for ArgKindXCom. + Key string + // Value is the literal value from the Dag file. Set only for ArgKindLiteral. + Value any + // DataType is the declared type to check the Go parameter against; empty + // is treated as DataTypeAny. + DataType DataType +} + +// paramKind classifies how a task-function parameter is filled at execution. +type paramKind int + +const ( + paramTIRunContext paramKind = iota + paramContext + paramLogger + paramClient + paramData +) + +// 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). + typ reflect.Type + // index is the parameter's position in the function signature, for error + // messages. + index int +} + +// 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 +} + +// NumData returns how many data parameters the analyzed function declares. +func (p *Plan) NumData() int { return p.numData } + +// 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())} + 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++ + } + p.params[i] = plan + } + return p, nil +} + +// Resolve builds the ordered argument values for one call. Injectable +// parameters receive values derived from ctx, logger, or client; data +// parameters consume args in declaration order. An error fails the task +// before its body runs. +func (p *Plan) Resolve( + ctx context.Context, + logger *slog.Logger, + client sdk.Client, + args []Arg, +) ([]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, + ) + } + out := make([]reflect.Value, len(p.params)) + argIdx := 0 + 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: + arg := args[argIdx] + v, err := p.resolveData(ctx, client, plan, arg, argIdx) + if err != nil { + return nil, err + } + out[i] = v + argIdx++ + } + } + return out, nil +} + +// resolveData produces the value for one data parameter from its argument +// spec: type-check against the declared Dag type, then decode a literal or +// pull-and-decode an XCom. +func (p *Plan) resolveData( + ctx context.Context, + c sdk.XComClient, + plan paramPlan, + arg Arg, + argIdx int, +) (reflect.Value, error) { + if err := checkDataType(arg.DataType, plan.typ); err != nil { + return reflect.Value{}, fmt.Errorf( + "task function %s: argument %d (parameter %d): %w", p.fnName, argIdx, plan.index, err, + ) + } + switch arg.Kind { + case ArgKindLiteral: + v, err := decodeValue(arg.Value, plan.typ) + if err != nil { + return reflect.Value{}, fmt.Errorf( + "task function %s: argument %d: decoding literal value into %s: %w", + p.fnName, argIdx, plan.typ, err, + ) + } + return v, nil + case ArgKindXCom: + workload, ok := ctx.Value(sdkcontext.WorkloadContextKey).(api.ExecuteTaskWorkload) + if !ok { + return reflect.Value{}, fmt.Errorf( + "task function %s: no workload in context, cannot resolve xcom argument %d", + p.fnName, argIdx, + ) + } + key := arg.Key + if key == "" { + key = api.XComReturnValueKey + } + // Pull from the upstream's unmapped instance (map_index nil); mapped + // upstream fan-in is out of scope for now. + raw, err := c.GetXCom(ctx, workload.TI.DagId, workload.TI.RunId, arg.TaskID, nil, key, nil) + if err != nil { + return reflect.Value{}, fmt.Errorf( + "task function %s: argument %d: pulling xcom from task %q (key %q): %w", + p.fnName, argIdx, arg.TaskID, key, err, + ) + } + v, err := decodeValue(raw, plan.typ) + if err != nil { + return reflect.Value{}, fmt.Errorf( + "task function %s: argument %d: decoding xcom from task %q into %s: %w", + p.fnName, argIdx, arg.TaskID, plan.typ, err, + ) + } + return v, nil + default: + return reflect.Value{}, fmt.Errorf( + "task function %s: argument %d: unknown argument kind %q", p.fnName, argIdx, arg.Kind, + ) + } +} + +// 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 +} + +// checkDataType verifies the Dag-declared type can bind to the Go parameter +// type. One pointer level is dereferenced first; DataTypeAny (or an empty +// declaration) skips the check, as does an `any` parameter. +func checkDataType(dt DataType, target reflect.Type) error { + if dt == "" || dt == DataTypeAny { + return nil + } + t := target + if t.Kind() == reflect.Pointer { + t = t.Elem() + } + if t.Kind() == reflect.Interface { + return nil + } + ok := false + switch dt { + case DataTypeString: + ok = t.Kind() == reflect.String + case DataTypeInteger: + switch t.Kind() { + case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, + reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64: + ok = true + } + case DataTypeNumber: + ok = t.Kind() == reflect.Float32 || t.Kind() == reflect.Float64 + case DataTypeBoolean: + ok = t.Kind() == reflect.Bool + case DataTypeObject: + ok = t.Kind() == reflect.Struct || t.Kind() == reflect.Map + case DataTypeArray: + ok = t.Kind() == reflect.Slice || t.Kind() == reflect.Array + default: + return fmt.Errorf("unknown declared type %q in the argument spec", dt) + } + if !ok { + return fmt.Errorf( + "the Dag declares type %q which cannot bind to Go parameter type %s", dt, 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..48992a7fc7372 --- /dev/null +++ b/go-sdk/pkg/binding/binding_test.go @@ -0,0 +1,393 @@ +// 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" + "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/sdkcontext" + "github.com/apache/airflow/go-sdk/sdk" +) + +type BindingSuite struct { + suite.Suite +} + +func TestBindingSuite(t *testing.T) { + suite.Run(t, &BindingSuite{}) +} + +// fakeXComClient records GetXCom calls and returns preconfigured values. +type fakeXComClient struct { + sdk.Client + + values map[string]any // "/" -> raw value + 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.calls = append(f.calls, fakeXComCall{dagID, runID, taskID, key, mapIndex}) + 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{{Kind: ArgKindLiteral, 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{ + {Kind: ArgKindLiteral, Value: "uk", DataType: DataTypeString}, + {Kind: ArgKindLiteral, Value: 3, DataType: DataTypeInteger}, + {Kind: ArgKindLiteral, Value: 1.5, DataType: DataTypeNumber}, + {Kind: ArgKindLiteral, Value: true, DataType: DataTypeBoolean}, + {Kind: ArgKindLiteral, Value: []any{"a", "b"}, DataType: DataTypeArray}, + {Kind: ArgKindLiteral, Value: map[string]any{"k": "v"}, DataType: DataTypeObject}, + }, &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{ + {Kind: ArgKindLiteral, Value: "uk", DataType: DataTypeString}, + {Kind: ArgKindLiteral, Value: map[string]any{"k": "v"}, DataType: DataTypeObject}, + }, &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) TestCheckDataTypeMatrix() { + cases := map[string]struct { + dt DataType + target reflect.Type + errContains string + }{ + "string-ok": {DataTypeString, reflect.TypeFor[string](), ""}, + "string-ptr-ok": {DataTypeString, reflect.TypeFor[*string](), ""}, + "string-vs-int": {DataTypeString, reflect.TypeFor[int](), "cannot bind"}, + "integer-ok": {DataTypeInteger, reflect.TypeFor[int64](), ""}, + "integer-uint-ok": {DataTypeInteger, reflect.TypeFor[uint32](), ""}, + "integer-vs-float": {DataTypeInteger, reflect.TypeFor[float64](), "cannot bind"}, + "number-ok": {DataTypeNumber, reflect.TypeFor[float32](), ""}, + "number-vs-int": {DataTypeNumber, reflect.TypeFor[int](), "cannot bind"}, + "boolean-ok": {DataTypeBoolean, reflect.TypeFor[bool](), ""}, + "boolean-vs-string": {DataTypeBoolean, reflect.TypeFor[string](), "cannot bind"}, + "object-map-ok": {DataTypeObject, reflect.TypeFor[map[string]int](), ""}, + "object-struct-ok": {DataTypeObject, reflect.TypeFor[struct{ A int }](), ""}, + "object-vs-slice": {DataTypeObject, reflect.TypeFor[[]int](), "cannot bind"}, + "array-slice-ok": {DataTypeArray, reflect.TypeFor[[]string](), ""}, + "array-array-ok": {DataTypeArray, reflect.TypeFor[[2]int](), ""}, + "array-vs-map": {DataTypeArray, reflect.TypeFor[map[string]any](), "cannot bind"}, + "any-skips": {DataTypeAny, reflect.TypeFor[chan int](), ""}, + "empty-skips": {DataType(""), reflect.TypeFor[string](), ""}, + "any-target-skips": {DataTypeString, reflect.TypeFor[any](), ""}, + "unknown-dt": {DataType("uuid"), reflect.TypeFor[string](), "unknown declared type"}, + } + for name, tt := range cases { + s.Run(name, func() { + err := checkDataType(tt.dt, 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{{Kind: ArgKindLiteral, Value: "uk", DataType: DataTypeString}}, + &fakeXComClient{}, + ) + if s.Assert().Error(err) { + s.Contains( + err.Error(), + `the Dag declares 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{{Kind: ArgKindLiteral, 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"` +} + +func (s *BindingSuite) TestResolveXComArgs() { + client := &fakeXComClient{values: map[string]any{ + "extract/return_value": map[string]any{"go_version": "go1.24", "timestamp": int64(42)}, + "extract/part": "part-value", + }} + + fn := func(res extractResult, part string) error { return nil } + got, err := s.resolve(fn, []Arg{ + {Kind: ArgKindXCom, TaskID: "extract", DataType: DataTypeObject}, + {Kind: ArgKindXCom, TaskID: "extract", Key: "part", DataType: DataTypeString}, + }, client) + s.Require().NoError(err) + s.Equal(extractResult{GoVersion: "go1.24", Timestamp: 42}, got[0].Interface()) + s.Equal("part-value", got[1].Interface()) + + s.Require().Len(client.calls, 2) + s.Equal("dag1", client.calls[0].dagID) + s.Equal("run1", client.calls[0].runID) + s.Equal("extract", client.calls[0].taskID) + s.Equal( + api.XComReturnValueKey, + client.calls[0].key, + "an empty key must default to the return-value key", + ) + s.Nil(client.calls[0].mapIndex, "v1 always pulls the unmapped upstream instance") + s.Equal("part", client.calls[1].key) +} + +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{{Kind: ArgKindXCom, 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{{Kind: ArgKindXCom, 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{{Kind: ArgKindXCom, 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{{Kind: ArgKindLiteral, Value: nil, DataType: DataTypeObject}}, + &fakeXComClient{}, + ) + s.Require().NoError(err) + s.Nil(got[0].Interface()) + + fnStr := func(country string) error { return nil } + _, err = s.resolve(fnStr, []Arg{{Kind: ArgKindLiteral, Value: nil}}, &fakeXComClient{}) + if s.Assert().Error(err) { + s.Contains(err.Error(), "not nilable") + } +} + +func (s *BindingSuite) TestResolveUnknownKind() { + fn := func(country string) error { return nil } + _, err := s.resolve(fn, []Arg{{Kind: ArgKind("template"), Value: "x"}}, &fakeXComClient{}) + if s.Assert().Error(err) { + s.Contains(err.Error(), `unknown argument kind "template"`) + } +} + +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{ + {Kind: ArgKindLiteral, Value: "uk", DataType: DataTypeString}, + }) + 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()) +} 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/defaults.gen.go b/go-sdk/pkg/execution/genmodels/defaults.gen.go index a4e9b1212823d..d411881bc8404 100644 --- a/go-sdk/pkg/execution/genmodels/defaults.gen.go +++ b/go-sdk/pkg/execution/genmodels/defaults.gen.go @@ -282,6 +282,17 @@ func (m *SucceedTask) DecodeMsgpack(dec *msgpack.Decoder) error { return nil } +// DecodeMsgpack applies TaskArgBinding's schema defaults that msgpack would otherwise skip. +func (m *TaskArgBinding) DecodeMsgpack(dec *msgpack.Decoder) error { + type alias TaskArgBinding + v := alias{DataType: TaskArgBindingDataType("any"), Key: "return_value"} + if err := dec.Decode(&v); err != nil { + return err + } + *m = TaskArgBinding(v) + return nil +} + // DecodeMsgpack applies TaskInstance's schema defaults that msgpack would otherwise skip. func (m *TaskInstance) DecodeMsgpack(dec *msgpack.Decoder) error { type alias TaskInstance diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index e6861d8c8add5..2e6caf47c3da3 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -20,6 +20,8 @@ package genmodels import "time" +type ArgBindings []TaskArgBinding + // Schema for AssetAliasModel used in AssetEventDagRunReference. type AssetAliasReferenceAssetEventDagRun struct { // Name corresponds to the JSON schema field "name". @@ -370,6 +372,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. @@ -629,6 +634,16 @@ const DagRunTypeScheduled DagRunType = "scheduled" type Data map[string]interface{} +type DataType string + +const DataTypeAny DataType = "any" +const DataTypeArray DataType = "array" +const DataTypeBoolean DataType = "boolean" +const DataTypeInteger DataType = "integer" +const DataTypeNumber DataType = "number" +const DataTypeObject DataType = "object" +const DataTypeString DataType = "string" + type Defaults []string // Update a task instance state to deferred. @@ -749,6 +764,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,6 +826,9 @@ 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"` @@ -834,6 +855,9 @@ 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"` @@ -1226,6 +1250,11 @@ type InactiveAssetsResult struct { type JsonValue interface{} +type Kind string + +const KindLiteral Kind = "literal" +const KindXcom Kind = "xcom" + // Lazily build information from the serialized DAG structure. // // An object that will present "enough" of the DAG like interface to update DAG db @@ -1564,6 +1593,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 +1628,44 @@ type TIRunContext struct { XcomKeysToClear []string `msgpack:"xcom_keys_to_clear,omitempty"` } +// One positional argument of a stub (foreign-runtime) task, in declaration order. +// +// A deliberately flat shape (“kind“ discriminates instead of a union) so the +// JSON schema +// generates a plain struct in the foreign-language SDKs consuming the supervisor +// schema. +type TaskArgBinding struct { + // DataType corresponds to the JSON schema field "data_type". + DataType TaskArgBindingDataType `msgpack:"data_type,omitempty"` + + // Key corresponds to the JSON schema field "key". + Key string `msgpack:"key,omitempty"` + + // Kind corresponds to the JSON schema field "kind". + Kind TaskArgBindingKind `msgpack:"kind"` + + // TaskID corresponds to the JSON schema field "task_id". + TaskID interface{} `msgpack:"task_id,omitempty"` + + // Value corresponds to the JSON schema field "value". + Value interface{} `msgpack:"value,omitempty"` +} + +type TaskArgBindingDataType string + +const TaskArgBindingDataTypeAny TaskArgBindingDataType = "any" +const TaskArgBindingDataTypeArray TaskArgBindingDataType = "array" +const TaskArgBindingDataTypeBoolean TaskArgBindingDataType = "boolean" +const TaskArgBindingDataTypeInteger TaskArgBindingDataType = "integer" +const TaskArgBindingDataTypeNumber TaskArgBindingDataType = "number" +const TaskArgBindingDataTypeObject TaskArgBindingDataType = "object" +const TaskArgBindingDataTypeString TaskArgBindingDataType = "string" + +type TaskArgBindingKind string + +const TaskArgBindingKindLiteral TaskArgBindingKind = "literal" +const TaskArgBindingKindXcom TaskArgBindingKind = "xcom" + type TaskBreadcrumbsResult struct { // Breadcrumbs corresponds to the JSON schema field "breadcrumbs". Breadcrumbs []TaskBreadcrumbsResultBreadcrumbsElem `msgpack:"breadcrumbs"` @@ -1636,6 +1706,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 +1850,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 +1865,19 @@ 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 Warnings []interface{} + +type VersionData map[string]interface{} + type ValidateInletsAndOutlets struct { // TIID corresponds to the JSON schema field "ti_id". TIID string `msgpack:"ti_id"` diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 2559359453866..5646a82e29206 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -230,6 +230,115 @@ 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{ + {Kind: "literal", DataType: "string", Value: "uk"}, + {Kind: "literal", DataType: "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{ + {Kind: "literal", DataType: "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") +} + +// 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{ + {Kind: "literal", DataType: "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) +} + 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..378cc931dda50 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-07-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..d6081d24ad4c3 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,33 @@ 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 := convertArgBindings(details.TIContext.ArgBindings) + + 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-neutral binding form. +func convertArgBindings(specsPtr *genmodels.ArgBindings) []binding.Arg { + if specsPtr == nil || len(*specsPtr) == 0 { + return nil + } + specs := *specsPtr + args := make([]binding.Arg, len(specs)) + for i, spec := range specs { + taskID := "" + if s, ok := spec.TaskID.(string); ok { + taskID = s + } + args[i] = binding.Arg{ + Kind: binding.ArgKind(spec.Kind), + TaskID: taskID, + Key: spec.Key, + Value: spec.Value, + DataType: binding.DataType(spec.DataType), + } + } + return args } // mapIndexPtr normalizes the supervisor's map_index into the optional form @@ -142,9 +169,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 +201,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/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index 08bcf163a56ad..72497f184e32e 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -18,8 +18,12 @@ from __future__ import annotations import ast -from collections.abc import Callable -from typing import TYPE_CHECKING, Any +import inspect +import json +import types +import typing +from collections.abc import Callable, Collection, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Union from airflow.providers.common.compat.sdk import ( DecoratedOperator, @@ -27,13 +31,146 @@ task_decorator_factory, ) +try: + from airflow.sdk.definitions.xcom_arg import PlainXComArg, XComArg +except ImportError: # Airflow 2 + from airflow.models.xcom_arg import PlainXComArg, XComArg # type: ignore[attr-defined,no-redef] + +try: + from airflow.sdk.definitions.context import KNOWN_CONTEXT_KEYS +except ImportError: # Airflow 2, and 3.0 where the SDK does not export it yet + from airflow.utils.context import KNOWN_CONTEXT_KEYS # type: ignore[attr-defined,no-redef] + if TYPE_CHECKING: from airflow.providers.common.compat.sdk import Context +def _data_type_from_annotation(annotation: Any) -> str: + """ + Map a stub function parameter annotation to the language-neutral arg-type vocabulary. + + The foreign runtime type-checks the bound value against the returned name; anything we + cannot classify confidently maps to ``"any"`` so binding falls back to a decode-only check. + """ + if annotation is inspect.Parameter.empty or annotation is None or annotation is Any: + return "any" + origin = typing.get_origin(annotation) + if origin is not None: + if origin is Union or origin is getattr(types, "UnionType", None): + members = [a for a in typing.get_args(annotation) if a is not type(None)] + if len(members) == 1: + return _data_type_from_annotation(members[0]) + return "any" + annotation = origin + if not isinstance(annotation, type): + return "any" + # bool subclasses int, and str/bytes are Sequences -- order matters. + if issubclass(annotation, bool): + return "boolean" + if issubclass(annotation, int): + return "integer" + if issubclass(annotation, float): + return "number" + if issubclass(annotation, str): + return "string" + if issubclass(annotation, bytes): + return "any" + if issubclass(annotation, (dict, Mapping)): + return "object" + if issubclass(annotation, (list, tuple, set, frozenset, Sequence)): + return "array" + return "any" + + +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 the execution API ``TaskArgBinding`` shape: an XCom + reference (``kind="xcom"``) for upstream TaskFlow outputs, or an inline value + (``kind="literal"``) for everything else. Returns ``None`` for parameterless stubs. + """ + signature = inspect.signature(python_callable) + + 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)" + ) + + if not signature.parameters: + return None + + bound = signature.bind(*op_args, **op_kwargs) + bound.apply_defaults() + + try: + hints = typing.get_type_hints(python_callable) + except Exception: + # Annotations that cannot be resolved at parse time (e.g. names behind + # TYPE_CHECKING with ``from __future__ import annotations``) degrade to "any". + hints = {} + + def annotation_for(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 + + spec: list[dict[str, Any]] = [] + for name, param in signature.parameters.items(): + value = bound.arguments[name] + data_type = _data_type_from_annotation(annotation_for(name, param)) + if isinstance(value, PlainXComArg): + spec.append( + { + "kind": "xcom", + "data_type": data_type, + "task_id": value.operator.task_id, + "key": value.key, + } + ) + continue + 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 (optionally " + "indexed by key) can cross the language boundary -- .map()/.zip()/.concat() " + "results are not supported" + ) + try: + json.dumps(value) + 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" + ) + spec.append({"kind": "literal", "data_type": data_type, "value": value}) + return spec + + class _StubOperator(DecoratedOperator): custom_operator_name: str = "@task.stub" + # Mapped stubs would need per-map-index arg specs, which the foreign runtime cannot + # receive yet; the task-sdk decorator machinery rejects .expand() at parse time for + # operator classes that opt out. + supports_expand: bool = False + def __init__( self, *, @@ -75,7 +212,14 @@ 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) + + @classmethod + def get_serialized_fields(cls): + return super().get_serialized_fields() | {"_arg_bindings"} def execute(self, context: Context) -> Any: raise RuntimeError( @@ -96,6 +240,9 @@ 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 positional-argument spec is delivered + to the foreign runtime, which binds the values onto the native task function. """ 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..3be7ea18e1032 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -17,12 +17,15 @@ from __future__ import annotations import contextlib +import typing +from typing import Any import pytest -from airflow.providers.standard.decorators.stub import stub +from airflow.providers.common.compat.sdk import DAG +from airflow.providers.standard.decorators.stub import _data_type_from_annotation, stub -from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS +from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS, AIRFLOW_V_3_4_PLUS def fn_ellipsis(): ... @@ -69,3 +72,129 @@ 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 == [ + {"kind": "literal", "data_type": "string", "value": "uk"}, + {"kind": "xcom", "data_type": "object", "task_id": "fn_extract", "key": "return_value"}, + {"kind": "literal", "data_type": "integer", "value": 3}, + ] + 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["part"], country="fr", retries_num=7) + + assert result.operator._arg_bindings == [ + {"kind": "literal", "data_type": "string", "value": "fr"}, + {"kind": "xcom", "data_type": "object", "task_id": "fn_extract", "key": "part"}, + {"kind": "literal", "data_type": "integer", "value": 7}, + ] + + def test_zero_param_stub_has_no_spec(self): + assert stub(fn_pass)().operator._arg_bindings is None + + def test_untyped_params_degrade_to_any(self): + with DAG(dag_id="d"): + result = stub(fn_untyped)(1, "x") + + assert result.operator._arg_bindings == [ + {"kind": "literal", "data_type": "any", "value": 1}, + {"kind": "literal", "data_type": "any", "value": "x"}, + ] + + def test_unresolvable_annotation_degrades_to_any(self): + def fn(x): ... + + fn.__annotations__ = {"x": "NotARealType"} + with DAG(dag_id="d"): + result = stub(fn)("v") + + assert result.operator._arg_bindings == [{"kind": "literal", "data_type": "any", "value": "v"}] + + def test_varargs_rejected(self): + with pytest.raises(ValueError, match="fixed number of parameters"): + stub(fn_varargs)() + + def test_varkw_rejected(self): + with pytest.raises(ValueError, match="fixed number of parameters"): + stub(fn_kwonly_varkw)() + + def test_context_key_param_rejected(self): + with pytest.raises(ValueError, match="is an Airflow context key"): + stub(fn_context_key)(1) + + 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_mapped_xcom_arg_rejected(self): + with DAG(dag_id="d"): + extracted = stub(fn_extract)() + with pytest.raises(ValueError, match="MapXComArg"): + stub(fn_transform)("uk", extracted.map(lambda v: v)) + + @pytest.mark.skipif( + not AIRFLOW_V_3_4_PLUS, reason="task-sdk honors the supports_expand opt-out from Airflow 3.4" + ) + def test_expand_rejected_at_parse_time(self): + with DAG(dag_id="d"): + with pytest.raises(TypeError, match="do not support dynamic task mapping"): + stub(fn_transform).expand(country=["uk", "fr"], extracted=[{}, {}]) + + +@pytest.mark.parametrize( + ("annotation", "expected"), + [ + pytest.param(str, "string", id="str"), + pytest.param(bool, "boolean", id="bool"), + pytest.param(int, "integer", id="int"), + pytest.param(float, "number", id="float"), + pytest.param(dict, "object", id="dict"), + pytest.param(dict[str, int], "object", id="dict-parameterized"), + pytest.param(typing.Mapping[str, int], "object", id="mapping"), + pytest.param(list, "array", id="list"), + pytest.param(list[int], "array", id="list-parameterized"), + pytest.param(tuple, "array", id="tuple"), + pytest.param(set, "array", id="set"), + pytest.param(typing.Sequence[int], "array", id="sequence"), + pytest.param(Any, "any", id="any"), + pytest.param(None, "any", id="none"), + pytest.param(bytes, "any", id="bytes"), + pytest.param(typing.Optional[str], "string", id="optional-str"), # noqa: UP045 -- legacy form on purpose + pytest.param(typing.Union[int, str], "any", id="union"), # noqa: UP007 -- legacy form on purpose + pytest.param(str | None, "string", id="pep604-optional"), + pytest.param(int | str, "any", id="pep604-union"), + pytest.param(contextlib.AbstractContextManager, "any", id="custom-class"), + ], +) +def test_data_type_from_annotation(annotation, expected): + assert _data_type_from_annotation(annotation) == expected diff --git a/task-sdk/pyproject.toml b/task-sdk/pyproject.toml index 67c7ca053562b..28e2812ec6a0b 100644 --- a/task-sdk/pyproject.toml +++ b/task-sdk/pyproject.toml @@ -253,6 +253,7 @@ disable-timestamp=true enable-version-header=true enum-field-as-literal='one' # When a single enum member, make it output a `Literal["..."]` input-file-type='openapi' +set-default-enum-member=true # `= DataType.ANY` not `= "any"`, keeping mypy happy with enum-typed fields output-model-type='pydantic_v2.BaseModel' output-datetime-class='AwareDatetime' target-python-version='3.10' diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 72aebeab3e49b..aa62452980d5c 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -371,6 +371,36 @@ class TITargetStatePayload(BaseModel): state: IntermediateTIState +class Kind(str, Enum): + XCOM = "xcom" + LITERAL = "literal" + + +class DataType(str, Enum): + STRING = "string" + INTEGER = "integer" + NUMBER = "number" + BOOLEAN = "boolean" + OBJECT = "object" + ARRAY = "array" + ANY = "any" + + +class TaskArgBinding(BaseModel): + """ + One positional argument of a stub (foreign-runtime) task, in declaration order. + + A deliberately flat shape (``kind`` discriminates instead of a union) so the JSON schema + generates a plain struct in the foreign-language SDKs consuming the supervisor schema. + """ + + kind: Annotated[Kind, Field(title="Kind")] + data_type: Annotated[DataType | None, Field(title="Data Type")] = DataType.ANY + task_id: Annotated[str | None, Field(title="Task Id")] = None + key: Annotated[str | None, Field(title="Key")] = "return_value" + value: JsonValue | None = None + + class TaskBreadcrumbsResponse(BaseModel): """ Response for task breadcrumbs. @@ -797,3 +827,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/bases/decorator.py b/task-sdk/src/airflow/sdk/bases/decorator.py index e45778725cc16..fc3a01cf95a53 100644 --- a/task-sdk/src/airflow/sdk/bases/decorator.py +++ b/task-sdk/src/airflow/sdk/bases/decorator.py @@ -588,6 +588,11 @@ def expand_kwargs(self, kwargs: OperatorExpandKwargsArgument, *, strict: bool = return self._expand(ListOfDictsExpandInput(kwargs), strict=strict) def _expand(self, expand_input: ExpandInput, *, strict: bool) -> XComArg: + if not getattr(self.operator_class, "supports_expand", True): + operator_name = ( + getattr(self.operator_class, "custom_operator_name", None) or self.operator_class.__name__ + ) + raise TypeError(f"{operator_name} tasks do not support dynamic task mapping (.expand())") ensure_xcomarg_return_value(expand_input.value) task_kwargs = self.kwargs.copy() 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 01e3c8970b5cb..66797d25d81bb 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -1,6 +1,6 @@ { "$schema": "https://json-schema.org/draft/2020-12/schema", - "api_version": "2026-06-16", + "api_version": "2026-07-30", "description": "Apache Airflow SDK Supervisor Schema", "$defs": { "AssetAliasReferenceAssetEventDagRun": { @@ -1395,6 +1395,19 @@ "title": "DagRunType", "type": "string" }, + "DataType": { + "enum": [ + "string", + "integer", + "number", + "boolean", + "object", + "array", + "any" + ], + "title": "DataType", + "type": "string" + }, "DeferTask": { "additionalProperties": false, "description": "Update a task instance state to deferred.", @@ -2998,6 +3011,14 @@ "type": "object" }, "JsonValue": {}, + "Kind": { + "enum": [ + "xcom", + "literal" + ], + "title": "Kind", + "type": "string" + }, "LazyDeserializedDAG": { "description": "Lazily build information from the serialized DAG structure.\n\nAn object that will present \"enough\" of the DAG like interface to update DAG db models etc, without having\nto deserialize the full DAG and Task hierarchy.", "properties": { @@ -4878,6 +4899,21 @@ ], "default": null, "title": "Start Date" + }, + "arg_bindings": { + "anyOf": [ + { + "items": { + "$ref": "#/$defs/TaskArgBinding" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Arg Bindings" } }, "required": [ @@ -4887,6 +4923,66 @@ "title": "TIRunContext", "type": "object" }, + "TaskArgBinding": { + "description": "One positional argument of a stub (foreign-runtime) task, in declaration order.\n\nA deliberately flat shape (``kind`` discriminates instead of a union) so the JSON schema\ngenerates a plain struct in the foreign-language SDKs consuming the supervisor schema.", + "properties": { + "kind": { + "enum": [ + "xcom", + "literal" + ], + "title": "Kind", + "type": "string" + }, + "data_type": { + "default": "any", + "enum": [ + "string", + "integer", + "number", + "boolean", + "object", + "array", + "any" + ], + "title": "Data Type", + "type": "string" + }, + "task_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Task Id" + }, + "key": { + "default": "return_value", + "title": "Key", + "type": "string" + }, + "value": { + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ], + "default": null + } + }, + "required": [ + "kind" + ], + "title": "TaskArgBinding", + "type": "object" + }, "TaskInstance": { "description": "Schema for TaskInstance model with minimal required fields needed for Runtime.", "properties": { 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..266aeb0c19736 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,11 @@ def get_bundle() -> VersionBundle: """ from cadwyn import HeadVersion, Version, VersionBundle + from airflow.sdk.execution_time.schema.versions.v2026_07_30 import AddArgBindingsToTIRunContext + return VersionBundle( HeadVersion(), + Version("2026-07-30", AddArgBindingsToTIRunContext), Version("2026-06-16"), ) diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_07_30.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_07_30.py new file mode 100644 index 0000000000000..fd25aef095bb8 --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_07_30.py @@ -0,0 +1,30 @@ +# 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 AddArgBindingsToTIRunContext(VersionChange): + """Add the ``arg_bindings`` positional-argument binding spec for stub (foreign-runtime) tasks.""" + + 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..4b146a62859ac 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,76 @@ 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. + + ``AddArgBindingsToTIRunContext`` 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=[ + {"kind": "literal", "data_type": "string", "value": "uk"}, + {"kind": "xcom", "data_type": "object", "task_id": "extract", "key": "return_value"}, + ], + ), + 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): + out = real_migrator.downgrade(startup_details, "2026-07-30") + assert out.ti_context.arg_bindings is not None + assert [a.kind for a in out.ti_context.arg_bindings] == ["literal", "xcom"] diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 4766c0e8a7949..0d3aac7e95a77 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -245,6 +245,11 @@ export type NextKwargs1 = export type XcomKeysToClear = string[]; export type ShouldRetry = boolean; export type StartDate2 = string | null; +export type ArgBindings = TaskArgBinding[] | null; +export type Kind = "xcom" | "literal"; +export type DataType = "string" | "integer" | "number" | "boolean" | "object" | "array" | "any"; +export type TaskId1 = string | null; +export type Key1 = string; export type Type13 = "TaskCallbackRequest"; export type Filepath2 = string; export type BundleName3 = string; @@ -294,6 +299,11 @@ export type Note1 = string | null; export type TeamName1 = string | null; export type Type18 = "DagRunResult"; export type Type19 = "DagRunStateResult"; +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "DataType". + */ +export type DataType1 = "string" | "integer" | "number" | "boolean" | "object" | "array" | "any"; export type State2 = "deferred" | null; export type Classpath = string; export type TriggerKwargs = @@ -311,20 +321,20 @@ export type NextKwargs2 = { export type RenderedMapIndex1 = string | null; export type Type20 = "DeferTask"; export type Name8 = string; -export type Key1 = string; +export type Key2 = string; export type Type21 = "DeleteAssetStateStoreByName"; export type Uri5 = string; -export type Key2 = string; +export type Key3 = string; export type Type22 = "DeleteAssetStateStoreByUri"; export type TiId2 = string; -export type Key3 = string; -export type Type23 = "DeleteTaskStateStore"; export type Key4 = string; -export type Type24 = "DeleteVariable"; +export type Type23 = "DeleteTaskStateStore"; export type Key5 = string; +export type Type24 = "DeleteVariable"; +export type Key6 = 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"; /** @@ -386,10 +396,10 @@ export type Extra8 = { } | null; export type Type30 = "GetAssetEventByAssetAlias"; export type Name11 = string; -export type Key6 = string; +export type Key7 = string; export type Type31 = "GetAssetStateStoreByName"; export type Uri8 = string; -export type Key7 = string; +export type Key8 = string; export type Type32 = "GetAssetStateStoreByUri"; export type AliasName1 = string; export type Type33 = "GetAssetsByAlias"; @@ -417,7 +427,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"; @@ -436,7 +446,7 @@ export type TiId5 = string; export type TryNumber1 = number; export type Type45 = "GetTaskRescheduleStartDate"; export type TiId6 = string; -export type Key8 = string; +export type Key9 = string; export type Type46 = "GetTaskStateStore"; export type DagId15 = string; export type MapIndex4 = number | null; @@ -445,34 +455,34 @@ export type TaskGroupId1 = string | null; export type LogicalDates2 = string[] | null; export type RunIds2 = string[] | null; export type Type47 = "GetTaskStates"; -export type Key9 = string; +export type Key10 = string; export type Type48 = "GetVariable"; export type Prefix = string | null; export type Limit2 = number; export type Offset = number; export type Type49 = "GetVariableKeys"; -export type Key10 = string; +export type Key11 = 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 Key12 = 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 Key13 = 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 Key14 = 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; @@ -494,6 +504,11 @@ export type AssignedUsers1 = HITLUser[] | null; export type Type54 = "HITLDetailRequestResult"; export type InactiveAssets = AssetProfile[] | null; export type Type55 = "InactiveAssetsResult"; +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "Kind". + */ +export type Kind1 = "xcom" | "literal"; export type Name12 = string | null; export type Type56 = "MaskSecret"; export type Ok = boolean; @@ -504,7 +519,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; @@ -515,7 +530,7 @@ export type TryNumber2 = number; export type MapIndex6 = number | null; export type Duration = number | null; export type Type60 = "PreviousTIResult"; -export type Key14 = string; +export type Key15 = string; export type Value1 = string | null; export type Description = string | null; export type Type61 = "PutVariable"; @@ -533,22 +548,22 @@ export type Type64 = "RetryTask"; export type Type65 = "SentFDs"; export type Fds = number[]; export type Name13 = string; -export type Key15 = string; +export type Key16 = string; export type Type66 = "SetAssetStateStoreByName"; export type Uri9 = string; -export type Key16 = string; +export type Key17 = string; export type Type67 = "SetAssetStateStoreByUri"; export type Type68 = "SetRenderedFields"; export type RenderedMapIndex3 = string; export type Type69 = "SetRenderedMapIndex"; export type TiId8 = string; -export type Key17 = string; +export type Key18 = string; export type ExpiresAt = string | null; export type Type70 = "SetTaskStateStore"; -export type Key18 = string; +export type Key19 = 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; @@ -608,12 +623,12 @@ export type Type83 = "ValidateInletsAndOutlets"; export type Keys = string[]; export type TotalEntries = number; export type Type84 = "VariableKeysResult"; -export type Key19 = string; +export type Key20 = string; export type Value2 = string | null; export type Type85 = "VariableResult"; export type Len = number; export type Type86 = "XComCountResponse"; -export type Key20 = string; +export type Key21 = string; export type Type87 = "XComResult"; export type Type88 = "XComSequenceIndexResult"; export type Root = JsonValue[]; @@ -999,6 +1014,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. @@ -1026,6 +1042,22 @@ export interface ConnectionResponse { port: Port1; extra: Extra6; } +/** + * One positional argument of a stub (foreign-runtime) task, in declaration order. + * + * A deliberately flat shape (``kind`` discriminates instead of a union) so the JSON schema + * generates a plain struct in the foreign-language SDKs consuming the supervisor schema. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskArgBinding". + */ +export interface TaskArgBinding { + kind: Kind; + data_type?: DataType; + task_id?: TaskId1; + key?: Key1; + value?: unknown; +} /** * Email notification request for task failures/retries. * @@ -1146,7 +1178,7 @@ export interface DeferTask { */ export interface DeleteAssetStateStoreByName { name: Name8; - key: Key1; + key: Key2; type?: Type21; } /** @@ -1155,7 +1187,7 @@ export interface DeleteAssetStateStoreByName { */ export interface DeleteAssetStateStoreByUri { uri: Uri5; - key: Key2; + key: Key3; type?: Type22; } /** @@ -1164,7 +1196,7 @@ export interface DeleteAssetStateStoreByUri { */ export interface DeleteTaskStateStore { ti_id: TiId2; - key: Key3; + key: Key4; type?: Type23; } /** @@ -1172,7 +1204,7 @@ export interface DeleteTaskStateStore { * via the `definition` "DeleteVariable". */ export interface DeleteVariable { - key: Key4; + key: Key5; type?: Type24; } /** @@ -1180,10 +1212,10 @@ export interface DeleteVariable { * via the `definition` "DeleteXCom". */ export interface DeleteXCom { - key: Key5; + key: Key6; dag_id: DagId6; run_id: RunId5; - task_id: TaskId1; + task_id: TaskId2; map_index?: MapIndex1; type?: Type25; } @@ -1245,7 +1277,7 @@ export interface GetAssetEventByAssetAlias { */ export interface GetAssetStateStoreByName { name: Name11; - key: Key6; + key: Key7; type?: Type31; } /** @@ -1254,7 +1286,7 @@ export interface GetAssetStateStoreByName { */ export interface GetAssetStateStoreByUri { uri: Uri8; - key: Key7; + key: Key8; type?: Type32; } /** @@ -1346,7 +1378,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; @@ -1390,7 +1422,7 @@ export interface GetTaskRescheduleStartDate { */ export interface GetTaskStateStore { ti_id: TiId6; - key: Key8; + key: Key9; type?: Type46; } /** @@ -1411,7 +1443,7 @@ export interface GetTaskStates { * via the `definition` "GetVariable". */ export interface GetVariable { - key: Key9; + key: Key10; type?: Type48; } /** @@ -1429,10 +1461,10 @@ export interface GetVariableKeys { * via the `definition` "GetXCom". */ export interface GetXCom { - key: Key10; + key: Key11; dag_id: DagId16; run_id: RunId9; - task_id: TaskId3; + task_id: TaskId4; map_index?: MapIndex5; include_prior_dates?: IncludePriorDates; type?: Type50; @@ -1444,10 +1476,10 @@ export interface GetXCom { * via the `definition` "GetXComCount". */ export interface GetXComCount { - key: Key11; + key: Key12; dag_id: DagId17; run_id: RunId10; - task_id: TaskId4; + task_id: TaskId5; type?: Type51; } /** @@ -1455,10 +1487,10 @@ export interface GetXComCount { * via the `definition` "GetXComSequenceItem". */ export interface GetXComSequenceItem { - key: Key12; + key: Key13; dag_id: DagId18; run_id: RunId11; - task_id: TaskId5; + task_id: TaskId6; offset: Offset1; type?: Type52; } @@ -1467,10 +1499,10 @@ export interface GetXComSequenceItem { * via the `definition` "GetXComSequenceSlice". */ export interface GetXComSequenceSlice { - key: Key13; + key: Key14; dag_id: DagId19; run_id: RunId12; - task_id: TaskId6; + task_id: TaskId7; start: Start; stop: Stop; step: Step; @@ -1551,7 +1583,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; @@ -1577,7 +1609,7 @@ export interface PreviousTIResult { * via the `definition` "PutVariable". */ export interface PutVariable { - key: Key14; + key: Key15; value: Value1; description: Description; type?: Type61; @@ -1629,7 +1661,7 @@ export interface SentFDs { */ export interface SetAssetStateStoreByName { name: Name13; - key: Key15; + key: Key16; value: JsonValue; type?: Type66; } @@ -1639,7 +1671,7 @@ export interface SetAssetStateStoreByName { */ export interface SetAssetStateStoreByUri { uri: Uri9; - key: Key16; + key: Key17; value: JsonValue; type?: Type67; } @@ -1672,7 +1704,7 @@ export interface SetRenderedMapIndex { */ export interface SetTaskStateStore { ti_id: TiId8; - key: Key17; + key: Key18; value: JsonValue; expires_at: ExpiresAt; type?: Type70; @@ -1682,11 +1714,11 @@ export interface SetTaskStateStore { * via the `definition` "SetXCom". */ export interface SetXCom { - key: Key18; + key: Key19; value: JsonValue; dag_id: DagId21; run_id: RunId14; - task_id: TaskId8; + task_id: TaskId9; map_index?: MapIndex7; dag_result?: DagResult1; mapped_length?: MappedLength; @@ -1843,7 +1875,7 @@ export interface VariableKeysResult { * via the `definition` "VariableResult". */ export interface VariableResult { - key: Key19; + key: Key20; value?: Value2; type?: Type85; } @@ -1862,7 +1894,7 @@ export interface XComCountResponse { * via the `definition` "XComResult". */ export interface XComResult { - key: Key20; + key: Key21; value: JsonValue; type?: Type87; } @@ -1888,4 +1920,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-07-30" as const;