From e5d95c4d954c8ebb04e899ab29c69178f2db9aac Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 10 Jul 2026 07:24:18 +0000 Subject: [PATCH 01/40] Support TaskFlow call syntax on stub tasks for the Go SDK Stub Dags could only declare argless tasks, so cross-language dataflow required hand-written GetXCom calls inside each Go task. Capturing the TaskFlow call's argument spec at parse time and delivering it through the Execution API and StartupDetails lets a Go task receive upstream outputs and Dag-file literals as plain typed parameters, with loud arity/type errors instead of silently zero-filled values. --- .../execution_api/datamodels/taskinstance.py | 35 ++ .../execution_api/routes/task_instances.py | 37 ++ .../execution_api/versions/__init__.py | 2 + .../execution_api/versions/v2026_06_30.py | 13 + .../versions/head/test_task_instances.py | 37 ++ .../v2026_04_17/test_task_instances.py | 43 ++ .../serialization/test_dag_serialization.py | 42 +- go-sdk/README.md | 19 +- .../0003-coordinator-protocol-msgpack-ipc.md | 11 +- go-sdk/bundle/bundlev1/task.go | 147 ++---- go-sdk/bundle/bundlev1/task_test.go | 87 +++- .../airflow-go-pack/pack_integration_test.go | 3 +- go-sdk/dags/go_examples.py | 11 +- go-sdk/example/bundle/main.go | 18 +- go-sdk/example/bundle/main_test.go | 6 +- go-sdk/pkg/binding/binding.go | 425 ++++++++++++++++++ go-sdk/pkg/binding/binding_test.go | 393 ++++++++++++++++ go-sdk/pkg/execution/frames.go | 7 + .../pkg/execution/genmodels/defaults.gen.go | 11 + go-sdk/pkg/execution/genmodels/models.gen.go | 220 ++++++--- go-sdk/pkg/execution/integration_test.go | 109 +++++ go-sdk/pkg/execution/messages.go | 2 +- go-sdk/pkg/execution/task_runner.go | 49 +- .../providers/standard/decorators/stub.py | 150 ++++++- .../unit/standard/decorators/test_stub.py | 141 +++++- task-sdk/pyproject.toml | 1 + .../airflow/sdk/api/datamodels/_generated.py | 31 ++ task-sdk/src/airflow/sdk/bases/decorator.py | 5 + .../sdk/execution_time/schema/schema.json | 98 +++- .../schema/versions/__init__.py | 3 + .../schema/versions/v2026_07_30.py | 30 ++ .../execution_time/schema/test_migrator.py | 82 +++- ts-sdk/src/generated/supervisor.ts | 146 +++--- 33 files changed, 2133 insertions(+), 281 deletions(-) create mode 100644 go-sdk/pkg/binding/binding.go create mode 100644 go-sdk/pkg/binding/binding_test.go create mode 100644 task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_07_30.py 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..2976ad26f905f 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 @@ -393,6 +393,34 @@ def safe_extract_from_orm(cls, data: Any) -> Any: return values +StubArgDataType = Literal["string", "integer", "number", "boolean", "object", "array", "any"] +"""Language-neutral value type a stub-task argument binds to in the foreign runtime.""" + + +class StubTaskArg(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: StubArgDataType = "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``.""" + + class TIRunContext(BaseModel): """Response schema for TaskInstance run context.""" @@ -435,6 +463,13 @@ class TIRunContext(BaseModel): always reflects when the task *first* started, not when it was rescheduled/resumed. """ + stub_args: list[StubTaskArg] | 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..bbec7f542f22a 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 @@ -53,6 +53,7 @@ InactiveAssetsResponse, PreviousTIResponse, PrevSuccessfulDagRunResponse, + StubTaskArg, TaskBreadcrumbsResponse, TaskStatesResponse, TIAwaitingInputStatePayload, @@ -80,6 +81,7 @@ from airflow.models.asset import AssetActive from airflow.models.base import ID_LEN from airflow.models.dag import DagModel +from airflow.models.dag_version import DagVersion from airflow.models.dagrun import DagRun as DR from airflow.models.hitl import HITLDetail from airflow.models.log import Log @@ -89,6 +91,8 @@ from airflow.models.trigger import Trigger, handle_event_submit from airflow.models.xcom import XComModel from airflow.serialization.definitions.assets import SerializedAsset, SerializedAssetUniqueKey +from airflow.serialization.enums import Encoding +from airflow.serialization.serialized_objects import BaseSerialization from airflow.state import get_state_backend from airflow.triggers.base import TriggerEvent from airflow.utils.sqlalchemy import get_dialect_name @@ -110,6 +114,30 @@ 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 ``stub_args`` so regular tasks never pay for it. +_STUB_TASK_TYPE = "_StubOperator" + + +def _get_stub_args(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.""" + 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("_stub_args"): + return BaseSerialization.deserialize(encoded) + return None + return None + @ti_id_router.patch( "/{task_instance_id}/run", @@ -163,6 +191,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 +340,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 ( + stub_args := _get_stub_args(ti.dag_version_id, ti.task_id, session=session) + ): + context.stub_args = [StubTaskArg.model_validate(arg) for arg in stub_args] + # 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..dc6a6bc3e38de 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 @@ -46,6 +46,7 @@ AddConnectionTestEndpoint, AddPartitionDateField, AddRetryPolicyFields, + AddStubArgsToTIRunContext, AddTaskAndAssetStateStoreEndpoints, AddTaskInstanceQueueField, AddTeamNameField, @@ -65,6 +66,7 @@ AddTaskAndAssetStateStoreEndpoints, AddAssetsByAliasEndpoint, AddPartitionDateField, + AddStubArgsToTIRunContext, ), 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..e58e8cf6b4a26 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 AddStubArgsToTIRunContext(VersionChange): + """Add the ``stub_args`` positional-argument binding spec for stub (foreign-runtime) tasks.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = (schema(TIRunContext).field("stub_args").didnt_exist,) + + @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type] + def remove_stub_args_field(response: ResponseInfo) -> None: # type: ignore[misc] + """Strip ``stub_args`` from the run context for older clients.""" + response.body.pop("stub_args", 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..c7cc42b6c10a5 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_stub_args_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_stub_args_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()["stub_args"] == [ + {"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 "stub_args" 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..c1cef8ab71d47 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 TestStubArgsFieldBackwardCompat: + @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_stub_args_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_stub_args_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 "stub_args" not in response.json() + + def test_head_version_includes_stub_args(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()["stub_args"] == [ + {"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..338e4857015f5 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 (``_stub_args``) survives Dag serialization.""" + from airflow.providers.standard.decorators.stub import stub + + def extract(): ... + + def transform(country: str, extracted: dict): ... + + with DAG(dag_id="stub_args_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 "_stub_args" not in encoded_tasks["extract"], "argless stubs must not serialize a spec" + assert encoded_tasks["transform"]["_stub_args"] == [ + { + 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"]._stub_args == [ + {"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"], "_stub_args") or ( + round_tripped.task_dict["extract"]._stub_args is None + ) + + def test_handle_v1_serdag(): v1 = { "__version": 1, 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..df66bb3888243 100644 --- a/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md +++ b/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md @@ -217,7 +217,10 @@ Supervisor Bundle binary (Go) │ │ ├── StartupDetails ────────────────────►│ │ (ti, dag_rel_path, bundle_info, │ - │ start_date, ti_context) │ + │ start_date, ti_context; the │ + │ ti_context carries stub_args, the │ + │ positional-argument spec captured │ + │ from the stub Dag's TaskFlow call) │ │ │ │ ├── lookup task: │ │ bundle.dags[ti.dag_id] @@ -225,6 +228,12 @@ Supervisor Bundle binary (Go) │ │ (returns TaskState{state:"removed"} │ │ if not found, mirroring Java) │ │ + │ ├── bind stub_args 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..84e9a1045f4e8 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: diff --git a/go-sdk/dags/go_examples.py b/go-sdk/dags/go_examples.py index 23e02dd5e49dc..0c54f2ed5c778 100644 --- a/go-sdk/dags/go_examples.py +++ b/go-sdk/dags/go_examples.py @@ -33,6 +33,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 +69,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 +90,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 diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 23e60bd1dd46e..c03d8c80a71d5 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -127,12 +127,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/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..642eb87f5adb9 100644 --- a/go-sdk/pkg/execution/genmodels/defaults.gen.go +++ b/go-sdk/pkg/execution/genmodels/defaults.gen.go @@ -258,6 +258,17 @@ func (m *RetryTask) DecodeMsgpack(dec *msgpack.Decoder) error { return nil } +// DecodeMsgpack applies StubTaskArg's schema defaults that msgpack would otherwise skip. +func (m *StubTaskArg) DecodeMsgpack(dec *msgpack.Decoder) error { + type alias StubTaskArg + v := alias{DataType: StubTaskArgDataType("any"), Key: "return_value"} + if err := dec.Decode(&v); err != nil { + return err + } + *m = StubTaskArg(v) + return nil +} + // DecodeMsgpack applies SucceedTask's schema defaults that msgpack would otherwise skip. func (m *SucceedTask) DecodeMsgpack(dec *msgpack.Decoder) error { type alias SucceedTask diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index e6861d8c8add5..974b87dcb8de1 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -370,6 +370,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 +632,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 +762,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 +824,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 +853,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 +1248,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 @@ -1531,6 +1558,46 @@ type StartupDetails struct { type States []string +type StubArgs []StubTaskArg + +// 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 StubTaskArg struct { + // DataType corresponds to the JSON schema field "data_type". + DataType StubTaskArgDataType `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 StubTaskArgKind `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 StubTaskArgDataType string + +const StubTaskArgDataTypeAny StubTaskArgDataType = "any" +const StubTaskArgDataTypeArray StubTaskArgDataType = "array" +const StubTaskArgDataTypeBoolean StubTaskArgDataType = "boolean" +const StubTaskArgDataTypeInteger StubTaskArgDataType = "integer" +const StubTaskArgDataTypeNumber StubTaskArgDataType = "number" +const StubTaskArgDataTypeObject StubTaskArgDataType = "object" +const StubTaskArgDataTypeString StubTaskArgDataType = "string" + +type StubTaskArgKind string + +const StubTaskArgKindLiteral StubTaskArgKind = "literal" +const StubTaskArgKindXcom StubTaskArgKind = "xcom" + // Update a task's state to success. Includes task_outlets and outlet_events for // registering asset events. type SucceedTask struct { @@ -1585,6 +1652,9 @@ type TIRunContext struct { // StartDate corresponds to the JSON schema field "start_date". StartDate interface{} `msgpack:"start_date,omitempty"` + // StubArgs corresponds to the JSON schema field "stub_args". + StubArgs *StubArgs `msgpack:"stub_args,omitempty"` + // TaskRescheduleCount corresponds to the JSON schema field // "task_reschedule_count". TaskRescheduleCount int `msgpack:"task_reschedule_count,omitempty"` @@ -1596,49 +1666,26 @@ type TIRunContext struct { XcomKeysToClear []string `msgpack:"xcom_keys_to_clear,omitempty"` } -type TaskBreadcrumbsResult struct { - // Breadcrumbs corresponds to the JSON schema field "breadcrumbs". - Breadcrumbs []TaskBreadcrumbsResultBreadcrumbsElem `msgpack:"breadcrumbs"` - - // Type corresponds to the JSON schema field "type". - Type string `msgpack:"type,omitempty"` -} - -type TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{} - -// Task callback status information. -// -// A Class with information about the success/failure TI callback to be executed. -// Currently, only failure -// callbacks when tasks are externally killed or experience heartbeat timeouts are -// run via DagFileProcessorProcess. -type TaskCallbackRequest struct { - // BundleName corresponds to the JSON schema field "bundle_name". - BundleName string `msgpack:"bundle_name"` - - // BundleVersion corresponds to the JSON schema field "bundle_version". - BundleVersion interface{} `msgpack:"bundle_version"` - - // ContextFromServer corresponds to the JSON schema field "context_from_server". - ContextFromServer *TIRunContext `msgpack:"context_from_server,omitempty"` - - // Filepath corresponds to the JSON schema field "filepath". - Filepath string `msgpack:"filepath"` - - // Msg corresponds to the JSON schema field "msg". - Msg interface{} `msgpack:"msg,omitempty"` +type VersionData map[string]interface{} - // TaskCallbackType corresponds to the JSON schema field "task_callback_type". - TaskCallbackType interface{} `msgpack:"task_callback_type,omitempty"` +type TaskInstanceState string - // TI corresponds to the JSON schema field "ti". - TI TaskInstance `msgpack:"ti"` +const TaskInstanceStateRemoved TaskInstanceState = "removed" +const TaskInstanceStateScheduled TaskInstanceState = "scheduled" +const TaskInstanceStateQueued TaskInstanceState = "queued" +const TaskInstanceStateRunning TaskInstanceState = "running" +const TaskInstanceStateSuccess TaskInstanceState = "success" +const TaskInstanceStateRestarting TaskInstanceState = "restarting" +const TaskInstanceStateFailed TaskInstanceState = "failed" - // Type corresponds to the JSON schema field "type". - Type string `msgpack:"type,omitempty"` -} +const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" +const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" +const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" +const TaskInstanceStateSkipped TaskInstanceState = "skipped" +const TaskInstanceStateDeferred TaskInstanceState = "deferred" +const TaskInstanceStateAwaitingInput TaskInstanceState = "awaiting_input" -type TaskIds []string +type Warnings []interface{} // Schema for TaskInstance model with minimal required fields needed for Runtime. type TaskInstance struct { @@ -1673,24 +1720,64 @@ type TaskInstance struct { TryNumber int `msgpack:"try_number"` } -type TaskInstanceState string +// 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"` -const TaskInstanceStateAwaitingInput TaskInstanceState = "awaiting_input" -const TaskInstanceStateDeferred TaskInstanceState = "deferred" -const TaskInstanceStateFailed TaskInstanceState = "failed" -const TaskInstanceStateQueued TaskInstanceState = "queued" -const TaskInstanceStateRemoved TaskInstanceState = "removed" -const TaskInstanceStateRestarting TaskInstanceState = "restarting" -const TaskInstanceStateRunning TaskInstanceState = "running" -const TaskInstanceStateScheduled TaskInstanceState = "scheduled" -const TaskInstanceStateSkipped TaskInstanceState = "skipped" -const TaskInstanceStateSuccess TaskInstanceState = "success" -const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" -const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" -const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" + // Value corresponds to the JSON schema field "value". + Value interface{} `msgpack:"value"` +} + +type TriggerKwargs map[string]JsonValue type TaskOutlets []AssetProfile +type TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{} + +type TaskBreadcrumbsResult struct { + // Breadcrumbs corresponds to the JSON schema field "breadcrumbs". + Breadcrumbs []TaskBreadcrumbsResultBreadcrumbsElem `msgpack:"breadcrumbs"` + + // Type corresponds to the JSON schema field "type". + Type string `msgpack:"type,omitempty"` +} + +// Task callback status information. +// +// A Class with information about the success/failure TI callback to be executed. +// Currently, only failure +// callbacks when tasks are externally killed or experience heartbeat timeouts are +// run via DagFileProcessorProcess. +type TaskCallbackRequest struct { + // BundleName corresponds to the JSON schema field "bundle_name". + BundleName string `msgpack:"bundle_name"` + + // BundleVersion corresponds to the JSON schema field "bundle_version". + BundleVersion interface{} `msgpack:"bundle_version"` + + // ContextFromServer corresponds to the JSON schema field "context_from_server". + ContextFromServer *TIRunContext `msgpack:"context_from_server,omitempty"` + + // Filepath corresponds to the JSON schema field "filepath". + Filepath string `msgpack:"filepath"` + + // Msg corresponds to the JSON schema field "msg". + Msg interface{} `msgpack:"msg,omitempty"` + + // TaskCallbackType corresponds to the JSON schema field "task_callback_type". + TaskCallbackType interface{} `msgpack:"task_callback_type,omitempty"` + + // TI corresponds to the JSON schema field "ti". + TI TaskInstance `msgpack:"ti"` + + // 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"` +} + // Response containing the first reschedule date for a task instance. type TaskRescheduleStartDate struct { // StartDate corresponds to the JSON schema field "start_date". @@ -1700,6 +1787,12 @@ type TaskRescheduleStartDate struct { Type string `msgpack:"type,omitempty"` } +type TaskStateState string + +const TaskStateStateFailed TaskStateState = "failed" +const TaskStateStateSkipped TaskStateState = "skipped" +const TaskStateStateRemoved TaskStateState = "removed" + // Update a task's state. // // If a process exits without sending one of these the state will be derived from @@ -1720,12 +1813,6 @@ type TaskState struct { Type string `msgpack:"type,omitempty"` } -type TaskStateState string - -const TaskStateStateFailed TaskStateState = "failed" -const TaskStateStateRemoved TaskStateState = "removed" -const TaskStateStateSkipped TaskStateState = "skipped" - // Response to GetTaskStateStore; wraps the generated API response for supervisor // to worker comms. type TaskStateStoreResult struct { @@ -1775,20 +1862,7 @@ type TriggerDagRun struct { Type string `msgpack:"type,omitempty"` } -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{} +type TaskIds []string // Update the response content part of an existing Human-in-the-loop response. type UpdateHITLDetail struct { diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 2559359453866..206cff6983079 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") } +// TestTaskRunnerBindsStubArgs covers the TaskFlow path through RunTask: the +// positional-argument spec in ti_context.stub_args binds literals onto the +// task function's data parameters. +func TestTaskRunnerBindsStubArgs(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{ + StubArgs: &genmodels.StubArgs{ + {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) +} + +// TestTaskRunnerStubArgsArityMismatch: an argument spec that does not match +// the function's data parameters fails the task loudly instead of running it +// with zero values. +func TestTaskRunnerStubArgsArityMismatch(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{ + StubArgs: &genmodels.StubArgs{ + {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") +} + +// TestTaskRunnerStubArgsTypeMismatch: a declared Dag type that cannot bind to +// the Go parameter type fails the task loudly before the body runs. +func TestTaskRunnerStubArgsTypeMismatch(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{ + StubArgs: &genmodels.StubArgs{ + {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..26d2117fc8c1c 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 := convertStubArgs(details.TIContext.StubArgs) + + return executeTask(ctx, task, args, details.TIContext.ShouldRetry, logger) +} + +// convertStubArgs maps the wire-model positional-argument spec (captured from +// the Python stub Dag's TaskFlow call) onto the runtime-neutral binding form. +func convertStubArgs(specsPtr *genmodels.StubArgs) []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..629fd18055d5c 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, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Union from airflow.providers.common.compat.sdk import ( DecoratedOperator, @@ -27,13 +31,143 @@ task_decorator_factory, ) +try: + from airflow.sdk.definitions.context import KNOWN_CONTEXT_KEYS + from airflow.sdk.definitions.xcom_arg import PlainXComArg, XComArg +except ImportError: # Airflow 2 + from airflow.models.xcom_arg import PlainXComArg, XComArg # type: ignore[no-redef] + from airflow.utils.context import KNOWN_CONTEXT_KEYS # type: ignore[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_stub_args( + python_callable: Callable, + op_args: Sequence[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 ``StubTaskArg`` 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; " + "context injection does not happen in a foreign runtime, so pass the value " + "explicitly under a different parameter name" + ) + + 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 +209,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._stub_args = _build_stub_args(python_callable, self.op_args, self.op_kwargs, self.task_id) + + @classmethod + def get_serialized_fields(cls): + return super().get_serialized_fields() | {"_stub_args"} def execute(self, context: Context) -> Any: raise RuntimeError( @@ -96,6 +237,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..bc7ddb2323bba 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -17,10 +17,12 @@ 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.standard.decorators.stub import _data_type_from_annotation, stub from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS @@ -69,3 +71,140 @@ 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 (``_stub_args``).""" + + def test_literal_and_xcom_spec(self): + from airflow.sdk import DAG + + with DAG(dag_id="d"): + extracted = stub(fn_extract)() + result = stub(fn_transform)("uk", extracted) + + op = result.operator + assert op._stub_args == [ + {"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): + from airflow.sdk import DAG + + with DAG(dag_id="d"): + extracted = stub(fn_extract)() + result = stub(fn_transform)(extracted=extracted["part"], country="fr", retries_num=7) + + assert result.operator._stub_args == [ + {"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._stub_args is None + + def test_untyped_params_degrade_to_any(self): + from airflow.sdk import DAG + + with DAG(dag_id="d"): + result = stub(fn_untyped)(1, "x") + + assert result.operator._stub_args == [ + {"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"} + from airflow.sdk import DAG + + with DAG(dag_id="d"): + result = stub(fn)("v") + + assert result.operator._stub_args == [{"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): + from airflow.sdk import DAG + + with DAG(dag_id="d"), pytest.raises(ValueError, match="not JSON-serializable"): + stub(fn_transform)("uk", object()) + + def test_mapped_xcom_arg_rejected(self): + from airflow.sdk import DAG + + with DAG(dag_id="d"): + extracted = stub(fn_extract)() + with pytest.raises(ValueError, match="MapXComArg"): + stub(fn_transform)("uk", extracted.map(lambda v: v)) + + def test_expand_rejected_at_parse_time(self): + from airflow.sdk import DAG + + 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 75eac99a93888..ea394d57b12c1 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..7319feb2181e0 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -246,6 +246,36 @@ class PreviousTIResponse(BaseModel): duration: Annotated[float | None, Field(title="Duration")] = None +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 StubTaskArg(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 TIAwaitingInputStatePayload(BaseModel): """ Schema for parking a TaskInstance in an awaiting_input state (Human-in-the-loop, no trigger). @@ -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 + stub_args: Annotated[list[StubTaskArg] | None, Field(title="Stub Args")] = 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 0ec8fe4e49a1c..0c1ffa2b51e91 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.", @@ -3046,6 +3059,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": { @@ -4845,6 +4866,66 @@ "title": "DagRun", "type": "object" }, + "StubTaskArg": { + "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": "StubTaskArg", + "type": "object" + }, "TIRunContext": { "description": "Response schema for TaskInstance run context.", "properties": { @@ -4926,6 +5007,21 @@ ], "default": null, "title": "Start Date" + }, + "stub_args": { + "anyOf": [ + { + "items": { + "$ref": "#/$defs/StubTaskArg" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Stub Args" } }, "required": [ diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py index 9491a8993fdc3..cd79c4625f57d 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 AddStubArgsToTIRunContext + return VersionBundle( HeadVersion(), + Version("2026-07-30", AddStubArgsToTIRunContext), 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..94a1e6ad81801 --- /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 AddStubArgsToTIRunContext(VersionChange): + """Add the ``stub_args`` positional-argument binding spec for stub (foreign-runtime) tasks.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = (schema(TIRunContext).field("stub_args").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..16de279b3962f 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 ``stub_args`` migration is covered by + :class:`TestRealBundleStubArgsDowngrade` 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 TestRealBundleStubArgsDowngrade: + """ + Drive the *real* supervisor bundle through the ``stub_args`` migration. + + ``AddStubArgsToTIRunContext`` 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, + stub_args=[ + {"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_stub_args_for_previous_version(self, real_migrator, startup_details): + out = real_migrator.downgrade(startup_details, "2026-06-16").model_dump() + assert "stub_args" not in out["ti_context"] + + def test_head_version_keeps_stub_args(self, real_migrator, startup_details): + out = real_migrator.downgrade(startup_details, "2026-07-30") + assert out.ti_context.stub_args is not None + assert [a.kind for a in out.ti_context.stub_args] == ["literal", "xcom"] diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index ab2632831ab95..de5a1eaab6456 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 StubArgs = StubTaskArg[] | 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"; /** @@ -390,10 +400,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"; @@ -421,7 +431,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"; @@ -440,7 +450,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; @@ -449,34 +459,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; @@ -498,6 +508,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; @@ -508,7 +523,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; @@ -519,7 +534,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"; @@ -537,22 +552,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; @@ -612,12 +627,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[]; @@ -1003,6 +1018,7 @@ export interface TIRunContext { xcom_keys_to_clear?: XcomKeysToClear; should_retry?: ShouldRetry; start_date?: StartDate2; + stub_args?: StubArgs; } /** * Variable schema for responses with fields that are needed for Runtime. @@ -1030,6 +1046,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` "StubTaskArg". + */ +export interface StubTaskArg { + kind: Kind; + data_type?: DataType; + task_id?: TaskId1; + key?: Key1; + value?: unknown; +} /** * Email notification request for task failures/retries. * @@ -1150,7 +1182,7 @@ export interface DeferTask { */ export interface DeleteAssetStateStoreByName { name: Name8; - key: Key1; + key: Key2; type?: Type21; } /** @@ -1159,7 +1191,7 @@ export interface DeleteAssetStateStoreByName { */ export interface DeleteAssetStateStoreByUri { uri: Uri5; - key: Key2; + key: Key3; type?: Type22; } /** @@ -1168,7 +1200,7 @@ export interface DeleteAssetStateStoreByUri { */ export interface DeleteTaskStateStore { ti_id: TiId2; - key: Key3; + key: Key4; type?: Type23; } /** @@ -1176,7 +1208,7 @@ export interface DeleteTaskStateStore { * via the `definition` "DeleteVariable". */ export interface DeleteVariable { - key: Key4; + key: Key5; type?: Type24; } /** @@ -1184,10 +1216,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; } @@ -1253,7 +1285,7 @@ export interface GetAssetEventByAssetAlias { */ export interface GetAssetStateStoreByName { name: Name11; - key: Key6; + key: Key7; type?: Type31; } /** @@ -1262,7 +1294,7 @@ export interface GetAssetStateStoreByName { */ export interface GetAssetStateStoreByUri { uri: Uri8; - key: Key7; + key: Key8; type?: Type32; } /** @@ -1354,7 +1386,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; @@ -1398,7 +1430,7 @@ export interface GetTaskRescheduleStartDate { */ export interface GetTaskStateStore { ti_id: TiId6; - key: Key8; + key: Key9; type?: Type46; } /** @@ -1419,7 +1451,7 @@ export interface GetTaskStates { * via the `definition` "GetVariable". */ export interface GetVariable { - key: Key9; + key: Key10; type?: Type48; } /** @@ -1437,10 +1469,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; @@ -1452,10 +1484,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; } /** @@ -1463,10 +1495,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; } @@ -1475,10 +1507,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; @@ -1559,7 +1591,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; @@ -1585,7 +1617,7 @@ export interface PreviousTIResult { * via the `definition` "PutVariable". */ export interface PutVariable { - key: Key14; + key: Key15; value: Value1; description: Description; type?: Type61; @@ -1637,7 +1669,7 @@ export interface SentFDs { */ export interface SetAssetStateStoreByName { name: Name13; - key: Key15; + key: Key16; value: JsonValue; type?: Type66; } @@ -1647,7 +1679,7 @@ export interface SetAssetStateStoreByName { */ export interface SetAssetStateStoreByUri { uri: Uri9; - key: Key16; + key: Key17; value: JsonValue; type?: Type67; } @@ -1680,7 +1712,7 @@ export interface SetRenderedMapIndex { */ export interface SetTaskStateStore { ti_id: TiId8; - key: Key17; + key: Key18; value: JsonValue; expires_at: ExpiresAt; type?: Type70; @@ -1690,11 +1722,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; @@ -1851,7 +1883,7 @@ export interface VariableKeysResult { * via the `definition` "VariableResult". */ export interface VariableResult { - key: Key19; + key: Key20; value?: Value2; type?: Type85; } @@ -1870,7 +1902,7 @@ export interface XComCountResponse { * via the `definition` "XComResult". */ export interface XComResult { - key: Key20; + key: Key21; value: JsonValue; type?: Type87; } @@ -1896,4 +1928,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; From c68b1ccd0bb70580fce467e9f57573eeb25c88e4 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Sat, 11 Jul 2026 12:15:51 +0000 Subject: [PATCH 02/40] Rename stub_args to arg_bindings across the TaskFlow stub pipeline "stub_args" leaked the _StubOperator implementation detail into the wire contract that foreign-language SDKs code-generate against; "arg bindings" names what the data actually is -- the ordered spec a runtime binds onto the task function. Renaming now, before the field ships in a released execution API or supervisor schema version, keeps the contract clean without any compatibility shims. --- .../execution_api/datamodels/taskinstance.py | 8 +- .../execution_api/routes/task_instances.py | 12 +- .../execution_api/versions/__init__.py | 4 +- .../execution_api/versions/v2026_06_30.py | 12 +- .../versions/head/test_task_instances.py | 8 +- .../v2026_04_17/test_task_instances.py | 12 +- .../serialization/test_dag_serialization.py | 14 +- .../0003-coordinator-protocol-msgpack-ipc.md | 9 +- .../pkg/execution/genmodels/defaults.gen.go | 22 +- go-sdk/pkg/execution/genmodels/models.gen.go | 215 +++++++++--------- go-sdk/pkg/execution/integration_test.go | 20 +- go-sdk/pkg/execution/task_runner.go | 6 +- .../providers/standard/decorators/stub.py | 8 +- .../unit/standard/decorators/test_stub.py | 12 +- .../airflow/sdk/api/datamodels/_generated.py | 62 ++--- .../sdk/execution_time/schema/schema.json | 126 +++++----- .../schema/versions/__init__.py | 4 +- .../schema/versions/v2026_07_30.py | 6 +- .../execution_time/schema/test_migrator.py | 22 +- ts-sdk/src/generated/supervisor.ts | 8 +- 20 files changed, 295 insertions(+), 295 deletions(-) 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 2976ad26f905f..66f6ec1dc824a 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 @@ -393,11 +393,11 @@ def safe_extract_from_orm(cls, data: Any) -> Any: return values -StubArgDataType = Literal["string", "integer", "number", "boolean", "object", "array", "any"] +ArgBindingDataType = Literal["string", "integer", "number", "boolean", "object", "array", "any"] """Language-neutral value type a stub-task argument binds to in the foreign runtime.""" -class StubTaskArg(BaseModel): +class TaskArgBinding(BaseModel): """ One positional argument of a stub (foreign-runtime) task, in declaration order. @@ -408,7 +408,7 @@ class StubTaskArg(BaseModel): kind: Literal["xcom", "literal"] """Whether the value comes from an upstream task's XCom or is a literal from the Dag file.""" - data_type: StubArgDataType = "any" + data_type: ArgBindingDataType = "any" """Declared type from the stub function's annotation; runtimes type-check against it.""" task_id: str | None = None @@ -463,7 +463,7 @@ class TIRunContext(BaseModel): always reflects when the task *first* started, not when it was rescheduled/resumed. """ - stub_args: list[StubTaskArg] | None = None + arg_bindings: list[TaskArgBinding] | None = None """ Ordered positional-argument binding spec for stub (foreign-runtime) tasks. 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 bbec7f542f22a..7573f34344ee7 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 @@ -53,7 +53,7 @@ InactiveAssetsResponse, PreviousTIResponse, PrevSuccessfulDagRunResponse, - StubTaskArg, + TaskArgBinding, TaskBreadcrumbsResponse, TaskStatesResponse, TIAwaitingInputStatePayload, @@ -116,11 +116,11 @@ # Task type recorded on the TI row (``TaskInstance.operator``) for # ``airflow.providers.standard.decorators.stub._StubOperator``. Used to gate the -# serialized-dag lookup for ``stub_args`` so regular tasks never pay for it. +# serialized-dag lookup for ``arg_bindings`` so regular tasks never pay for it. _STUB_TASK_TYPE = "_StubOperator" -def _get_stub_args(dag_version_id: UUID | None, task_id: str, *, session) -> list[dict] | None: +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.""" if dag_version_id is None: return None @@ -133,7 +133,7 @@ def _get_stub_args(dag_version_id: UUID | None, task_id: str, *, session) -> lis 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("_stub_args"): + if encoded := var.get("_arg_bindings"): return BaseSerialization.deserialize(encoded) return None return None @@ -343,9 +343,9 @@ def ti_run( # 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 ( - stub_args := _get_stub_args(ti.dag_version_id, ti.task_id, session=session) + arg_bindings := _get_arg_bindings(ti.dag_version_id, ti.task_id, session=session) ): - context.stub_args = [StubTaskArg.model_validate(arg) for arg in stub_args] + context.arg_bindings = [TaskArgBinding.model_validate(arg) for arg in arg_bindings] # Only set if they are non-null if 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 dc6a6bc3e38de..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,12 +41,12 @@ RemoveUpstreamMapIndexesField, ) from airflow.api_fastapi.execution_api.versions.v2026_06_30 import ( + AddArgBindingsToTIRunContext, AddAssetsByAliasEndpoint, AddAwaitingInputStatePayload, AddConnectionTestEndpoint, AddPartitionDateField, AddRetryPolicyFields, - AddStubArgsToTIRunContext, AddTaskAndAssetStateStoreEndpoints, AddTaskInstanceQueueField, AddTeamNameField, @@ -66,7 +66,7 @@ AddTaskAndAssetStateStoreEndpoints, AddAssetsByAliasEndpoint, AddPartitionDateField, - AddStubArgsToTIRunContext, + 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 e58e8cf6b4a26..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 @@ -143,14 +143,14 @@ def remove_partition_date_from_dag_run(response: ResponseInfo) -> None: # type: response.body["dag_run"].pop("partition_date", None) -class AddStubArgsToTIRunContext(VersionChange): - """Add the ``stub_args`` positional-argument binding spec for stub (foreign-runtime) tasks.""" +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("stub_args").didnt_exist,) + 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_stub_args_field(response: ResponseInfo) -> None: # type: ignore[misc] - """Strip ``stub_args`` from the run context for older clients.""" - response.body.pop("stub_args", None) + 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 c7cc42b6c10a5..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,7 +370,7 @@ async def workload_token(request: Request) -> TIToken: assert extras["scope"] == "execution" assert extras["sub"] == str(ti.id) - def test_ti_run_returns_stub_args_for_stub_task(self, client, dag_maker): + 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 @@ -378,7 +378,7 @@ def extract(): ... def transform(country: str, extracted: dict): ... - with dag_maker("test_stub_args_dag", serialized=True): + with dag_maker("test_arg_bindings_dag", serialized=True): stub(transform)("uk", stub(extract)()) dr = dag_maker.create_dagrun() @@ -397,7 +397,7 @@ def transform(country: str, extracted: dict): ... response = client.patch(f"/execution/task-instances/{tis['transform'].id}/run", json=payload) assert response.status_code == 200 - assert response.json()["stub_args"] == [ + assert response.json()["arg_bindings"] == [ {"kind": "literal", "data_type": "string", "value": "uk"}, {"kind": "xcom", "data_type": "object", "task_id": "extract", "key": "return_value"}, ] @@ -405,7 +405,7 @@ def transform(country: str, extracted: dict): ... # 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 "stub_args" not in response.json() + 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.""" 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 c1cef8ab71d47..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 @@ -86,7 +86,7 @@ def test_head_version_includes_team_name_field(self, client, session, create_tas assert response.json()["dag_run"]["team_name"] is None -class TestStubArgsFieldBackwardCompat: +class TestArgBindingsFieldBackwardCompat: @pytest.fixture(autouse=True) def _freeze_time(self, time_machine): time_machine.move_to(TIMESTAMP_STR, tick=False) @@ -105,7 +105,7 @@ def extract(): ... def transform(country: str, extracted: dict): ... - with dag_maker("test_stub_args_compat_dag", serialized=True): + with dag_maker("test_arg_bindings_compat_dag", serialized=True): stub(transform)("uk", stub(extract)()) dr = dag_maker.create_dagrun() @@ -115,15 +115,15 @@ def transform(country: str, extracted: dict): ... dag_maker.session.flush() return tis["transform"] - def test_old_version_strips_stub_args_even_when_set(self, old_ver_client, stub_ti): + 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 "stub_args" not in response.json() + assert "arg_bindings" not in response.json() - def test_head_version_includes_stub_args(self, client, stub_ti): + 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()["stub_args"] == [ + 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 338e4857015f5..f03ac2d64fc96 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -3406,20 +3406,20 @@ def inner(): def test_stub_task_args_round_trip(): - """The stub task's TaskFlow arg spec (``_stub_args``) survives Dag serialization.""" + """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="stub_args_dag", schedule=None) as dag: + 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 "_stub_args" not in encoded_tasks["extract"], "argless stubs must not serialize a spec" - assert encoded_tasks["transform"]["_stub_args"] == [ + 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"}, @@ -3436,12 +3436,12 @@ def transform(country: str, extracted: dict): ... ] round_tripped = DagSerialization.from_dict(ser_dag) - assert round_tripped.task_dict["transform"]._stub_args == [ + 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"], "_stub_args") or ( - round_tripped.task_dict["extract"]._stub_args is None + assert not hasattr(round_tripped.task_dict["extract"], "_arg_bindings") or ( + round_tripped.task_dict["extract"]._arg_bindings is None ) diff --git a/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md b/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md index df66bb3888243..e7535ab6bea75 100644 --- a/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md +++ b/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md @@ -218,9 +218,10 @@ Supervisor Bundle binary (Go) ├── StartupDetails ────────────────────►│ │ (ti, dag_rel_path, bundle_info, │ │ start_date, ti_context; the │ - │ ti_context carries stub_args, the │ - │ positional-argument spec captured │ - │ from the stub Dag's TaskFlow call) │ + │ ti_context carries arg_bindings, │ + │ the positional-argument spec │ + │ captured from the stub Dag's │ + │ TaskFlow call) │ │ │ │ ├── lookup task: │ │ bundle.dags[ti.dag_id] @@ -228,7 +229,7 @@ Supervisor Bundle binary (Go) │ │ (returns TaskState{state:"removed"} │ │ if not found, mirroring Java) │ │ - │ ├── bind stub_args onto the task + │ ├── bind arg_bindings onto the task │ │ fn's data parameters (literals │ │ decode directly; xcom refs pull │ │ below); arity/type mismatch diff --git a/go-sdk/pkg/execution/genmodels/defaults.gen.go b/go-sdk/pkg/execution/genmodels/defaults.gen.go index 642eb87f5adb9..d411881bc8404 100644 --- a/go-sdk/pkg/execution/genmodels/defaults.gen.go +++ b/go-sdk/pkg/execution/genmodels/defaults.gen.go @@ -258,17 +258,6 @@ func (m *RetryTask) DecodeMsgpack(dec *msgpack.Decoder) error { return nil } -// DecodeMsgpack applies StubTaskArg's schema defaults that msgpack would otherwise skip. -func (m *StubTaskArg) DecodeMsgpack(dec *msgpack.Decoder) error { - type alias StubTaskArg - v := alias{DataType: StubTaskArgDataType("any"), Key: "return_value"} - if err := dec.Decode(&v); err != nil { - return err - } - *m = StubTaskArg(v) - return nil -} - // DecodeMsgpack applies SucceedTask's schema defaults that msgpack would otherwise skip. func (m *SucceedTask) DecodeMsgpack(dec *msgpack.Decoder) error { type alias SucceedTask @@ -293,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 974b87dcb8de1..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". @@ -1558,46 +1560,6 @@ type StartupDetails struct { type States []string -type StubArgs []StubTaskArg - -// 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 StubTaskArg struct { - // DataType corresponds to the JSON schema field "data_type". - DataType StubTaskArgDataType `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 StubTaskArgKind `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 StubTaskArgDataType string - -const StubTaskArgDataTypeAny StubTaskArgDataType = "any" -const StubTaskArgDataTypeArray StubTaskArgDataType = "array" -const StubTaskArgDataTypeBoolean StubTaskArgDataType = "boolean" -const StubTaskArgDataTypeInteger StubTaskArgDataType = "integer" -const StubTaskArgDataTypeNumber StubTaskArgDataType = "number" -const StubTaskArgDataTypeObject StubTaskArgDataType = "object" -const StubTaskArgDataTypeString StubTaskArgDataType = "string" - -type StubTaskArgKind string - -const StubTaskArgKindLiteral StubTaskArgKind = "literal" -const StubTaskArgKindXcom StubTaskArgKind = "xcom" - // Update a task's state to success. Includes task_outlets and outlet_events for // registering asset events. type SucceedTask struct { @@ -1631,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"` @@ -1652,9 +1617,6 @@ type TIRunContext struct { // StartDate corresponds to the JSON schema field "start_date". StartDate interface{} `msgpack:"start_date,omitempty"` - // StubArgs corresponds to the JSON schema field "stub_args". - StubArgs *StubArgs `msgpack:"stub_args,omitempty"` - // TaskRescheduleCount corresponds to the JSON schema field // "task_reschedule_count". TaskRescheduleCount int `msgpack:"task_reschedule_count,omitempty"` @@ -1666,74 +1628,43 @@ type TIRunContext struct { XcomKeysToClear []string `msgpack:"xcom_keys_to_clear,omitempty"` } -type VersionData map[string]interface{} - -type TaskInstanceState string - -const TaskInstanceStateRemoved TaskInstanceState = "removed" -const TaskInstanceStateScheduled TaskInstanceState = "scheduled" -const TaskInstanceStateQueued TaskInstanceState = "queued" -const TaskInstanceStateRunning TaskInstanceState = "running" -const TaskInstanceStateSuccess TaskInstanceState = "success" -const TaskInstanceStateRestarting TaskInstanceState = "restarting" -const TaskInstanceStateFailed TaskInstanceState = "failed" - -const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" -const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" -const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" -const TaskInstanceStateSkipped TaskInstanceState = "skipped" -const TaskInstanceStateDeferred TaskInstanceState = "deferred" -const TaskInstanceStateAwaitingInput TaskInstanceState = "awaiting_input" - -type Warnings []interface{} - -// Schema for TaskInstance model with minimal required fields needed for Runtime. -type TaskInstance struct { - // ContextCarrier corresponds to the JSON schema field "context_carrier". - ContextCarrier *ContextCarrier `msgpack:"context_carrier,omitempty"` - - // DagID corresponds to the JSON schema field "dag_id". - DagID string `msgpack:"dag_id"` - - // DagVersionID corresponds to the JSON schema field "dag_version_id". - DagVersionID string `msgpack:"dag_version_id"` - - // Hostname corresponds to the JSON schema field "hostname". - Hostname interface{} `msgpack:"hostname,omitempty"` - - // ID corresponds to the JSON schema field "id". - ID string `msgpack:"id"` - - // MapIndex corresponds to the JSON schema field "map_index". - MapIndex *int `msgpack:"map_index,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"` - // Queue corresponds to the JSON schema field "queue". - Queue string `msgpack:"queue,omitempty"` + // Key corresponds to the JSON schema field "key". + Key string `msgpack:"key,omitempty"` - // RunID corresponds to the JSON schema field "run_id". - RunID string `msgpack:"run_id"` + // Kind corresponds to the JSON schema field "kind". + Kind TaskArgBindingKind `msgpack:"kind"` // TaskID corresponds to the JSON schema field "task_id". - TaskID string `msgpack:"task_id"` - - // TryNumber corresponds to the JSON schema field "try_number". - TryNumber int `msgpack:"try_number"` -} - -// 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"` + TaskID interface{} `msgpack:"task_id,omitempty"` // Value corresponds to the JSON schema field "value". - Value interface{} `msgpack:"value"` + Value interface{} `msgpack:"value,omitempty"` } -type TriggerKwargs map[string]JsonValue +type TaskArgBindingDataType string -type TaskOutlets []AssetProfile +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 TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{} +type TaskArgBindingKind string + +const TaskArgBindingKindLiteral TaskArgBindingKind = "literal" +const TaskArgBindingKindXcom TaskArgBindingKind = "xcom" type TaskBreadcrumbsResult struct { // Breadcrumbs corresponds to the JSON schema field "breadcrumbs". @@ -1743,6 +1674,8 @@ type TaskBreadcrumbsResult struct { Type string `msgpack:"type,omitempty"` } +type TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{} + // Task callback status information. // // A Class with information about the success/failure TI callback to be executed. @@ -1778,6 +1711,59 @@ type TaskCallbackRequest struct { VersionData *VersionData `msgpack:"version_data,omitempty"` } +type TaskIds []string + +// Schema for TaskInstance model with minimal required fields needed for Runtime. +type TaskInstance struct { + // ContextCarrier corresponds to the JSON schema field "context_carrier". + ContextCarrier *ContextCarrier `msgpack:"context_carrier,omitempty"` + + // DagID corresponds to the JSON schema field "dag_id". + DagID string `msgpack:"dag_id"` + + // DagVersionID corresponds to the JSON schema field "dag_version_id". + DagVersionID string `msgpack:"dag_version_id"` + + // Hostname corresponds to the JSON schema field "hostname". + Hostname interface{} `msgpack:"hostname,omitempty"` + + // ID corresponds to the JSON schema field "id". + ID string `msgpack:"id"` + + // MapIndex corresponds to the JSON schema field "map_index". + MapIndex *int `msgpack:"map_index,omitempty"` + + // Queue corresponds to the JSON schema field "queue". + Queue string `msgpack:"queue,omitempty"` + + // RunID corresponds to the JSON schema field "run_id". + RunID string `msgpack:"run_id"` + + // TaskID corresponds to the JSON schema field "task_id". + TaskID string `msgpack:"task_id"` + + // TryNumber corresponds to the JSON schema field "try_number". + TryNumber int `msgpack:"try_number"` +} + +type TaskInstanceState string + +const TaskInstanceStateAwaitingInput TaskInstanceState = "awaiting_input" +const TaskInstanceStateDeferred TaskInstanceState = "deferred" +const TaskInstanceStateFailed TaskInstanceState = "failed" +const TaskInstanceStateQueued TaskInstanceState = "queued" +const TaskInstanceStateRemoved TaskInstanceState = "removed" +const TaskInstanceStateRestarting TaskInstanceState = "restarting" +const TaskInstanceStateRunning TaskInstanceState = "running" +const TaskInstanceStateScheduled TaskInstanceState = "scheduled" +const TaskInstanceStateSkipped TaskInstanceState = "skipped" +const TaskInstanceStateSuccess TaskInstanceState = "success" +const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" +const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" +const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" + +type TaskOutlets []AssetProfile + // Response containing the first reschedule date for a task instance. type TaskRescheduleStartDate struct { // StartDate corresponds to the JSON schema field "start_date". @@ -1787,12 +1773,6 @@ type TaskRescheduleStartDate struct { Type string `msgpack:"type,omitempty"` } -type TaskStateState string - -const TaskStateStateFailed TaskStateState = "failed" -const TaskStateStateSkipped TaskStateState = "skipped" -const TaskStateStateRemoved TaskStateState = "removed" - // Update a task's state. // // If a process exits without sending one of these the state will be derived from @@ -1813,6 +1793,12 @@ type TaskState struct { Type string `msgpack:"type,omitempty"` } +type TaskStateState string + +const TaskStateStateFailed TaskStateState = "failed" +const TaskStateStateRemoved TaskStateState = "removed" +const TaskStateStateSkipped TaskStateState = "skipped" + // Response to GetTaskStateStore; wraps the generated API response for supervisor // to worker comms. type TaskStateStoreResult struct { @@ -1862,7 +1848,7 @@ type TriggerDagRun struct { Type string `msgpack:"type,omitempty"` } -type TaskIds []string +type TriggerKwargs map[string]JsonValue // Update the response content part of an existing Human-in-the-loop response. type UpdateHITLDetail struct { @@ -1879,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 206cff6983079..5646a82e29206 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -230,10 +230,10 @@ func TestTaskRunnerPanicRetry(t *testing.T) { assertRetryTask(t, result, "panic: something went wrong") } -// TestTaskRunnerBindsStubArgs covers the TaskFlow path through RunTask: the -// positional-argument spec in ti_context.stub_args binds literals onto the +// 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 TestTaskRunnerBindsStubArgs(t *testing.T) { +func TestTaskRunnerBindsArgs(t *testing.T) { var gotCountry string var gotMeta map[string]any bundle := buildBundle(t, func(r bundlev1.Registry) { @@ -255,7 +255,7 @@ func TestTaskRunnerBindsStubArgs(t *testing.T) { }, BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, TIContext: genmodels.TIRunContext{ - StubArgs: &genmodels.StubArgs{ + ArgBindings: &genmodels.ArgBindings{ {Kind: "literal", DataType: "string", Value: "uk"}, {Kind: "literal", DataType: "object", Value: map[string]any{"k": "v"}}, }, @@ -271,10 +271,10 @@ func TestTaskRunnerBindsStubArgs(t *testing.T) { assert.Equal(t, map[string]any{"k": "v"}, gotMeta) } -// TestTaskRunnerStubArgsArityMismatch: an argument spec that does not match +// 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 TestTaskRunnerStubArgsArityMismatch(t *testing.T) { +func TestTaskRunnerArgBindingsArityMismatch(t *testing.T) { ran := false bundle := buildBundle(t, func(r bundlev1.Registry) { r.AddDag("test_dag").AddTaskWithName("transform", @@ -294,7 +294,7 @@ func TestTaskRunnerStubArgsArityMismatch(t *testing.T) { }, BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, TIContext: genmodels.TIRunContext{ - StubArgs: &genmodels.StubArgs{ + ArgBindings: &genmodels.ArgBindings{ {Kind: "literal", DataType: "string", Value: "uk"}, }, }, @@ -308,9 +308,9 @@ func TestTaskRunnerStubArgsArityMismatch(t *testing.T) { assert.False(t, ran, "the task body must not run on an arity mismatch") } -// TestTaskRunnerStubArgsTypeMismatch: a declared Dag type that cannot bind to +// TestTaskRunnerArgBindingsTypeMismatch: a declared Dag type that cannot bind to // the Go parameter type fails the task loudly before the body runs. -func TestTaskRunnerStubArgsTypeMismatch(t *testing.T) { +func TestTaskRunnerArgBindingsTypeMismatch(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { r.AddDag("test_dag").AddTaskWithName("transform", func(count int) error { return nil }) @@ -326,7 +326,7 @@ func TestTaskRunnerStubArgsTypeMismatch(t *testing.T) { }, BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, TIContext: genmodels.TIRunContext{ - StubArgs: &genmodels.StubArgs{ + ArgBindings: &genmodels.ArgBindings{ {Kind: "literal", DataType: "string", Value: "uk"}, }, }, diff --git a/go-sdk/pkg/execution/task_runner.go b/go-sdk/pkg/execution/task_runner.go index 26d2117fc8c1c..d6081d24ad4c3 100644 --- a/go-sdk/pkg/execution/task_runner.go +++ b/go-sdk/pkg/execution/task_runner.go @@ -125,14 +125,14 @@ func RunTask( ctx = context.WithValue(ctx, sdkcontext.SdkClientContextKey, sdk.Client(client)) ctx = context.WithValue(ctx, sdkcontext.RuntimeContextKey, runtimeContext) - args := convertStubArgs(details.TIContext.StubArgs) + args := convertArgBindings(details.TIContext.ArgBindings) return executeTask(ctx, task, args, details.TIContext.ShouldRetry, logger) } -// convertStubArgs maps the wire-model positional-argument spec (captured from +// convertArgBindings maps the wire-model positional-argument spec (captured from // the Python stub Dag's TaskFlow call) onto the runtime-neutral binding form. -func convertStubArgs(specsPtr *genmodels.StubArgs) []binding.Arg { +func convertArgBindings(specsPtr *genmodels.ArgBindings) []binding.Arg { if specsPtr == nil || len(*specsPtr) == 0 { return nil } diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index 629fd18055d5c..8c3f1398444fb 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -79,7 +79,7 @@ def _data_type_from_annotation(annotation: Any) -> str: return "any" -def _build_stub_args( +def _build_arg_bindings( python_callable: Callable, op_args: Sequence[Any], op_kwargs: Mapping[str, Any], @@ -88,7 +88,7 @@ def _build_stub_args( """ 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 ``StubTaskArg`` shape: an XCom + 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. """ @@ -212,11 +212,11 @@ def __init__( # 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._stub_args = _build_stub_args(python_callable, self.op_args, self.op_kwargs, self.task_id) + 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() | {"_stub_args"} + return super().get_serialized_fields() | {"_arg_bindings"} def execute(self, context: Context) -> Any: raise RuntimeError( diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index bc7ddb2323bba..84d350dc2981a 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -92,7 +92,7 @@ def fn_context_key(ti): ... class TestStubTaskflowArgs: - """The TaskFlow call on a stub captures the ordered positional-arg spec (``_stub_args``).""" + """The TaskFlow call on a stub captures the ordered positional-arg spec (``_arg_bindings``).""" def test_literal_and_xcom_spec(self): from airflow.sdk import DAG @@ -102,7 +102,7 @@ def test_literal_and_xcom_spec(self): result = stub(fn_transform)("uk", extracted) op = result.operator - assert op._stub_args == [ + 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}, @@ -116,14 +116,14 @@ def test_kwargs_normalize_to_declaration_order(self): extracted = stub(fn_extract)() result = stub(fn_transform)(extracted=extracted["part"], country="fr", retries_num=7) - assert result.operator._stub_args == [ + 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._stub_args is None + assert stub(fn_pass)().operator._arg_bindings is None def test_untyped_params_degrade_to_any(self): from airflow.sdk import DAG @@ -131,7 +131,7 @@ def test_untyped_params_degrade_to_any(self): with DAG(dag_id="d"): result = stub(fn_untyped)(1, "x") - assert result.operator._stub_args == [ + assert result.operator._arg_bindings == [ {"kind": "literal", "data_type": "any", "value": 1}, {"kind": "literal", "data_type": "any", "value": "x"}, ] @@ -145,7 +145,7 @@ def fn(x): ... with DAG(dag_id="d"): result = stub(fn)("v") - assert result.operator._stub_args == [{"kind": "literal", "data_type": "any", "value": "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"): diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 7319feb2181e0..aa62452980d5c 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -246,36 +246,6 @@ class PreviousTIResponse(BaseModel): duration: Annotated[float | None, Field(title="Duration")] = None -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 StubTaskArg(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 TIAwaitingInputStatePayload(BaseModel): """ Schema for parking a TaskInstance in an awaiting_input state (Human-in-the-loop, no trigger). @@ -401,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. @@ -827,4 +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 - stub_args: Annotated[list[StubTaskArg] | None, Field(title="Stub Args")] = None + arg_bindings: Annotated[list[TaskArgBinding] | None, Field(title="Arg Bindings")] = None diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 0c1ffa2b51e91..4dbfc70dead32 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -4866,66 +4866,6 @@ "title": "DagRun", "type": "object" }, - "StubTaskArg": { - "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": "StubTaskArg", - "type": "object" - }, "TIRunContext": { "description": "Response schema for TaskInstance run context.", "properties": { @@ -5008,11 +4948,11 @@ "default": null, "title": "Start Date" }, - "stub_args": { + "arg_bindings": { "anyOf": [ { "items": { - "$ref": "#/$defs/StubTaskArg" + "$ref": "#/$defs/TaskArgBinding" }, "type": "array" }, @@ -5021,7 +4961,7 @@ } ], "default": null, - "title": "Stub Args" + "title": "Arg Bindings" } }, "required": [ @@ -5031,6 +4971,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 cd79c4625f57d..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,11 +37,11 @@ def get_bundle() -> VersionBundle: """ from cadwyn import HeadVersion, Version, VersionBundle - from airflow.sdk.execution_time.schema.versions.v2026_07_30 import AddStubArgsToTIRunContext + from airflow.sdk.execution_time.schema.versions.v2026_07_30 import AddArgBindingsToTIRunContext return VersionBundle( HeadVersion(), - Version("2026-07-30", AddStubArgsToTIRunContext), + 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 index 94a1e6ad81801..fd25aef095bb8 100644 --- 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 @@ -22,9 +22,9 @@ from airflow.sdk.api.datamodels._generated import TIRunContext -class AddStubArgsToTIRunContext(VersionChange): - """Add the ``stub_args`` positional-argument binding spec for stub (foreign-runtime) tasks.""" +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("stub_args").didnt_exist,) + 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 16de279b3962f..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 @@ -106,8 +106,8 @@ class TestSchemaVersionMigratorDowngrade: """ Drive the downgrade direction against a mock bundle so we can pin *field-level* migration behaviour independent of the real bundle's - contents. The real bundle's ``stub_args`` migration is covered by - :class:`TestRealBundleStubArgsDowngrade` below. + contents. The real bundle's ``arg_bindings`` migration is covered by + :class:`TestRealBundleArgBindingsDowngrade` below. """ @pytest.fixture @@ -368,11 +368,11 @@ def test_accessing_bundle_loads_cadwyn(self): subprocess.run([sys.executable, "-c", code], check=True, capture_output=True, text=True) -class TestRealBundleStubArgsDowngrade: +class TestRealBundleArgBindingsDowngrade: """ - Drive the *real* supervisor bundle through the ``stub_args`` migration. + Drive the *real* supervisor bundle through the ``arg_bindings`` migration. - ``AddStubArgsToTIRunContext`` is the bundle's first ``schema(...)`` + ``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 @@ -419,7 +419,7 @@ def startup_details(self): consumed_asset_events=[], ), max_tries=1, - stub_args=[ + arg_bindings=[ {"kind": "literal", "data_type": "string", "value": "uk"}, {"kind": "xcom", "data_type": "object", "task_id": "extract", "key": "return_value"}, ], @@ -431,11 +431,11 @@ def startup_details(self): def real_migrator(self) -> SchemaVersionMigrator: return get_schema_version_migrator() - def test_downgrade_strips_stub_args_for_previous_version(self, real_migrator, startup_details): + 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 "stub_args" not in out["ti_context"] + assert "arg_bindings" not in out["ti_context"] - def test_head_version_keeps_stub_args(self, real_migrator, startup_details): + 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.stub_args is not None - assert [a.kind for a in out.ti_context.stub_args] == ["literal", "xcom"] + 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 de5a1eaab6456..528f45cb8d23b 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -245,7 +245,7 @@ export type NextKwargs1 = export type XcomKeysToClear = string[]; export type ShouldRetry = boolean; export type StartDate2 = string | null; -export type StubArgs = StubTaskArg[] | 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; @@ -1018,7 +1018,7 @@ export interface TIRunContext { xcom_keys_to_clear?: XcomKeysToClear; should_retry?: ShouldRetry; start_date?: StartDate2; - stub_args?: StubArgs; + arg_bindings?: ArgBindings; } /** * Variable schema for responses with fields that are needed for Runtime. @@ -1053,9 +1053,9 @@ export interface ConnectionResponse { * 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` "StubTaskArg". + * via the `definition` "TaskArgBinding". */ -export interface StubTaskArg { +export interface TaskArgBinding { kind: Kind; data_type?: DataType; task_id?: TaskId1; From a3f1c9585323ad2df5cf23911480fe29351167bd Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Sun, 12 Jul 2026 05:42:46 +0000 Subject: [PATCH 03/40] Fix stub decorator compatibility with released Airflow versions The single try/except made Airflow 3.0 (whose SDK predates KNOWN_CONTEXT_KEYS) fall back to the Airflow 2 import paths and fail; the arg-capture tests imported airflow.sdk directly, which does not exist on 2.11; and the .expand() rejection relies on the supports_expand opt-out that only ships with Airflow 3.4. Also reword the context-key rejection to stop implying foreign runtimes have no task context -- the lang SDKs inject their own natively; stub signatures just must not declare Airflow context parameters. --- .../tests_common/test_utils/version_compat.py | 1 + .../providers/standard/decorators/stub.py | 17 +++++++++------- .../unit/standard/decorators/test_stub.py | 20 +++++-------------- 3 files changed, 16 insertions(+), 22 deletions(-) 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/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index 8c3f1398444fb..72497f184e32e 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -22,7 +22,7 @@ import json import types import typing -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Collection, Mapping, Sequence from typing import TYPE_CHECKING, Any, Union from airflow.providers.common.compat.sdk import ( @@ -32,11 +32,14 @@ ) try: - from airflow.sdk.definitions.context import KNOWN_CONTEXT_KEYS from airflow.sdk.definitions.xcom_arg import PlainXComArg, XComArg except ImportError: # Airflow 2 - from airflow.models.xcom_arg import PlainXComArg, XComArg # type: ignore[no-redef] - from airflow.utils.context import KNOWN_CONTEXT_KEYS # type: ignore[no-redef] + 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 @@ -81,7 +84,7 @@ def _data_type_from_annotation(annotation: Any) -> str: def _build_arg_bindings( python_callable: Callable, - op_args: Sequence[Any], + op_args: Collection[Any], op_kwargs: Mapping[str, Any], task_id: str, ) -> list[dict[str, Any]] | None: @@ -103,8 +106,8 @@ def _build_arg_bindings( 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; " - "context injection does not happen in a foreign runtime, so pass the value " - "explicitly under a different parameter name" + "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: diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index 84d350dc2981a..3be7ea18e1032 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -22,9 +22,10 @@ import pytest +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(): ... @@ -95,8 +96,6 @@ class TestStubTaskflowArgs: """The TaskFlow call on a stub captures the ordered positional-arg spec (``_arg_bindings``).""" def test_literal_and_xcom_spec(self): - from airflow.sdk import DAG - with DAG(dag_id="d"): extracted = stub(fn_extract)() result = stub(fn_transform)("uk", extracted) @@ -110,8 +109,6 @@ def test_literal_and_xcom_spec(self): assert op.upstream_task_ids == {"fn_extract"} def test_kwargs_normalize_to_declaration_order(self): - from airflow.sdk import DAG - with DAG(dag_id="d"): extracted = stub(fn_extract)() result = stub(fn_transform)(extracted=extracted["part"], country="fr", retries_num=7) @@ -126,8 +123,6 @@ 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): - from airflow.sdk import DAG - with DAG(dag_id="d"): result = stub(fn_untyped)(1, "x") @@ -140,8 +135,6 @@ def test_unresolvable_annotation_degrades_to_any(self): def fn(x): ... fn.__annotations__ = {"x": "NotARealType"} - from airflow.sdk import DAG - with DAG(dag_id="d"): result = stub(fn)("v") @@ -160,22 +153,19 @@ def test_context_key_param_rejected(self): stub(fn_context_key)(1) def test_non_json_literal_rejected(self): - from airflow.sdk import DAG - with DAG(dag_id="d"), pytest.raises(ValueError, match="not JSON-serializable"): stub(fn_transform)("uk", object()) def test_mapped_xcom_arg_rejected(self): - from airflow.sdk import DAG - 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): - from airflow.sdk import DAG - with DAG(dag_id="d"): with pytest.raises(TypeError, match="do not support dynamic task mapping"): stub(fn_transform).expand(country=["uk", "fr"], extracted=[{}, {}]) From 9ddf063edcb83291d65c3cb89750158a01e3d34e Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 13 Jul 2026 04:18:03 +0000 Subject: [PATCH 04/40] Move stub arg-binding support out of shared execution API paths Only the Multi-Lang stub-task path needs the serialized-dag machinery and the arg-binding models, so regular task-run requests should not pay for them: the TaskArgBinding datamodels move to a dedicated module and the serialized-dag imports become local to the stub lookup. The OpenAPI schema is unchanged (component names stay the same), which is why no execution API version bump accompanies this commit. --- .../datamodels/task_arg_binding.py | 58 +++++++++++++++++++ .../execution_api/datamodels/taskinstance.py | 29 +--------- .../execution_api/routes/task_instances.py | 11 ++-- 3 files changed, 66 insertions(+), 32 deletions(-) create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py 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 66f6ec1dc824a..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, @@ -393,34 +394,6 @@ def safe_extract_from_orm(cls, data: Any) -> Any: return values -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``.""" - - class TIRunContext(BaseModel): """Response schema for TaskInstance run 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 7573f34344ee7..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,11 +49,11 @@ 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, PrevSuccessfulDagRunResponse, - TaskArgBinding, TaskBreadcrumbsResponse, TaskStatesResponse, TIAwaitingInputStatePayload, @@ -81,7 +81,6 @@ from airflow.models.asset import AssetActive from airflow.models.base import ID_LEN from airflow.models.dag import DagModel -from airflow.models.dag_version import DagVersion from airflow.models.dagrun import DagRun as DR from airflow.models.hitl import HITLDetail from airflow.models.log import Log @@ -91,8 +90,6 @@ from airflow.models.trigger import Trigger, handle_event_submit from airflow.models.xcom import XComModel from airflow.serialization.definitions.assets import SerializedAsset, SerializedAssetUniqueKey -from airflow.serialization.enums import Encoding -from airflow.serialization.serialized_objects import BaseSerialization from airflow.state import get_state_backend from airflow.triggers.base import TriggerEvent from airflow.utils.sqlalchemy import get_dialect_name @@ -122,6 +119,12 @@ 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) From 3700c175e74e6ca4777a9af48ae27a36707268fb Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 13 Jul 2026 04:18:56 +0000 Subject: [PATCH 05/40] Add a Go SDK example Dag covering the full TaskFlow binding surface simple_dag only exercises the minimal binding: one literal and one XCom argument. The new taskflow_binding_dag locks in the rest of the surface end to end -- scalar and array literals, keyword arguments, a defaulted None, and XCom fan-in from two upstream Go tasks bound onto a strict struct and a typed slice -- with the Go task verifying every bound value so binding regressions fail the example run loudly. --- .../test_go_sdk_taskflow_binding.py | 102 ++++++++++++++ .../airflow-go-pack/pack_integration_test.go | 5 + go-sdk/dags/go_examples.py | 55 +++++++- go-sdk/example/bundle/main.go | 6 + .../bundle/taskflowbinding/taskflowbinding.go | 129 ++++++++++++++++++ .../taskflowbinding/taskflowbinding_test.go | 60 ++++++++ 6 files changed, 354 insertions(+), 3 deletions(-) create mode 100644 airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_taskflow_binding.py create mode 100644 go-sdk/example/bundle/taskflowbinding/taskflowbinding.go create mode 100644 go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go 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/go-sdk/cmd/airflow-go-pack/pack_integration_test.go b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go index 84e9a1045f4e8..f6432595a0458 100644 --- a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go +++ b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go @@ -154,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 0c54f2ed5c778..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 @@ -114,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 c03d8c80a71d5..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 } 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") +} From 79ce1d39c3da0a3ba2ae93752ee09435395ce664 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Sun, 19 Jul 2026 13:15:49 +0000 Subject: [PATCH 06/40] Support sdk.TaskInput struct-field injection for Go SDK stub tasks Naming every stub argument as a separate flat Go parameter gets unwieldy as the argument count grows, and there was no way for a Go task to pull an XCom that the Python TaskFlow call itself never passed. A struct that embeds sdk.TaskInput lets a task bind many arguments by name (or an explicit ad hoc XCom pull) onto one parameter instead, while the existing flat/positional binding keeps working unchanged for functions that don't opt in. This required adding a name to the wire-level TaskArgBinding spec so a struct field can look itself up by the Dag's TaskFlow argument name regardless of declaration order on either side, since Go cannot recover a plain function parameter's name via reflection the way it can for a struct's fields. --- .../datamodels/task_arg_binding.py | 3 + .../execution_api/versions/v2026_06_30.py | 6 +- .../versions/head/test_task_instances.py | 10 +- .../v2026_04_17/test_task_instances.py | 10 +- .../serialization/test_dag_serialization.py | 13 +- .../test_go_sdk_taskflow_binding.py | 15 +- go-sdk/README.md | 45 ++ .../airflow-go-pack/pack_integration_test.go | 1 + go-sdk/dags/go_examples.py | 19 +- go-sdk/example/bundle/main.go | 1 + .../bundle/taskflowbinding/taskflowbinding.go | 48 +++ .../taskflowbinding/taskflowbinding_test.go | 24 ++ go-sdk/pkg/binding/binding.go | 397 ++++++++++++++++-- go-sdk/pkg/binding/binding_test.go | 217 ++++++++++ go-sdk/pkg/execution/genmodels/models.gen.go | 3 + go-sdk/pkg/execution/integration_test.go | 45 ++ go-sdk/pkg/execution/task_runner.go | 1 + go-sdk/sdk/context.go | 24 ++ .../providers/standard/decorators/stub.py | 8 +- .../unit/standard/decorators/test_stub.py | 32 +- .../airflow/sdk/api/datamodels/_generated.py | 1 + .../sdk/execution_time/schema/schema.json | 125 +++--- .../schema/versions/v2026_07_30.py | 7 +- .../execution_time/schema/test_migrator.py | 11 +- 24 files changed, 945 insertions(+), 121 deletions(-) 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 index 432c7ce3a4e20..a696f25edc62b 100644 --- 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 @@ -42,6 +42,9 @@ class TaskArgBinding(BaseModel): generates a plain struct in the foreign-language SDKs consuming the supervisor schema. """ + name: str + """The stub function's parameter name this binding fills, in declaration order.""" + kind: Literal["xcom", "literal"] """Whether the value comes from an upstream task's XCom or is a literal from the Dag file.""" 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 15b04fc1db8d2..a85b163dcba9c 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 @@ -25,6 +25,7 @@ schema, ) +from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import TaskArgBinding from airflow.api_fastapi.execution_api.datamodels.taskinstance import ( DagRun, TaskInstance, @@ -148,7 +149,10 @@ class AddArgBindingsToTIRunContext(VersionChange): description = __doc__ - instructions_to_migrate_to_previous_version = (schema(TIRunContext).field("arg_bindings").didnt_exist,) + instructions_to_migrate_to_previous_version = ( + schema(TIRunContext).field("arg_bindings").didnt_exist, + schema(TaskArgBinding).field("name").didnt_exist, + ) @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type] def remove_arg_bindings_field(response: ResponseInfo) -> None: # type: ignore[misc] 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 ff6678a08ff4b..dedb25c2f1083 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 @@ -398,8 +398,14 @@ def transform(country: str, extracted: dict): ... 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"}, + {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, + { + "name": "extracted", + "kind": "xcom", + "data_type": "object", + "task_id": "extract", + "key": "return_value", + }, ] # An argless stub has no captured spec, so the field stays unset. 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 e51444d239b86..7f648f1134039 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 @@ -124,6 +124,12 @@ 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"}, + {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, + { + "name": "extracted", + "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 f03ac2d64fc96..bb0354e5741bf 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -3422,11 +3422,12 @@ def transform(country: str, extracted: dict): ... assert encoded_tasks["transform"]["_arg_bindings"] == [ { Encoding.TYPE: DAT.DICT, - Encoding.VAR: {"kind": "literal", "data_type": "string", "value": "uk"}, + Encoding.VAR: {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, }, { Encoding.TYPE: DAT.DICT, Encoding.VAR: { + "name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract", @@ -3437,8 +3438,14 @@ def transform(country: str, extracted: dict): ... 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"}, + {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, + { + "name": "extracted", + "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 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 index 2a6065801f105..ab087eb7f7bd8 100644 --- 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 @@ -72,7 +72,7 @@ def test_all_tasks_succeeded(completed_run: _CompletedRun): 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"): + for task_id in ("make_config", "make_numbers", "combine", "combine_via_task_input"): assert completed_run.ti_states.get(task_id) == "success", completed_run.ti_states @@ -100,3 +100,16 @@ def test_combine_summary_reflects_bound_arguments(completed_run: _CompletedRun): "sum": 20, "note_was_null": True, } + + +def test_combine_via_task_input_summary_reflects_bound_arguments(completed_run: _CompletedRun): + """``combine_via_task_input`` demonstrates the Go SDK's ``sdk.TaskInput`` + struct-field injection mode: ``region_code``/``threshold`` bind by name onto + the struct exactly like ``combine``'s flat parameters do, while the struct's + third field is an ad hoc XCom pull of ``make_config``'s return value declared + purely in Go, with no corresponding TaskFlow call argument here.""" + assert completed_run.xcom("combine_via_task_input") == { + "region": "eu-west-1", + "threshold": 0.75, + "environment": "production", + } diff --git a/go-sdk/README.md b/go-sdk/README.md index 3f3b14ad44c8f..6b06cfe744155 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -141,6 +141,51 @@ Asking for the narrowest interface a task needs (e.g. `sdk.VariableClient` inste unit testing easier and documents which Airflow features the task touches. `RegisterDags` is the single source of truth for which `dag_id`s and `task_id`s a bundle can run. +### TaskInput structs + +A struct that anonymously embeds `sdk.TaskInput` opts into **per-field, name-based** binding instead +of a long flat parameter list — at most one such parameter is allowed per function, and it can be +mixed with plain flat parameters: + +```go +type CombineInput struct { + sdk.TaskInput // one-line opt-in, zero runtime cost + Region string `arg:"region_code"` // named lookup against the TaskFlow call argument "region_code" + Threshold float64 // no tag -> falls back to the snake_cased field name "threshold" + Config Config `xcom:"make_config"` // ad hoc pull of make_config's return-value XCom, independent + // of the TaskFlow call -- there is no "config" argument at all +} + +func Combine(ctx sdk.TIRunContext, log *slog.Logger, input CombineInput) (any, error) { + // input.Region, input.Threshold, input.Config are all populated. + return nil, nil +} +``` + +Each exported field supports three all-optional tags: + +- `arg:""` — bind from the TaskFlow call argument with this name (matched against the stub + function's Python parameter name, independent of declaration order on either side). With no tag, + the field's own Go name, snake_cased (`RatioValue` → `ratio_value`, `TaskID` → `task_id`), is used. +- `xcom:""` — an explicit, ad hoc XCom pull from the named upstream task, fully independent + of the TaskFlow call: the field need not correspond to any argument the Dag file passes at all. This + pull is unchecked — the runtime does not verify `` is an actual upstream dependency, the + same trust model as calling `sdk.Client.GetXCom` by hand. +- `xcom-key:""` — the XCom key for an `xcom:`-tagged field; defaults to the return-value key. + Setting it without `xcom:` is a registration-time error (a key with no task id is meaningless), as is + setting both `arg:` and `xcom:` on the same field. + +When a `TaskInput` struct and plain flat parameters coexist in the same function, the struct's fields +claim entries out of the TaskFlow call's argument spec by name first; the *remaining, unclaimed* +entries are then distributed, in their original relative order, onto the flat parameters in +declaration order. With no `TaskInput` struct present, this is exactly today's positional-only +behaviour. A plain custom struct type *without* the `sdk.TaskInput` embed is unaffected by any of +this — it keeps working as a single flat data parameter, JSON-decoded whole from one TaskFlow +argument (see `Config` in +[`example/bundle/taskflowbinding/taskflowbinding.go`](./example/bundle/taskflowbinding/taskflowbinding.go)), +which is a different mechanism from per-field `TaskInput` binding. See +[`CombineViaTaskInput`](./example/bundle/taskflowbinding/taskflowbinding.go) for a full worked example. + ### Reading the task runtime context Declare an `sdk.TIRunContext` parameter on a task to read the identifiers and scheduling timestamps of the diff --git a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go index f6432595a0458..09d4143d1793c 100644 --- a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go +++ b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go @@ -159,6 +159,7 @@ dags: - "make_config" - "make_numbers" - "combine" + - "combine_via_task_input" ` assert.Equal(t, expectedManifest, string(metadata)) diff --git a/go-sdk/dags/go_examples.py b/go-sdk/dags/go_examples.py index 39328747283c3..ae91b1512d43c 100644 --- a/go-sdk/dags/go_examples.py +++ b/go-sdk/dags/go_examples.py @@ -20,7 +20,9 @@ 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). +(stressing the TaskFlow argument-binding surface -- both the flat parameter +list ``combine`` binds onto and the ``sdk.TaskInput`` struct ``combine_via_task_input`` +binds onto instead, 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 @@ -138,6 +140,10 @@ def combine( ): ... +@task.stub(queue="golang") +def combine_via_task_input(region_code: str, threshold: float): ... + + @dag(dag_id="taskflow_binding_dag") def taskflow_binding_dag(): """ @@ -150,16 +156,25 @@ def taskflow_binding_dag(): 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_via_task_input`` demonstrates the Go SDK's ``sdk.TaskInput`` struct + injection mode: ``region_code``/``threshold`` bind by name onto the struct's + fields exactly like ``combine``'s flat parameters do, but the struct's third + field is an ad hoc XCom pull of ``make_config``'s return value declared purely + in Go (an ``xcom:`` struct tag) -- it is never passed as a TaskFlow argument + here, so the explicit ``>>`` below is what orders it after ``make_config``. """ + config = make_config() combine( "summary", 3, 2.5, True, ["metrics", "hourly"], - config=make_config(), + config=config, numbers=make_numbers(), ) + config >> combine_via_task_input(region_code="eu-west-1", threshold=0.75) taskflow_binding_dag() diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 1eaf213666913..581044ce3aca4 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -60,6 +60,7 @@ func (m *myBundle) RegisterDags(dagbag v1.Registry) error { bindingDag.AddTaskWithName("make_config", taskflowbinding.MakeConfig) bindingDag.AddTaskWithName("make_numbers", taskflowbinding.MakeNumbers) bindingDag.AddTaskWithName("combine", taskflowbinding.Combine) + bindingDag.AddTaskWithName("combine_via_task_input", taskflowbinding.CombineViaTaskInput) return nil } diff --git a/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go b/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go index a23b470ed93c6..07fd12c11a527 100644 --- a/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go +++ b/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go @@ -20,6 +20,9 @@ // 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. +// CombineViaTaskInput additionally shows the sdk.TaskInput struct-field +// injection mode: the same binding surface collapsed into one struct +// parameter instead of a long flat list. package taskflowbinding import ( @@ -127,3 +130,48 @@ func Combine( "note_was_null": note == nil, }, nil } + +// CombineInput demonstrates the sdk.TaskInput struct-field injection mode: an +// ergonomic alternative to Combine's long flat parameter list. Region binds +// by an explicit arg: tag; Threshold has no tag, so it falls back to its Go +// field name snake_cased ("threshold"); Config is an ad hoc XCom pull of +// make_config's return value, independent of the Python call's TaskFlow +// arguments entirely. +type CombineInput struct { + sdk.TaskInput + Region string `arg:"region_code"` + Threshold float64 + Config Config ` xcom:"make_config"` +} + +// CombineViaTaskInput is the TaskInput-struct sibling of Combine: the same +// kind of binding surface -- a named literal, a snake_case-fallback literal, +// and an ad hoc XCom pull -- collapsed into one struct parameter instead of +// many flat ones. The Python side calls it as +// +// combine_via_task_input(region_code="eu-west-1", threshold=0.75) +func CombineViaTaskInput(ctx sdk.TIRunContext, log *slog.Logger, input CombineInput) (any, error) { + if input.Region != "eu-west-1" || input.Threshold != 0.75 { + return nil, fmt.Errorf( + "TaskInput fields bound incorrectly: region=%q threshold=%v", + input.Region, + input.Threshold, + ) + } + if want := (Config{Environment: "production", Region: "eu-west-1", Debug: true}); input.Config != want { + return nil, fmt.Errorf( + "ad hoc xcom field bound incorrectly: config=%+v, want %+v", input.Config, want, + ) + } + + log.InfoContext(ctx, "Bound TaskInput struct", + "region", input.Region, + "threshold", input.Threshold, + "environment", input.Config.Environment, + ) + return map[string]any{ + "region": input.Region, + "threshold": input.Threshold, + "environment": input.Config.Environment, + }, nil +} diff --git a/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go b/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go index 51168778ce9b8..298ec8508c9aa 100644 --- a/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go +++ b/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go @@ -58,3 +58,27 @@ func TestCombineRejectsWrongBinding(t *testing.T) { ) assert.ErrorContains(t, err, "object XCom bound incorrectly") } + +func TestCombineViaTaskInput(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + got, err := CombineViaTaskInput(ctx, slog.Default(), CombineInput{ + Region: "eu-west-1", + Threshold: 0.75, + Config: Config{Environment: "production", Region: "eu-west-1", Debug: true}, + }) + require.NoError(t, err) + + summary, ok := got.(map[string]any) + require.True(t, ok, "CombineViaTaskInput should return a map summary, got %T", got) + assert.Equal(t, "production", summary["environment"]) +} + +func TestCombineViaTaskInputRejectsWrongBinding(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + _, err := CombineViaTaskInput(ctx, slog.Default(), CombineInput{ + Region: "eu-west-1", + Threshold: 0.75, + Config: Config{}, + }) + assert.ErrorContains(t, err, "ad hoc xcom field bound incorrectly") +} diff --git a/go-sdk/pkg/binding/binding.go b/go-sdk/pkg/binding/binding.go index c8a621835e017..e5fe6af725d8b 100644 --- a/go-sdk/pkg/binding/binding.go +++ b/go-sdk/pkg/binding/binding.go @@ -18,16 +18,33 @@ // 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: +// Three 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. +// - Data parameters: everything else (except TaskInput structs, below), 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. +// - TaskInput structs: a struct that anonymously embeds sdk.TaskInput opts +// into per-field, name-based binding instead of consuming one positional +// slot as a whole-value decode target. Each exported field binds by name +// (an `arg:""` tag, or its Go field name snake_cased) against the +// Dag's TaskFlow call arguments, or by an explicit, ad hoc XCom pull (an +// `xcom:""` tag, with an optional `xcom-key:""`) that never +// consults the positional argument spec at all. At most one such +// parameter is allowed per function. +// +// A TaskInput struct's fields are resolved first, by name, claiming entries +// out of the argument spec; the remaining unclaimed entries are then +// distributed, in their original relative order, onto the plain flat data +// parameters in declaration order -- so flat parameters and a TaskInput +// struct can coexist in the same function signature regardless of where each +// sits, and (with no TaskInput struct present) this reduces to exactly +// today's positional-only behaviour. // // Analyze inspects a function once at registration and returns a Plan; Resolve // builds the call arguments for each execution from that Plan and the @@ -46,6 +63,8 @@ import ( "fmt" "log/slog" "reflect" + "strings" + "unicode" "github.com/apache/airflow/go-sdk/pkg/api" "github.com/apache/airflow/go-sdk/pkg/sdkcontext" @@ -78,6 +97,10 @@ const ( // 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 { + // Name is the stub function's parameter name this binding fills. Always + // populated; used to match a TaskInput struct field's `arg:` tag (or its + // snake_cased field-name fallback). + Name string Kind ArgKind // TaskID is the upstream task to pull from. Set only for ArgKindXCom. TaskID string @@ -100,16 +123,54 @@ const ( paramLogger paramClient paramData + // paramTaskInput is a struct that anonymously embeds sdk.TaskInput, + // opting into per-field, name-based binding instead of consuming one + // positional slot as a whole-value decode target. + paramTaskInput +) + +// taskInputFieldSource discriminates how a TaskInput struct field is filled. +type taskInputFieldSource int + +const ( + // taskInputFieldFromArg claims a named entry from the argument spec. + taskInputFieldFromArg taskInputFieldSource = iota + // taskInputFieldFromXCom pulls directly via an `xcom:` tag, independent + // of the argument spec. + taskInputFieldFromXCom ) +// taskInputField describes how Resolve fills one exported field of a +// TaskInput-embedding struct. Precomputed once by Analyze. +type taskInputField struct { + // structIndex is the field's index within the struct, for + // reflect.Value.Field. + structIndex int + // goName is the Go field name, for error messages. + goName string + fieldType reflect.Type + source taskInputFieldSource + // argName is the name to claim from the argument spec. Set only for + // taskInputFieldFromArg. + argName string + // xcomTaskID/xcomKey identify the ad hoc pull. Set only for + // taskInputFieldFromXCom. + xcomTaskID string + xcomKey string +} + // paramPlan describes how Resolve fills a single task-function parameter. type paramPlan struct { kind paramKind - // typ is the declared Go type of a data parameter (kind == paramData). + // typ is the declared Go type of a data parameter (kind == paramData) or + // the struct type (kind == paramTaskInput). typ reflect.Type // index is the parameter's position in the function signature, for error // messages. index int + // fields describes each exported field's binding. Set only for + // kind == paramTaskInput. + fields []taskInputField } // Plan is the precomputed recipe for filling a task function's parameters. It @@ -129,11 +190,24 @@ func (p *Plan) NumData() int { return p.numData } // anything else is a registration error. func Analyze(fnType reflect.Type, fnName string) (*Plan, error) { p := &Plan{fnName: fnName, params: make([]paramPlan, fnType.NumIn())} + seenTaskInput := -1 for i := range fnType.NumIn() { plan, err := classifyParam(fnName, fnType.In(i), i) if err != nil { return nil, err } + if plan.kind == paramTaskInput { + if seenTaskInput >= 0 { + return nil, fmt.Errorf( + "task function %s: parameter %d: only one TaskInput struct parameter is allowed "+ + "per function (parameter %d already is one)", + fnName, + i, + seenTaskInput, + ) + } + seenTaskInput = i + } if plan.kind == paramData { p.numData++ } @@ -143,24 +217,65 @@ func Analyze(fnType reflect.Type, fnName string) (*Plan, error) { } // 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. +// parameters receive values derived from ctx, logger, or client. A TaskInput +// struct's fields are resolved first, by name (or an explicit ad hoc xcom +// pull), claiming entries out of args; the remaining unclaimed entries are +// then distributed, in their original relative order, onto the plain flat +// data parameters 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 { + byName := make(map[string]int, len(args)) + for i, a := range args { + byName[a.Name] = i + } + claimed := make([]bool, len(args)) + + out := make([]reflect.Value, len(p.params)) + for i, plan := range p.params { + if plan.kind != paramTaskInput { + continue + } + v, err := p.resolveTaskInput(ctx, client, plan, args, byName, claimed) + if err != nil { + return nil, err + } + out[i] = v + } + + remaining := 0 + for _, c := range claimed { + if !c { + remaining++ + } + } + if remaining != 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, + p.fnName, remaining, p.numData, ) } - out := make([]reflect.Value, len(p.params)) - argIdx := 0 + + argCursor := 0 + nextUnclaimed := func() Arg { + for argCursor < len(args) { + i := argCursor + argCursor++ + if !claimed[i] { + return args[i] + } + } + // Unreachable: the remaining/numData check above guarantees enough + // unclaimed args exist for every flat parameter still to be filled. + panic("binding: exhausted unclaimed args despite a passing arity check") + } + + flatIdx := 0 for i, plan := range p.params { switch plan.kind { case paramTIRunContext: @@ -178,22 +293,22 @@ func (p *Plan) Resolve( out[i] = reflect.ValueOf(logger) case paramClient: out[i] = reflect.ValueOf(client) + case paramTaskInput: + // Already resolved above. case paramData: - arg := args[argIdx] - v, err := p.resolveData(ctx, client, plan, arg, argIdx) + v, err := p.resolveData(ctx, client, plan, nextUnclaimed(), flatIdx) if err != nil { return nil, err } out[i] = v - argIdx++ + flatIdx++ } } 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. +// resolveData produces the value for one flat data parameter from its +// argument spec. func (p *Plan) resolveData( ctx context.Context, c sdk.XComClient, @@ -201,18 +316,100 @@ func (p *Plan) resolveData( 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, - ) + return p.resolveOne( + ctx, c, plan.typ, arg, + fmt.Sprintf("argument %d (parameter %d)", argIdx, plan.index), + fmt.Sprintf("argument %d", argIdx), + ) +} + +// resolveTaskInput builds the struct value for one TaskInput parameter. A +// field tagged `xcom:` pulls directly and never touches byName/claimed; a +// field claiming an argument spec entry by name marks it claimed so the +// later flat-parameter cursor skips it. +func (p *Plan) resolveTaskInput( + ctx context.Context, + client sdk.Client, + plan paramPlan, + args []Arg, + byName map[string]int, + claimed []bool, +) (reflect.Value, error) { + structType := plan.typ + isPtr := structType.Kind() == reflect.Pointer + if isPtr { + structType = structType.Elem() + } + structVal := reflect.New(structType).Elem() + + for _, tif := range plan.fields { + typeCheckCtx := fmt.Sprintf("TaskInput field %s (parameter %d)", tif.goName, plan.index) + generalCtx := fmt.Sprintf("TaskInput field %s", tif.goName) + + var arg Arg + switch tif.source { + case taskInputFieldFromXCom: + // DataTypeAny: an ad hoc pull has no Dag-declared type to check + // against, so decoding is driven entirely by the Go field's type. + arg = Arg{ + Kind: ArgKindXCom, + TaskID: tif.xcomTaskID, + Key: tif.xcomKey, + DataType: DataTypeAny, + } + case taskInputFieldFromArg: + idx, ok := byName[tif.argName] + if !ok { + return reflect.Value{}, fmt.Errorf( + "task function %s: %s: arg name %q not found among the Dag's TaskFlow call arguments", + p.fnName, + typeCheckCtx, + tif.argName, + ) + } + claimed[idx] = true + arg = args[idx] + } + + v, err := p.resolveOne(ctx, client, tif.fieldType, arg, typeCheckCtx, generalCtx) + if err != nil { + return reflect.Value{}, err + } + structVal.Field(tif.structIndex).Set(v) + } + + if isPtr { + return structVal.Addr(), nil + } + return structVal, nil +} + +// resolveOne decodes one argument-spec entry into a value assignable to +// targetType: type-check against the declared Dag type, then decode a +// literal or pull-and-decode an XCom. Shared by resolveData (one flat +// parameter) and resolveTaskInput (one TaskInput struct field) so a literal- +// or xcom-kind entry resolves identically regardless of which parameter +// shape it fills. typeCheckCtx/generalCtx are error-message prefixes: the +// former (used only for the type-check error) additionally names the +// parameter index, matching this package's existing error conventions. +func (p *Plan) resolveOne( + ctx context.Context, + c sdk.XComClient, + targetType reflect.Type, + arg Arg, + typeCheckCtx string, + generalCtx string, +) (reflect.Value, error) { + if err := checkDataType(arg.DataType, targetType); err != nil { + return reflect.Value{}, fmt.Errorf("task function %s: %s: %w", p.fnName, typeCheckCtx, err) } switch arg.Kind { case ArgKindLiteral: - v, err := decodeValue(arg.Value, plan.typ) + v, err := decodeValue(arg.Value, targetType) 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, + "task function %s: %s: decoding literal value into %s: %w", + p.fnName, generalCtx, targetType, err, ) } return v, nil @@ -220,8 +417,8 @@ func (p *Plan) resolveData( 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, + "task function %s: %s: no workload in context, cannot resolve xcom argument", + p.fnName, generalCtx, ) } key := arg.Key @@ -233,21 +430,21 @@ func (p *Plan) resolveData( 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, + "task function %s: %s: pulling xcom from task %q (key %q): %w", + p.fnName, generalCtx, arg.TaskID, key, err, ) } - v, err := decodeValue(raw, plan.typ) + v, err := decodeValue(raw, targetType) 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, + "task function %s: %s: decoding xcom from task %q into %s: %w", + p.fnName, generalCtx, arg.TaskID, targetType, err, ) } return v, nil default: return reflect.Value{}, fmt.Errorf( - "task function %s: argument %d: unknown argument kind %q", p.fnName, argIdx, arg.Kind, + "task function %s: %s: unknown argument kind %q", p.fnName, generalCtx, arg.Kind, ) } } @@ -283,6 +480,13 @@ func classifyParam(fnName string, in reflect.Type, index int) (paramPlan, error) fnName, index, in, explainClientMismatch(in), ) } + if structType := taskInputStructType(in); structType != nil { + fields, err := buildTaskInputFields(fnName, structType, index) + if err != nil { + return paramPlan{}, err + } + return paramPlan{kind: paramTaskInput, typ: in, index: index, fields: fields}, nil + } if !isDecodableType(in) { return paramPlan{}, fmt.Errorf( "task function %s: parameter %d: type %s cannot receive a task argument "+ @@ -293,6 +497,128 @@ func classifyParam(fnName string, in reflect.Type, index int) (paramPlan, error) return paramPlan{kind: paramData, typ: in, index: index}, nil } +// taskInputStructType reports whether in (after dereferencing one pointer +// level, matching checkDataType's convention) is a struct that anonymously +// embeds sdk.TaskInput, returning that struct type or nil. +func taskInputStructType(in reflect.Type) reflect.Type { + t := in + if t.Kind() == reflect.Pointer { + t = t.Elem() + } + if t.Kind() != reflect.Struct { + return nil + } + for i := range t.NumField() { + f := t.Field(i) + if f.Anonymous && f.Type == taskInputType { + return t + } + } + return nil +} + +// buildTaskInputFields validates and precomputes the field-binding plan for a +// TaskInput-embedding struct parameter. It runs once at registration time so +// a misconfigured struct fails loudly before any task ever executes. +func buildTaskInputFields( + fnName string, + structType reflect.Type, + paramIndex int, +) ([]taskInputField, error) { + var fields []taskInputField + seenArgNames := make(map[string]string) // resolved arg name -> Go field name that claims it + + for i := range structType.NumField() { + f := structType.Field(i) + if f.Anonymous && f.Type == taskInputType { + continue + } + if !f.IsExported() { + continue + } + + argTag, hasArg := f.Tag.Lookup("arg") + xcomTag, hasXCom := f.Tag.Lookup("xcom") + xcomKeyTag, hasXComKey := f.Tag.Lookup("xcom-key") + + if hasArg && hasXCom { + return nil, fmt.Errorf( + "task function %s: parameter %d: TaskInput field %s: cannot set both `arg` and `xcom` tags", + fnName, + paramIndex, + f.Name, + ) + } + if hasXComKey && !hasXCom { + return nil, fmt.Errorf( + "task function %s: parameter %d: TaskInput field %s: `xcom-key` requires `xcom` to also "+ + "be set (a key with no task id is meaningless)", + fnName, + paramIndex, + f.Name, + ) + } + if !isDecodableType(f.Type) { + return nil, fmt.Errorf( + "task function %s: parameter %d: TaskInput field %s: type %s cannot receive a task "+ + "argument (func/chan/unsafe-pointer values cannot be decoded)", + fnName, + paramIndex, + f.Name, + f.Type, + ) + } + + tif := taskInputField{structIndex: i, goName: f.Name, fieldType: f.Type} + if hasXCom { + tif.source = taskInputFieldFromXCom + tif.xcomTaskID = xcomTag + tif.xcomKey = xcomKeyTag + if tif.xcomKey == "" { + tif.xcomKey = api.XComReturnValueKey + } + } else { + tif.source = taskInputFieldFromArg + tif.argName = argTag + if tif.argName == "" { + tif.argName = snakeCase(f.Name) + } + if existing, ok := seenArgNames[tif.argName]; ok { + return nil, fmt.Errorf( + "task function %s: parameter %d: TaskInput fields %s and %s both bind arg name %q", + fnName, paramIndex, existing, f.Name, tif.argName, + ) + } + seenArgNames[tif.argName] = f.Name + } + fields = append(fields, tif) + } + return fields, nil +} + +// snakeCase converts a Go exported field name (UpperCamelCase, acronyms +// preserved as a run) to the wire's snake_case convention, e.g. +// "RatioValue" -> "ratio_value", "TaskID" -> "task_id". Used as the fallback +// arg name for a TaskInput struct field with no explicit `arg:` tag. +func snakeCase(name string) string { + runes := []rune(name) + var b strings.Builder + for i, r := range runes { + if unicode.IsUpper(r) { + prevLower := i > 0 && unicode.IsLower(runes[i-1]) + prevUpper := i > 0 && unicode.IsUpper(runes[i-1]) + nextLower := i+1 < len(runes) && unicode.IsLower(runes[i+1]) + if i > 0 && (prevLower || (prevUpper && nextLower)) { + b.WriteByte('_') + } + b.WriteRune(unicode.ToLower(r)) + } else { + b.WriteRune(r) + } + } + return b.String() +} + // 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. @@ -385,6 +711,7 @@ var ( tiRunContextType = reflect.TypeFor[sdk.TIRunContext]() slogLoggerType = reflect.TypeFor[*slog.Logger]() clientType = reflect.TypeFor[sdk.Client]() + taskInputType = reflect.TypeFor[sdk.TaskInput]() ) func isContext(inType reflect.Type) bool { diff --git a/go-sdk/pkg/binding/binding_test.go b/go-sdk/pkg/binding/binding_test.go index 48992a7fc7372..d7c573f683c2e 100644 --- a/go-sdk/pkg/binding/binding_test.go +++ b/go-sdk/pkg/binding/binding_test.go @@ -287,6 +287,45 @@ type extractResult struct { Timestamp int64 `json:"timestamp"` } +// simpleTaskInput is the minimal TaskInput struct: one field, no tags, so it +// falls back to matching its own (lowercased) field name. +type simpleTaskInput struct { + sdk.TaskInput + Name string +} + +// nonEmbeddingStruct has no sdk.TaskInput sentinel, so it must keep resolving +// as today's whole-value decode target, not per-field TaskInput binding. +type nonEmbeddingStruct struct { + Name string +} + +// combineInput exercises every TaskInput field-tag combination: Name falls +// back to its field name, Count is explicitly named, Note is an ad hoc XCom +// pull at the default key, and Debug is an ad hoc XCom pull at a custom key. +type combineInput struct { + sdk.TaskInput + Name string + Count int `arg:"count"` + Note *string ` xcom:"make_note"` + Debug bool ` xcom:"make_config" xcom-key:"debug_flag"` +} + +// reportInput deliberately declares Ratio before Region, the reverse of the +// wire order those names appear in, to prove field declaration order is +// irrelevant to by-name claiming. +type reportInput struct { + sdk.TaskInput + Ratio float64 + Region string `arg:"region"` +} + +// xcomOnlyInput's sole field never touches the argument spec at all. +type xcomOnlyInput struct { + sdk.TaskInput + Extra string `xcom:"make_config"` +} + func (s *BindingSuite) TestResolveXComArgs() { client := &fakeXComClient{values: map[string]any{ "extract/return_value": map[string]any{"go_version": "go1.24", "timestamp": int64(42)}, @@ -391,3 +430,181 @@ func (s *BindingSuite) TestResolveTIRunContextRebuild() { s.Equal(dagRun, rc.DagRun()) s.Equal("uk", got[1].Interface()) } + +func (s *BindingSuite) TestAnalyzeTaskInputClassification() { + plan := analyze(s, func(input simpleTaskInput) error { return nil }) + s.Zero(plan.NumData(), "a TaskInput struct claims by name, not by position") + + ptrPlan := analyze(s, func(input *simpleTaskInput) error { return nil }) + s.Zero(ptrPlan.NumData(), "a pointer to a TaskInput struct is detected the same way") + + plainPlan := analyze(s, func(cfg nonEmbeddingStruct) error { return nil }) + s.Equal( + 1, plainPlan.NumData(), + "a plain struct without the TaskInput sentinel stays a whole-value data parameter", + ) +} + +func (s *BindingSuite) TestAnalyzeTaskInputTagValidation() { + type conflictingTags struct { + sdk.TaskInput + Field string `arg:"a" xcom:"b"` + } + type orphanXComKey struct { + sdk.TaskInput + Field string `xcom-key:"k"` + } + type duplicateArgNames struct { + sdk.TaskInput + A string + B string `arg:"a"` + } + type nonDecodableXComField struct { + sdk.TaskInput + Bad chan int `xcom:"t"` + } + + cases := map[string]struct { + fn any + errContains string + }{ + "conflicting-arg-and-xcom-tags": { + func(input conflictingTags) error { return nil }, + "cannot set both `arg` and `xcom` tags", + }, + "orphan-xcom-key": { + func(input orphanXComKey) error { return nil }, + "`xcom-key` requires `xcom` to also be set", + }, + "duplicate-arg-names": { + func(input duplicateArgNames) error { return nil }, + `fields A and B both bind arg name "a"`, + }, + "non-decodable-xcom-field": { + func(input nonDecodableXComField) error { return nil }, + "cannot receive a task argument", + }, + "two-taskinput-params": { + func(a simpleTaskInput, b simpleTaskInput) error { return nil }, + "only one TaskInput struct parameter is allowed", + }, + } + for name, tt := range cases { + s.Run(name, func() { + _, err := Analyze(reflect.TypeOf(tt.fn), "testFn") + if s.Assert().Error(err) { + s.Assert().Contains(err.Error(), tt.errContains) + } + }) + } +} + +func (s *BindingSuite) TestResolveTaskInputAllStruct() { + client := &fakeXComClient{values: map[string]any{ + "make_note/return_value": "hello", + "make_config/debug_flag": true, + }} + fn := func(input combineInput) error { return nil } + got, err := s.resolve(fn, []Arg{ + {Name: "name", Kind: ArgKindLiteral, Value: "widget", DataType: DataTypeString}, + {Name: "count", Kind: ArgKindLiteral, Value: 7, DataType: DataTypeInteger}, + }, client) + s.Require().NoError(err) + + input := got[0].Interface().(combineInput) + s.Equal("widget", input.Name) + s.Equal(7, input.Count) + s.Require().NotNil(input.Note) + s.Equal("hello", *input.Note) + s.True(input.Debug) + + s.Require().Len(client.calls, 2, "the ad hoc xcom: fields, in struct declaration order") + s.Equal("make_note", client.calls[0].taskID) + s.Equal( + api.XComReturnValueKey, + client.calls[0].key, + "no xcom-key tag defaults to the return-value key", + ) + s.Equal("make_config", client.calls[1].taskID) + s.Equal("debug_flag", client.calls[1].key) +} + +func (s *BindingSuite) TestResolveTaskInputMixedWithFlat() { + fn := func(prefix string, input reportInput, suffix string) error { return nil } + got, err := s.resolve(fn, []Arg{ + {Name: "prefix", Kind: ArgKindLiteral, Value: "head", DataType: DataTypeString}, + {Name: "region", Kind: ArgKindXCom, TaskID: "make_region", DataType: DataTypeString}, + {Name: "ratio", Kind: ArgKindLiteral, Value: 0.5, DataType: DataTypeNumber}, + {Name: "suffix", Kind: ArgKindLiteral, Value: "footer", DataType: DataTypeString}, + }, &fakeXComClient{values: map[string]any{"make_region/return_value": "east"}}) + s.Require().NoError(err) + + s.Equal( + "head", + got[0].Interface(), + "the leading flat parameter claims the first unclaimed wire entry", + ) + input := got[1].Interface().(reportInput) + s.Equal("east", input.Region, "Region resolves by name despite being declared after Ratio") + s.Equal(0.5, input.Ratio) + s.Equal( + "footer", + got[2].Interface(), + "the trailing flat parameter claims the last unclaimed wire entry", + ) +} + +func (s *BindingSuite) TestResolveTaskInputXComTagIndependentOfArgs() { + client := &fakeXComClient{values: map[string]any{"make_config/return_value": "cfg-value"}} + fn := func(input xcomOnlyInput) error { return nil } + got, err := s.resolve(fn, nil, client) + s.Require().NoError(err) + + input := got[0].Interface().(xcomOnlyInput) + s.Equal("cfg-value", input.Extra) + s.Require().Len(client.calls, 1) + s.Equal("make_config", client.calls[0].taskID) + s.Equal(api.XComReturnValueKey, client.calls[0].key) +} + +func (s *BindingSuite) TestResolveTaskInputLiteralThroughArgName() { + fn := func(input simpleTaskInput) error { return nil } + got, err := s.resolve(fn, []Arg{ + {Name: "name", Kind: ArgKindLiteral, Value: "widget", DataType: DataTypeString}, + }, &fakeXComClient{}) + s.Require().NoError(err) + s.Equal("widget", got[0].Interface().(simpleTaskInput).Name) +} + +func (s *BindingSuite) TestResolveTaskInputPointerStruct() { + fn := func(input *simpleTaskInput) error { return nil } + got, err := s.resolve(fn, []Arg{ + {Name: "name", Kind: ArgKindLiteral, Value: "widget", DataType: DataTypeString}, + }, &fakeXComClient{}) + s.Require().NoError(err) + input := got[0].Interface().(*simpleTaskInput) + s.Require().NotNil(input) + s.Equal("widget", input.Name) +} + +func (s *BindingSuite) TestResolveTaskInputUnmatchedArgName() { + fn := func(input simpleTaskInput) error { return nil } + _, err := s.resolve(fn, []Arg{ + {Name: "different_name", Kind: ArgKindLiteral, Value: "x", DataType: DataTypeString}, + }, &fakeXComClient{}) + if s.Assert().Error(err) { + s.Contains(err.Error(), `arg name "name" not found`) + } +} + +func (s *BindingSuite) TestResolveTaskInputArityMismatchForLeftoverArgs() { + fn := func(input simpleTaskInput, extra string) error { return nil } + _, err := s.resolve(fn, []Arg{ + {Name: "name", Kind: ArgKindLiteral, Value: "widget", DataType: DataTypeString}, + }, &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)") + } +} diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index 2e6caf47c3da3..8b94e03847a07 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -1644,6 +1644,9 @@ type TaskArgBinding struct { // Kind corresponds to the JSON schema field "kind". Kind TaskArgBindingKind `msgpack:"kind"` + // Name corresponds to the JSON schema field "name". + Name string `msgpack:"name"` + // TaskID corresponds to the JSON schema field "task_id". TaskID interface{} `msgpack:"task_id,omitempty"` diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 5646a82e29206..3454c58f910a4 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -308,6 +308,51 @@ func TestTaskRunnerArgBindingsArityMismatch(t *testing.T) { assert.False(t, ran, "the task body must not run on an arity mismatch") } +// combineInput is a TaskInput struct whose sole field claims a named entry +// out of ti_context.arg_bindings. +type combineInput struct { + sdk.TaskInput + Region string `arg:"region"` +} + +// TestTaskRunnerBindsTaskInputStructArgs covers the TaskFlow path through +// RunTask for a TaskInput struct parameter: convertArgBindings must propagate +// each spec's Name through to binding.Arg so the struct's `arg:"region"` field +// can claim it by name. +func TestTaskRunnerBindsTaskInputStructArgs(t *testing.T) { + var got combineInput + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTaskWithName("transform", + func(input combineInput) error { + got = input + return nil + }) + }) + + details := &genmodels.StartupDetails{ + TI: genmodels.TaskInstance{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "transform", + RunID: "run1", + MapIndex: ptr(-1), + }, + BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, + TIContext: genmodels.TIRunContext{ + ArgBindings: &genmodels.ArgBindings{ + {Name: "region", Kind: "literal", DataType: "string", Value: "eu-west-1"}, + }, + }, + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(context.Background(), bundle, details, comm, logger) + assertSucceedTask(t, result) + assert.Equal(t, "eu-west-1", got.Region) +} + // 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) { diff --git a/go-sdk/pkg/execution/task_runner.go b/go-sdk/pkg/execution/task_runner.go index d6081d24ad4c3..bb09726aa9584 100644 --- a/go-sdk/pkg/execution/task_runner.go +++ b/go-sdk/pkg/execution/task_runner.go @@ -144,6 +144,7 @@ func convertArgBindings(specsPtr *genmodels.ArgBindings) []binding.Arg { taskID = s } args[i] = binding.Arg{ + Name: spec.Name, Kind: binding.ArgKind(spec.Kind), TaskID: taskID, Key: spec.Key, diff --git a/go-sdk/sdk/context.go b/go-sdk/sdk/context.go index afd0dd3705051..8400ff8430191 100644 --- a/go-sdk/sdk/context.go +++ b/go-sdk/sdk/context.go @@ -108,3 +108,27 @@ type DagRun struct { DataIntervalStart *time.Time DataIntervalEnd *time.Time } + +// TaskInput is a zero-size marker embedded anonymously in a struct to opt +// that struct into per-field, name-based TaskFlow argument binding -- an +// ergonomic alternative to a long flat parameter list. Each exported field of +// such a struct may carry an `arg:""` tag naming the stub's TaskFlow +// argument to bind (falling back to the field's own name, snake_cased, when +// omitted), or an `xcom:""` tag (with an optional companion +// `xcom-key:""`, defaulting to the return-value key) for an explicit, +// ad hoc XCom pull independent of the TaskFlow call: +// +// type CombineInput struct { +// sdk.TaskInput +// Name string +// Count int `arg:"count"` +// Extra string `xcom:"make_config" xcom-key:"environment"` +// } +// +// func Combine(ctx sdk.TIRunContext, log *slog.Logger, input CombineInput) (any, error) +// +// Embedding costs nothing at runtime: the marker occupies zero bytes and is +// never read; its only purpose is for the binding package to detect it +// reflectively at task-registration time. See the go-sdk README's "TaskInput +// structs" section for the full field-tag reference. +type TaskInput struct{} diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index 72497f184e32e..bef2cc5b52442 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -93,7 +93,10 @@ def _build_arg_bindings( 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. + (``kind="literal"``) for everything else. ``name`` is always the stub function's parameter + name, so a foreign runtime can bind by name (e.g. the Go SDK's ``sdk.TaskInput`` struct + fields) in addition to the existing positional order. Returns ``None`` for parameterless + stubs. """ signature = inspect.signature(python_callable) @@ -137,6 +140,7 @@ def annotation_for(name: str, param: inspect.Parameter) -> Any: if isinstance(value, PlainXComArg): spec.append( { + "name": name, "kind": "xcom", "data_type": data_type, "task_id": value.operator.task_id, @@ -159,7 +163,7 @@ def annotation_for(name: str, param: inspect.Parameter) -> Any: 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}) + spec.append({"name": name, "kind": "literal", "data_type": data_type, "value": value}) return spec diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index 3be7ea18e1032..b8ba9a8ef20da 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -102,9 +102,15 @@ def test_literal_and_xcom_spec(self): 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}, + {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, + { + "name": "extracted", + "kind": "xcom", + "data_type": "object", + "task_id": "fn_extract", + "key": "return_value", + }, + {"name": "retries_num", "kind": "literal", "data_type": "integer", "value": 3}, ] assert op.upstream_task_ids == {"fn_extract"} @@ -114,9 +120,15 @@ def test_kwargs_normalize_to_declaration_order(self): 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}, + {"name": "country", "kind": "literal", "data_type": "string", "value": "fr"}, + { + "name": "extracted", + "kind": "xcom", + "data_type": "object", + "task_id": "fn_extract", + "key": "part", + }, + {"name": "retries_num", "kind": "literal", "data_type": "integer", "value": 7}, ] def test_zero_param_stub_has_no_spec(self): @@ -127,8 +139,8 @@ def test_untyped_params_degrade_to_any(self): 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"}, + {"name": "a", "kind": "literal", "data_type": "any", "value": 1}, + {"name": "b", "kind": "literal", "data_type": "any", "value": "x"}, ] def test_unresolvable_annotation_degrades_to_any(self): @@ -138,7 +150,9 @@ def fn(x): ... with DAG(dag_id="d"): result = stub(fn)("v") - assert result.operator._arg_bindings == [{"kind": "literal", "data_type": "any", "value": "v"}] + assert result.operator._arg_bindings == [ + {"name": "x", "kind": "literal", "data_type": "any", "value": "v"} + ] def test_varargs_rejected(self): with pytest.raises(ValueError, match="fixed number of parameters"): diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index aa62452980d5c..9366039637a35 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -394,6 +394,7 @@ class TaskArgBinding(BaseModel): generates a plain struct in the foreign-language SDKs consuming the supervisor schema. """ + name: Annotated[str, Field(title="Name")] 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 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 4dbfc70dead32..065129868185c 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -4584,6 +4584,71 @@ "title": "ConnectionResponse", "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": { + "name": { + "title": "Name", + "type": "string" + }, + "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": [ + "name", + "kind" + ], + "title": "TaskArgBinding", + "type": "object" + }, "AssetEventDagRunReference": { "additionalProperties": false, "description": "Schema for AssetEvent model used in DagRun.", @@ -4971,66 +5036,6 @@ "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/v2026_07_30.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_07_30.py index fd25aef095bb8..e69fe456f33c8 100644 --- 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 @@ -19,7 +19,7 @@ from cadwyn import VersionChange, schema -from airflow.sdk.api.datamodels._generated import TIRunContext +from airflow.sdk.api.datamodels._generated import TaskArgBinding, TIRunContext class AddArgBindingsToTIRunContext(VersionChange): @@ -27,4 +27,7 @@ class AddArgBindingsToTIRunContext(VersionChange): description = __doc__ - instructions_to_migrate_to_previous_version = (schema(TIRunContext).field("arg_bindings").didnt_exist,) + instructions_to_migrate_to_previous_version = ( + schema(TIRunContext).field("arg_bindings").didnt_exist, + schema(TaskArgBinding).field("name").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 4b146a62859ac..44ab182a41421 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 @@ -420,8 +420,14 @@ def startup_details(self): ), max_tries=1, arg_bindings=[ - {"kind": "literal", "data_type": "string", "value": "uk"}, - {"kind": "xcom", "data_type": "object", "task_id": "extract", "key": "return_value"}, + {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, + { + "name": "extracted", + "kind": "xcom", + "data_type": "object", + "task_id": "extract", + "key": "return_value", + }, ], ), sentry_integration="", @@ -439,3 +445,4 @@ 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"] + assert [a.name for a in out.ti_context.arg_bindings] == ["country", "extracted"] From 755eb105529308e078f49efc0045250a4419458c Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 20 Jul 2026 01:47:15 +0000 Subject: [PATCH 07/40] Use the @task.stub decorator in stub arg-binding regression tests These tests built stub tasks with the raw stub(fn)(...) call instead of the @task.stub decorator every real Dag (Go/TS/Java examples) already uses, so a reader comparing the tests to real usage saw a syntax the feature doesn't actually ship. --- .../versions/head/test_task_instances.py | 11 ++++++----- .../versions/v2026_04_17/test_task_instances.py | 12 +++++++----- .../unit/serialization/test_dag_serialization.py | 13 ++++++++----- 3 files changed, 21 insertions(+), 15 deletions(-) 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 dedb25c2f1083..2abac5bd68172 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 @@ -372,14 +372,15 @@ async def workload_token(request: Request) -> TIToken: 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 + with dag_maker("test_arg_bindings_dag", serialized=True): - def extract(): ... + @task.stub + def extract(): ... - def transform(country: str, extracted: dict): ... + @task.stub + def transform(country: str, extracted: dict): ... - with dag_maker("test_arg_bindings_dag", serialized=True): - stub(transform)("uk", stub(extract)()) + transform("uk", extract()) dr = dag_maker.create_dagrun() tis = {ti.task_id: ti for ti in dr.get_task_instances()} 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 7f648f1134039..2b55a6e2ae3f4 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 @@ -20,6 +20,7 @@ import pytest from airflow._shared.timezones import timezone +from airflow.sdk import task from airflow.utils.state import DagRunState, State from tests_common.test_utils.config import conf_vars @@ -99,14 +100,15 @@ def teardown_method(self): @pytest.fixture def stub_ti(self, dag_maker): - from airflow.providers.standard.decorators.stub import stub + with dag_maker("test_arg_bindings_compat_dag", serialized=True): - def extract(): ... + @task.stub + def extract(): ... - def transform(country: str, extracted: dict): ... + @task.stub + def transform(country: str, extracted: dict): ... - with dag_maker("test_arg_bindings_compat_dag", serialized=True): - stub(transform)("uk", stub(extract)()) + transform("uk", extract()) dr = dag_maker.create_dagrun() tis = {ti.task_id: ti for ti in dr.get_task_instances()} diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index bb0354e5741bf..6e23878691a8e 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -3407,14 +3407,17 @@ def inner(): 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 + from airflow.sdk import task + + with DAG(dag_id="arg_bindings_dag", schedule=None) as dag: - def extract(): ... + @task.stub + def extract(): ... - def transform(country: str, extracted: dict): ... + @task.stub + def transform(country: str, extracted: dict): ... - with DAG(dag_id="arg_bindings_dag", schedule=None) as dag: - stub(transform)("uk", stub(extract)()) + transform("uk", extract()) ser_dag = DagSerialization.to_dict(dag) encoded_tasks = {t[Encoding.VAR]["task_id"]: t[Encoding.VAR] for t in ser_dag["dag"]["tasks"]} From b53d8462c2d68a5acbe871fdf0639fff2ad055cd Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 20 Jul 2026 01:53:36 +0000 Subject: [PATCH 08/40] Promote ArgBindingDataType to a real enum for stable codegen naming As a plain Literal type alias, the field's generated model came out under a generic, field-derived name (DataType) in both the task-sdk client model and the Go SDK's generated types, rather than the ArgBindingDataType name declared in the source. A real Enum class carries its own name through codegen, so providers/standard can import it directly instead of re-deriving the same string vocabulary by hand, with a hand-written fallback for Airflow 2 where the execution-API generated models aren't importable. --- .../datamodels/task_arg_binding.py | 16 ++++- .../execution_api/versions/v2026_06_30.py | 9 ++- .../pkg/execution/genmodels/defaults.gen.go | 2 +- go-sdk/pkg/execution/genmodels/models.gen.go | 58 ++++++++----------- .../providers/standard/decorators/stub.py | 43 +++++++++----- .../airflow/sdk/api/datamodels/_generated.py | 26 +++++---- .../sdk/execution_time/schema/schema.json | 41 +++++-------- 7 files changed, 107 insertions(+), 88 deletions(-) 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 index a696f25edc62b..171801f0837ad 100644 --- 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 @@ -24,14 +24,24 @@ from __future__ import annotations +from enum import Enum 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 ArgBindingDataType(str, Enum): + """Language-neutral value type a stub-task argument binds to in the foreign runtime.""" + + STRING = "string" + INTEGER = "integer" + NUMBER = "number" + BOOLEAN = "boolean" + OBJECT = "object" + ARRAY = "array" + ANY = "any" class TaskArgBinding(BaseModel): @@ -48,7 +58,7 @@ class TaskArgBinding(BaseModel): 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" + data_type: ArgBindingDataType = ArgBindingDataType.ANY """Declared type from the stub function's annotation; runtimes type-check against it.""" task_id: str | None = None 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 a85b163dcba9c..0deb96d0940f2 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 @@ -145,7 +145,14 @@ def remove_partition_date_from_dag_run(response: ResponseInfo) -> None: # type: class AddArgBindingsToTIRunContext(VersionChange): - """Add the ``arg_bindings`` positional-argument binding spec for stub (foreign-runtime) tasks.""" + """ + Add the ``arg_bindings`` positional-argument binding spec for stub (foreign-runtime) tasks. + + ``TaskArgBinding.data_type`` is declared as the ``ArgBindingDataType`` enum rather than an + inline ``Literal``; the wire representation (a JSON string) is unchanged, so no migration + instruction is needed -- this version has not been released with the ``arg_bindings`` field + in any other shape. + """ description = __doc__ diff --git a/go-sdk/pkg/execution/genmodels/defaults.gen.go b/go-sdk/pkg/execution/genmodels/defaults.gen.go index d411881bc8404..a6193f2b9af99 100644 --- a/go-sdk/pkg/execution/genmodels/defaults.gen.go +++ b/go-sdk/pkg/execution/genmodels/defaults.gen.go @@ -285,7 +285,7 @@ func (m *SucceedTask) DecodeMsgpack(dec *msgpack.Decoder) error { // 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"} + v := alias{DataType: ArgBindingDataType("any"), Key: "return_value"} if err := dec.Decode(&v); err != nil { return err } diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index 8b94e03847a07..d6ce08114a491 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -20,6 +20,16 @@ package genmodels import "time" +type ArgBindingDataType string + +const ArgBindingDataTypeAny ArgBindingDataType = "any" +const ArgBindingDataTypeArray ArgBindingDataType = "array" +const ArgBindingDataTypeBoolean ArgBindingDataType = "boolean" +const ArgBindingDataTypeInteger ArgBindingDataType = "integer" +const ArgBindingDataTypeNumber ArgBindingDataType = "number" +const ArgBindingDataTypeObject ArgBindingDataType = "object" +const ArgBindingDataTypeString ArgBindingDataType = "string" + type ArgBindings []TaskArgBinding // Schema for AssetAliasModel used in AssetEventDagRunReference. @@ -634,16 +644,6 @@ 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. @@ -1636,7 +1636,7 @@ type TIRunContext struct { // schema. type TaskArgBinding struct { // DataType corresponds to the JSON schema field "data_type". - DataType TaskArgBindingDataType `msgpack:"data_type,omitempty"` + DataType ArgBindingDataType `msgpack:"data_type,omitempty"` // Key corresponds to the JSON schema field "key". Key string `msgpack:"key,omitempty"` @@ -1654,16 +1654,6 @@ type TaskArgBinding struct { 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" @@ -1868,19 +1858,6 @@ 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"` @@ -1900,6 +1877,15 @@ type VariableKeysResult 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 VariableResult struct { // Key corresponds to the JSON schema field "key". Key string `msgpack:"key"` @@ -1911,6 +1897,10 @@ type VariableResult struct { Value interface{} `msgpack:"value,omitempty"` } +type VersionData map[string]interface{} + +type Warnings []interface{} + type XComCountResponse struct { // Len corresponds to the JSON schema field "len". Len int `msgpack:"len"` diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index bef2cc5b52442..b1ad5cc19295e 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -18,6 +18,7 @@ from __future__ import annotations import ast +import enum import inspect import json import types @@ -41,45 +42,61 @@ 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] +try: + from airflow.sdk.api.datamodels._generated import ArgBindingDataType +except ImportError: # Airflow 2 -- no task-sdk execution-API generated models + + class ArgBindingDataType(str, enum.Enum): # type: ignore[no-redef] + """Language-neutral value type a stub-task argument binds to in the foreign runtime.""" + + STRING = "string" + INTEGER = "integer" + NUMBER = "number" + BOOLEAN = "boolean" + OBJECT = "object" + ARRAY = "array" + ANY = "any" + + if TYPE_CHECKING: from airflow.providers.common.compat.sdk import Context -def _data_type_from_annotation(annotation: Any) -> str: +def _data_type_from_annotation(annotation: Any) -> ArgBindingDataType: """ 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. + 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" + return ArgBindingDataType.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" + return ArgBindingDataType.ANY annotation = origin if not isinstance(annotation, type): - return "any" + return ArgBindingDataType.ANY # bool subclasses int, and str/bytes are Sequences -- order matters. if issubclass(annotation, bool): - return "boolean" + return ArgBindingDataType.BOOLEAN if issubclass(annotation, int): - return "integer" + return ArgBindingDataType.INTEGER if issubclass(annotation, float): - return "number" + return ArgBindingDataType.NUMBER if issubclass(annotation, str): - return "string" + return ArgBindingDataType.STRING if issubclass(annotation, bytes): - return "any" + return ArgBindingDataType.ANY if issubclass(annotation, (dict, Mapping)): - return "object" + return ArgBindingDataType.OBJECT if issubclass(annotation, (list, tuple, set, frozenset, Sequence)): - return "array" - return "any" + return ArgBindingDataType.ARRAY + return ArgBindingDataType.ANY def _build_arg_bindings( diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 9366039637a35..fe12557698456 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -30,6 +30,20 @@ API_VERSION: Final[str] = "2026-06-30" +class ArgBindingDataType(str, Enum): + """ + Language-neutral value type a stub-task argument binds to in the foreign runtime. + """ + + STRING = "string" + INTEGER = "integer" + NUMBER = "number" + BOOLEAN = "boolean" + OBJECT = "object" + ARRAY = "array" + ANY = "any" + + class AssetAliasReferenceAssetEventDagRun(BaseModel): """ Schema for AssetAliasModel used in AssetEventDagRunReference. @@ -376,16 +390,6 @@ class Kind(str, Enum): 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. @@ -396,7 +400,7 @@ class TaskArgBinding(BaseModel): name: Annotated[str, Field(title="Name")] kind: Annotated[Kind, Field(title="Kind")] - data_type: Annotated[DataType | None, Field(title="Data Type")] = DataType.ANY + data_type: ArgBindingDataType | None = ArgBindingDataType.ANY task_id: Annotated[str | None, Field(title="Task Id")] = None key: Annotated[str | None, Field(title="Key")] = "return_value" value: JsonValue | None = None diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 065129868185c..4cb56479e80a6 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -3,6 +3,20 @@ "api_version": "2026-07-30", "description": "Apache Airflow SDK Supervisor Schema", "$defs": { + "ArgBindingDataType": { + "description": "Language-neutral value type a stub-task argument binds to in the foreign runtime.", + "enum": [ + "string", + "integer", + "number", + "boolean", + "object", + "array", + "any" + ], + "title": "ArgBindingDataType", + "type": "string" + }, "AssetAliasReferenceAssetEventDagRun": { "additionalProperties": false, "description": "Schema for AssetAliasModel used in AssetEventDagRunReference.", @@ -1395,19 +1409,6 @@ "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.", @@ -4600,18 +4601,8 @@ "type": "string" }, "data_type": { - "default": "any", - "enum": [ - "string", - "integer", - "number", - "boolean", - "object", - "array", - "any" - ], - "title": "Data Type", - "type": "string" + "$ref": "#/$defs/ArgBindingDataType", + "default": "any" }, "task_id": { "anyOf": [ From 854479badf2dd7a5862bd32adceb9f5a1a6e4685 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 20 Jul 2026 01:55:10 +0000 Subject: [PATCH 09/40] Rework the Go SDK TaskInput binding example set The combined TaskInput example mixed all three field-binding modes (arg: tag, no tag, xcom: tag) into one struct, so no single task demonstrated any one mode in isolation. Split it into via_struct_no_tags, via_struct_arg_tag, and via_struct_xcom_tag, and renamed combine to via_flat_args to make the positional/keyword-style split between flat and struct binding legible at the call-site naming level. A TaskInput struct field whose name has no matching TaskFlow call argument now stays at its Go zero value instead of failing the task -- keyword-argument semantics (an unpassed name falls back to its default) rather than the strict arity check flat, positional parameters get. via_struct_unmatched_arg exercises this directly. --- .../test_go_sdk_taskflow_binding.py | 64 +++++-- go-sdk/README.md | 14 +- .../airflow-go-pack/pack_integration_test.go | 7 +- go-sdk/dags/go_examples.py | 75 +++++--- go-sdk/example/bundle/main.go | 7 +- .../bundle/taskflowbinding/taskflowbinding.go | 178 ++++++++++++++---- .../taskflowbinding/taskflowbinding_test.go | 89 +++++++-- go-sdk/pkg/binding/binding.go | 30 ++- go-sdk/pkg/binding/binding_test.go | 27 ++- 9 files changed, 393 insertions(+), 98 deletions(-) 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 index ab087eb7f7bd8..c874b72de26cd 100644 --- 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 @@ -19,9 +19,9 @@ 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. +bound onto ``[]int``). The Go ``via_flat_args`` task verifies every bound +value and errors on any mismatch, so a green run *is* the binding assertion; +the tests here check the run outcome and the summary XCom it pushes. """ from __future__ import annotations @@ -67,12 +67,20 @@ def completed_run() -> _CompletedRun: def test_all_tasks_succeeded(completed_run: _CompletedRun): - """The Go ``combine`` task errors on any mis-bound argument, so success here + """The Go ``via_flat_args`` task errors on any mis-bound argument, so success here proves every literal, XCom, keyword, and defaulted-None binding was correct.""" assert completed_run.state == "success", ( f"expected the run to succeed; got {completed_run.state!r}. task states: {completed_run.ti_states}" ) - for task_id in ("make_config", "make_numbers", "combine", "combine_via_task_input"): + for task_id in ( + "make_config", + "make_numbers", + "via_flat_args", + "via_struct_no_tags", + "via_struct_arg_tag", + "via_struct_xcom_tag", + "via_struct_unmatched_arg", + ): assert completed_run.ti_states.get(task_id) == "success", completed_run.ti_states @@ -86,10 +94,10 @@ def test_upstream_xcoms_keep_their_shapes(completed_run: _CompletedRun): 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 +def test_via_flat_args_summary_reflects_bound_arguments(completed_run: _CompletedRun): + """``via_flat_args`` re-emits every bound value, confirming types survived the Python literal / XCom -> Go parameter -> XCom round trip.""" - assert completed_run.xcom("combine") == { + assert completed_run.xcom("via_flat_args") == { "name": "summary", "count": 3, "ratio": 2.5, @@ -102,14 +110,40 @@ def test_combine_summary_reflects_bound_arguments(completed_run: _CompletedRun): } -def test_combine_via_task_input_summary_reflects_bound_arguments(completed_run: _CompletedRun): - """``combine_via_task_input`` demonstrates the Go SDK's ``sdk.TaskInput`` - struct-field injection mode: ``region_code``/``threshold`` bind by name onto - the struct exactly like ``combine``'s flat parameters do, while the struct's - third field is an ad hoc XCom pull of ``make_config``'s return value declared - purely in Go, with no corresponding TaskFlow call argument here.""" - assert completed_run.xcom("combine_via_task_input") == { +def test_via_struct_no_tags_reflects_bound_arguments(completed_run: _CompletedRun): + """``via_struct_no_tags`` demonstrates the Go SDK's ``sdk.TaskInput`` struct-field + injection mode with no field tags at all: both fields bind by their Go field name + snake_cased.""" + assert completed_run.xcom("via_struct_no_tags") == { + "region_code": "eu-west-1", + "threshold": 0.75, + } + + +def test_via_struct_arg_tag_reflects_bound_arguments(completed_run: _CompletedRun): + """``via_struct_arg_tag`` demonstrates the ``arg:`` tag renaming a struct field + away from its snake_cased default.""" + assert completed_run.xcom("via_struct_arg_tag") == { "region": "eu-west-1", "threshold": 0.75, + } + + +def test_via_struct_xcom_tag_reflects_bound_arguments(completed_run: _CompletedRun): + """``via_struct_xcom_tag`` demonstrates the ``xcom:`` tag: its ``Config`` field is + an ad hoc pull of ``make_config``'s return value, with no corresponding TaskFlow + call argument here.""" + assert completed_run.xcom("via_struct_xcom_tag") == { + "threshold": 0.75, "environment": "production", } + + +def test_via_struct_unmatched_arg_reflects_zero_valued_field(completed_run: _CompletedRun): + """``via_struct_unmatched_arg`` demonstrates that a struct field whose name has + no corresponding TaskFlow call argument stays at its Go zero value instead of + failing the task -- kwarg-style, an unpassed name simply isn't bound.""" + assert completed_run.xcom("via_struct_unmatched_arg") == { + "region": "eu-west-1", + "missing_was_empty": True, + } diff --git a/go-sdk/README.md b/go-sdk/README.md index 6b06cfe744155..e8afd899adb28 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -145,7 +145,14 @@ source of truth for which `dag_id`s and `task_id`s a bundle can run. A struct that anonymously embeds `sdk.TaskInput` opts into **per-field, name-based** binding instead of a long flat parameter list — at most one such parameter is allowed per function, and it can be -mixed with plain flat parameters: +mixed with plain flat parameters. + +Conceptually, a plain flat parameter list is **positional-argument** binding: order matters, and +every parameter must be filled or the task fails before its body runs. A `TaskInput` struct is +closer to **keyword-argument** binding: fields match by name instead of position, and (see the +`arg:` bullet below) a field whose name has no corresponding TaskFlow call argument is simply left +at its Go zero value rather than failing the task — the same way an unpassed keyword argument falls +back to a caller-side default in a kwargs-style call. ```go type CombineInput struct { @@ -167,6 +174,8 @@ Each exported field supports three all-optional tags: - `arg:""` — bind from the TaskFlow call argument with this name (matched against the stub function's Python parameter name, independent of declaration order on either side). With no tag, the field's own Go name, snake_cased (`RatioValue` → `ratio_value`, `TaskID` → `task_id`), is used. + If no TaskFlow call argument carries that name, the field is simply left at its Go zero value — + it does not fail the task, kwarg-style (see `ViaStructUnmatchedArg` below). - `xcom:""` — an explicit, ad hoc XCom pull from the named upstream task, fully independent of the TaskFlow call: the field need not correspond to any argument the Dag file passes at all. This pull is unchecked — the runtime does not verify `` is an actual upstream dependency, the @@ -184,7 +193,8 @@ this — it keeps working as a single flat data parameter, JSON-decoded whole fr argument (see `Config` in [`example/bundle/taskflowbinding/taskflowbinding.go`](./example/bundle/taskflowbinding/taskflowbinding.go)), which is a different mechanism from per-field `TaskInput` binding. See -[`CombineViaTaskInput`](./example/bundle/taskflowbinding/taskflowbinding.go) for a full worked example. +[`ViaStructNoTags`, `ViaStructArgTag`, `ViaStructXComTag`, and `ViaStructUnmatchedArg`](./example/bundle/taskflowbinding/taskflowbinding.go) +for a full worked example of each tag mode — and the unmatched-field case — in isolation. ### Reading the task runtime context 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 09d4143d1793c..3e29c5b5641af 100644 --- a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go +++ b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go @@ -158,8 +158,11 @@ dags: tasks: - "make_config" - "make_numbers" - - "combine" - - "combine_via_task_input" + - "via_flat_args" + - "via_struct_no_tags" + - "via_struct_arg_tag" + - "via_struct_xcom_tag" + - "via_struct_unmatched_arg" ` assert.Equal(t, expectedManifest, string(metadata)) diff --git a/go-sdk/dags/go_examples.py b/go-sdk/dags/go_examples.py index ae91b1512d43c..dab624452e5ef 100644 --- a/go-sdk/dags/go_examples.py +++ b/go-sdk/dags/go_examples.py @@ -20,9 +20,11 @@ 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 -- both the flat parameter -list ``combine`` binds onto and the ``sdk.TaskInput`` struct ``combine_via_task_input`` -binds onto instead, see its Dag function below). +(stressing the TaskFlow argument-binding surface -- the flat, positional +parameter list ``via_flat_args`` binds onto, plus four ``sdk.TaskInput`` +(keyword-style) struct examples, ``via_struct_no_tags``/``via_struct_arg_tag``/ +``via_struct_xcom_tag``/``via_struct_unmatched_arg``, each isolating one +field-binding mode; see its Dag function below). ``simple_dag`` sandwiches the Go tasks between two native Python tasks so the run exercises XCom across the language boundary, the same way @@ -128,7 +130,7 @@ def make_numbers(): ... @task.stub(queue="golang") -def combine( +def via_flat_args( name: str, count: int, ratio: float, @@ -141,7 +143,19 @@ def combine( @task.stub(queue="golang") -def combine_via_task_input(region_code: str, threshold: float): ... +def via_struct_no_tags(region_code: str, threshold: float): ... + + +@task.stub(queue="golang") +def via_struct_arg_tag(region_code: str, threshold: float): ... + + +@task.stub(queue="golang") +def via_struct_xcom_tag(threshold: float): ... + + +@task.stub(queue="golang") +def via_struct_unmatched_arg(region_code: str): ... @dag(dag_id="taskflow_binding_dag") @@ -149,23 +163,39 @@ 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_via_task_input`` demonstrates the Go SDK's ``sdk.TaskInput`` struct - injection mode: ``region_code``/``threshold`` bind by name onto the struct's - fields exactly like ``combine``'s flat parameters do, but the struct's third - field is an ad hoc XCom pull of ``make_config``'s return value declared purely - in Go (an ``xcom:`` struct tag) -- it is never passed as a TaskFlow argument - here, so the explicit ``>>`` below is what orders it after ``make_config``. + Conceptually, the flat parameter list is *positional-argument* binding: order + matters, and every parameter must be filled or the task fails before it runs. + ``sdk.TaskInput`` structs are closer to *keyword-argument* binding: fields match + by name, and (see ``via_struct_unmatched_arg`` below) a field whose name has no + corresponding TaskFlow call argument simply stays at its zero value instead of + failing the task -- the same way an unpassed keyword argument falls back to a + default in kwargs-style calls. + + ``via_flat_args``'s one mixed positional/keyword call carries literals of every + scalar type plus an array literal, and fans in XComs from *two* upstream Go + tasks: ``make_config`` returns an object that binds onto a strictly-decoded Go + struct, ``make_numbers`` an array that binds onto ``[]int``. ``note`` is not + passed, so its ``None`` default is captured and arrives in Go as a nil + ``*string``. The Go ``via_flat_args`` (``go-sdk/example/bundle/taskflowbinding``) + verifies every bound value and fails the task on any mismatch. + + Four further tasks demonstrate the Go SDK's ``sdk.TaskInput`` struct injection + mode, one field-binding mode at a time: + + * ``via_struct_no_tags``: both struct fields fall back to their Go field name + snake_cased -- no ``arg:``/``xcom:`` tags at all. + * ``via_struct_arg_tag``: one field is renamed via an explicit ``arg:`` tag, + proving the tag remaps the name rather than coincidentally matching it. + * ``via_struct_xcom_tag``: one field is an ad hoc XCom pull of ``make_config``'s + return value declared purely in Go (an ``xcom:`` struct tag) -- it is never + passed as a TaskFlow argument here, so the explicit ``>>`` below is what + orders it after ``make_config``. + * ``via_struct_unmatched_arg``: the Go struct declares a field with no + corresponding argument in this TaskFlow call at all -- it stays at its Go + zero value rather than failing the task. """ config = make_config() - combine( + via_flat_args( "summary", 3, 2.5, @@ -174,7 +204,10 @@ def taskflow_binding_dag(): config=config, numbers=make_numbers(), ) - config >> combine_via_task_input(region_code="eu-west-1", threshold=0.75) + via_struct_no_tags(region_code="eu-west-1", threshold=0.75) + via_struct_arg_tag(region_code="eu-west-1", threshold=0.75) + config >> via_struct_xcom_tag(threshold=0.75) + via_struct_unmatched_arg(region_code="eu-west-1") taskflow_binding_dag() diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 581044ce3aca4..7f020f8c2c429 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -59,8 +59,11 @@ func (m *myBundle) RegisterDags(dagbag v1.Registry) error { bindingDag := dagbag.AddDag("taskflow_binding_dag") bindingDag.AddTaskWithName("make_config", taskflowbinding.MakeConfig) bindingDag.AddTaskWithName("make_numbers", taskflowbinding.MakeNumbers) - bindingDag.AddTaskWithName("combine", taskflowbinding.Combine) - bindingDag.AddTaskWithName("combine_via_task_input", taskflowbinding.CombineViaTaskInput) + bindingDag.AddTaskWithName("via_flat_args", taskflowbinding.ViaFlatArgs) + bindingDag.AddTaskWithName("via_struct_no_tags", taskflowbinding.ViaStructNoTags) + bindingDag.AddTaskWithName("via_struct_arg_tag", taskflowbinding.ViaStructArgTag) + bindingDag.AddTaskWithName("via_struct_xcom_tag", taskflowbinding.ViaStructXComTag) + bindingDag.AddTaskWithName("via_struct_unmatched_arg", taskflowbinding.ViaStructUnmatchedArg) return nil } diff --git a/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go b/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go index 07fd12c11a527..475ab8da8ceb5 100644 --- a/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go +++ b/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go @@ -15,14 +15,19 @@ // 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 +// Package taskflowbinding holds the taskflow_binding_dag tasks. ViaFlatArgs is +// positional-argument binding pushed to its limit: literals of every scalar // type, an array literal, keyword arguments, a defaulted null, and XCom fan-in -// from two upstream Go tasks decoded into a strict struct and a typed slice. -// CombineViaTaskInput additionally shows the sdk.TaskInput struct-field -// injection mode: the same binding surface collapsed into one struct -// parameter instead of a long flat list. +// from two upstream Go tasks decoded into a strict struct and a typed slice -- +// where simple_dag's transform shows the minimal case (one literal, one +// XCom), this shows the full argument surface. The ViaStruct* functions +// instead show the sdk.TaskInput struct-field injection mode -- conceptually +// keyword-argument binding, where fields match by name and an unmatched name +// is left at its zero value rather than failing the task -- one field-binding +// tag at a time: ViaStructNoTags (plain snake_case name fallback), +// ViaStructArgTag (an explicit `arg:` rename), ViaStructXComTag (an ad hoc +// `xcom:` pull), and ViaStructUnmatchedArg (a field whose name has no +// corresponding TaskFlow call argument at all). package taskflowbinding import ( @@ -33,8 +38,8 @@ import ( "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 +// Config is the object make_config returns as its XCom; via_flat_args declares +// the same struct as a parameter, so the round trip exercises strict struct // decoding (an unknown or renamed key fails the task rather than silently // zeroing a field). type Config struct { @@ -43,7 +48,7 @@ type Config struct { Debug bool `json:"debug"` } -// MakeConfig pushes an object XCom that combine binds onto its Config parameter. +// MakeConfig pushes an object XCom that via_flat_args binds onto its Config parameter. func MakeConfig(log *slog.Logger) (any, error) { cfg := Config{Environment: "production", Region: "eu-west-1", Debug: true} log.Info( @@ -58,23 +63,23 @@ func MakeConfig(log *slog.Logger) (any, error) { return cfg, nil } -// MakeNumbers pushes an array XCom that combine binds onto its []int parameter. +// MakeNumbers pushes an array XCom that via_flat_args binds onto its []int parameter. func MakeNumbers(log *slog.Logger) (any, error) { numbers := []int{1, 1, 2, 3, 5, 8} log.Info("Pushing numbers", "numbers", fmt.Sprint(numbers)) return numbers, nil } -// Combine receives every argument shape the stub Dag can express. The Python -// side calls it as +// ViaFlatArgs receives every argument shape the stub Dag can express as plain, +// positional data parameters. The Python side calls it as // -// combine("summary", 3, 2.5, True, ["metrics", "hourly"], -// config=make_config(), numbers=make_numbers()) +// via_flat_args("summary", 3, 2.5, True, ["metrics", "hourly"], +// config=make_config(), numbers=make_numbers()) // // so the bound values are fixed; any mismatch below is a binding regression // and fails the task loudly. note is never passed and falls back to the stub's // None default, arriving as a nil *string. -func Combine( +func ViaFlatArgs( ctx sdk.TIRunContext, log *slog.Logger, name string, @@ -131,26 +136,60 @@ func Combine( }, nil } -// CombineInput demonstrates the sdk.TaskInput struct-field injection mode: an -// ergonomic alternative to Combine's long flat parameter list. Region binds -// by an explicit arg: tag; Threshold has no tag, so it falls back to its Go -// field name snake_cased ("threshold"); Config is an ad hoc XCom pull of -// make_config's return value, independent of the Python call's TaskFlow -// arguments entirely. -type CombineInput struct { +// ViaStructNoTagsInput demonstrates the sdk.TaskInput struct-field injection +// mode with no field tags at all: both fields fall back to their Go field +// name snake_cased ("region_code", "threshold"). +type ViaStructNoTagsInput struct { + sdk.TaskInput + RegionCode string + Threshold float64 +} + +// ViaStructNoTags is called as +// +// via_struct_no_tags(region_code="eu-west-1", threshold=0.75) +func ViaStructNoTags( + ctx sdk.TIRunContext, + log *slog.Logger, + input ViaStructNoTagsInput, +) (any, error) { + if input.RegionCode != "eu-west-1" || input.Threshold != 0.75 { + return nil, fmt.Errorf( + "TaskInput fields bound incorrectly: region_code=%q threshold=%v", + input.RegionCode, + input.Threshold, + ) + } + + log.InfoContext(ctx, "Bound TaskInput struct (no tags)", + "region_code", input.RegionCode, + "threshold", input.Threshold, + ) + return map[string]any{ + "region_code": input.RegionCode, + "threshold": input.Threshold, + }, nil +} + +// ViaStructArgTagInput demonstrates the sdk.TaskInput struct-field injection +// mode with an explicit arg: tag: Region binds to the "region_code" TaskFlow +// argument under a renamed Go field, proving the tag remaps the name rather +// than coincidentally matching it; Threshold has no tag and falls back to its +// snake_cased field name. +type ViaStructArgTagInput struct { sdk.TaskInput Region string `arg:"region_code"` Threshold float64 - Config Config ` xcom:"make_config"` } -// CombineViaTaskInput is the TaskInput-struct sibling of Combine: the same -// kind of binding surface -- a named literal, a snake_case-fallback literal, -// and an ad hoc XCom pull -- collapsed into one struct parameter instead of -// many flat ones. The Python side calls it as +// ViaStructArgTag is called as // -// combine_via_task_input(region_code="eu-west-1", threshold=0.75) -func CombineViaTaskInput(ctx sdk.TIRunContext, log *slog.Logger, input CombineInput) (any, error) { +// via_struct_arg_tag(region_code="eu-west-1", threshold=0.75) +func ViaStructArgTag( + ctx sdk.TIRunContext, + log *slog.Logger, + input ViaStructArgTagInput, +) (any, error) { if input.Region != "eu-west-1" || input.Threshold != 0.75 { return nil, fmt.Errorf( "TaskInput fields bound incorrectly: region=%q threshold=%v", @@ -158,20 +197,93 @@ func CombineViaTaskInput(ctx sdk.TIRunContext, log *slog.Logger, input CombineIn input.Threshold, ) } + + log.InfoContext(ctx, "Bound TaskInput struct (arg: tag)", + "region", input.Region, + "threshold", input.Threshold, + ) + return map[string]any{ + "region": input.Region, + "threshold": input.Threshold, + }, nil +} + +// ViaStructXComTagInput demonstrates the sdk.TaskInput struct-field injection +// mode with an xcom: tag: Config is an ad hoc pull of make_config's return +// value, independent of the Python call's TaskFlow arguments entirely; +// Threshold has no tag and falls back to its snake_cased field name. +type ViaStructXComTagInput struct { + sdk.TaskInput + Threshold float64 + Config Config `xcom:"make_config"` +} + +// ViaStructXComTag is called as +// +// via_struct_xcom_tag(threshold=0.75) +// +// with the Python Dag ordering it after make_config explicitly (the xcom: +// pull is not a TaskFlow argument, so there is no implicit dependency). +func ViaStructXComTag( + ctx sdk.TIRunContext, + log *slog.Logger, + input ViaStructXComTagInput, +) (any, error) { + if input.Threshold != 0.75 { + return nil, fmt.Errorf("TaskInput field bound incorrectly: threshold=%v", input.Threshold) + } if want := (Config{Environment: "production", Region: "eu-west-1", Debug: true}); input.Config != want { return nil, fmt.Errorf( "ad hoc xcom field bound incorrectly: config=%+v, want %+v", input.Config, want, ) } - log.InfoContext(ctx, "Bound TaskInput struct", - "region", input.Region, + log.InfoContext(ctx, "Bound TaskInput struct (xcom: tag)", "threshold", input.Threshold, "environment", input.Config.Environment, ) return map[string]any{ - "region": input.Region, "threshold": input.Threshold, "environment": input.Config.Environment, }, nil } + +// ViaStructUnmatchedArgInput demonstrates that a field whose name has no +// corresponding TaskFlow call argument at all is left at its Go zero value +// rather than failing the task: Region binds normally, but Missing's arg +// name is never among this call's arguments -- conceptually, an unpassed +// keyword argument falling back to its default in a kwargs-style call. +type ViaStructUnmatchedArgInput struct { + sdk.TaskInput + Region string `arg:"region_code"` + Missing string `arg:"does_not_exist"` +} + +// ViaStructUnmatchedArg is called as +// +// via_struct_unmatched_arg(region_code="eu-west-1") +// +// -- the stub only declares region_code, so Missing's arg name never appears +// among the call's arguments and stays at its Go zero value (""). +func ViaStructUnmatchedArg( + ctx sdk.TIRunContext, log *slog.Logger, input ViaStructUnmatchedArgInput, +) (any, error) { + if input.Region != "eu-west-1" { + return nil, fmt.Errorf("TaskInput field bound incorrectly: region=%q", input.Region) + } + if input.Missing != "" { + return nil, fmt.Errorf( + "expected the unmatched field to stay at its Go zero value, got missing=%q", + input.Missing, + ) + } + + log.InfoContext(ctx, "Bound TaskInput struct (unmatched arg)", + "region", input.Region, + "missing_was_empty", input.Missing == "", + ) + return map[string]any{ + "region": input.Region, + "missing_was_empty": input.Missing == "", + }, nil +} diff --git a/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go b/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go index 298ec8508c9aa..6bd28dd99459a 100644 --- a/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go +++ b/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go @@ -30,9 +30,9 @@ import ( // 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) { +func TestViaFlatArgs(t *testing.T) { ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) - got, err := Combine(ctx, slog.Default(), + got, err := ViaFlatArgs(ctx, slog.Default(), "summary", 3, 2.5, true, []string{"metrics", "hourly"}, Config{Environment: "production", Region: "eu-west-1", Debug: true}, @@ -42,14 +42,14 @@ func TestCombine(t *testing.T) { require.NoError(t, err) summary, ok := got.(map[string]any) - require.True(t, ok, "Combine should return a map summary, got %T", got) + require.True(t, ok, "ViaFlatArgs should return a map summary, got %T", got) assert.Equal(t, 20, summary["sum"]) assert.Equal(t, true, summary["note_was_null"]) } -func TestCombineRejectsWrongBinding(t *testing.T) { +func TestViaFlatArgsRejectsWrongBinding(t *testing.T) { ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) - _, err := Combine(ctx, slog.Default(), + _, err := ViaFlatArgs(ctx, slog.Default(), "summary", 3, 2.5, true, []string{"metrics", "hourly"}, Config{}, @@ -59,26 +59,93 @@ func TestCombineRejectsWrongBinding(t *testing.T) { assert.ErrorContains(t, err, "object XCom bound incorrectly") } -func TestCombineViaTaskInput(t *testing.T) { +func TestViaStructNoTags(t *testing.T) { ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) - got, err := CombineViaTaskInput(ctx, slog.Default(), CombineInput{ + got, err := ViaStructNoTags(ctx, slog.Default(), ViaStructNoTagsInput{ + RegionCode: "eu-west-1", + Threshold: 0.75, + }) + require.NoError(t, err) + + summary, ok := got.(map[string]any) + require.True(t, ok, "ViaStructNoTags should return a map summary, got %T", got) + assert.Equal(t, "eu-west-1", summary["region_code"]) +} + +func TestViaStructNoTagsRejectsWrongBinding(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + _, err := ViaStructNoTags(ctx, slog.Default(), ViaStructNoTagsInput{ + RegionCode: "wrong-region", + Threshold: 0.75, + }) + assert.ErrorContains(t, err, "TaskInput fields bound incorrectly") +} + +func TestViaStructArgTag(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + got, err := ViaStructArgTag(ctx, slog.Default(), ViaStructArgTagInput{ Region: "eu-west-1", Threshold: 0.75, + }) + require.NoError(t, err) + + summary, ok := got.(map[string]any) + require.True(t, ok, "ViaStructArgTag should return a map summary, got %T", got) + assert.Equal(t, "eu-west-1", summary["region"]) +} + +func TestViaStructArgTagRejectsWrongBinding(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + _, err := ViaStructArgTag(ctx, slog.Default(), ViaStructArgTagInput{ + Region: "wrong-region", + Threshold: 0.75, + }) + assert.ErrorContains(t, err, "TaskInput fields bound incorrectly") +} + +func TestViaStructXComTag(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + got, err := ViaStructXComTag(ctx, slog.Default(), ViaStructXComTagInput{ + Threshold: 0.75, Config: Config{Environment: "production", Region: "eu-west-1", Debug: true}, }) require.NoError(t, err) summary, ok := got.(map[string]any) - require.True(t, ok, "CombineViaTaskInput should return a map summary, got %T", got) + require.True(t, ok, "ViaStructXComTag should return a map summary, got %T", got) assert.Equal(t, "production", summary["environment"]) } -func TestCombineViaTaskInputRejectsWrongBinding(t *testing.T) { +func TestViaStructXComTagRejectsWrongBinding(t *testing.T) { ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) - _, err := CombineViaTaskInput(ctx, slog.Default(), CombineInput{ - Region: "eu-west-1", + _, err := ViaStructXComTag(ctx, slog.Default(), ViaStructXComTagInput{ Threshold: 0.75, Config: Config{}, }) assert.ErrorContains(t, err, "ad hoc xcom field bound incorrectly") } + +func TestViaStructUnmatchedArg(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + // Missing is left at its Go zero value, exactly as binding.Resolve leaves an + // unmatched TaskInput field -- this task fn is unit-testable independent of + // the binding package precisely because it declares that expectation itself. + got, err := ViaStructUnmatchedArg(ctx, slog.Default(), ViaStructUnmatchedArgInput{ + Region: "eu-west-1", + Missing: "", + }) + require.NoError(t, err) + + summary, ok := got.(map[string]any) + require.True(t, ok, "ViaStructUnmatchedArg should return a map summary, got %T", got) + assert.Equal(t, true, summary["missing_was_empty"]) +} + +func TestViaStructUnmatchedArgRejectsNonZeroMissingField(t *testing.T) { + ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) + _, err := ViaStructUnmatchedArg(ctx, slog.Default(), ViaStructUnmatchedArgInput{ + Region: "eu-west-1", + Missing: "unexpected", + }) + assert.ErrorContains(t, err, "expected the unmatched field to stay at its Go zero value") +} diff --git a/go-sdk/pkg/binding/binding.go b/go-sdk/pkg/binding/binding.go index e5fe6af725d8b..73343842bd533 100644 --- a/go-sdk/pkg/binding/binding.go +++ b/go-sdk/pkg/binding/binding.go @@ -46,11 +46,20 @@ // sits, and (with no TaskInput struct present) this reduces to exactly // today's positional-only behaviour. // +// Conceptually, flat data parameters are positional-argument binding: order +// matters, and every parameter must be filled or Resolve fails the task +// before its body runs. A TaskInput struct is closer to keyword-argument +// binding: fields match by name, and a field whose name has no corresponding +// TaskFlow call argument is simply left at its Go zero value instead of +// failing the task -- the same way an unpassed keyword argument falls back +// to its default in a kwargs-style call. +// // 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. +// per-execution argument spec, failing loudly on arity or type mismatches of +// flat data parameters. 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). @@ -326,7 +335,9 @@ func (p *Plan) resolveData( // resolveTaskInput builds the struct value for one TaskInput parameter. A // field tagged `xcom:` pulls directly and never touches byName/claimed; a // field claiming an argument spec entry by name marks it claimed so the -// later flat-parameter cursor skips it. +// later flat-parameter cursor skips it. A field whose name claims nothing +// (kwarg-style: it was never "passed") is left unset at its Go zero value +// rather than failing the task. func (p *Plan) resolveTaskInput( ctx context.Context, client sdk.Client, @@ -360,12 +371,11 @@ func (p *Plan) resolveTaskInput( case taskInputFieldFromArg: idx, ok := byName[tif.argName] if !ok { - return reflect.Value{}, fmt.Errorf( - "task function %s: %s: arg name %q not found among the Dag's TaskFlow call arguments", - p.fnName, - typeCheckCtx, - tif.argName, - ) + // No TaskFlow call argument carries this name -- kwarg-style, an + // unpassed name leaves the field at its Go zero value rather than + // failing the task (unlike a flat data parameter, where arity is + // checked strictly; see the package doc comment). + continue } claimed[idx] = true arg = args[idx] diff --git a/go-sdk/pkg/binding/binding_test.go b/go-sdk/pkg/binding/binding_test.go index d7c573f683c2e..a0a070bd56e02 100644 --- a/go-sdk/pkg/binding/binding_test.go +++ b/go-sdk/pkg/binding/binding_test.go @@ -294,6 +294,15 @@ type simpleTaskInput struct { Name string } +// twoFieldTaskInput has one field the args always match (Name) and one whose +// arg name is never present in the tests that use it (Missing), to prove an +// unmatched field is left at its Go zero value instead of failing the task. +type twoFieldTaskInput struct { + sdk.TaskInput + Name string + Missing string `arg:"missing"` +} + // nonEmbeddingStruct has no sdk.TaskInput sentinel, so it must keep resolving // as today's whole-value decode target, not per-field TaskInput binding. type nonEmbeddingStruct struct { @@ -587,16 +596,30 @@ func (s *BindingSuite) TestResolveTaskInputPointerStruct() { s.Equal("widget", input.Name) } -func (s *BindingSuite) TestResolveTaskInputUnmatchedArgName() { +func (s *BindingSuite) TestResolveTaskInputUnmatchedArgNameLeavesFieldZeroValued() { fn := func(input simpleTaskInput) error { return nil } _, err := s.resolve(fn, []Arg{ {Name: "different_name", Kind: ArgKindLiteral, Value: "x", DataType: DataTypeString}, }, &fakeXComClient{}) + // simpleTaskInput has no other flat parameter to absorb "different_name", + // so the arity check (0 unclaimed args expected) still fails the task -- + // this asserts the unmatched TaskInput field itself is not what errors. if s.Assert().Error(err) { - s.Contains(err.Error(), `arg name "name" not found`) + s.Contains(err.Error(), "argument count mismatch") } } +func (s *BindingSuite) TestResolveTaskInputUnmatchedArgNameZeroValuedAlongsideMatch() { + fn := func(input twoFieldTaskInput) error { return nil } + got, err := s.resolve(fn, []Arg{ + {Name: "name", Kind: ArgKindLiteral, Value: "widget", DataType: DataTypeString}, + }, &fakeXComClient{}) + s.Require().NoError(err) + input := got[0].Interface().(twoFieldTaskInput) + s.Equal("widget", input.Name, "the matched field binds normally") + s.Equal("", input.Missing, "the unmatched field is left at its Go zero value, not an error") +} + func (s *BindingSuite) TestResolveTaskInputArityMismatchForLeftoverArgs() { fn := func(input simpleTaskInput, extra string) error { return nil } _, err := s.resolve(fn, []Arg{ From 9f96f509fabde7081a7fa7ceae42a7ac9d348282 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 20 Jul 2026 15:59:37 +0000 Subject: [PATCH 10/40] Split TaskArgBinding into XComArgBinding and LiteralArgBinding The flat kind-discriminated shape kept foreign-language codegen simple but left each variant's contract implicit: task_id was nullable even though every xcom binding has one, and value/key were dead weight on the opposite kind. Modelling arg_bindings as a kind-discriminated union makes the contracts explicit on every wire (OpenAPI, supervisor schema, Go, TypeScript) - xcom bindings now require task_id - and lets the Go runtime mirror the split as a sealed sum type instead of branching on a string field, so malformed specs fail the task before its body runs. --- .../datamodels/task_arg_binding.py | 58 +++++++--- .../execution_api/routes/task_instances.py | 8 +- .../execution_api/versions/v2026_06_30.py | 11 +- .../versions/head/test_task_instances.py | 8 ++ go-sdk/bundle/bundlev1/task_test.go | 9 +- go-sdk/pkg/binding/binding.go | 95 +++++++++------ go-sdk/pkg/binding/binding_test.go | 83 ++++++++------ .../pkg/execution/genmodels/defaults.gen.go | 33 ++++-- go-sdk/pkg/execution/genmodels/models.gen.go | 73 ++++++------ go-sdk/pkg/execution/integration_test.go | 105 ++++++++++++++++- go-sdk/pkg/execution/task_runner.go | 56 ++++++--- .../providers/standard/decorators/stub.py | 12 +- .../airflow/sdk/api/datamodels/_generated.py | 48 ++++---- .../sdk/execution_time/schema/schema.json | 94 +++++++++------ .../schema/versions/v2026_07_30.py | 12 +- .../execution_time/schema/test_migrator.py | 11 +- ts-sdk/src/generated/supervisor.ts | 108 ++++++++++++------ 17 files changed, 563 insertions(+), 261 deletions(-) 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 index 171801f0837ad..87aec3b829db5 100644 --- 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 @@ -20,14 +20,21 @@ 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. + +Each binding is one variant of a union discriminated on ``kind``: an ``XComArgBinding`` +pulls the value from an upstream task's XCom, a ``LiteralArgBinding`` carries an inline +value from the Dag file. Both variants still emit plain named structs +(``$defs/XComArgBinding``, ``$defs/LiteralArgBinding``) for the foreign-language SDKs +consuming the supervisor schema. """ from __future__ import annotations from enum import Enum -from typing import Literal +from typing import Annotated, Literal -from pydantic import JsonValue +from pydantic import Field, JsonValue +from typing_extensions import TypeAliasType from airflow.api_fastapi.core_api.base import BaseModel @@ -44,28 +51,49 @@ class ArgBindingDataType(str, Enum): ANY = "any" -class TaskArgBinding(BaseModel): - """ - One positional argument of a stub (foreign-runtime) task, in declaration order. +class XComArgBinding(BaseModel): + """One positional stub-task argument pulled from an upstream task's XCom.""" - 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. - """ + # No default on purpose: a required ``kind`` stays non-nullable through the OpenAPI + # round trip, which discriminated-union codegen needs (a defaulted field turns + # ``Literal`` into ``Literal | None`` in the generated task-sdk models). + kind: Literal["xcom"] name: str """The stub function's parameter name this binding fills, in declaration order.""" - 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 = 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``.""" + task_id: str + """Upstream task id to pull the XCom from.""" key: str = "return_value" - """XCom key to pull. Only meaningful when ``kind`` is ``xcom``.""" + """XCom key to pull.""" + + +class LiteralArgBinding(BaseModel): + """One positional stub-task argument carrying an inline literal from the Dag file.""" + + kind: Literal["literal"] + """Required like ``XComArgBinding.kind``; see the note there.""" + + name: str + """The stub function's parameter name this binding fills, in declaration order.""" + + data_type: ArgBindingDataType = ArgBindingDataType.ANY + """Declared type from the stub function's annotation; runtimes type-check against it.""" value: JsonValue | None = None - """The literal value from the Dag file. Only set when ``kind`` is ``literal``.""" + """The literal value from the Dag file.""" + + +# A named alias (TypeAliasType, not a bare Annotated) so the union lands in every +# schema as its own named definition instead of an anonymous field-title-derived one. +# The explicit title lets the supervisor-schema dump merge this def with the +# task-sdk-generated twin (its core/SDK dedup keys on titles). +TaskArgBinding = TypeAliasType( + "TaskArgBinding", + Annotated[XComArgBinding | LiteralArgBinding, Field(discriminator="kind", title="TaskArgBinding")], +) +"""One positional argument of a stub (foreign-runtime) task, in declaration order.""" 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 e3061e68b28c4..7ddbdf6dcd5ce 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -32,7 +32,7 @@ from opentelemetry import trace from opentelemetry.trace import StatusCode from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter from sqlalchemy import and_, func, or_, tuple_, update from sqlalchemy.engine import CursorResult from sqlalchemy.exc import DataError, NoResultFound, SQLAlchemyError @@ -116,6 +116,10 @@ # serialized-dag lookup for ``arg_bindings`` so regular tasks never pay for it. _STUB_TASK_TYPE = "_StubOperator" +# Validates the serialized-dag arg-binding dicts into the kind-discriminated +# TaskArgBinding union; built once at import, not per request. +_arg_bindings_adapter: TypeAdapter[list[TaskArgBinding]] = TypeAdapter(list[TaskArgBinding]) + 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.""" @@ -348,7 +352,7 @@ def ti_run( 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] + context.arg_bindings = _arg_bindings_adapter.validate_python(arg_bindings) # Only set if they are non-null if ti.next_method: 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 0deb96d0940f2..b9e5a7b8206d9 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 @@ -25,7 +25,10 @@ schema, ) -from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import TaskArgBinding +from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import ( + LiteralArgBinding, + XComArgBinding, +) from airflow.api_fastapi.execution_api.datamodels.taskinstance import ( DagRun, TaskInstance, @@ -148,7 +151,8 @@ class AddArgBindingsToTIRunContext(VersionChange): """ Add the ``arg_bindings`` positional-argument binding spec for stub (foreign-runtime) tasks. - ``TaskArgBinding.data_type`` is declared as the ``ArgBindingDataType`` enum rather than an + Each entry is a discriminated union of ``XComArgBinding`` and ``LiteralArgBinding`` keyed + on ``kind``. ``data_type`` is declared as the ``ArgBindingDataType`` enum rather than an inline ``Literal``; the wire representation (a JSON string) is unchanged, so no migration instruction is needed -- this version has not been released with the ``arg_bindings`` field in any other shape. @@ -158,7 +162,8 @@ class AddArgBindingsToTIRunContext(VersionChange): instructions_to_migrate_to_previous_version = ( schema(TIRunContext).field("arg_bindings").didnt_exist, - schema(TaskArgBinding).field("name").didnt_exist, + schema(XComArgBinding).field("name").didnt_exist, + schema(LiteralArgBinding).field("name").didnt_exist, ) @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type] 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 2abac5bd68172..cdff621c7d2aa 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 @@ -31,6 +31,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import StatusCode from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator +from pydantic import ValidationError from sqlalchemy import select, update from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import Session @@ -414,6 +415,13 @@ def transform(country: str, extracted: dict): ... assert response.status_code == 200 assert "arg_bindings" not in response.json() + def test_arg_bindings_adapter_rejects_unknown_kind(self): + """The discriminated union refuses serialized specs with an unrecognised kind.""" + from airflow.api_fastapi.execution_api.routes.task_instances import _arg_bindings_adapter + + with pytest.raises(ValidationError, match="does not match any of the expected tags"): + _arg_bindings_adapter.validate_python([{"name": "country", "kind": "template", "value": "x"}]) + 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/go-sdk/bundle/bundlev1/task_test.go b/go-sdk/bundle/bundlev1/task_test.go index 1535fe5eb097c..6136529a0806d 100644 --- a/go-sdk/bundle/bundlev1/task_test.go +++ b/go-sdk/bundle/bundlev1/task_test.go @@ -206,9 +206,8 @@ func (s *TaskSuite) TestExecuteArgsBindsDataParameters() { 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, + binding.LiteralArg{Value: "uk", DataType: binding.DataTypeString}, + binding.LiteralArg{ Value: map[string]any{"k": "v"}, DataType: binding.DataTypeObject, }, @@ -242,8 +241,8 @@ func (s *TaskSuite) TestExecuteArgsArityMismatch() { context.Background(), slog.New(logging.NewTeeLogger()), []binding.Arg{ - {Kind: binding.ArgKindLiteral, Value: "uk"}, - {Kind: binding.ArgKindLiteral, Value: "de"}, + binding.LiteralArg{Value: "uk"}, + binding.LiteralArg{Value: "de"}, }, ) if s.Assert().Error(err) { diff --git a/go-sdk/pkg/binding/binding.go b/go-sdk/pkg/binding/binding.go index 73343842bd533..1db9b0651399d 100644 --- a/go-sdk/pkg/binding/binding.go +++ b/go-sdk/pkg/binding/binding.go @@ -80,14 +80,6 @@ import ( "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 @@ -103,26 +95,53 @@ const ( ) // 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 { - // Name is the stub function's parameter name this binding fills. Always - // populated; used to match a TaskInput struct field's `arg:` tag (or its - // snake_cased field-name fallback). +// declaration order: an XComArg or a LiteralArg. A sealed sum type mirroring +// the wire model's XComArgBinding/LiteralArgBinding split, kept runtime-neutral +// so this package stays decoupled from the generated coordinator schema types. +type Arg interface { + // ArgName is the stub function's parameter name this binding fills; used + // to match a TaskInput struct field's `arg:` tag (or its snake_cased + // field-name fallback). + ArgName() string + // DeclaredType is the Dag-declared language-neutral type for the argument; + // empty is treated as DataTypeAny. + DeclaredType() DataType + // sealedArg restricts implementations to this package, keeping + // resolveOne's type switch the single exhaustive consumer. + sealedArg() +} + +// XComArg sources the argument from an upstream task's XCom. +type XComArg struct { + // Name is the stub function's parameter name this binding fills. Name string - Kind ArgKind - // TaskID is the upstream task to pull from. Set only for ArgKindXCom. + // TaskID is the upstream task to pull from. TaskID string - // Key is the XCom key to pull; empty means the return-value key. Set only - // for ArgKindXCom. + // Key is the XCom key to pull; empty means the return-value key. Key string - // Value is the literal value from the Dag file. Set only for ArgKindLiteral. + // DataType is the declared type to check the Go parameter against. + DataType DataType +} + +// LiteralArg carries an inline value from the Dag file. +type LiteralArg struct { + // Name is the stub function's parameter name this binding fills. + Name string + // Value is the literal value from the Dag file. Value any - // DataType is the declared type to check the Go parameter against; empty - // is treated as DataTypeAny. + // DataType is the declared type to check the Go parameter against. DataType DataType } +func (a XComArg) ArgName() string { return a.Name } +func (a LiteralArg) ArgName() string { return a.Name } + +func (a XComArg) DeclaredType() DataType { return a.DataType } +func (a LiteralArg) DeclaredType() DataType { return a.DataType } + +func (XComArg) sealedArg() {} +func (LiteralArg) sealedArg() {} + // paramKind classifies how a task-function parameter is filled at execution. type paramKind int @@ -240,7 +259,11 @@ func (p *Plan) Resolve( ) ([]reflect.Value, error) { byName := make(map[string]int, len(args)) for i, a := range args { - byName[a.Name] = i + // A nil entry can claim no name; it stays unclaimed and fails loudly in + // resolveOne when the flat-parameter cursor reaches it. + if a != nil { + byName[a.ArgName()] = i + } } claimed := make([]bool, len(args)) @@ -362,8 +385,7 @@ func (p *Plan) resolveTaskInput( case taskInputFieldFromXCom: // DataTypeAny: an ad hoc pull has no Dag-declared type to check // against, so decoding is driven entirely by the Go field's type. - arg = Arg{ - Kind: ArgKindXCom, + arg = XComArg{ TaskID: tif.xcomTaskID, Key: tif.xcomKey, DataType: DataTypeAny, @@ -410,12 +432,17 @@ func (p *Plan) resolveOne( typeCheckCtx string, generalCtx string, ) (reflect.Value, error) { - if err := checkDataType(arg.DataType, targetType); err != nil { + if arg == nil { + return reflect.Value{}, fmt.Errorf( + "task function %s: %s: nil argument binding", p.fnName, generalCtx, + ) + } + if err := checkDataType(arg.DeclaredType(), targetType); err != nil { return reflect.Value{}, fmt.Errorf("task function %s: %s: %w", p.fnName, typeCheckCtx, err) } - switch arg.Kind { - case ArgKindLiteral: - v, err := decodeValue(arg.Value, targetType) + switch a := arg.(type) { + case LiteralArg: + v, err := decodeValue(a.Value, targetType) if err != nil { return reflect.Value{}, fmt.Errorf( "task function %s: %s: decoding literal value into %s: %w", @@ -423,7 +450,7 @@ func (p *Plan) resolveOne( ) } return v, nil - case ArgKindXCom: + case XComArg: workload, ok := ctx.Value(sdkcontext.WorkloadContextKey).(api.ExecuteTaskWorkload) if !ok { return reflect.Value{}, fmt.Errorf( @@ -431,30 +458,30 @@ func (p *Plan) resolveOne( p.fnName, generalCtx, ) } - key := arg.Key + key := a.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) + raw, err := c.GetXCom(ctx, workload.TI.DagId, workload.TI.RunId, a.TaskID, nil, key, nil) if err != nil { return reflect.Value{}, fmt.Errorf( "task function %s: %s: pulling xcom from task %q (key %q): %w", - p.fnName, generalCtx, arg.TaskID, key, err, + p.fnName, generalCtx, a.TaskID, key, err, ) } v, err := decodeValue(raw, targetType) if err != nil { return reflect.Value{}, fmt.Errorf( "task function %s: %s: decoding xcom from task %q into %s: %w", - p.fnName, generalCtx, arg.TaskID, targetType, err, + p.fnName, generalCtx, a.TaskID, targetType, err, ) } return v, nil default: return reflect.Value{}, fmt.Errorf( - "task function %s: %s: unknown argument kind %q", p.fnName, generalCtx, arg.Kind, + "task function %s: %s: unsupported argument binding %T", p.fnName, generalCtx, arg, ) } } diff --git a/go-sdk/pkg/binding/binding_test.go b/go-sdk/pkg/binding/binding_test.go index a0a070bd56e02..f73b32074ec19 100644 --- a/go-sdk/pkg/binding/binding_test.go +++ b/go-sdk/pkg/binding/binding_test.go @@ -174,7 +174,7 @@ func (s *BindingSuite) TestResolveArityMismatch() { _, err = s.resolve( func() error { return nil }, - []Arg{{Kind: ArgKindLiteral, Value: "uk"}}, + []Arg{LiteralArg{Value: "uk"}}, &fakeXComClient{}, ) if s.Assert().Error(err) { @@ -187,12 +187,12 @@ func (s *BindingSuite) TestResolveLiterals() { 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}, + LiteralArg{Value: "uk", DataType: DataTypeString}, + LiteralArg{Value: 3, DataType: DataTypeInteger}, + LiteralArg{Value: 1.5, DataType: DataTypeNumber}, + LiteralArg{Value: true, DataType: DataTypeBoolean}, + LiteralArg{Value: []any{"a", "b"}, DataType: DataTypeArray}, + LiteralArg{Value: map[string]any{"k": "v"}, DataType: DataTypeObject}, }, &fakeXComClient{}) s.Require().NoError(err) s.Equal("uk", got[0].Interface()) @@ -208,8 +208,8 @@ func (s *BindingSuite) TestResolveInterleavedInjectables() { return nil } got, err := s.resolve(fn, []Arg{ - {Kind: ArgKindLiteral, Value: "uk", DataType: DataTypeString}, - {Kind: ArgKindLiteral, Value: map[string]any{"k": "v"}, DataType: DataTypeObject}, + LiteralArg{Value: "uk", DataType: DataTypeString}, + LiteralArg{Value: map[string]any{"k": "v"}, DataType: DataTypeObject}, }, &fakeXComClient{}) s.Require().NoError(err) s.NotNil(got[0].Interface().(*slog.Logger)) @@ -261,7 +261,7 @@ func (s *BindingSuite) TestResolveTypeMismatchFailsLoudly() { fn := func(count int) error { return nil } _, err := s.resolve( fn, - []Arg{{Kind: ArgKindLiteral, Value: "uk", DataType: DataTypeString}}, + []Arg{LiteralArg{Value: "uk", DataType: DataTypeString}}, &fakeXComClient{}, ) if s.Assert().Error(err) { @@ -276,7 +276,7 @@ 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{}) + _, err := s.resolve(fn, []Arg{LiteralArg{Value: "uk"}}, &fakeXComClient{}) if s.Assert().Error(err) { s.Contains(err.Error(), "decoding literal value into int") } @@ -343,8 +343,8 @@ func (s *BindingSuite) TestResolveXComArgs() { 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}, + XComArg{TaskID: "extract", DataType: DataTypeObject}, + XComArg{TaskID: "extract", Key: "part", DataType: DataTypeString}, }, client) s.Require().NoError(err) s.Equal(extractResult{GoVersion: "go1.24", Timestamp: 42}, got[0].Interface()) @@ -368,7 +368,7 @@ func (s *BindingSuite) TestResolveXComStrictStructDecode() { "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) + _, err := s.resolve(fn, []Arg{XComArg{TaskID: "extract"}}, client) if s.Assert().Error(err) { s.Contains(err.Error(), `decoding xcom from task "extract"`) s.Contains(err.Error(), "unknown field") @@ -378,7 +378,7 @@ func (s *BindingSuite) TestResolveXComStrictStructDecode() { 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) + _, err := s.resolve(fn, []Arg{XComArg{TaskID: "extract"}}, client) if s.Assert().Error(err) { s.Contains(err.Error(), `pulling xcom from task "extract"`) } @@ -388,7 +388,7 @@ 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"}}, + []Arg{XComArg{TaskID: "extract"}}, ) if s.Assert().Error(err) { s.Contains(err.Error(), "no workload in context") @@ -399,24 +399,41 @@ 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}}, + []Arg{LiteralArg{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{}) + _, err = s.resolve(fnStr, []Arg{LiteralArg{Value: nil}}, &fakeXComClient{}) if s.Assert().Error(err) { s.Contains(err.Error(), "not nilable") } } -func (s *BindingSuite) TestResolveUnknownKind() { +// fakeArg is an out-of-catalogue Arg variant: the compiler seals the sum type +// to this package, so the defensive default branch can only be reached from +// inside it. +type fakeArg struct{} + +func (fakeArg) ArgName() string { return "fake" } +func (fakeArg) DeclaredType() DataType { return DataTypeAny } +func (fakeArg) sealedArg() {} + +func (s *BindingSuite) TestResolveUnsupportedVariant() { + fn := func(country string) error { return nil } + _, err := s.resolve(fn, []Arg{fakeArg{}}, &fakeXComClient{}) + if s.Assert().Error(err) { + s.Contains(err.Error(), "unsupported argument binding binding.fakeArg") + } +} + +func (s *BindingSuite) TestResolveNilArg() { fn := func(country string) error { return nil } - _, err := s.resolve(fn, []Arg{{Kind: ArgKind("template"), Value: "x"}}, &fakeXComClient{}) + _, err := s.resolve(fn, []Arg{nil}, &fakeXComClient{}) if s.Assert().Error(err) { - s.Contains(err.Error(), `unknown argument kind "template"`) + s.Contains(err.Error(), "nil argument binding") } } @@ -431,7 +448,7 @@ func (s *BindingSuite) TestResolveTIRunContextRebuild() { 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}, + LiteralArg{Value: "uk", DataType: DataTypeString}, }) s.Require().NoError(err) rc := got[0].Interface().(sdk.TIRunContext) @@ -515,8 +532,8 @@ func (s *BindingSuite) TestResolveTaskInputAllStruct() { }} fn := func(input combineInput) error { return nil } got, err := s.resolve(fn, []Arg{ - {Name: "name", Kind: ArgKindLiteral, Value: "widget", DataType: DataTypeString}, - {Name: "count", Kind: ArgKindLiteral, Value: 7, DataType: DataTypeInteger}, + LiteralArg{Name: "name", Value: "widget", DataType: DataTypeString}, + LiteralArg{Name: "count", Value: 7, DataType: DataTypeInteger}, }, client) s.Require().NoError(err) @@ -541,10 +558,10 @@ func (s *BindingSuite) TestResolveTaskInputAllStruct() { func (s *BindingSuite) TestResolveTaskInputMixedWithFlat() { fn := func(prefix string, input reportInput, suffix string) error { return nil } got, err := s.resolve(fn, []Arg{ - {Name: "prefix", Kind: ArgKindLiteral, Value: "head", DataType: DataTypeString}, - {Name: "region", Kind: ArgKindXCom, TaskID: "make_region", DataType: DataTypeString}, - {Name: "ratio", Kind: ArgKindLiteral, Value: 0.5, DataType: DataTypeNumber}, - {Name: "suffix", Kind: ArgKindLiteral, Value: "footer", DataType: DataTypeString}, + LiteralArg{Name: "prefix", Value: "head", DataType: DataTypeString}, + XComArg{Name: "region", TaskID: "make_region", DataType: DataTypeString}, + LiteralArg{Name: "ratio", Value: 0.5, DataType: DataTypeNumber}, + LiteralArg{Name: "suffix", Value: "footer", DataType: DataTypeString}, }, &fakeXComClient{values: map[string]any{"make_region/return_value": "east"}}) s.Require().NoError(err) @@ -579,7 +596,7 @@ func (s *BindingSuite) TestResolveTaskInputXComTagIndependentOfArgs() { func (s *BindingSuite) TestResolveTaskInputLiteralThroughArgName() { fn := func(input simpleTaskInput) error { return nil } got, err := s.resolve(fn, []Arg{ - {Name: "name", Kind: ArgKindLiteral, Value: "widget", DataType: DataTypeString}, + LiteralArg{Name: "name", Value: "widget", DataType: DataTypeString}, }, &fakeXComClient{}) s.Require().NoError(err) s.Equal("widget", got[0].Interface().(simpleTaskInput).Name) @@ -588,7 +605,7 @@ func (s *BindingSuite) TestResolveTaskInputLiteralThroughArgName() { func (s *BindingSuite) TestResolveTaskInputPointerStruct() { fn := func(input *simpleTaskInput) error { return nil } got, err := s.resolve(fn, []Arg{ - {Name: "name", Kind: ArgKindLiteral, Value: "widget", DataType: DataTypeString}, + LiteralArg{Name: "name", Value: "widget", DataType: DataTypeString}, }, &fakeXComClient{}) s.Require().NoError(err) input := got[0].Interface().(*simpleTaskInput) @@ -599,7 +616,7 @@ func (s *BindingSuite) TestResolveTaskInputPointerStruct() { func (s *BindingSuite) TestResolveTaskInputUnmatchedArgNameLeavesFieldZeroValued() { fn := func(input simpleTaskInput) error { return nil } _, err := s.resolve(fn, []Arg{ - {Name: "different_name", Kind: ArgKindLiteral, Value: "x", DataType: DataTypeString}, + LiteralArg{Name: "different_name", Value: "x", DataType: DataTypeString}, }, &fakeXComClient{}) // simpleTaskInput has no other flat parameter to absorb "different_name", // so the arity check (0 unclaimed args expected) still fails the task -- @@ -612,7 +629,7 @@ func (s *BindingSuite) TestResolveTaskInputUnmatchedArgNameLeavesFieldZeroValued func (s *BindingSuite) TestResolveTaskInputUnmatchedArgNameZeroValuedAlongsideMatch() { fn := func(input twoFieldTaskInput) error { return nil } got, err := s.resolve(fn, []Arg{ - {Name: "name", Kind: ArgKindLiteral, Value: "widget", DataType: DataTypeString}, + LiteralArg{Name: "name", Value: "widget", DataType: DataTypeString}, }, &fakeXComClient{}) s.Require().NoError(err) input := got[0].Interface().(twoFieldTaskInput) @@ -623,7 +640,7 @@ func (s *BindingSuite) TestResolveTaskInputUnmatchedArgNameZeroValuedAlongsideMa func (s *BindingSuite) TestResolveTaskInputArityMismatchForLeftoverArgs() { fn := func(input simpleTaskInput, extra string) error { return nil } _, err := s.resolve(fn, []Arg{ - {Name: "name", Kind: ArgKindLiteral, Value: "widget", DataType: DataTypeString}, + LiteralArg{Name: "name", Value: "widget", DataType: DataTypeString}, }, &fakeXComClient{}) if s.Assert().Error(err) { s.Contains(err.Error(), "argument count mismatch") diff --git a/go-sdk/pkg/execution/genmodels/defaults.gen.go b/go-sdk/pkg/execution/genmodels/defaults.gen.go index a6193f2b9af99..ba7f35b75431a 100644 --- a/go-sdk/pkg/execution/genmodels/defaults.gen.go +++ b/go-sdk/pkg/execution/genmodels/defaults.gen.go @@ -186,6 +186,17 @@ func (m *HITLDetailRequestResult) DecodeMsgpack(dec *msgpack.Decoder) error { return nil } +// DecodeMsgpack applies LiteralArgBinding's schema defaults that msgpack would otherwise skip. +func (m *LiteralArgBinding) DecodeMsgpack(dec *msgpack.Decoder) error { + type alias LiteralArgBinding + v := alias{DataType: ArgBindingDataType("any")} + if err := dec.Decode(&v); err != nil { + return err + } + *m = LiteralArgBinding(v) + return nil +} + // DecodeMsgpack applies PreviousTIResponse's schema defaults that msgpack would otherwise skip. func (m *PreviousTIResponse) DecodeMsgpack(dec *msgpack.Decoder) error { type alias PreviousTIResponse @@ -282,17 +293,6 @@ 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: ArgBindingDataType("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 @@ -327,3 +327,14 @@ func (m *TriggerDagRun) DecodeMsgpack(dec *msgpack.Decoder) error { *m = TriggerDagRun(v) return nil } + +// DecodeMsgpack applies XComArgBinding's schema defaults that msgpack would otherwise skip. +func (m *XComArgBinding) DecodeMsgpack(dec *msgpack.Decoder) error { + type alias XComArgBinding + v := alias{DataType: ArgBindingDataType("any"), Key: "return_value"} + if err := dec.Decode(&v); err != nil { + return err + } + *m = XComArgBinding(v) + return nil +} diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index d6ce08114a491..38d305ea88ca7 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -1250,11 +1250,6 @@ 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 @@ -1268,6 +1263,21 @@ type LazyDeserializedDAG struct { LastLoaded interface{} `msgpack:"last_loaded,omitempty"` } +// One positional stub-task argument carrying an inline literal from the Dag file. +type LiteralArgBinding struct { + // DataType corresponds to the JSON schema field "data_type". + DataType ArgBindingDataType `msgpack:"data_type,omitempty"` + + // Kind corresponds to the JSON schema field "kind". + Kind string `msgpack:"kind"` + + // Name corresponds to the JSON schema field "name". + Name string `msgpack:"name"` + + // Value corresponds to the JSON schema field "value". + Value interface{} `msgpack:"value,omitempty"` +} + type LogicalDates []time.Time // Add a new value to be redacted in task logs. @@ -1628,36 +1638,7 @@ 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 ArgBindingDataType `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"` - - // Name corresponds to the JSON schema field "name". - Name string `msgpack:"name"` - - // 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 TaskArgBindingKind string - -const TaskArgBindingKindLiteral TaskArgBindingKind = "literal" -const TaskArgBindingKindXcom TaskArgBindingKind = "xcom" +type TaskArgBinding interface{} type TaskBreadcrumbsResult struct { // Breadcrumbs corresponds to the JSON schema field "breadcrumbs". @@ -1886,6 +1867,10 @@ type VariableResponse struct { Value interface{} `msgpack:"value"` } +type VersionData map[string]interface{} + +type Warnings []interface{} + type VariableResult struct { // Key corresponds to the JSON schema field "key". Key string `msgpack:"key"` @@ -1897,9 +1882,23 @@ type VariableResult struct { Value interface{} `msgpack:"value,omitempty"` } -type VersionData map[string]interface{} +// One positional stub-task argument pulled from an upstream task's XCom. +type XComArgBinding struct { + // DataType corresponds to the JSON schema field "data_type". + DataType ArgBindingDataType `msgpack:"data_type,omitempty"` -type Warnings []interface{} + // Key corresponds to the JSON schema field "key". + Key string `msgpack:"key,omitempty"` + + // Kind corresponds to the JSON schema field "kind". + Kind string `msgpack:"kind"` + + // Name corresponds to the JSON schema field "name". + Name string `msgpack:"name"` + + // TaskID corresponds to the JSON schema field "task_id". + TaskID string `msgpack:"task_id"` +} type XComCountResponse struct { // Len corresponds to the JSON schema field "len". diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 3454c58f910a4..1567c5f65adde 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -256,8 +256,18 @@ func TestTaskRunnerBindsArgs(t *testing.T) { 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"}}, + map[string]any{ + "name": "country", + "kind": "literal", + "data_type": "string", + "value": "uk", + }, + map[string]any{ + "name": "meta", + "kind": "literal", + "data_type": "object", + "value": map[string]any{"k": "v"}, + }, }, }, } @@ -295,7 +305,12 @@ func TestTaskRunnerArgBindingsArityMismatch(t *testing.T) { BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, TIContext: genmodels.TIRunContext{ ArgBindings: &genmodels.ArgBindings{ - {Kind: "literal", DataType: "string", Value: "uk"}, + map[string]any{ + "name": "country", + "kind": "literal", + "data_type": "string", + "value": "uk", + }, }, }, } @@ -340,7 +355,12 @@ func TestTaskRunnerBindsTaskInputStructArgs(t *testing.T) { BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, TIContext: genmodels.TIRunContext{ ArgBindings: &genmodels.ArgBindings{ - {Name: "region", Kind: "literal", DataType: "string", Value: "eu-west-1"}, + map[string]any{ + "name": "region", + "kind": "literal", + "data_type": "string", + "value": "eu-west-1", + }, }, }, } @@ -372,7 +392,12 @@ func TestTaskRunnerArgBindingsTypeMismatch(t *testing.T) { BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, TIContext: genmodels.TIRunContext{ ArgBindings: &genmodels.ArgBindings{ - {Kind: "literal", DataType: "string", Value: "uk"}, + map[string]any{ + "name": "count", + "kind": "literal", + "data_type": "string", + "value": "uk", + }, }, }, } @@ -384,6 +409,76 @@ func TestTaskRunnerArgBindingsTypeMismatch(t *testing.T) { assertTaskState(t, result, genmodels.TaskStateStateFailed) } +// TestTaskRunnerArgBindingsUnknownKind: a wire spec whose kind is neither xcom +// nor literal fails the task before the body runs. +func TestTaskRunnerArgBindingsUnknownKind(t *testing.T) { + ran := false + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTaskWithName("transform", + func(country string) error { + ran = true + return nil + }) + }) + + details := &genmodels.StartupDetails{ + TI: genmodels.TaskInstance{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "transform", + RunID: "run1", + MapIndex: ptr(-1), + }, + BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, + TIContext: genmodels.TIRunContext{ + ArgBindings: &genmodels.ArgBindings{ + map[string]any{"name": "country", "kind": "template", "value": "x"}, + }, + }, + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(context.Background(), bundle, details, comm, logger) + assertTaskState(t, result, genmodels.TaskStateStateFailed) + assert.False(t, ran, "the task body must not run on an unknown binding kind") +} + +// TestTaskRunnerArgBindingsMalformedElement: a wire spec element that is not a +// map at all fails the task before the body runs. +func TestTaskRunnerArgBindingsMalformedElement(t *testing.T) { + ran := false + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTaskWithName("transform", + func(country string) error { + ran = true + return nil + }) + }) + + details := &genmodels.StartupDetails{ + TI: genmodels.TaskInstance{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "transform", + RunID: "run1", + MapIndex: ptr(-1), + }, + BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, + TIContext: genmodels.TIRunContext{ + ArgBindings: &genmodels.ArgBindings{"bogus"}, + }, + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(context.Background(), bundle, details, comm, logger) + assertTaskState(t, result, genmodels.TaskStateStateFailed) + assert.False(t, ran, "the task body must not run on a malformed binding element") +} + 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/task_runner.go b/go-sdk/pkg/execution/task_runner.go index bb09726aa9584..0ec3d94b4d44e 100644 --- a/go-sdk/pkg/execution/task_runner.go +++ b/go-sdk/pkg/execution/task_runner.go @@ -125,34 +125,58 @@ func RunTask( ctx = context.WithValue(ctx, sdkcontext.SdkClientContextKey, sdk.Client(client)) ctx = context.WithValue(ctx, sdkcontext.RuntimeContextKey, runtimeContext) - args := convertArgBindings(details.TIContext.ArgBindings) + args, err := convertArgBindings(details.TIContext.ArgBindings) + if err != nil { + logger.Error("Invalid arg_bindings spec from supervisor", + "dag_id", details.TI.DagID, + "task_id", details.TI.TaskID, + "error", err, + ) + return genmodels.TaskState{ + State: genmodels.TaskStateStateFailed, + EndDate: time.Now().UTC(), + } + } return executeTask(ctx, task, args, details.TIContext.ShouldRetry, logger) } // convertArgBindings maps the wire-model positional-argument spec (captured from -// the Python stub Dag's TaskFlow call) onto the runtime-neutral binding form. -func convertArgBindings(specsPtr *genmodels.ArgBindings) []binding.Arg { +// the Python stub Dag's TaskFlow call) onto the runtime binding sum type. The +// wire union generates untyped items (msgpack delivers each XComArgBinding / +// LiteralArgBinding as a plain map), so the kind dispatch and the schema +// defaults (data_type "any", xcom key "return_value") are applied here. +func convertArgBindings(specsPtr *genmodels.ArgBindings) ([]binding.Arg, error) { if specsPtr == nil || len(*specsPtr) == 0 { - return nil + return nil, nil } specs := *specsPtr args := make([]binding.Arg, len(specs)) - for i, spec := range specs { - taskID := "" - if s, ok := spec.TaskID.(string); ok { - taskID = s + for i, raw := range specs { + m, ok := raw.(map[string]any) + if !ok { + return nil, fmt.Errorf("arg_bindings[%d]: unexpected wire shape %T", i, raw) + } + name, _ := m["name"].(string) + dataType := binding.DataTypeAny + if s, ok := m["data_type"].(string); ok && s != "" { + dataType = binding.DataType(s) } - args[i] = binding.Arg{ - Name: spec.Name, - Kind: binding.ArgKind(spec.Kind), - TaskID: taskID, - Key: spec.Key, - Value: spec.Value, - DataType: binding.DataType(spec.DataType), + switch kind, _ := m["kind"].(string); kind { + case "xcom": + taskID, _ := m["task_id"].(string) + key := "return_value" + if s, ok := m["key"].(string); ok && s != "" { + key = s + } + args[i] = binding.XComArg{Name: name, TaskID: taskID, Key: key, DataType: dataType} + case "literal": + args[i] = binding.LiteralArg{Name: name, Value: m["value"], DataType: dataType} + default: + return nil, fmt.Errorf("arg_bindings[%d]: unknown kind %q", i, kind) } } - return args + return args, nil } // mapIndexPtr normalizes the supervisor's map_index into the optional form diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index b1ad5cc19295e..31268c2d17c8b 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -108,12 +108,12 @@ def _build_arg_bindings( """ 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. ``name`` is always the stub function's parameter - name, so a foreign runtime can bind by name (e.g. the Go SDK's ``sdk.TaskInput`` struct - fields) in addition to the existing positional order. Returns ``None`` for parameterless - stubs. + Each spec entry is a plain dict matching one variant of the execution API's + ``TaskArgBinding`` union: an ``XComArgBinding`` (``kind="xcom"``) for upstream TaskFlow + outputs, or a ``LiteralArgBinding`` (``kind="literal"``) for everything else. ``name`` is + always the stub function's parameter name, so a foreign runtime can bind by name (e.g. the + Go SDK's ``sdk.TaskInput`` struct fields) in addition to the existing positional order. + Returns ``None`` for parameterless stubs. """ signature = inspect.signature(python_callable) diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index fe12557698456..cc68b1486c79d 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -232,6 +232,17 @@ class IntermediateTIState(str, Enum): AWAITING_INPUT = "awaiting_input" +class LiteralArgBinding(BaseModel): + """ + One positional stub-task argument carrying an inline literal from the Dag file. + """ + + kind: Annotated[Literal["literal"], Field(title="Kind")] + name: Annotated[str, Field(title="Name")] + data_type: ArgBindingDataType | None = ArgBindingDataType.ANY + value: JsonValue | None = None + + class PrevSuccessfulDagRunResponse(BaseModel): """ Schema for response with previous successful DagRun information for Task Template Context. @@ -385,27 +396,6 @@ class TITargetStatePayload(BaseModel): state: IntermediateTIState -class Kind(str, Enum): - XCOM = "xcom" - LITERAL = "literal" - - -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. - """ - - name: Annotated[str, Field(title="Name")] - kind: Annotated[Kind, Field(title="Kind")] - data_type: ArgBindingDataType | None = ArgBindingDataType.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. @@ -548,6 +538,18 @@ class VariableResponse(BaseModel): value: Annotated[str | None, Field(title="Value")] = None +class XComArgBinding(BaseModel): + """ + One positional stub-task argument pulled from an upstream task's XCom. + """ + + kind: Annotated[Literal["xcom"], Field(title="Kind")] + name: Annotated[str, Field(title="Name")] + data_type: ArgBindingDataType | None = ArgBindingDataType.ANY + task_id: Annotated[str, Field(title="Task Id")] + key: Annotated[str | None, Field(title="Key")] = "return_value" + + class XComResponse(BaseModel): """ XCom schema for responses with fields that are needed for Runtime. @@ -745,6 +747,10 @@ class TITerminalStatePayload(BaseModel): rendered_map_index: Annotated[str | None, Field(title="Rendered Map Index")] = None +class TaskArgBinding(RootModel[XComArgBinding | LiteralArgBinding]): + root: Annotated[XComArgBinding | LiteralArgBinding, Field(discriminator="kind", title="TaskArgBinding")] + + class AssetEventDagRunReference(BaseModel): """ Schema for AssetEvent model used in DagRun. 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 4cb56479e80a6..c9bec70012a67 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -3060,14 +3060,6 @@ "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": { @@ -4585,59 +4577,91 @@ "title": "ConnectionResponse", "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.", + "LiteralArgBinding": { + "description": "One positional stub-task argument carrying an inline literal from the Dag file.", "properties": { - "name": { - "title": "Name", - "type": "string" - }, "kind": { - "enum": [ - "xcom", - "literal" - ], + "const": "literal", "title": "Kind", "type": "string" }, + "name": { + "title": "Name", + "type": "string" + }, "data_type": { "$ref": "#/$defs/ArgBindingDataType", "default": "any" }, - "task_id": { + "value": { "anyOf": [ { - "type": "string" + "$ref": "#/$defs/JsonValue" }, { "type": "null" } ], - "default": null, - "title": "Task Id" + "default": null + } + }, + "required": [ + "kind", + "name" + ], + "title": "LiteralArgBinding", + "type": "object" + }, + "TaskArgBinding": { + "discriminator": { + "mapping": { + "literal": "#/$defs/LiteralArgBinding", + "xcom": "#/$defs/XComArgBinding" + }, + "propertyName": "kind" + }, + "oneOf": [ + { + "$ref": "#/$defs/XComArgBinding" + }, + { + "$ref": "#/$defs/LiteralArgBinding" + } + ], + "title": "TaskArgBinding" + }, + "XComArgBinding": { + "description": "One positional stub-task argument pulled from an upstream task's XCom.", + "properties": { + "kind": { + "const": "xcom", + "title": "Kind", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "data_type": { + "$ref": "#/$defs/ArgBindingDataType", + "default": "any" + }, + "task_id": { + "title": "Task Id", + "type": "string" }, "key": { "default": "return_value", "title": "Key", "type": "string" - }, - "value": { - "anyOf": [ - { - "$ref": "#/$defs/JsonValue" - }, - { - "type": "null" - } - ], - "default": null } }, "required": [ + "kind", "name", - "kind" + "task_id" ], - "title": "TaskArgBinding", + "title": "XComArgBinding", "type": "object" }, "AssetEventDagRunReference": { 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 index e69fe456f33c8..119963de7e855 100644 --- 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 @@ -19,15 +19,21 @@ from cadwyn import VersionChange, schema -from airflow.sdk.api.datamodels._generated import TaskArgBinding, TIRunContext +from airflow.sdk.api.datamodels._generated import LiteralArgBinding, TIRunContext, XComArgBinding class AddArgBindingsToTIRunContext(VersionChange): - """Add the ``arg_bindings`` positional-argument binding spec for stub (foreign-runtime) tasks.""" + """ + Add the ``arg_bindings`` positional-argument binding spec for stub (foreign-runtime) tasks. + + Each entry is a discriminated union of ``XComArgBinding`` and ``LiteralArgBinding`` + keyed on ``kind``. + """ description = __doc__ instructions_to_migrate_to_previous_version = ( schema(TIRunContext).field("arg_bindings").didnt_exist, - schema(TaskArgBinding).field("name").didnt_exist, + schema(XComArgBinding).field("name").didnt_exist, + schema(LiteralArgBinding).field("name").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 44ab182a41421..8f803dcefb951 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 @@ -442,7 +442,14 @@ def test_downgrade_strips_arg_bindings_for_previous_version(self, real_migrator, assert "arg_bindings" not in out["ti_context"] def test_head_version_keeps_arg_bindings(self, real_migrator, startup_details): + from airflow.sdk.api.datamodels._generated import LiteralArgBinding, XComArgBinding + out = real_migrator.downgrade(startup_details, "2026-07-30") assert out.ti_context.arg_bindings is not None - assert [a.kind for a in out.ti_context.arg_bindings] == ["literal", "xcom"] - assert [a.name for a in out.ti_context.arg_bindings] == ["country", "extracted"] + literal, xcom = (a.root for a in out.ti_context.arg_bindings) + assert isinstance(literal, LiteralArgBinding) + assert literal.value == "uk" + assert literal.name == "country" + assert isinstance(xcom, XComArgBinding) + assert xcom.task_id == "extract" + assert xcom.name == "extracted" diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 528f45cb8d23b..748a4c6e20499 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -22,6 +22,20 @@ // // Re-run with: pnpm run generate:supervisor +/** + * Language-neutral value type a stub-task argument binds to in the foreign runtime. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "ArgBindingDataType". + */ +export type ArgBindingDataType = + | "string" + | "integer" + | "number" + | "boolean" + | "object" + | "array" + | "any"; export type Name = string; export type Id = number; export type Timestamp = string; @@ -246,10 +260,39 @@ 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; +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskArgBinding". + */ +export type TaskArgBinding = XComArgBinding | LiteralArgBinding; +export type Kind = "xcom"; +export type Name8 = string; +/** + * Language-neutral value type a stub-task argument binds to in the foreign runtime. + */ +export type ArgBindingDataType1 = + | "string" + | "integer" + | "number" + | "boolean" + | "object" + | "array" + | "any"; +export type TaskId1 = string; export type Key1 = string; +export type Kind1 = "literal"; +export type Name9 = string; +/** + * Language-neutral value type a stub-task argument binds to in the foreign runtime. + */ +export type ArgBindingDataType2 = + | "string" + | "integer" + | "number" + | "boolean" + | "object" + | "array" + | "any"; export type Type13 = "TaskCallbackRequest"; export type Filepath2 = string; export type BundleName3 = string; @@ -299,11 +342,6 @@ 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 = @@ -320,7 +358,7 @@ export type NextKwargs2 = { } | null; export type RenderedMapIndex1 = string | null; export type Type20 = "DeferTask"; -export type Name8 = string; +export type Name10 = string; export type Key2 = string; export type Type21 = "DeleteAssetStateStoreByName"; export type Uri5 = string; @@ -372,11 +410,11 @@ export type ErrorType1 = | "PERMISSION_DENIED" | "GENERIC_ERROR" | "API_SERVER_ERROR"; -export type Name9 = string; +export type Name11 = string; export type Type27 = "GetAssetByName"; export type Uri6 = string; export type Type28 = "GetAssetByUri"; -export type Name10 = string | null; +export type Name12 = string | null; export type Uri7 = string | null; export type After = string | null; export type Before = string | null; @@ -399,7 +437,7 @@ export type Extra8 = { [k: string]: string; } | null; export type Type30 = "GetAssetEventByAssetAlias"; -export type Name11 = string; +export type Name13 = string; export type Key7 = string; export type Type31 = "GetAssetStateStoreByName"; export type Uri8 = string; @@ -508,12 +546,7 @@ 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 Name14 = string | null; export type Type56 = "MaskSecret"; export type Ok = boolean; export type Type57 = "OKResponse"; @@ -551,7 +584,7 @@ export type RetryReason = string | null; export type Type64 = "RetryTask"; export type Type65 = "SentFDs"; export type Fds = number[]; -export type Name13 = string; +export type Name15 = string; export type Key16 = string; export type Type66 = "SetAssetStateStoreByName"; export type Uri9 = string; @@ -1047,19 +1080,28 @@ export interface ConnectionResponse { 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. + * One positional stub-task argument pulled from an upstream task's XCom. * * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema - * via the `definition` "TaskArgBinding". + * via the `definition` "XComArgBinding". */ -export interface TaskArgBinding { +export interface XComArgBinding { kind: Kind; - data_type?: DataType; - task_id?: TaskId1; + name: Name8; + data_type?: ArgBindingDataType1; + task_id: TaskId1; key?: Key1; +} +/** + * One positional stub-task argument carrying an inline literal from the Dag file. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "LiteralArgBinding". + */ +export interface LiteralArgBinding { + kind: Kind1; + name: Name9; + data_type?: ArgBindingDataType2; value?: unknown; } /** @@ -1181,7 +1223,7 @@ export interface DeferTask { * via the `definition` "DeleteAssetStateStoreByName". */ export interface DeleteAssetStateStoreByName { - name: Name8; + name: Name10; key: Key2; type?: Type21; } @@ -1237,7 +1279,7 @@ export interface ErrorResponse { * via the `definition` "GetAssetByName". */ export interface GetAssetByName { - name: Name9; + name: Name11; type?: Type27; } /** @@ -1253,7 +1295,7 @@ export interface GetAssetByUri { * via the `definition` "GetAssetEventByAsset". */ export interface GetAssetEventByAsset { - name: Name10; + name: Name12; uri: Uri7; after?: After; before?: Before; @@ -1284,7 +1326,7 @@ export interface GetAssetEventByAssetAlias { * via the `definition` "GetAssetStateStoreByName". */ export interface GetAssetStateStoreByName { - name: Name11; + name: Name13; key: Key7; type?: Type31; } @@ -1552,7 +1594,7 @@ export interface InactiveAssetsResult { */ export interface MaskSecret { value: JsonValue; - name?: Name12; + name?: Name14; type?: Type56; } /** @@ -1668,7 +1710,7 @@ export interface SentFDs { * via the `definition` "SetAssetStateStoreByName". */ export interface SetAssetStateStoreByName { - name: Name13; + name: Name15; key: Key16; value: JsonValue; type?: Type66; From a3e4b9f7d721e5f17a11a3de2a141b84a94d8041 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 20 Jul 2026 16:05:22 +0000 Subject: [PATCH 11/40] Drop the xcom struct tag from Go SDK TaskInput field binding An ad hoc `xcom:""` pull baked the upstream task id into the compiled Go binary, hiding a data dependency from the Dag file that owns task wiring on the Python side (the example even needed a manual >> to order the pull's upstream). Fields now bind exclusively by argument name -- an `arg:""` tag, or the snake_cased field name when the tag is omitted -- so every value a task consumes stays visible in its TaskFlow call, and a task that needs an extra XCom can still ask for it explicitly through the injected client. --- go-sdk/README.md | 28 ++-- .../airflow-go-pack/pack_integration_test.go | 1 - go-sdk/dags/go_examples.py | 22 +-- go-sdk/example/bundle/main.go | 1 - .../bundle/taskflowbinding/taskflowbinding.go | 47 +------ .../taskflowbinding/taskflowbinding_test.go | 22 --- go-sdk/pkg/binding/binding.go | 128 +++++------------- go-sdk/pkg/binding/binding_test.go | 78 ++--------- go-sdk/sdk/context.go | 9 +- 9 files changed, 66 insertions(+), 270 deletions(-) diff --git a/go-sdk/README.md b/go-sdk/README.md index e8afd899adb28..0db21a1a0ef62 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -159,30 +159,20 @@ type CombineInput struct { sdk.TaskInput // one-line opt-in, zero runtime cost Region string `arg:"region_code"` // named lookup against the TaskFlow call argument "region_code" Threshold float64 // no tag -> falls back to the snake_cased field name "threshold" - Config Config `xcom:"make_config"` // ad hoc pull of make_config's return-value XCom, independent - // of the TaskFlow call -- there is no "config" argument at all } func Combine(ctx sdk.TIRunContext, log *slog.Logger, input CombineInput) (any, error) { - // input.Region, input.Threshold, input.Config are all populated. + // input.Region and input.Threshold are both populated. return nil, nil } ``` -Each exported field supports three all-optional tags: - -- `arg:""` — bind from the TaskFlow call argument with this name (matched against the stub - function's Python parameter name, independent of declaration order on either side). With no tag, - the field's own Go name, snake_cased (`RatioValue` → `ratio_value`, `TaskID` → `task_id`), is used. - If no TaskFlow call argument carries that name, the field is simply left at its Go zero value — - it does not fail the task, kwarg-style (see `ViaStructUnmatchedArg` below). -- `xcom:""` — an explicit, ad hoc XCom pull from the named upstream task, fully independent - of the TaskFlow call: the field need not correspond to any argument the Dag file passes at all. This - pull is unchecked — the runtime does not verify `` is an actual upstream dependency, the - same trust model as calling `sdk.Client.GetXCom` by hand. -- `xcom-key:""` — the XCom key for an `xcom:`-tagged field; defaults to the return-value key. - Setting it without `xcom:` is a registration-time error (a key with no task id is meaningless), as is - setting both `arg:` and `xcom:` on the same field. +Each exported field binds from the TaskFlow call argument named by its optional `arg:""` tag +(matched against the stub function's Python parameter name, independent of declaration order on +either side). With no tag, the field's own Go name, snake_cased (`RatioValue` → `ratio_value`, +`TaskID` → `task_id`), is used. If no TaskFlow call argument carries that name, the field is simply +left at its Go zero value — it does not fail the task, kwarg-style (see `ViaStructUnmatchedArg` +below). When a `TaskInput` struct and plain flat parameters coexist in the same function, the struct's fields claim entries out of the TaskFlow call's argument spec by name first; the *remaining, unclaimed* @@ -193,8 +183,8 @@ this — it keeps working as a single flat data parameter, JSON-decoded whole fr argument (see `Config` in [`example/bundle/taskflowbinding/taskflowbinding.go`](./example/bundle/taskflowbinding/taskflowbinding.go)), which is a different mechanism from per-field `TaskInput` binding. See -[`ViaStructNoTags`, `ViaStructArgTag`, `ViaStructXComTag`, and `ViaStructUnmatchedArg`](./example/bundle/taskflowbinding/taskflowbinding.go) -for a full worked example of each tag mode — and the unmatched-field case — in isolation. +[`ViaStructNoTags`, `ViaStructArgTag`, and `ViaStructUnmatchedArg`](./example/bundle/taskflowbinding/taskflowbinding.go) +for a full worked example of each field-binding mode — and the unmatched-field case — in isolation. ### Reading the task runtime context 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 3e29c5b5641af..652e943ba63e8 100644 --- a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go +++ b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go @@ -161,7 +161,6 @@ dags: - "via_flat_args" - "via_struct_no_tags" - "via_struct_arg_tag" - - "via_struct_xcom_tag" - "via_struct_unmatched_arg" ` assert.Equal(t, expectedManifest, string(metadata)) diff --git a/go-sdk/dags/go_examples.py b/go-sdk/dags/go_examples.py index dab624452e5ef..6d78261257e1c 100644 --- a/go-sdk/dags/go_examples.py +++ b/go-sdk/dags/go_examples.py @@ -21,10 +21,10 @@ load, below), ``concurrent_xcom_dag`` (one ``pull_xcoms_concurrently`` task timing sequential vs goroutine XCom pulls), and ``taskflow_binding_dag`` (stressing the TaskFlow argument-binding surface -- the flat, positional -parameter list ``via_flat_args`` binds onto, plus four ``sdk.TaskInput`` +parameter list ``via_flat_args`` binds onto, plus three ``sdk.TaskInput`` (keyword-style) struct examples, ``via_struct_no_tags``/``via_struct_arg_tag``/ -``via_struct_xcom_tag``/``via_struct_unmatched_arg``, each isolating one -field-binding mode; see its Dag function below). +``via_struct_unmatched_arg``, each isolating one field-binding mode; see its +Dag function below). ``simple_dag`` sandwiches the Go tasks between two native Python tasks so the run exercises XCom across the language boundary, the same way @@ -150,10 +150,6 @@ def via_struct_no_tags(region_code: str, threshold: float): ... def via_struct_arg_tag(region_code: str, threshold: float): ... -@task.stub(queue="golang") -def via_struct_xcom_tag(threshold: float): ... - - @task.stub(queue="golang") def via_struct_unmatched_arg(region_code: str): ... @@ -179,34 +175,28 @@ def taskflow_binding_dag(): ``*string``. The Go ``via_flat_args`` (``go-sdk/example/bundle/taskflowbinding``) verifies every bound value and fails the task on any mismatch. - Four further tasks demonstrate the Go SDK's ``sdk.TaskInput`` struct injection + Three further tasks demonstrate the Go SDK's ``sdk.TaskInput`` struct injection mode, one field-binding mode at a time: * ``via_struct_no_tags``: both struct fields fall back to their Go field name - snake_cased -- no ``arg:``/``xcom:`` tags at all. + snake_cased -- no ``arg:`` tags at all. * ``via_struct_arg_tag``: one field is renamed via an explicit ``arg:`` tag, proving the tag remaps the name rather than coincidentally matching it. - * ``via_struct_xcom_tag``: one field is an ad hoc XCom pull of ``make_config``'s - return value declared purely in Go (an ``xcom:`` struct tag) -- it is never - passed as a TaskFlow argument here, so the explicit ``>>`` below is what - orders it after ``make_config``. * ``via_struct_unmatched_arg``: the Go struct declares a field with no corresponding argument in this TaskFlow call at all -- it stays at its Go zero value rather than failing the task. """ - config = make_config() via_flat_args( "summary", 3, 2.5, True, ["metrics", "hourly"], - config=config, + config=make_config(), numbers=make_numbers(), ) via_struct_no_tags(region_code="eu-west-1", threshold=0.75) via_struct_arg_tag(region_code="eu-west-1", threshold=0.75) - config >> via_struct_xcom_tag(threshold=0.75) via_struct_unmatched_arg(region_code="eu-west-1") diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 7f020f8c2c429..88917fcf183cc 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -62,7 +62,6 @@ func (m *myBundle) RegisterDags(dagbag v1.Registry) error { bindingDag.AddTaskWithName("via_flat_args", taskflowbinding.ViaFlatArgs) bindingDag.AddTaskWithName("via_struct_no_tags", taskflowbinding.ViaStructNoTags) bindingDag.AddTaskWithName("via_struct_arg_tag", taskflowbinding.ViaStructArgTag) - bindingDag.AddTaskWithName("via_struct_xcom_tag", taskflowbinding.ViaStructXComTag) bindingDag.AddTaskWithName("via_struct_unmatched_arg", taskflowbinding.ViaStructUnmatchedArg) return nil diff --git a/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go b/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go index 475ab8da8ceb5..c5f5bd150c10b 100644 --- a/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go +++ b/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go @@ -24,10 +24,9 @@ // instead show the sdk.TaskInput struct-field injection mode -- conceptually // keyword-argument binding, where fields match by name and an unmatched name // is left at its zero value rather than failing the task -- one field-binding -// tag at a time: ViaStructNoTags (plain snake_case name fallback), -// ViaStructArgTag (an explicit `arg:` rename), ViaStructXComTag (an ad hoc -// `xcom:` pull), and ViaStructUnmatchedArg (a field whose name has no -// corresponding TaskFlow call argument at all). +// mode at a time: ViaStructNoTags (plain snake_case name fallback), +// ViaStructArgTag (an explicit `arg:` rename), and ViaStructUnmatchedArg (a +// field whose name has no corresponding TaskFlow call argument at all). package taskflowbinding import ( @@ -208,46 +207,6 @@ func ViaStructArgTag( }, nil } -// ViaStructXComTagInput demonstrates the sdk.TaskInput struct-field injection -// mode with an xcom: tag: Config is an ad hoc pull of make_config's return -// value, independent of the Python call's TaskFlow arguments entirely; -// Threshold has no tag and falls back to its snake_cased field name. -type ViaStructXComTagInput struct { - sdk.TaskInput - Threshold float64 - Config Config `xcom:"make_config"` -} - -// ViaStructXComTag is called as -// -// via_struct_xcom_tag(threshold=0.75) -// -// with the Python Dag ordering it after make_config explicitly (the xcom: -// pull is not a TaskFlow argument, so there is no implicit dependency). -func ViaStructXComTag( - ctx sdk.TIRunContext, - log *slog.Logger, - input ViaStructXComTagInput, -) (any, error) { - if input.Threshold != 0.75 { - return nil, fmt.Errorf("TaskInput field bound incorrectly: threshold=%v", input.Threshold) - } - if want := (Config{Environment: "production", Region: "eu-west-1", Debug: true}); input.Config != want { - return nil, fmt.Errorf( - "ad hoc xcom field bound incorrectly: config=%+v, want %+v", input.Config, want, - ) - } - - log.InfoContext(ctx, "Bound TaskInput struct (xcom: tag)", - "threshold", input.Threshold, - "environment", input.Config.Environment, - ) - return map[string]any{ - "threshold": input.Threshold, - "environment": input.Config.Environment, - }, nil -} - // ViaStructUnmatchedArgInput demonstrates that a field whose name has no // corresponding TaskFlow call argument at all is left at its Go zero value // rather than failing the task: Region binds normally, but Missing's arg diff --git a/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go b/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go index 6bd28dd99459a..648dec1473fde 100644 --- a/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go +++ b/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go @@ -103,28 +103,6 @@ func TestViaStructArgTagRejectsWrongBinding(t *testing.T) { assert.ErrorContains(t, err, "TaskInput fields bound incorrectly") } -func TestViaStructXComTag(t *testing.T) { - ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) - got, err := ViaStructXComTag(ctx, slog.Default(), ViaStructXComTagInput{ - Threshold: 0.75, - Config: Config{Environment: "production", Region: "eu-west-1", Debug: true}, - }) - require.NoError(t, err) - - summary, ok := got.(map[string]any) - require.True(t, ok, "ViaStructXComTag should return a map summary, got %T", got) - assert.Equal(t, "production", summary["environment"]) -} - -func TestViaStructXComTagRejectsWrongBinding(t *testing.T) { - ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) - _, err := ViaStructXComTag(ctx, slog.Default(), ViaStructXComTagInput{ - Threshold: 0.75, - Config: Config{}, - }) - assert.ErrorContains(t, err, "ad hoc xcom field bound incorrectly") -} - func TestViaStructUnmatchedArg(t *testing.T) { ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) // Missing is left at its Go zero value, exactly as binding.Resolve leaves an diff --git a/go-sdk/pkg/binding/binding.go b/go-sdk/pkg/binding/binding.go index 1db9b0651399d..99febaf776507 100644 --- a/go-sdk/pkg/binding/binding.go +++ b/go-sdk/pkg/binding/binding.go @@ -32,11 +32,10 @@ // - TaskInput structs: a struct that anonymously embeds sdk.TaskInput opts // into per-field, name-based binding instead of consuming one positional // slot as a whole-value decode target. Each exported field binds by name -// (an `arg:""` tag, or its Go field name snake_cased) against the -// Dag's TaskFlow call arguments, or by an explicit, ad hoc XCom pull (an -// `xcom:""` tag, with an optional `xcom-key:""`) that never -// consults the positional argument spec at all. At most one such -// parameter is allowed per function. +// against the Dag's TaskFlow call arguments: an `arg:""` tag names +// the argument to claim, and a field with no tag falls back to its own Go +// field name, snake_cased. At most one such parameter is allowed per +// function. // // A TaskInput struct's fields are resolved first, by name, claiming entries // out of the argument spec; the remaining unclaimed entries are then @@ -157,17 +156,6 @@ const ( paramTaskInput ) -// taskInputFieldSource discriminates how a TaskInput struct field is filled. -type taskInputFieldSource int - -const ( - // taskInputFieldFromArg claims a named entry from the argument spec. - taskInputFieldFromArg taskInputFieldSource = iota - // taskInputFieldFromXCom pulls directly via an `xcom:` tag, independent - // of the argument spec. - taskInputFieldFromXCom -) - // taskInputField describes how Resolve fills one exported field of a // TaskInput-embedding struct. Precomputed once by Analyze. type taskInputField struct { @@ -177,14 +165,9 @@ type taskInputField struct { // goName is the Go field name, for error messages. goName string fieldType reflect.Type - source taskInputFieldSource - // argName is the name to claim from the argument spec. Set only for - // taskInputFieldFromArg. + // argName is the name to claim from the argument spec: the field's `arg:` + // tag, or its snake_cased Go name when the tag is omitted. argName string - // xcomTaskID/xcomKey identify the ad hoc pull. Set only for - // taskInputFieldFromXCom. - xcomTaskID string - xcomKey string } // paramPlan describes how Resolve fills a single task-function parameter. @@ -246,11 +229,10 @@ func Analyze(fnType reflect.Type, fnName string) (*Plan, error) { // Resolve builds the ordered argument values for one call. Injectable // parameters receive values derived from ctx, logger, or client. A TaskInput -// struct's fields are resolved first, by name (or an explicit ad hoc xcom -// pull), claiming entries out of args; the remaining unclaimed entries are -// then distributed, in their original relative order, onto the plain flat -// data parameters in declaration order. An error fails the task before its -// body runs. +// struct's fields are resolved first, by name, claiming entries out of args; +// the remaining unclaimed entries are then distributed, in their original +// relative order, onto the plain flat data parameters in declaration order. +// An error fails the task before its body runs. func (p *Plan) Resolve( ctx context.Context, logger *slog.Logger, @@ -356,7 +338,6 @@ func (p *Plan) resolveData( } // resolveTaskInput builds the struct value for one TaskInput parameter. A -// field tagged `xcom:` pulls directly and never touches byName/claimed; a // field claiming an argument spec entry by name marks it claimed so the // later flat-parameter cursor skips it. A field whose name claims nothing // (kwarg-style: it was never "passed") is left unset at its Go zero value @@ -377,33 +358,21 @@ func (p *Plan) resolveTaskInput( structVal := reflect.New(structType).Elem() for _, tif := range plan.fields { - typeCheckCtx := fmt.Sprintf("TaskInput field %s (parameter %d)", tif.goName, plan.index) - generalCtx := fmt.Sprintf("TaskInput field %s", tif.goName) - - var arg Arg - switch tif.source { - case taskInputFieldFromXCom: - // DataTypeAny: an ad hoc pull has no Dag-declared type to check - // against, so decoding is driven entirely by the Go field's type. - arg = XComArg{ - TaskID: tif.xcomTaskID, - Key: tif.xcomKey, - DataType: DataTypeAny, - } - case taskInputFieldFromArg: - idx, ok := byName[tif.argName] - if !ok { - // No TaskFlow call argument carries this name -- kwarg-style, an - // unpassed name leaves the field at its Go zero value rather than - // failing the task (unlike a flat data parameter, where arity is - // checked strictly; see the package doc comment). - continue - } - claimed[idx] = true - arg = args[idx] + idx, ok := byName[tif.argName] + if !ok { + // No TaskFlow call argument carries this name -- kwarg-style, an + // unpassed name leaves the field at its Go zero value rather than + // failing the task (unlike a flat data parameter, where arity is + // checked strictly; see the package doc comment). + continue } + claimed[idx] = true - v, err := p.resolveOne(ctx, client, tif.fieldType, arg, typeCheckCtx, generalCtx) + v, err := p.resolveOne( + ctx, client, tif.fieldType, args[idx], + fmt.Sprintf("TaskInput field %s (parameter %d)", tif.goName, plan.index), + fmt.Sprintf("TaskInput field %s", tif.goName), + ) if err != nil { return reflect.Value{}, err } @@ -574,27 +543,6 @@ func buildTaskInputFields( continue } - argTag, hasArg := f.Tag.Lookup("arg") - xcomTag, hasXCom := f.Tag.Lookup("xcom") - xcomKeyTag, hasXComKey := f.Tag.Lookup("xcom-key") - - if hasArg && hasXCom { - return nil, fmt.Errorf( - "task function %s: parameter %d: TaskInput field %s: cannot set both `arg` and `xcom` tags", - fnName, - paramIndex, - f.Name, - ) - } - if hasXComKey && !hasXCom { - return nil, fmt.Errorf( - "task function %s: parameter %d: TaskInput field %s: `xcom-key` requires `xcom` to also "+ - "be set (a key with no task id is meaningless)", - fnName, - paramIndex, - f.Name, - ) - } if !isDecodableType(f.Type) { return nil, fmt.Errorf( "task function %s: parameter %d: TaskInput field %s: type %s cannot receive a task "+ @@ -607,27 +555,17 @@ func buildTaskInputFields( } tif := taskInputField{structIndex: i, goName: f.Name, fieldType: f.Type} - if hasXCom { - tif.source = taskInputFieldFromXCom - tif.xcomTaskID = xcomTag - tif.xcomKey = xcomKeyTag - if tif.xcomKey == "" { - tif.xcomKey = api.XComReturnValueKey - } - } else { - tif.source = taskInputFieldFromArg - tif.argName = argTag - if tif.argName == "" { - tif.argName = snakeCase(f.Name) - } - if existing, ok := seenArgNames[tif.argName]; ok { - return nil, fmt.Errorf( - "task function %s: parameter %d: TaskInput fields %s and %s both bind arg name %q", - fnName, paramIndex, existing, f.Name, tif.argName, - ) - } - seenArgNames[tif.argName] = f.Name + tif.argName = f.Tag.Get("arg") + if tif.argName == "" { + tif.argName = snakeCase(f.Name) + } + if existing, ok := seenArgNames[tif.argName]; ok { + return nil, fmt.Errorf( + "task function %s: parameter %d: TaskInput fields %s and %s both bind arg name %q", + fnName, paramIndex, existing, f.Name, tif.argName, + ) } + seenArgNames[tif.argName] = f.Name fields = append(fields, tif) } return fields, nil diff --git a/go-sdk/pkg/binding/binding_test.go b/go-sdk/pkg/binding/binding_test.go index f73b32074ec19..beeccb5affc25 100644 --- a/go-sdk/pkg/binding/binding_test.go +++ b/go-sdk/pkg/binding/binding_test.go @@ -309,15 +309,13 @@ type nonEmbeddingStruct struct { Name string } -// combineInput exercises every TaskInput field-tag combination: Name falls -// back to its field name, Count is explicitly named, Note is an ad hoc XCom -// pull at the default key, and Debug is an ad hoc XCom pull at a custom key. +// combineInput exercises both TaskInput field-binding modes side by side: +// Name falls back to its snake_cased field name, Count is explicitly named +// via its `arg:` tag. type combineInput struct { sdk.TaskInput Name string - Count int `arg:"count"` - Note *string ` xcom:"make_note"` - Debug bool ` xcom:"make_config" xcom-key:"debug_flag"` + Count int `arg:"count"` } // reportInput deliberately declares Ratio before Region, the reverse of the @@ -329,12 +327,6 @@ type reportInput struct { Region string `arg:"region"` } -// xcomOnlyInput's sole field never touches the argument spec at all. -type xcomOnlyInput struct { - sdk.TaskInput - Extra string `xcom:"make_config"` -} - func (s *BindingSuite) TestResolveXComArgs() { client := &fakeXComClient{values: map[string]any{ "extract/return_value": map[string]any{"go_version": "go1.24", "timestamp": int64(42)}, @@ -471,43 +463,27 @@ func (s *BindingSuite) TestAnalyzeTaskInputClassification() { ) } -func (s *BindingSuite) TestAnalyzeTaskInputTagValidation() { - type conflictingTags struct { - sdk.TaskInput - Field string `arg:"a" xcom:"b"` - } - type orphanXComKey struct { - sdk.TaskInput - Field string `xcom-key:"k"` - } +func (s *BindingSuite) TestAnalyzeTaskInputValidation() { type duplicateArgNames struct { sdk.TaskInput A string B string `arg:"a"` } - type nonDecodableXComField struct { + type nonDecodableField struct { sdk.TaskInput - Bad chan int `xcom:"t"` + Bad chan int } cases := map[string]struct { fn any errContains string }{ - "conflicting-arg-and-xcom-tags": { - func(input conflictingTags) error { return nil }, - "cannot set both `arg` and `xcom` tags", - }, - "orphan-xcom-key": { - func(input orphanXComKey) error { return nil }, - "`xcom-key` requires `xcom` to also be set", - }, "duplicate-arg-names": { func(input duplicateArgNames) error { return nil }, `fields A and B both bind arg name "a"`, }, - "non-decodable-xcom-field": { - func(input nonDecodableXComField) error { return nil }, + "non-decodable-field": { + func(input nonDecodableField) error { return nil }, "cannot receive a task argument", }, "two-taskinput-params": { @@ -526,33 +502,16 @@ func (s *BindingSuite) TestAnalyzeTaskInputTagValidation() { } func (s *BindingSuite) TestResolveTaskInputAllStruct() { - client := &fakeXComClient{values: map[string]any{ - "make_note/return_value": "hello", - "make_config/debug_flag": true, - }} fn := func(input combineInput) error { return nil } got, err := s.resolve(fn, []Arg{ LiteralArg{Name: "name", Value: "widget", DataType: DataTypeString}, LiteralArg{Name: "count", Value: 7, DataType: DataTypeInteger}, - }, client) + }, &fakeXComClient{}) s.Require().NoError(err) input := got[0].Interface().(combineInput) - s.Equal("widget", input.Name) - s.Equal(7, input.Count) - s.Require().NotNil(input.Note) - s.Equal("hello", *input.Note) - s.True(input.Debug) - - s.Require().Len(client.calls, 2, "the ad hoc xcom: fields, in struct declaration order") - s.Equal("make_note", client.calls[0].taskID) - s.Equal( - api.XComReturnValueKey, - client.calls[0].key, - "no xcom-key tag defaults to the return-value key", - ) - s.Equal("make_config", client.calls[1].taskID) - s.Equal("debug_flag", client.calls[1].key) + s.Equal("widget", input.Name, "the untagged field claims its snake_cased field name") + s.Equal(7, input.Count, "the `arg:` tag claims its named entry") } func (s *BindingSuite) TestResolveTaskInputMixedWithFlat() { @@ -580,19 +539,6 @@ func (s *BindingSuite) TestResolveTaskInputMixedWithFlat() { ) } -func (s *BindingSuite) TestResolveTaskInputXComTagIndependentOfArgs() { - client := &fakeXComClient{values: map[string]any{"make_config/return_value": "cfg-value"}} - fn := func(input xcomOnlyInput) error { return nil } - got, err := s.resolve(fn, nil, client) - s.Require().NoError(err) - - input := got[0].Interface().(xcomOnlyInput) - s.Equal("cfg-value", input.Extra) - s.Require().Len(client.calls, 1) - s.Equal("make_config", client.calls[0].taskID) - s.Equal(api.XComReturnValueKey, client.calls[0].key) -} - func (s *BindingSuite) TestResolveTaskInputLiteralThroughArgName() { fn := func(input simpleTaskInput) error { return nil } got, err := s.resolve(fn, []Arg{ diff --git a/go-sdk/sdk/context.go b/go-sdk/sdk/context.go index 8400ff8430191..6050ed52dae2b 100644 --- a/go-sdk/sdk/context.go +++ b/go-sdk/sdk/context.go @@ -113,16 +113,13 @@ type DagRun struct { // that struct into per-field, name-based TaskFlow argument binding -- an // ergonomic alternative to a long flat parameter list. Each exported field of // such a struct may carry an `arg:""` tag naming the stub's TaskFlow -// argument to bind (falling back to the field's own name, snake_cased, when -// omitted), or an `xcom:""` tag (with an optional companion -// `xcom-key:""`, defaulting to the return-value key) for an explicit, -// ad hoc XCom pull independent of the TaskFlow call: +// argument to bind, falling back to the field's own name, snake_cased, when +// omitted: // // type CombineInput struct { // sdk.TaskInput // Name string -// Count int `arg:"count"` -// Extra string `xcom:"make_config" xcom-key:"environment"` +// Count int `arg:"count"` // } // // func Combine(ctx sdk.TIRunContext, log *slog.Logger, input CombineInput) (any, error) From 79bed136be6d07249e62b5366e16f003f5ecdb70 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 21 Jul 2026 03:30:09 +0000 Subject: [PATCH 12/40] Bind untagged Go SDK TaskInput fields by their verbatim field name The snake_cased fallback silently rewrote Go field names into wire argument names, hiding the cross-language mapping from the reader; and because an unmatched TaskInput field kwarg-style falls back to its zero value, a wrong guess about the conversion never failed loudly. Matching the field name verbatim removes that magic: every snake_case Python parameter a field binds is now spelled out as an explicit `arg:` tag in the Go source. The e2e module also still referenced the via_struct_xcom_tag task removed with the xcom struct tag, which would have failed the suite against the current Dag. --- .../test_go_sdk_taskflow_binding.py | 20 +++-------- go-sdk/README.md | 10 +++--- go-sdk/dags/go_examples.py | 17 +++++---- .../bundle/taskflowbinding/taskflowbinding.go | 22 ++++++------ go-sdk/pkg/binding/binding.go | 36 +++---------------- go-sdk/pkg/binding/binding_test.go | 22 ++++++------ go-sdk/sdk/context.go | 5 +-- 7 files changed, 52 insertions(+), 80 deletions(-) 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 index c874b72de26cd..1dd793fbfa91b 100644 --- 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 @@ -78,7 +78,6 @@ def test_all_tasks_succeeded(completed_run: _CompletedRun): "via_flat_args", "via_struct_no_tags", "via_struct_arg_tag", - "via_struct_xcom_tag", "via_struct_unmatched_arg", ): assert completed_run.ti_states.get(task_id) == "success", completed_run.ti_states @@ -112,8 +111,8 @@ def test_via_flat_args_summary_reflects_bound_arguments(completed_run: _Complete def test_via_struct_no_tags_reflects_bound_arguments(completed_run: _CompletedRun): """``via_struct_no_tags`` demonstrates the Go SDK's ``sdk.TaskInput`` struct-field - injection mode with no field tags at all: both fields bind by their Go field name - snake_cased.""" + injection mode with no field tags at all: each field binds the TaskFlow argument + spelled exactly like its Go field name (``RegionCode``, ``Threshold``).""" assert completed_run.xcom("via_struct_no_tags") == { "region_code": "eu-west-1", "threshold": 0.75, @@ -121,24 +120,15 @@ def test_via_struct_no_tags_reflects_bound_arguments(completed_run: _CompletedRu def test_via_struct_arg_tag_reflects_bound_arguments(completed_run: _CompletedRun): - """``via_struct_arg_tag`` demonstrates the ``arg:`` tag renaming a struct field - away from its snake_cased default.""" + """``via_struct_arg_tag`` demonstrates explicit ``arg:`` tags: ``Region`` is + genuinely renamed to ``region_code``, and ``Threshold`` is tagged ``threshold`` + to pull the snake_case argument its verbatim field name would miss.""" assert completed_run.xcom("via_struct_arg_tag") == { "region": "eu-west-1", "threshold": 0.75, } -def test_via_struct_xcom_tag_reflects_bound_arguments(completed_run: _CompletedRun): - """``via_struct_xcom_tag`` demonstrates the ``xcom:`` tag: its ``Config`` field is - an ad hoc pull of ``make_config``'s return value, with no corresponding TaskFlow - call argument here.""" - assert completed_run.xcom("via_struct_xcom_tag") == { - "threshold": 0.75, - "environment": "production", - } - - def test_via_struct_unmatched_arg_reflects_zero_valued_field(completed_run: _CompletedRun): """``via_struct_unmatched_arg`` demonstrates that a struct field whose name has no corresponding TaskFlow call argument stays at its Go zero value instead of diff --git a/go-sdk/README.md b/go-sdk/README.md index 0db21a1a0ef62..9830de9758b3b 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -158,7 +158,7 @@ back to a caller-side default in a kwargs-style call. type CombineInput struct { sdk.TaskInput // one-line opt-in, zero runtime cost Region string `arg:"region_code"` // named lookup against the TaskFlow call argument "region_code" - Threshold float64 // no tag -> falls back to the snake_cased field name "threshold" + Threshold float64 `arg:"threshold"` // tags also bridge Go's UpperCamelCase to a snake_case argument } func Combine(ctx sdk.TIRunContext, log *slog.Logger, input CombineInput) (any, error) { @@ -169,10 +169,10 @@ func Combine(ctx sdk.TIRunContext, log *slog.Logger, input CombineInput) (any, e Each exported field binds from the TaskFlow call argument named by its optional `arg:""` tag (matched against the stub function's Python parameter name, independent of declaration order on -either side). With no tag, the field's own Go name, snake_cased (`RatioValue` → `ratio_value`, -`TaskID` → `task_id`), is used. If no TaskFlow call argument carries that name, the field is simply -left at its Go zero value — it does not fail the task, kwarg-style (see `ViaStructUnmatchedArg` -below). +either side). With no tag, the field's own Go name is matched verbatim — an untagged `Threshold` +only binds an argument literally spelled `Threshold`, so a snake_case Python parameter needs an +explicit tag. If no TaskFlow call argument carries that name, the field is simply left at its Go +zero value — it does not fail the task, kwarg-style (see `ViaStructUnmatchedArg` below). When a `TaskInput` struct and plain flat parameters coexist in the same function, the struct's fields claim entries out of the TaskFlow call's argument spec by name first; the *remaining, unclaimed* diff --git a/go-sdk/dags/go_examples.py b/go-sdk/dags/go_examples.py index 6d78261257e1c..dc9143bc16436 100644 --- a/go-sdk/dags/go_examples.py +++ b/go-sdk/dags/go_examples.py @@ -142,8 +142,10 @@ def via_flat_args( ): ... +# Capitalized parameters on purpose: with no ``arg:`` tags on the Go side, each +# struct field binds the argument spelled exactly like its Go field name. @task.stub(queue="golang") -def via_struct_no_tags(region_code: str, threshold: float): ... +def via_struct_no_tags(RegionCode: str, Threshold: float): ... @task.stub(queue="golang") @@ -178,10 +180,13 @@ def taskflow_binding_dag(): Three further tasks demonstrate the Go SDK's ``sdk.TaskInput`` struct injection mode, one field-binding mode at a time: - * ``via_struct_no_tags``: both struct fields fall back to their Go field name - snake_cased -- no ``arg:`` tags at all. - * ``via_struct_arg_tag``: one field is renamed via an explicit ``arg:`` tag, - proving the tag remaps the name rather than coincidentally matching it. + * ``via_struct_no_tags``: no ``arg:`` tags at all -- each struct field binds + the argument spelled exactly like its Go field name, hence this stub's + capitalized ``RegionCode``/``Threshold`` parameters. + * ``via_struct_arg_tag``: every field names its argument via an explicit + ``arg:`` tag -- ``Region`` is genuinely renamed to ``region_code``, and + ``Threshold`` is tagged ``threshold`` to pull the snake_case argument its + verbatim field name would miss. * ``via_struct_unmatched_arg``: the Go struct declares a field with no corresponding argument in this TaskFlow call at all -- it stays at its Go zero value rather than failing the task. @@ -195,7 +200,7 @@ def taskflow_binding_dag(): config=make_config(), numbers=make_numbers(), ) - via_struct_no_tags(region_code="eu-west-1", threshold=0.75) + via_struct_no_tags(RegionCode="eu-west-1", Threshold=0.75) via_struct_arg_tag(region_code="eu-west-1", threshold=0.75) via_struct_unmatched_arg(region_code="eu-west-1") diff --git a/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go b/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go index c5f5bd150c10b..7346c7160cb34 100644 --- a/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go +++ b/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go @@ -24,8 +24,8 @@ // instead show the sdk.TaskInput struct-field injection mode -- conceptually // keyword-argument binding, where fields match by name and an unmatched name // is left at its zero value rather than failing the task -- one field-binding -// mode at a time: ViaStructNoTags (plain snake_case name fallback), -// ViaStructArgTag (an explicit `arg:` rename), and ViaStructUnmatchedArg (a +// mode at a time: ViaStructNoTags (verbatim field-name fallback), +// ViaStructArgTag (explicit `arg:` naming), and ViaStructUnmatchedArg (a // field whose name has no corresponding TaskFlow call argument at all). package taskflowbinding @@ -136,8 +136,9 @@ func ViaFlatArgs( } // ViaStructNoTagsInput demonstrates the sdk.TaskInput struct-field injection -// mode with no field tags at all: both fields fall back to their Go field -// name snake_cased ("region_code", "threshold"). +// mode with no field tags at all: each field binds the TaskFlow call argument +// spelled exactly like its Go field name ("RegionCode", "Threshold"), which +// is why the stub declares capitalized parameters. type ViaStructNoTagsInput struct { sdk.TaskInput RegionCode string @@ -146,7 +147,7 @@ type ViaStructNoTagsInput struct { // ViaStructNoTags is called as // -// via_struct_no_tags(region_code="eu-west-1", threshold=0.75) +// via_struct_no_tags(RegionCode="eu-west-1", Threshold=0.75) func ViaStructNoTags( ctx sdk.TIRunContext, log *slog.Logger, @@ -171,14 +172,15 @@ func ViaStructNoTags( } // ViaStructArgTagInput demonstrates the sdk.TaskInput struct-field injection -// mode with an explicit arg: tag: Region binds to the "region_code" TaskFlow +// mode with explicit arg: tags: Region binds to the "region_code" TaskFlow // argument under a renamed Go field, proving the tag remaps the name rather -// than coincidentally matching it; Threshold has no tag and falls back to its -// snake_cased field name. +// than coincidentally matching it; Threshold is intentionally tagged +// "threshold" because an untagged field would only match an argument spelled +// exactly "Threshold". type ViaStructArgTagInput struct { sdk.TaskInput - Region string `arg:"region_code"` - Threshold float64 + Region string `arg:"region_code"` + Threshold float64 `arg:"threshold"` } // ViaStructArgTag is called as diff --git a/go-sdk/pkg/binding/binding.go b/go-sdk/pkg/binding/binding.go index 99febaf776507..0e79af1156847 100644 --- a/go-sdk/pkg/binding/binding.go +++ b/go-sdk/pkg/binding/binding.go @@ -33,9 +33,8 @@ // into per-field, name-based binding instead of consuming one positional // slot as a whole-value decode target. Each exported field binds by name // against the Dag's TaskFlow call arguments: an `arg:""` tag names -// the argument to claim, and a field with no tag falls back to its own Go -// field name, snake_cased. At most one such parameter is allowed per -// function. +// the argument to claim, and a field with no tag claims its own Go field +// name, verbatim. At most one such parameter is allowed per function. // // A TaskInput struct's fields are resolved first, by name, claiming entries // out of the argument spec; the remaining unclaimed entries are then @@ -71,8 +70,6 @@ import ( "fmt" "log/slog" "reflect" - "strings" - "unicode" "github.com/apache/airflow/go-sdk/pkg/api" "github.com/apache/airflow/go-sdk/pkg/sdkcontext" @@ -99,7 +96,7 @@ const ( // so this package stays decoupled from the generated coordinator schema types. type Arg interface { // ArgName is the stub function's parameter name this binding fills; used - // to match a TaskInput struct field's `arg:` tag (or its snake_cased + // to match a TaskInput struct field's `arg:` tag (or its verbatim // field-name fallback). ArgName() string // DeclaredType is the Dag-declared language-neutral type for the argument; @@ -166,7 +163,7 @@ type taskInputField struct { goName string fieldType reflect.Type // argName is the name to claim from the argument spec: the field's `arg:` - // tag, or its snake_cased Go name when the tag is omitted. + // tag, or its Go field name verbatim when the tag is omitted. argName string } @@ -557,7 +554,7 @@ func buildTaskInputFields( tif := taskInputField{structIndex: i, goName: f.Name, fieldType: f.Type} tif.argName = f.Tag.Get("arg") if tif.argName == "" { - tif.argName = snakeCase(f.Name) + tif.argName = f.Name } if existing, ok := seenArgNames[tif.argName]; ok { return nil, fmt.Errorf( @@ -571,29 +568,6 @@ func buildTaskInputFields( return fields, nil } -// snakeCase converts a Go exported field name (UpperCamelCase, acronyms -// preserved as a run) to the wire's snake_case convention, e.g. -// "RatioValue" -> "ratio_value", "TaskID" -> "task_id". Used as the fallback -// arg name for a TaskInput struct field with no explicit `arg:` tag. -func snakeCase(name string) string { - runes := []rune(name) - var b strings.Builder - for i, r := range runes { - if unicode.IsUpper(r) { - prevLower := i > 0 && unicode.IsLower(runes[i-1]) - prevUpper := i > 0 && unicode.IsUpper(runes[i-1]) - nextLower := i+1 < len(runes) && unicode.IsLower(runes[i+1]) - if i > 0 && (prevLower || (prevUpper && nextLower)) { - b.WriteByte('_') - } - b.WriteRune(unicode.ToLower(r)) - } else { - b.WriteRune(r) - } - } - return b.String() -} - // 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. diff --git a/go-sdk/pkg/binding/binding_test.go b/go-sdk/pkg/binding/binding_test.go index beeccb5affc25..b9abd6ab422f7 100644 --- a/go-sdk/pkg/binding/binding_test.go +++ b/go-sdk/pkg/binding/binding_test.go @@ -288,7 +288,7 @@ type extractResult struct { } // simpleTaskInput is the minimal TaskInput struct: one field, no tags, so it -// falls back to matching its own (lowercased) field name. +// falls back to matching its own field name, verbatim. type simpleTaskInput struct { sdk.TaskInput Name string @@ -310,7 +310,7 @@ type nonEmbeddingStruct struct { } // combineInput exercises both TaskInput field-binding modes side by side: -// Name falls back to its snake_cased field name, Count is explicitly named +// Name falls back to its verbatim field name, Count is explicitly named // via its `arg:` tag. type combineInput struct { sdk.TaskInput @@ -467,7 +467,7 @@ func (s *BindingSuite) TestAnalyzeTaskInputValidation() { type duplicateArgNames struct { sdk.TaskInput A string - B string `arg:"a"` + B string `arg:"A"` } type nonDecodableField struct { sdk.TaskInput @@ -480,7 +480,7 @@ func (s *BindingSuite) TestAnalyzeTaskInputValidation() { }{ "duplicate-arg-names": { func(input duplicateArgNames) error { return nil }, - `fields A and B both bind arg name "a"`, + `fields A and B both bind arg name "A"`, }, "non-decodable-field": { func(input nonDecodableField) error { return nil }, @@ -504,13 +504,13 @@ func (s *BindingSuite) TestAnalyzeTaskInputValidation() { func (s *BindingSuite) TestResolveTaskInputAllStruct() { fn := func(input combineInput) error { return nil } got, err := s.resolve(fn, []Arg{ - LiteralArg{Name: "name", Value: "widget", DataType: DataTypeString}, + LiteralArg{Name: "Name", Value: "widget", DataType: DataTypeString}, LiteralArg{Name: "count", Value: 7, DataType: DataTypeInteger}, }, &fakeXComClient{}) s.Require().NoError(err) input := got[0].Interface().(combineInput) - s.Equal("widget", input.Name, "the untagged field claims its snake_cased field name") + s.Equal("widget", input.Name, "the untagged field claims its verbatim field name") s.Equal(7, input.Count, "the `arg:` tag claims its named entry") } @@ -519,7 +519,7 @@ func (s *BindingSuite) TestResolveTaskInputMixedWithFlat() { got, err := s.resolve(fn, []Arg{ LiteralArg{Name: "prefix", Value: "head", DataType: DataTypeString}, XComArg{Name: "region", TaskID: "make_region", DataType: DataTypeString}, - LiteralArg{Name: "ratio", Value: 0.5, DataType: DataTypeNumber}, + LiteralArg{Name: "Ratio", Value: 0.5, DataType: DataTypeNumber}, LiteralArg{Name: "suffix", Value: "footer", DataType: DataTypeString}, }, &fakeXComClient{values: map[string]any{"make_region/return_value": "east"}}) s.Require().NoError(err) @@ -542,7 +542,7 @@ func (s *BindingSuite) TestResolveTaskInputMixedWithFlat() { func (s *BindingSuite) TestResolveTaskInputLiteralThroughArgName() { fn := func(input simpleTaskInput) error { return nil } got, err := s.resolve(fn, []Arg{ - LiteralArg{Name: "name", Value: "widget", DataType: DataTypeString}, + LiteralArg{Name: "Name", Value: "widget", DataType: DataTypeString}, }, &fakeXComClient{}) s.Require().NoError(err) s.Equal("widget", got[0].Interface().(simpleTaskInput).Name) @@ -551,7 +551,7 @@ func (s *BindingSuite) TestResolveTaskInputLiteralThroughArgName() { func (s *BindingSuite) TestResolveTaskInputPointerStruct() { fn := func(input *simpleTaskInput) error { return nil } got, err := s.resolve(fn, []Arg{ - LiteralArg{Name: "name", Value: "widget", DataType: DataTypeString}, + LiteralArg{Name: "Name", Value: "widget", DataType: DataTypeString}, }, &fakeXComClient{}) s.Require().NoError(err) input := got[0].Interface().(*simpleTaskInput) @@ -575,7 +575,7 @@ func (s *BindingSuite) TestResolveTaskInputUnmatchedArgNameLeavesFieldZeroValued func (s *BindingSuite) TestResolveTaskInputUnmatchedArgNameZeroValuedAlongsideMatch() { fn := func(input twoFieldTaskInput) error { return nil } got, err := s.resolve(fn, []Arg{ - LiteralArg{Name: "name", Value: "widget", DataType: DataTypeString}, + LiteralArg{Name: "Name", Value: "widget", DataType: DataTypeString}, }, &fakeXComClient{}) s.Require().NoError(err) input := got[0].Interface().(twoFieldTaskInput) @@ -586,7 +586,7 @@ func (s *BindingSuite) TestResolveTaskInputUnmatchedArgNameZeroValuedAlongsideMa func (s *BindingSuite) TestResolveTaskInputArityMismatchForLeftoverArgs() { fn := func(input simpleTaskInput, extra string) error { return nil } _, err := s.resolve(fn, []Arg{ - LiteralArg{Name: "name", Value: "widget", DataType: DataTypeString}, + LiteralArg{Name: "Name", Value: "widget", DataType: DataTypeString}, }, &fakeXComClient{}) if s.Assert().Error(err) { s.Contains(err.Error(), "argument count mismatch") diff --git a/go-sdk/sdk/context.go b/go-sdk/sdk/context.go index 6050ed52dae2b..0db323b15b602 100644 --- a/go-sdk/sdk/context.go +++ b/go-sdk/sdk/context.go @@ -113,8 +113,9 @@ type DagRun struct { // that struct into per-field, name-based TaskFlow argument binding -- an // ergonomic alternative to a long flat parameter list. Each exported field of // such a struct may carry an `arg:""` tag naming the stub's TaskFlow -// argument to bind, falling back to the field's own name, snake_cased, when -// omitted: +// argument to bind, falling back to the field's own name, verbatim, when +// omitted -- so an untagged Name binds the argument "Name", and a snake_case +// argument like "count" needs an explicit tag: // // type CombineInput struct { // sdk.TaskInput From b68f0bb9ab507eb19bcf884e3a10d51371f86419 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 21 Jul 2026 03:33:59 +0000 Subject: [PATCH 13/40] Build the arg-bindings TypeAdapter lazily in the execution API Most workloads are not stub operators, so constructing the discriminated-union adapter at module import made every execution API process pay for it up front. Moving it next to the TaskArgBinding models behind a cached getter defers the cost to the first stub-task run and leaves the _STUB_TASK_TYPE gate as the only stub-specific module-level state in the route. --- .../execution_api/datamodels/task_arg_binding.py | 15 ++++++++++++++- .../execution_api/routes/task_instances.py | 10 +++------- .../versions/head/test_task_instances.py | 6 ++++-- 3 files changed, 21 insertions(+), 10 deletions(-) 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 index 87aec3b829db5..24dced0762e33 100644 --- 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 @@ -31,9 +31,10 @@ from __future__ import annotations from enum import Enum +from functools import cache from typing import Annotated, Literal -from pydantic import Field, JsonValue +from pydantic import Field, JsonValue, TypeAdapter from typing_extensions import TypeAliasType from airflow.api_fastapi.core_api.base import BaseModel @@ -97,3 +98,15 @@ class LiteralArgBinding(BaseModel): Annotated[XComArgBinding | LiteralArgBinding, Field(discriminator="kind", title="TaskArgBinding")], ) """One positional argument of a stub (foreign-runtime) task, in declaration order.""" + + +@cache +def get_arg_bindings_adapter() -> TypeAdapter[list[TaskArgBinding]]: + """ + Validate serialized arg-binding dicts into the kind-discriminated ``TaskArgBinding`` union. + + Constructed lazily on first use (then cached): only the stub-task path in the + execution API ever needs the adapter, and most workloads are not stub operators, + so regular task runs never pay for building it. + """ + return TypeAdapter(list[TaskArgBinding]) 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 7ddbdf6dcd5ce..d49578f1bc483 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -32,7 +32,7 @@ from opentelemetry import trace from opentelemetry.trace import StatusCode from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator -from pydantic import JsonValue, TypeAdapter +from pydantic import JsonValue from sqlalchemy import and_, func, or_, tuple_, update from sqlalchemy.engine import CursorResult from sqlalchemy.exc import DataError, NoResultFound, SQLAlchemyError @@ -49,7 +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.task_arg_binding import get_arg_bindings_adapter from airflow.api_fastapi.execution_api.datamodels.taskinstance import ( InactiveAssetsResponse, PreviousTIResponse, @@ -116,10 +116,6 @@ # serialized-dag lookup for ``arg_bindings`` so regular tasks never pay for it. _STUB_TASK_TYPE = "_StubOperator" -# Validates the serialized-dag arg-binding dicts into the kind-discriminated -# TaskArgBinding union; built once at import, not per request. -_arg_bindings_adapter: TypeAdapter[list[TaskArgBinding]] = TypeAdapter(list[TaskArgBinding]) - 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.""" @@ -352,7 +348,7 @@ def ti_run( if ti.operator == _STUB_TASK_TYPE and ( arg_bindings := _get_arg_bindings(ti.dag_version_id, ti.task_id, session=session) ): - context.arg_bindings = _arg_bindings_adapter.validate_python(arg_bindings) + context.arg_bindings = get_arg_bindings_adapter().validate_python(arg_bindings) # Only set if they are non-null if ti.next_method: 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 cdff621c7d2aa..c88fccf0f71e1 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 @@ -417,10 +417,12 @@ def transform(country: str, extracted: dict): ... def test_arg_bindings_adapter_rejects_unknown_kind(self): """The discriminated union refuses serialized specs with an unrecognised kind.""" - from airflow.api_fastapi.execution_api.routes.task_instances import _arg_bindings_adapter + from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import get_arg_bindings_adapter with pytest.raises(ValidationError, match="does not match any of the expected tags"): - _arg_bindings_adapter.validate_python([{"name": "country", "kind": "template", "value": "x"}]) + get_arg_bindings_adapter().validate_python( + [{"name": "country", "kind": "template", "value": "x"}] + ) def test_dynamic_task_mapping_with_parse_time_value(self, client, dag_maker): """Test that dynamic task mapping works correctly with parse-time values.""" From 5cef708a079b5f1be3aa354c0a1bae1869db9a71 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 21 Jul 2026 07:03:20 +0000 Subject: [PATCH 14/40] Tighten the stub-task arg-binding contract after self-review The XCom key was always return_value for a TaskFlow call, so the key field carried no information; it is removed end to end (datamodel, serialized spec, supervisor schema, generated task-sdk and Go models) and indexing a stub argument by a custom key now fails at parse time instead of being silently representable. Mixing flat positional data parameters with a TaskInput struct in one Go task signature was too ambiguous to reason about, so Analyze now rejects it: a function declares one binding shape or the other. The Go binding sum type and its DataType vocabulary are now defined in terms of the generated supervisor-schema models rather than hand-written mirrors, so they cannot drift from the wire contract, and every via_struct_* example task now binds an XCom-sourced argument (make_region) alongside a literal so struct-field binding is exercised with both sources end to end. The TypeScript supervisor model bump is left out of this PR on purpose. --- .../datamodels/task_arg_binding.py | 15 +- .../execution_api/routes/task_instances.py | 3 +- .../execution_api/versions/v2026_06_30.py | 10 +- .../versions/head/test_task_instances.py | 8 +- .../v2026_04_17/test_task_instances.py | 8 +- .../serialization/test_dag_serialization.py | 16 +- .../test_go_sdk_taskflow_binding.py | 12 +- go-sdk/README.md | 11 +- go-sdk/dags/go_examples.py | 16 +- go-sdk/example/bundle/main.go | 1 + .../bundle/taskflowbinding/taskflowbinding.go | 28 ++- go-sdk/pkg/binding/binding.go | 128 ++++++----- go-sdk/pkg/binding/binding_test.go | 64 +++--- .../pkg/execution/genmodels/defaults.gen.go | 2 +- go-sdk/pkg/execution/genmodels/models.gen.go | 3 - go-sdk/pkg/execution/task_runner.go | 8 +- .../providers/standard/decorators/stub.py | 12 +- .../unit/standard/decorators/test_stub.py | 24 +- .../airflow/sdk/api/datamodels/_generated.py | 1 - .../sdk/execution_time/schema/schema.json | 5 - .../execution_time/schema/test_migrator.py | 8 +- ts-sdk/src/generated/supervisor.ts | 212 ++++++------------ 22 files changed, 244 insertions(+), 351 deletions(-) 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 index 24dced0762e33..1171a5fc6f1ce 100644 --- 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 @@ -20,12 +20,6 @@ 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. - -Each binding is one variant of a union discriminated on ``kind``: an ``XComArgBinding`` -pulls the value from an upstream task's XCom, a ``LiteralArgBinding`` carries an inline -value from the Dag file. Both variants still emit plain named structs -(``$defs/XComArgBinding``, ``$defs/LiteralArgBinding``) for the foreign-language SDKs -consuming the supervisor schema. """ from __future__ import annotations @@ -55,9 +49,6 @@ class ArgBindingDataType(str, Enum): class XComArgBinding(BaseModel): """One positional stub-task argument pulled from an upstream task's XCom.""" - # No default on purpose: a required ``kind`` stays non-nullable through the OpenAPI - # round trip, which discriminated-union codegen needs (a defaulted field turns - # ``Literal`` into ``Literal | None`` in the generated task-sdk models). kind: Literal["xcom"] name: str @@ -67,17 +58,13 @@ class XComArgBinding(BaseModel): """Declared type from the stub function's annotation; runtimes type-check against it.""" task_id: str - """Upstream task id to pull the XCom from.""" - - key: str = "return_value" - """XCom key to pull.""" + """Upstream task id to pull the XCom from; the ``return_value`` XCom is always the one pulled.""" class LiteralArgBinding(BaseModel): """One positional stub-task argument carrying an inline literal from the Dag file.""" kind: Literal["literal"] - """Required like ``XComArgBinding.kind``; see the note there.""" name: str """The stub function's parameter name this binding fills, in declaration order.""" 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 d49578f1bc483..eb27a0a214a80 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 @@ -130,8 +130,7 @@ def _get_arg_bindings(dag_version_id: UUID | None, task_id: str, *, session) -> 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: + if not (data := dag_version.serialized_dag.data): return None for task in data.get("dag", {}).get("tasks", []): var = task.get(Encoding.VAR) or {} 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 b9e5a7b8206d9..15a097fac5f70 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 @@ -148,15 +148,7 @@ def remove_partition_date_from_dag_run(response: ResponseInfo) -> None: # type: class AddArgBindingsToTIRunContext(VersionChange): - """ - Add the ``arg_bindings`` positional-argument binding spec for stub (foreign-runtime) tasks. - - Each entry is a discriminated union of ``XComArgBinding`` and ``LiteralArgBinding`` keyed - on ``kind``. ``data_type`` is declared as the ``ArgBindingDataType`` enum rather than an - inline ``Literal``; the wire representation (a JSON string) is unchanged, so no migration - instruction is needed -- this version has not been released with the ``arg_bindings`` field - in any other shape. - """ + """Add the ``arg_bindings`` positional-argument binding spec for stub (foreign-runtime) tasks.""" description = __doc__ 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 c88fccf0f71e1..85d25ea2a5a0b 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 @@ -401,13 +401,7 @@ def transform(country: str, extracted: dict): ... assert response.status_code == 200 assert response.json()["arg_bindings"] == [ {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, - { - "name": "extracted", - "kind": "xcom", - "data_type": "object", - "task_id": "extract", - "key": "return_value", - }, + {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract"}, ] # An argless stub has no captured spec, so the field stays unset. 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 2b55a6e2ae3f4..f88d22ff10e1c 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 @@ -127,11 +127,5 @@ def test_head_version_includes_arg_bindings(self, client, stub_ti): assert response.status_code == 200 assert response.json()["arg_bindings"] == [ {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, - { - "name": "extracted", - "kind": "xcom", - "data_type": "object", - "task_id": "extract", - "key": "return_value", - }, + {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract"}, ] diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index 6e23878691a8e..15b98987be7a4 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -3429,26 +3429,14 @@ def transform(country: str, extracted: dict): ... }, { Encoding.TYPE: DAT.DICT, - Encoding.VAR: { - "name": "extracted", - "kind": "xcom", - "data_type": "object", - "task_id": "extract", - "key": "return_value", - }, + Encoding.VAR: {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract"}, }, ] round_tripped = DagSerialization.from_dict(ser_dag) assert round_tripped.task_dict["transform"]._arg_bindings == [ {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, - { - "name": "extracted", - "kind": "xcom", - "data_type": "object", - "task_id": "extract", - "key": "return_value", - }, + {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract"}, ] assert not hasattr(round_tripped.task_dict["extract"], "_arg_bindings") or ( round_tripped.task_dict["extract"]._arg_bindings is None 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 index 1dd793fbfa91b..f6f8033a23118 100644 --- 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 @@ -75,6 +75,7 @@ def test_all_tasks_succeeded(completed_run: _CompletedRun): for task_id in ( "make_config", "make_numbers", + "make_region", "via_flat_args", "via_struct_no_tags", "via_struct_arg_tag", @@ -84,13 +85,14 @@ def test_all_tasks_succeeded(completed_run: _CompletedRun): def test_upstream_xcoms_keep_their_shapes(completed_run: _CompletedRun): - """The Go struct arrives as an object XCom and the ``[]int`` as an array.""" + """The Go struct arrives as an object XCom, the ``[]int`` as an array, the region as a string.""" assert completed_run.xcom("make_config") == { "environment": "production", "region": "eu-west-1", "debug": True, } assert completed_run.xcom("make_numbers") == [1, 1, 2, 3, 5, 8] + assert completed_run.xcom("make_region") == "eu-west-1" def test_via_flat_args_summary_reflects_bound_arguments(completed_run: _CompletedRun): @@ -112,7 +114,8 @@ def test_via_flat_args_summary_reflects_bound_arguments(completed_run: _Complete def test_via_struct_no_tags_reflects_bound_arguments(completed_run: _CompletedRun): """``via_struct_no_tags`` demonstrates the Go SDK's ``sdk.TaskInput`` struct-field injection mode with no field tags at all: each field binds the TaskFlow argument - spelled exactly like its Go field name (``RegionCode``, ``Threshold``).""" + spelled exactly like its Go field name (``RegionCode``, ``Threshold``). The region + is ``make_region``'s XCom, so a struct field binds an XCom-sourced value here.""" assert completed_run.xcom("via_struct_no_tags") == { "region_code": "eu-west-1", "threshold": 0.75, @@ -121,8 +124,9 @@ def test_via_struct_no_tags_reflects_bound_arguments(completed_run: _CompletedRu def test_via_struct_arg_tag_reflects_bound_arguments(completed_run: _CompletedRun): """``via_struct_arg_tag`` demonstrates explicit ``arg:`` tags: ``Region`` is - genuinely renamed to ``region_code``, and ``Threshold`` is tagged ``threshold`` - to pull the snake_case argument its verbatim field name would miss.""" + genuinely renamed to ``region_code`` (bound from ``make_region``'s XCom), and + ``Threshold`` is tagged ``threshold`` to pull the snake_case literal its + verbatim field name would miss.""" assert completed_run.xcom("via_struct_arg_tag") == { "region": "eu-west-1", "threshold": 0.75, diff --git a/go-sdk/README.md b/go-sdk/README.md index 9830de9758b3b..34c9a9dcc2475 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -144,8 +144,9 @@ source of truth for which `dag_id`s and `task_id`s a bundle can run. ### TaskInput structs A struct that anonymously embeds `sdk.TaskInput` opts into **per-field, name-based** binding instead -of a long flat parameter list — at most one such parameter is allowed per function, and it can be -mixed with plain flat parameters. +of a long flat parameter list. At most one such parameter is allowed per function, and it cannot be +combined with plain flat data parameters — a task function declares one shape or the other, and +registration fails on a signature that mixes them. Conceptually, a plain flat parameter list is **positional-argument** binding: order matters, and every parameter must be filled or the task fails before its body runs. A `TaskInput` struct is @@ -174,11 +175,7 @@ only binds an argument literally spelled `Threshold`, so a snake_case Python par explicit tag. If no TaskFlow call argument carries that name, the field is simply left at its Go zero value — it does not fail the task, kwarg-style (see `ViaStructUnmatchedArg` below). -When a `TaskInput` struct and plain flat parameters coexist in the same function, the struct's fields -claim entries out of the TaskFlow call's argument spec by name first; the *remaining, unclaimed* -entries are then distributed, in their original relative order, onto the flat parameters in -declaration order. With no `TaskInput` struct present, this is exactly today's positional-only -behaviour. A plain custom struct type *without* the `sdk.TaskInput` embed is unaffected by any of +A plain custom struct type *without* the `sdk.TaskInput` embed is unaffected by any of this — it keeps working as a single flat data parameter, JSON-decoded whole from one TaskFlow argument (see `Config` in [`example/bundle/taskflowbinding/taskflowbinding.go`](./example/bundle/taskflowbinding/taskflowbinding.go)), diff --git a/go-sdk/dags/go_examples.py b/go-sdk/dags/go_examples.py index dc9143bc16436..121a59f2cb9b3 100644 --- a/go-sdk/dags/go_examples.py +++ b/go-sdk/dags/go_examples.py @@ -129,6 +129,10 @@ def make_config(): ... def make_numbers(): ... +@task.stub(queue="golang") +def make_region(): ... + + @task.stub(queue="golang") def via_flat_args( name: str, @@ -178,7 +182,10 @@ def taskflow_binding_dag(): verifies every bound value and fails the task on any mismatch. Three further tasks demonstrate the Go SDK's ``sdk.TaskInput`` struct injection - mode, one field-binding mode at a time: + mode. Each call mixes a literal (``threshold``) with an XCom reference: the + ``region_code`` argument is ``make_region``'s output, so every struct example + also proves an XCom-sourced value binds onto a struct field. One field-binding + mode at a time: * ``via_struct_no_tags``: no ``arg:`` tags at all -- each struct field binds the argument spelled exactly like its Go field name, hence this stub's @@ -200,9 +207,10 @@ def taskflow_binding_dag(): config=make_config(), numbers=make_numbers(), ) - via_struct_no_tags(RegionCode="eu-west-1", Threshold=0.75) - via_struct_arg_tag(region_code="eu-west-1", threshold=0.75) - via_struct_unmatched_arg(region_code="eu-west-1") + region = make_region() + via_struct_no_tags(RegionCode=region, Threshold=0.75) + via_struct_arg_tag(region_code=region, threshold=0.75) + via_struct_unmatched_arg(region_code=region) taskflow_binding_dag() diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index 88917fcf183cc..f4d4696c4ed42 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -59,6 +59,7 @@ func (m *myBundle) RegisterDags(dagbag v1.Registry) error { bindingDag := dagbag.AddDag("taskflow_binding_dag") bindingDag.AddTaskWithName("make_config", taskflowbinding.MakeConfig) bindingDag.AddTaskWithName("make_numbers", taskflowbinding.MakeNumbers) + bindingDag.AddTaskWithName("make_region", taskflowbinding.MakeRegion) bindingDag.AddTaskWithName("via_flat_args", taskflowbinding.ViaFlatArgs) bindingDag.AddTaskWithName("via_struct_no_tags", taskflowbinding.ViaStructNoTags) bindingDag.AddTaskWithName("via_struct_arg_tag", taskflowbinding.ViaStructArgTag) diff --git a/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go b/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go index 7346c7160cb34..47c70fe365828 100644 --- a/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go +++ b/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go @@ -26,7 +26,9 @@ // is left at its zero value rather than failing the task -- one field-binding // mode at a time: ViaStructNoTags (verbatim field-name fallback), // ViaStructArgTag (explicit `arg:` naming), and ViaStructUnmatchedArg (a -// field whose name has no corresponding TaskFlow call argument at all). +// field whose name has no corresponding TaskFlow call argument at all). Each +// ViaStruct* call binds MakeRegion's XCom onto its region field alongside a +// literal, so struct fields are exercised with both argument sources. package taskflowbinding import ( @@ -69,6 +71,15 @@ func MakeNumbers(log *slog.Logger) (any, error) { return numbers, nil } +// MakeRegion pushes a string XCom that every ViaStruct* task binds onto a +// struct field, so each field-binding mode is exercised with an XCom-sourced +// argument and not just literals. +func MakeRegion(log *slog.Logger) (any, error) { + region := "eu-west-1" + log.Info("Pushing region", "region", region) + return region, nil +} + // ViaFlatArgs receives every argument shape the stub Dag can express as plain, // positional data parameters. The Python side calls it as // @@ -147,7 +158,9 @@ type ViaStructNoTagsInput struct { // ViaStructNoTags is called as // -// via_struct_no_tags(RegionCode="eu-west-1", Threshold=0.75) +// via_struct_no_tags(RegionCode=make_region(), Threshold=0.75) +// +// so RegionCode arrives via make_region's XCom and Threshold as a literal. func ViaStructNoTags( ctx sdk.TIRunContext, log *slog.Logger, @@ -185,7 +198,9 @@ type ViaStructArgTagInput struct { // ViaStructArgTag is called as // -// via_struct_arg_tag(region_code="eu-west-1", threshold=0.75) +// via_struct_arg_tag(region_code=make_region(), threshold=0.75) +// +// so Region arrives via make_region's XCom and Threshold as a literal. func ViaStructArgTag( ctx sdk.TIRunContext, log *slog.Logger, @@ -222,10 +237,11 @@ type ViaStructUnmatchedArgInput struct { // ViaStructUnmatchedArg is called as // -// via_struct_unmatched_arg(region_code="eu-west-1") +// via_struct_unmatched_arg(region_code=make_region()) // -// -- the stub only declares region_code, so Missing's arg name never appears -// among the call's arguments and stays at its Go zero value (""). +// -- the stub only declares region_code (bound from make_region's XCom), so +// Missing's arg name never appears among the call's arguments and stays at +// its Go zero value (""). func ViaStructUnmatchedArg( ctx sdk.TIRunContext, log *slog.Logger, input ViaStructUnmatchedArgInput, ) (any, error) { diff --git a/go-sdk/pkg/binding/binding.go b/go-sdk/pkg/binding/binding.go index 0e79af1156847..399d3ebb2b420 100644 --- a/go-sdk/pkg/binding/binding.go +++ b/go-sdk/pkg/binding/binding.go @@ -36,13 +36,11 @@ // the argument to claim, and a field with no tag claims its own Go field // name, verbatim. At most one such parameter is allowed per function. // -// A TaskInput struct's fields are resolved first, by name, claiming entries -// out of the argument spec; the remaining unclaimed entries are then -// distributed, in their original relative order, onto the plain flat data -// parameters in declaration order -- so flat parameters and a TaskInput -// struct can coexist in the same function signature regardless of where each -// sits, and (with no TaskInput struct present) this reduces to exactly -// today's positional-only behaviour. +// The two data-parameter shapes are mutually exclusive: a function declares +// either plain flat data parameters or one TaskInput struct, never both. +// Analyze rejects a signature that mixes them -- splitting one TaskFlow +// call's arguments between by-name claiming and positional order is too +// ambiguous to reason about. // // Conceptually, flat data parameters are positional-argument binding: order // matters, and every parameter must be filled or Resolve fails the task @@ -70,30 +68,35 @@ import ( "fmt" "log/slog" "reflect" + "strings" "github.com/apache/airflow/go-sdk/pkg/api" + "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" "github.com/apache/airflow/go-sdk/pkg/sdkcontext" "github.com/apache/airflow/go-sdk/sdk" ) // 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 +// argument (from the stub function's annotation on the Python side). It +// aliases the enum generated from the supervisor schema so the vocabulary +// cannot drift from the wire model. +type DataType = genmodels.ArgBindingDataType const ( - DataTypeString DataType = "string" - DataTypeInteger DataType = "integer" - DataTypeNumber DataType = "number" - DataTypeBoolean DataType = "boolean" - DataTypeObject DataType = "object" - DataTypeArray DataType = "array" - DataTypeAny DataType = "any" + DataTypeString = genmodels.ArgBindingDataTypeString + DataTypeInteger = genmodels.ArgBindingDataTypeInteger + DataTypeNumber = genmodels.ArgBindingDataTypeNumber + DataTypeBoolean = genmodels.ArgBindingDataTypeBoolean + DataTypeObject = genmodels.ArgBindingDataTypeObject + DataTypeArray = genmodels.ArgBindingDataTypeArray + DataTypeAny = genmodels.ArgBindingDataTypeAny ) // Arg is one positional argument for a task function's data parameters, in -// declaration order: an XComArg or a LiteralArg. A sealed sum type mirroring -// the wire model's XComArgBinding/LiteralArgBinding split, kept runtime-neutral -// so this package stays decoupled from the generated coordinator schema types. +// declaration order: an XComArg or a LiteralArg. A sealed sum type over the +// wire model's XComArgBinding/LiteralArgBinding split; each variant is +// defined in terms of its generated schema struct so the fields cannot drift +// from the coordinator protocol. type Arg interface { // ArgName is the stub function's parameter name this binding fills; used // to match a TaskInput struct field's `arg:` tag (or its verbatim @@ -107,27 +110,15 @@ type Arg interface { sealedArg() } -// XComArg sources the argument from an upstream task's XCom. -type XComArg struct { - // Name is the stub function's parameter name this binding fills. - Name string - // TaskID is the upstream task to pull from. - TaskID string - // Key is the XCom key to pull; empty means the return-value key. - Key string - // DataType is the declared type to check the Go parameter against. - DataType DataType -} +// XComArg sources the argument from an upstream task's return-value XCom. +// Kind is carried by the generated shape but unused here: the Go type itself +// is the discriminant. +type XComArg genmodels.XComArgBinding -// LiteralArg carries an inline value from the Dag file. -type LiteralArg struct { - // Name is the stub function's parameter name this binding fills. - Name string - // Value is the literal value from the Dag file. - Value any - // DataType is the declared type to check the Go parameter against. - DataType DataType -} +// LiteralArg carries an inline value from the Dag file. Kind is carried by +// the generated shape but unused here: the Go type itself is the +// discriminant. +type LiteralArg genmodels.LiteralArgBinding func (a XComArg) ArgName() string { return a.Name } func (a LiteralArg) ArgName() string { return a.Name } @@ -184,9 +175,10 @@ type paramPlan struct { // 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 + fnName string + params []paramPlan + numData int + hasTaskInput bool } // NumData returns how many data parameters the analyzed function declares. @@ -221,15 +213,24 @@ func Analyze(fnType reflect.Type, fnName string) (*Plan, error) { } p.params[i] = plan } + if seenTaskInput >= 0 { + p.hasTaskInput = true + if p.numData > 0 { + return nil, fmt.Errorf( + "task function %s: cannot mix a TaskInput struct parameter (parameter %d) with "+ + "plain data parameters; declare either flat data parameters or one TaskInput struct", + fnName, seenTaskInput, + ) + } + } return p, nil } // Resolve builds the ordered argument values for one call. Injectable // parameters receive values derived from ctx, logger, or client. A TaskInput -// struct's fields are resolved first, by name, claiming entries out of args; -// the remaining unclaimed entries are then distributed, in their original -// relative order, onto the plain flat data parameters in declaration order. -// An error fails the task before its body runs. +// struct's fields claim entries out of args by name; plain flat data +// parameters consume args in declaration order (the two shapes are mutually +// exclusive; see Analyze). An error fails the task before its body runs. func (p *Plan) Resolve( ctx context.Context, logger *slog.Logger, @@ -265,6 +266,24 @@ func (p *Plan) Resolve( } } if remaining != p.numData { + if p.hasTaskInput { + names := make([]string, 0, remaining) + for i, c := range claimed { + if c { + continue + } + name := "" + if args[i] != nil { + name = fmt.Sprintf("%q", args[i].ArgName()) + } + names = append(names, name) + } + return nil, fmt.Errorf( + "task function %s: %d TaskFlow call argument(s) not claimed by any TaskInput "+ + "field: %s", + p.fnName, remaining, strings.Join(names, ", "), + ) + } 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)", @@ -424,17 +443,16 @@ func (p *Plan) resolveOne( p.fnName, generalCtx, ) } - key := a.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, a.TaskID, nil, key, nil) + // Always the return-value XCom -- a stub Dag cannot reference any other + // key. Pull from the upstream's unmapped instance (map_index nil); + // mapped upstream fan-in is out of scope for now. + raw, err := c.GetXCom( + ctx, workload.TI.DagId, workload.TI.RunId, a.TaskID, nil, api.XComReturnValueKey, nil, + ) if err != nil { return reflect.Value{}, fmt.Errorf( - "task function %s: %s: pulling xcom from task %q (key %q): %w", - p.fnName, generalCtx, a.TaskID, key, err, + "task function %s: %s: pulling xcom from task %q: %w", + p.fnName, generalCtx, a.TaskID, err, ) } v, err := decodeValue(raw, targetType) diff --git a/go-sdk/pkg/binding/binding_test.go b/go-sdk/pkg/binding/binding_test.go index b9abd6ab422f7..b82324f19f1a3 100644 --- a/go-sdk/pkg/binding/binding_test.go +++ b/go-sdk/pkg/binding/binding_test.go @@ -327,20 +327,27 @@ type reportInput struct { Region string `arg:"region"` } +// mixedInput pairs a TaskInput struct with a plain data parameter in the +// functions that assert Analyze rejects that combination. +type mixedInput struct { + sdk.TaskInput + Name string +} + 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", + "probe/return_value": "probe-value", }} - fn := func(res extractResult, part string) error { return nil } + fn := func(res extractResult, probe string) error { return nil } got, err := s.resolve(fn, []Arg{ XComArg{TaskID: "extract", DataType: DataTypeObject}, - XComArg{TaskID: "extract", Key: "part", DataType: DataTypeString}, + XComArg{TaskID: "probe", 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.Equal("probe-value", got[1].Interface()) s.Require().Len(client.calls, 2) s.Equal("dag1", client.calls[0].dagID) @@ -349,10 +356,10 @@ func (s *BindingSuite) TestResolveXComArgs() { s.Equal( api.XComReturnValueKey, client.calls[0].key, - "an empty key must default to the return-value key", + "an XCom argument always pulls the return-value key", ) s.Nil(client.calls[0].mapIndex, "v1 always pulls the unmapped upstream instance") - s.Equal("part", client.calls[1].key) + s.Equal(api.XComReturnValueKey, client.calls[1].key) } func (s *BindingSuite) TestResolveXComStrictStructDecode() { @@ -490,6 +497,14 @@ func (s *BindingSuite) TestAnalyzeTaskInputValidation() { func(a simpleTaskInput, b simpleTaskInput) error { return nil }, "only one TaskInput struct parameter is allowed", }, + "mixed-with-flat-data-param": { + func(prefix string, input mixedInput) error { return nil }, + "cannot mix a TaskInput struct parameter", + }, + "mixed-with-trailing-flat-data-param": { + func(input mixedInput, suffix string) error { return nil }, + "cannot mix a TaskInput struct parameter", + }, } for name, tt := range cases { s.Run(name, func() { @@ -514,29 +529,17 @@ func (s *BindingSuite) TestResolveTaskInputAllStruct() { s.Equal(7, input.Count, "the `arg:` tag claims its named entry") } -func (s *BindingSuite) TestResolveTaskInputMixedWithFlat() { - fn := func(prefix string, input reportInput, suffix string) error { return nil } +func (s *BindingSuite) TestResolveTaskInputXComArg() { + fn := func(log *slog.Logger, input reportInput) error { return nil } got, err := s.resolve(fn, []Arg{ - LiteralArg{Name: "prefix", Value: "head", DataType: DataTypeString}, XComArg{Name: "region", TaskID: "make_region", DataType: DataTypeString}, LiteralArg{Name: "Ratio", Value: 0.5, DataType: DataTypeNumber}, - LiteralArg{Name: "suffix", Value: "footer", DataType: DataTypeString}, }, &fakeXComClient{values: map[string]any{"make_region/return_value": "east"}}) s.Require().NoError(err) - s.Equal( - "head", - got[0].Interface(), - "the leading flat parameter claims the first unclaimed wire entry", - ) input := got[1].Interface().(reportInput) s.Equal("east", input.Region, "Region resolves by name despite being declared after Ratio") s.Equal(0.5, input.Ratio) - s.Equal( - "footer", - got[2].Interface(), - "the trailing flat parameter claims the last unclaimed wire entry", - ) } func (s *BindingSuite) TestResolveTaskInputLiteralThroughArgName() { @@ -559,16 +562,15 @@ func (s *BindingSuite) TestResolveTaskInputPointerStruct() { s.Equal("widget", input.Name) } -func (s *BindingSuite) TestResolveTaskInputUnmatchedArgNameLeavesFieldZeroValued() { +func (s *BindingSuite) TestResolveTaskInputUnclaimedArgFailsLoudly() { fn := func(input simpleTaskInput) error { return nil } _, err := s.resolve(fn, []Arg{ LiteralArg{Name: "different_name", Value: "x", DataType: DataTypeString}, }, &fakeXComClient{}) - // simpleTaskInput has no other flat parameter to absorb "different_name", - // so the arity check (0 unclaimed args expected) still fails the task -- - // this asserts the unmatched TaskInput field itself is not what errors. + // No TaskInput field claims "different_name"; the leftover argument fails + // the task rather than being dropped silently. if s.Assert().Error(err) { - s.Contains(err.Error(), "argument count mismatch") + s.Contains(err.Error(), `not claimed by any TaskInput field: "different_name"`) } } @@ -582,15 +584,3 @@ func (s *BindingSuite) TestResolveTaskInputUnmatchedArgNameZeroValuedAlongsideMa s.Equal("widget", input.Name, "the matched field binds normally") s.Equal("", input.Missing, "the unmatched field is left at its Go zero value, not an error") } - -func (s *BindingSuite) TestResolveTaskInputArityMismatchForLeftoverArgs() { - fn := func(input simpleTaskInput, extra string) error { return nil } - _, err := s.resolve(fn, []Arg{ - LiteralArg{Name: "Name", Value: "widget", DataType: DataTypeString}, - }, &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)") - } -} diff --git a/go-sdk/pkg/execution/genmodels/defaults.gen.go b/go-sdk/pkg/execution/genmodels/defaults.gen.go index ba7f35b75431a..0b885e9d6e446 100644 --- a/go-sdk/pkg/execution/genmodels/defaults.gen.go +++ b/go-sdk/pkg/execution/genmodels/defaults.gen.go @@ -331,7 +331,7 @@ func (m *TriggerDagRun) DecodeMsgpack(dec *msgpack.Decoder) error { // DecodeMsgpack applies XComArgBinding's schema defaults that msgpack would otherwise skip. func (m *XComArgBinding) DecodeMsgpack(dec *msgpack.Decoder) error { type alias XComArgBinding - v := alias{DataType: ArgBindingDataType("any"), Key: "return_value"} + v := alias{DataType: ArgBindingDataType("any")} if err := dec.Decode(&v); err != nil { return err } diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index 38d305ea88ca7..33909abeecf4a 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -1887,9 +1887,6 @@ type XComArgBinding struct { // DataType corresponds to the JSON schema field "data_type". DataType ArgBindingDataType `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 string `msgpack:"kind"` diff --git a/go-sdk/pkg/execution/task_runner.go b/go-sdk/pkg/execution/task_runner.go index 0ec3d94b4d44e..3d18aae33157e 100644 --- a/go-sdk/pkg/execution/task_runner.go +++ b/go-sdk/pkg/execution/task_runner.go @@ -145,7 +145,7 @@ func RunTask( // the Python stub Dag's TaskFlow call) onto the runtime binding sum type. The // wire union generates untyped items (msgpack delivers each XComArgBinding / // LiteralArgBinding as a plain map), so the kind dispatch and the schema -// defaults (data_type "any", xcom key "return_value") are applied here. +// default (data_type "any") are applied here. func convertArgBindings(specsPtr *genmodels.ArgBindings) ([]binding.Arg, error) { if specsPtr == nil || len(*specsPtr) == 0 { return nil, nil @@ -165,11 +165,7 @@ func convertArgBindings(specsPtr *genmodels.ArgBindings) ([]binding.Arg, error) switch kind, _ := m["kind"].(string); kind { case "xcom": taskID, _ := m["task_id"].(string) - key := "return_value" - if s, ok := m["key"].(string); ok && s != "" { - key = s - } - args[i] = binding.XComArg{Name: name, TaskID: taskID, Key: key, DataType: dataType} + args[i] = binding.XComArg{Name: name, TaskID: taskID, DataType: dataType} case "literal": args[i] = binding.LiteralArg{Name: name, Value: m["value"], DataType: dataType} default: diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index 31268c2d17c8b..ddfd64fd31eb0 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -155,22 +155,26 @@ def annotation_for(name: str, param: inspect.Parameter) -> Any: value = bound.arguments[name] data_type = _data_type_from_annotation(annotation_for(name, param)) if isinstance(value, PlainXComArg): + if value.key != "return_value": + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} references the XCom key " + f"{value.key!r}; only an upstream task's return value can cross the language " + "boundary -- indexing an output by a custom key is not supported" + ) spec.append( { "name": name, "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" + f"{type(value).__name__}; only direct upstream task outputs can cross the " + "language boundary -- .map()/.zip()/.concat() results are not supported" ) try: json.dumps(value) diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index b8ba9a8ef20da..02bf3c7c58e54 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -103,13 +103,7 @@ def test_literal_and_xcom_spec(self): op = result.operator assert op._arg_bindings == [ {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, - { - "name": "extracted", - "kind": "xcom", - "data_type": "object", - "task_id": "fn_extract", - "key": "return_value", - }, + {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "fn_extract"}, {"name": "retries_num", "kind": "literal", "data_type": "integer", "value": 3}, ] assert op.upstream_task_ids == {"fn_extract"} @@ -117,20 +111,20 @@ def test_literal_and_xcom_spec(self): 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) + result = stub(fn_transform)(extracted=extracted, country="fr", retries_num=7) assert result.operator._arg_bindings == [ {"name": "country", "kind": "literal", "data_type": "string", "value": "fr"}, - { - "name": "extracted", - "kind": "xcom", - "data_type": "object", - "task_id": "fn_extract", - "key": "part", - }, + {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "fn_extract"}, {"name": "retries_num", "kind": "literal", "data_type": "integer", "value": 7}, ] + def test_custom_xcom_key_rejected(self): + with DAG(dag_id="d"): + extracted = stub(fn_extract)() + with pytest.raises(ValueError, match="indexing an output by a custom key"): + stub(fn_transform)("uk", extracted["part"]) + def test_zero_param_stub_has_no_spec(self): assert stub(fn_pass)().operator._arg_bindings is None diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index cc68b1486c79d..48b598d90915e 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -547,7 +547,6 @@ class XComArgBinding(BaseModel): name: Annotated[str, Field(title="Name")] data_type: ArgBindingDataType | None = ArgBindingDataType.ANY task_id: Annotated[str, Field(title="Task Id")] - key: Annotated[str | None, Field(title="Key")] = "return_value" class XComResponse(BaseModel): 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 c9bec70012a67..c1298dff6c743 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -4649,11 +4649,6 @@ "task_id": { "title": "Task Id", "type": "string" - }, - "key": { - "default": "return_value", - "title": "Key", - "type": "string" } }, "required": [ 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 8f803dcefb951..44d22a255e218 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 @@ -421,13 +421,7 @@ def startup_details(self): max_tries=1, arg_bindings=[ {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, - { - "name": "extracted", - "kind": "xcom", - "data_type": "object", - "task_id": "extract", - "key": "return_value", - }, + {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract"}, ], ), sentry_integration="", diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 748a4c6e20499..ab2632831ab95 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -22,20 +22,6 @@ // // Re-run with: pnpm run generate:supervisor -/** - * Language-neutral value type a stub-task argument binds to in the foreign runtime. - * - * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema - * via the `definition` "ArgBindingDataType". - */ -export type ArgBindingDataType = - | "string" - | "integer" - | "number" - | "boolean" - | "object" - | "array" - | "any"; export type Name = string; export type Id = number; export type Timestamp = string; @@ -259,40 +245,6 @@ export type NextKwargs1 = export type XcomKeysToClear = string[]; export type ShouldRetry = boolean; export type StartDate2 = string | null; -export type ArgBindings = TaskArgBinding[] | null; -/** - * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema - * via the `definition` "TaskArgBinding". - */ -export type TaskArgBinding = XComArgBinding | LiteralArgBinding; -export type Kind = "xcom"; -export type Name8 = string; -/** - * Language-neutral value type a stub-task argument binds to in the foreign runtime. - */ -export type ArgBindingDataType1 = - | "string" - | "integer" - | "number" - | "boolean" - | "object" - | "array" - | "any"; -export type TaskId1 = string; -export type Key1 = string; -export type Kind1 = "literal"; -export type Name9 = string; -/** - * Language-neutral value type a stub-task argument binds to in the foreign runtime. - */ -export type ArgBindingDataType2 = - | "string" - | "integer" - | "number" - | "boolean" - | "object" - | "array" - | "any"; export type Type13 = "TaskCallbackRequest"; export type Filepath2 = string; export type BundleName3 = string; @@ -358,21 +310,21 @@ export type NextKwargs2 = { } | null; export type RenderedMapIndex1 = string | null; export type Type20 = "DeferTask"; -export type Name10 = string; -export type Key2 = string; +export type Name8 = string; +export type Key1 = string; export type Type21 = "DeleteAssetStateStoreByName"; export type Uri5 = string; -export type Key3 = string; +export type Key2 = string; export type Type22 = "DeleteAssetStateStoreByUri"; export type TiId2 = string; -export type Key4 = string; +export type Key3 = string; export type Type23 = "DeleteTaskStateStore"; -export type Key5 = string; +export type Key4 = string; export type Type24 = "DeleteVariable"; -export type Key6 = string; +export type Key5 = string; export type DagId6 = string; export type RunId5 = string; -export type TaskId2 = string; +export type TaskId1 = string; export type MapIndex1 = number | null; export type Type25 = "DeleteXCom"; /** @@ -410,11 +362,11 @@ export type ErrorType1 = | "PERMISSION_DENIED" | "GENERIC_ERROR" | "API_SERVER_ERROR"; -export type Name11 = string; +export type Name9 = string; export type Type27 = "GetAssetByName"; export type Uri6 = string; export type Type28 = "GetAssetByUri"; -export type Name12 = string | null; +export type Name10 = string | null; export type Uri7 = string | null; export type After = string | null; export type Before = string | null; @@ -437,11 +389,11 @@ export type Extra8 = { [k: string]: string; } | null; export type Type30 = "GetAssetEventByAssetAlias"; -export type Name13 = string; -export type Key7 = string; +export type Name11 = string; +export type Key6 = string; export type Type31 = "GetAssetStateStoreByName"; export type Uri8 = string; -export type Key8 = string; +export type Key7 = string; export type Type32 = "GetAssetStateStoreByUri"; export type AliasName1 = string; export type Type33 = "GetAssetsByAlias"; @@ -469,7 +421,7 @@ export type LogicalDate3 = string; export type State3 = string | null; export type Type41 = "GetPreviousDagRun"; export type DagId12 = string; -export type TaskId3 = string; +export type TaskId2 = string; export type LogicalDate4 = string | null; export type MapIndex2 = number; export type Type42 = "GetPreviousTI"; @@ -488,7 +440,7 @@ export type TiId5 = string; export type TryNumber1 = number; export type Type45 = "GetTaskRescheduleStartDate"; export type TiId6 = string; -export type Key9 = string; +export type Key8 = string; export type Type46 = "GetTaskStateStore"; export type DagId15 = string; export type MapIndex4 = number | null; @@ -497,34 +449,34 @@ export type TaskGroupId1 = string | null; export type LogicalDates2 = string[] | null; export type RunIds2 = string[] | null; export type Type47 = "GetTaskStates"; -export type Key10 = string; +export type Key9 = string; export type Type48 = "GetVariable"; export type Prefix = string | null; export type Limit2 = number; export type Offset = number; export type Type49 = "GetVariableKeys"; -export type Key11 = string; +export type Key10 = string; export type DagId16 = string; export type RunId9 = string; -export type TaskId4 = string; +export type TaskId3 = string; export type MapIndex5 = number | null; export type IncludePriorDates = boolean; export type Type50 = "GetXCom"; -export type Key12 = string; +export type Key11 = string; export type DagId17 = string; export type RunId10 = string; -export type TaskId5 = string; +export type TaskId4 = string; export type Type51 = "GetXComCount"; -export type Key13 = string; +export type Key12 = string; export type DagId18 = string; export type RunId11 = string; -export type TaskId6 = string; +export type TaskId5 = string; export type Offset1 = number; export type Type52 = "GetXComSequenceItem"; -export type Key14 = string; +export type Key13 = string; export type DagId19 = string; export type RunId12 = string; -export type TaskId7 = string; +export type TaskId6 = string; export type Start = number | null; export type Stop = number | null; export type Step = number | null; @@ -546,7 +498,7 @@ export type AssignedUsers1 = HITLUser[] | null; export type Type54 = "HITLDetailRequestResult"; export type InactiveAssets = AssetProfile[] | null; export type Type55 = "InactiveAssetsResult"; -export type Name14 = string | null; +export type Name12 = string | null; export type Type56 = "MaskSecret"; export type Ok = boolean; export type Type57 = "OKResponse"; @@ -556,7 +508,7 @@ export type StartDate4 = string | null; export type EndDate3 = string | null; export type Type58 = "PrevSuccessfulDagRunResult"; export type Type59 = "PreviousDagRunResult"; -export type TaskId8 = string; +export type TaskId7 = string; export type DagId20 = string; export type RunId13 = string; export type LogicalDate5 = string | null; @@ -567,7 +519,7 @@ export type TryNumber2 = number; export type MapIndex6 = number | null; export type Duration = number | null; export type Type60 = "PreviousTIResult"; -export type Key15 = string; +export type Key14 = string; export type Value1 = string | null; export type Description = string | null; export type Type61 = "PutVariable"; @@ -584,23 +536,23 @@ export type RetryReason = string | null; export type Type64 = "RetryTask"; export type Type65 = "SentFDs"; export type Fds = number[]; -export type Name15 = string; -export type Key16 = string; +export type Name13 = string; +export type Key15 = string; export type Type66 = "SetAssetStateStoreByName"; export type Uri9 = string; -export type Key17 = string; +export type Key16 = string; export type Type67 = "SetAssetStateStoreByUri"; export type Type68 = "SetRenderedFields"; export type RenderedMapIndex3 = string; export type Type69 = "SetRenderedMapIndex"; export type TiId8 = string; -export type Key18 = string; +export type Key17 = string; export type ExpiresAt = string | null; export type Type70 = "SetTaskStateStore"; -export type Key19 = string; +export type Key18 = string; export type DagId21 = string; export type RunId14 = string; -export type TaskId9 = string; +export type TaskId8 = string; export type MapIndex7 = number | null; export type DagResult1 = boolean; export type MappedLength = number | null; @@ -660,12 +612,12 @@ export type Type83 = "ValidateInletsAndOutlets"; export type Keys = string[]; export type TotalEntries = number; export type Type84 = "VariableKeysResult"; -export type Key20 = string; +export type Key19 = string; export type Value2 = string | null; export type Type85 = "VariableResult"; export type Len = number; export type Type86 = "XComCountResponse"; -export type Key21 = string; +export type Key20 = string; export type Type87 = "XComResult"; export type Type88 = "XComSequenceIndexResult"; export type Root = JsonValue[]; @@ -1051,7 +1003,6 @@ 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. @@ -1079,31 +1030,6 @@ export interface ConnectionResponse { port: Port1; extra: Extra6; } -/** - * One positional stub-task argument pulled from an upstream task's XCom. - * - * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema - * via the `definition` "XComArgBinding". - */ -export interface XComArgBinding { - kind: Kind; - name: Name8; - data_type?: ArgBindingDataType1; - task_id: TaskId1; - key?: Key1; -} -/** - * One positional stub-task argument carrying an inline literal from the Dag file. - * - * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema - * via the `definition` "LiteralArgBinding". - */ -export interface LiteralArgBinding { - kind: Kind1; - name: Name9; - data_type?: ArgBindingDataType2; - value?: unknown; -} /** * Email notification request for task failures/retries. * @@ -1223,8 +1149,8 @@ export interface DeferTask { * via the `definition` "DeleteAssetStateStoreByName". */ export interface DeleteAssetStateStoreByName { - name: Name10; - key: Key2; + name: Name8; + key: Key1; type?: Type21; } /** @@ -1233,7 +1159,7 @@ export interface DeleteAssetStateStoreByName { */ export interface DeleteAssetStateStoreByUri { uri: Uri5; - key: Key3; + key: Key2; type?: Type22; } /** @@ -1242,7 +1168,7 @@ export interface DeleteAssetStateStoreByUri { */ export interface DeleteTaskStateStore { ti_id: TiId2; - key: Key4; + key: Key3; type?: Type23; } /** @@ -1250,7 +1176,7 @@ export interface DeleteTaskStateStore { * via the `definition` "DeleteVariable". */ export interface DeleteVariable { - key: Key5; + key: Key4; type?: Type24; } /** @@ -1258,10 +1184,10 @@ export interface DeleteVariable { * via the `definition` "DeleteXCom". */ export interface DeleteXCom { - key: Key6; + key: Key5; dag_id: DagId6; run_id: RunId5; - task_id: TaskId2; + task_id: TaskId1; map_index?: MapIndex1; type?: Type25; } @@ -1279,7 +1205,7 @@ export interface ErrorResponse { * via the `definition` "GetAssetByName". */ export interface GetAssetByName { - name: Name11; + name: Name9; type?: Type27; } /** @@ -1295,7 +1221,7 @@ export interface GetAssetByUri { * via the `definition` "GetAssetEventByAsset". */ export interface GetAssetEventByAsset { - name: Name12; + name: Name10; uri: Uri7; after?: After; before?: Before; @@ -1326,8 +1252,8 @@ export interface GetAssetEventByAssetAlias { * via the `definition` "GetAssetStateStoreByName". */ export interface GetAssetStateStoreByName { - name: Name13; - key: Key7; + name: Name11; + key: Key6; type?: Type31; } /** @@ -1336,7 +1262,7 @@ export interface GetAssetStateStoreByName { */ export interface GetAssetStateStoreByUri { uri: Uri8; - key: Key8; + key: Key7; type?: Type32; } /** @@ -1428,7 +1354,7 @@ export interface GetPreviousDagRun { */ export interface GetPreviousTI { dag_id: DagId12; - task_id: TaskId3; + task_id: TaskId2; logical_date?: LogicalDate4; map_index?: MapIndex2; state?: TaskInstanceState | null; @@ -1472,7 +1398,7 @@ export interface GetTaskRescheduleStartDate { */ export interface GetTaskStateStore { ti_id: TiId6; - key: Key9; + key: Key8; type?: Type46; } /** @@ -1493,7 +1419,7 @@ export interface GetTaskStates { * via the `definition` "GetVariable". */ export interface GetVariable { - key: Key10; + key: Key9; type?: Type48; } /** @@ -1511,10 +1437,10 @@ export interface GetVariableKeys { * via the `definition` "GetXCom". */ export interface GetXCom { - key: Key11; + key: Key10; dag_id: DagId16; run_id: RunId9; - task_id: TaskId4; + task_id: TaskId3; map_index?: MapIndex5; include_prior_dates?: IncludePriorDates; type?: Type50; @@ -1526,10 +1452,10 @@ export interface GetXCom { * via the `definition` "GetXComCount". */ export interface GetXComCount { - key: Key12; + key: Key11; dag_id: DagId17; run_id: RunId10; - task_id: TaskId5; + task_id: TaskId4; type?: Type51; } /** @@ -1537,10 +1463,10 @@ export interface GetXComCount { * via the `definition` "GetXComSequenceItem". */ export interface GetXComSequenceItem { - key: Key13; + key: Key12; dag_id: DagId18; run_id: RunId11; - task_id: TaskId6; + task_id: TaskId5; offset: Offset1; type?: Type52; } @@ -1549,10 +1475,10 @@ export interface GetXComSequenceItem { * via the `definition` "GetXComSequenceSlice". */ export interface GetXComSequenceSlice { - key: Key14; + key: Key13; dag_id: DagId19; run_id: RunId12; - task_id: TaskId7; + task_id: TaskId6; start: Start; stop: Stop; step: Step; @@ -1594,7 +1520,7 @@ export interface InactiveAssetsResult { */ export interface MaskSecret { value: JsonValue; - name?: Name14; + name?: Name12; type?: Type56; } /** @@ -1633,7 +1559,7 @@ export interface PreviousDagRunResult { * via the `definition` "PreviousTIResponse". */ export interface PreviousTIResponse { - task_id: TaskId8; + task_id: TaskId7; dag_id: DagId20; run_id: RunId13; logical_date?: LogicalDate5; @@ -1659,7 +1585,7 @@ export interface PreviousTIResult { * via the `definition` "PutVariable". */ export interface PutVariable { - key: Key15; + key: Key14; value: Value1; description: Description; type?: Type61; @@ -1710,8 +1636,8 @@ export interface SentFDs { * via the `definition` "SetAssetStateStoreByName". */ export interface SetAssetStateStoreByName { - name: Name15; - key: Key16; + name: Name13; + key: Key15; value: JsonValue; type?: Type66; } @@ -1721,7 +1647,7 @@ export interface SetAssetStateStoreByName { */ export interface SetAssetStateStoreByUri { uri: Uri9; - key: Key17; + key: Key16; value: JsonValue; type?: Type67; } @@ -1754,7 +1680,7 @@ export interface SetRenderedMapIndex { */ export interface SetTaskStateStore { ti_id: TiId8; - key: Key18; + key: Key17; value: JsonValue; expires_at: ExpiresAt; type?: Type70; @@ -1764,11 +1690,11 @@ export interface SetTaskStateStore { * via the `definition` "SetXCom". */ export interface SetXCom { - key: Key19; + key: Key18; value: JsonValue; dag_id: DagId21; run_id: RunId14; - task_id: TaskId9; + task_id: TaskId8; map_index?: MapIndex7; dag_result?: DagResult1; mapped_length?: MappedLength; @@ -1925,7 +1851,7 @@ export interface VariableKeysResult { * via the `definition` "VariableResult". */ export interface VariableResult { - key: Key20; + key: Key19; value?: Value2; type?: Type85; } @@ -1944,7 +1870,7 @@ export interface XComCountResponse { * via the `definition` "XComResult". */ export interface XComResult { - key: Key21; + key: Key20; value: JsonValue; type?: Type87; } @@ -1970,4 +1896,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-07-30" as const; +export const SUPERVISOR_API_VERSION = "2026-06-16" as const; From 7e2d0f229dd96fc98e997da99619652f36e2eabd Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 21 Jul 2026 09:26:10 +0000 Subject: [PATCH 15/40] Ship arg_bindings in a new 2026-07-30 execution API version The 2026-06-30 execution API version already shipped in Airflow 3.3.0, so appending the arg_bindings migration to it would mutate a released version, which the execution API versioning policy forbids; the change now opens version 2026-07-30, matching the supervisor-schema date. The rest addresses a local multi-reviewer audit of the branch: - The airflow-go-pack integration test's expected manifest was missing the make_region task added to the example bundle, failing go test. - The per-field cadwyn didnt_exist instructions on XComArgBinding and LiteralArgBinding name fields could never apply (arg_bindings is stripped wholesale on downgrade) and are dropped on both the execution API and the supervisor schema side; the supervisor-schema change class is renamed so the two same-named migrations cannot be confused. - The stub decorator's hand-rolled version-split imports now go through the common.compat sdk seam (new PlainXComArg and KNOWN_CONTEXT_KEYS exports). - The Go runtime validates required wire-spec fields (name, xcom task_id) instead of silently binding empty strings, populates the carried Kind discriminant, and reports a binding bookkeeping bug as a task error instead of panicking the worker. - The supports_expand opt-out is now covered by task-sdk-level tests, the ti_run serialized-dag scan moved onto LazyDeserializedDAG next to its sibling accessors, and assorted review nits (exception types, stale wording, enum comparisons) are fixed. --- .../execution_api/routes/task_instances.py | 13 +-- .../execution_api/versions/__init__.py | 4 +- .../execution_api/versions/v2026_06_30.py | 21 ----- .../execution_api/versions/v2026_07_30.py | 40 +++++++++ .../serialization/serialized_objects.py | 15 ++++ .../v2026_04_17/test_task_instances.py | 45 ---------- .../versions/v2026_07_30/__init__.py | 16 ++++ .../v2026_07_30/test_task_instances.py | 88 +++++++++++++++++++ .../airflow-go-pack/pack_integration_test.go | 1 + go-sdk/pkg/binding/binding.go | 39 +++++--- go-sdk/pkg/execution/integration_test.go | 51 +++++++++++ go-sdk/pkg/execution/task_runner.go | 23 ++++- .../airflow/providers/common/compat/sdk.py | 8 +- providers/standard/pyproject.toml | 2 +- .../providers/standard/decorators/stub.py | 28 +++--- .../unit/standard/decorators/test_stub.py | 72 ++++++++++----- .../airflow/sdk/api/datamodels/_generated.py | 2 +- .../schema/versions/__init__.py | 6 +- .../schema/versions/v2026_07_30.py | 15 ++-- .../tests/task_sdk/bases/test_decorator.py | 26 ++++++ .../execution_time/schema/test_migrator.py | 2 +- 21 files changed, 367 insertions(+), 150 deletions(-) create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_07_30.py create mode 100644 airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_07_30/__init__.py create mode 100644 airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_07_30/test_task_instances.py 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 eb27a0a214a80..668f21520e171 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 @@ -118,12 +118,11 @@ 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.""" + """Extract the stub task's serialized arg spec from its Dag version's serialized 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 + from airflow.serialization.serialized_objects import LazyDeserializedDAG if dag_version_id is None: return None @@ -132,13 +131,7 @@ def _get_arg_bindings(dag_version_id: UUID | None, task_id: str, *, session) -> return None if not (data := dag_version.serialized_dag.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 + return LazyDeserializedDAG(data=data).get_task_arg_bindings(task_id) @ti_id_router.patch( 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 f4a08db4dc7fa..03d50957e8a84 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,7 +41,6 @@ RemoveUpstreamMapIndexesField, ) from airflow.api_fastapi.execution_api.versions.v2026_06_30 import ( - AddArgBindingsToTIRunContext, AddAssetsByAliasEndpoint, AddAwaitingInputStatePayload, AddConnectionTestEndpoint, @@ -52,9 +51,11 @@ AddTeamNameField, AddVariableKeysEndpoint, ) +from airflow.api_fastapi.execution_api.versions.v2026_07_30 import AddArgBindingsToTIRunContext bundle = VersionBundle( HeadVersion(), + Version("2026-07-30", AddArgBindingsToTIRunContext), Version( "2026-06-30", AddVariableKeysEndpoint, @@ -66,7 +67,6 @@ 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 15a097fac5f70..e89e2ed04cc5d 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 @@ -25,10 +25,6 @@ schema, ) -from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import ( - LiteralArgBinding, - XComArgBinding, -) from airflow.api_fastapi.execution_api.datamodels.taskinstance import ( DagRun, TaskInstance, @@ -145,20 +141,3 @@ 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, - schema(XComArgBinding).field("name").didnt_exist, - schema(LiteralArgBinding).field("name").didnt_exist, - ) - - @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type] - def remove_arg_bindings_field(response: ResponseInfo) -> None: # type: ignore[misc] - """Strip ``arg_bindings`` from the run context for older clients.""" - response.body.pop("arg_bindings", None) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_07_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_07_30.py new file mode 100644 index 0000000000000..2b456eae4da1f --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_07_30.py @@ -0,0 +1,40 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from cadwyn import ( + ResponseInfo, + VersionChange, + convert_response_to_previous_version_for, + schema, +) + +from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIRunContext + + +class AddArgBindingsToTIRunContext(VersionChange): + """Add the ``arg_bindings`` argument-binding spec for stub (foreign-runtime) tasks.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = (schema(TIRunContext).field("arg_bindings").didnt_exist,) + + @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type] + def remove_arg_bindings_field(response: ResponseInfo) -> None: # type: ignore[misc] + """Strip ``arg_bindings`` from the run context for older clients.""" + response.body.pop("arg_bindings", None) diff --git a/airflow-core/src/airflow/serialization/serialized_objects.py b/airflow-core/src/airflow/serialization/serialized_objects.py index 54bc3389c64ce..fd154f00b96ee 100644 --- a/airflow-core/src/airflow/serialization/serialized_objects.py +++ b/airflow-core/src/airflow/serialization/serialized_objects.py @@ -2277,6 +2277,21 @@ def __getattr__(self, name: str, /) -> Any: def timetable(self) -> Timetable: return decode_timetable(self.data["dag"]["timetable"]) + def get_task_arg_bindings(self, task_id: str) -> list | None: + """ + Extract one task's serialized ``_arg_bindings`` spec without deserializing the Dag. + + The spec is captured at parse time from a stub task's TaskFlow call; ``None`` for + tasks that carry no spec (regular tasks, and stub tasks with no parameters). + """ + for task in self.data["dag"]["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 + @property def has_task_concurrency_limits(self) -> bool: return any( 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 f88d22ff10e1c..71775cb50590a 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 @@ -20,7 +20,6 @@ import pytest from airflow._shared.timezones import timezone -from airflow.sdk import task from airflow.utils.state import DagRunState, State from tests_common.test_utils.config import conf_vars @@ -85,47 +84,3 @@ 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): - with dag_maker("test_arg_bindings_compat_dag", serialized=True): - - @task.stub - def extract(): ... - - @task.stub - def transform(country: str, extracted: dict): ... - - transform("uk", extract()) - - dr = dag_maker.create_dagrun() - tis = {ti.task_id: ti for ti in dr.get_task_instances()} - for ti in tis.values(): - ti.set_state(State.QUEUED) - dag_maker.session.flush() - return tis["transform"] - - def test_old_version_strips_arg_bindings_even_when_set(self, old_ver_client, stub_ti): - response = old_ver_client.patch(f"/execution/task-instances/{stub_ti.id}/run", json=RUN_PATCH_BODY) - assert response.status_code == 200 - assert "arg_bindings" not in response.json() - - def test_head_version_includes_arg_bindings(self, client, stub_ti): - response = client.patch(f"/execution/task-instances/{stub_ti.id}/run", json=RUN_PATCH_BODY) - assert response.status_code == 200 - assert response.json()["arg_bindings"] == [ - {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, - {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract"}, - ] diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_07_30/__init__.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_07_30/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_07_30/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_07_30/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_07_30/test_task_instances.py new file mode 100644 index 0000000000000..9821baf162d90 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_07_30/test_task_instances.py @@ -0,0 +1,88 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import pytest + +from airflow.sdk import task +from airflow.utils.state import State + +from tests_common.test_utils.db import clear_db_runs + +pytestmark = pytest.mark.db_test + +TIMESTAMP_STR = "2024-09-30T12:00:00Z" + +RUN_PATCH_BODY = { + "state": "running", + "hostname": "h", + "unixname": "u", + "pid": 1, + "start_date": TIMESTAMP_STR, +} + + +@pytest.fixture +def old_ver_client(client): + """Execution API version immediately before ``arg_bindings`` was added.""" + client.headers["Airflow-API-Version"] = "2026-06-30" + return client + + +class TestArgBindingsFieldBackwardCompat: + @pytest.fixture(autouse=True) + def _freeze_time(self, time_machine): + time_machine.move_to(TIMESTAMP_STR, tick=False) + + def setup_method(self): + clear_db_runs() + + def teardown_method(self): + clear_db_runs() + + @pytest.fixture + def stub_ti(self, dag_maker): + with dag_maker("test_arg_bindings_compat_dag", serialized=True): + + @task.stub + def extract(): ... + + @task.stub + def transform(country: str, extracted: dict): ... + + transform("uk", extract()) + + dr = dag_maker.create_dagrun() + tis = {ti.task_id: ti for ti in dr.get_task_instances()} + for ti in tis.values(): + ti.set_state(State.QUEUED) + dag_maker.session.flush() + return tis["transform"] + + def test_old_version_strips_arg_bindings_even_when_set(self, old_ver_client, stub_ti): + response = old_ver_client.patch(f"/execution/task-instances/{stub_ti.id}/run", json=RUN_PATCH_BODY) + assert response.status_code == 200 + assert "arg_bindings" not in response.json() + + def test_head_version_includes_arg_bindings(self, client, stub_ti): + response = client.patch(f"/execution/task-instances/{stub_ti.id}/run", json=RUN_PATCH_BODY) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, + {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract"}, + ] 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 652e943ba63e8..3be151660a8a3 100644 --- a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go +++ b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go @@ -158,6 +158,7 @@ dags: tasks: - "make_config" - "make_numbers" + - "make_region" - "via_flat_args" - "via_struct_no_tags" - "via_struct_arg_tag" diff --git a/go-sdk/pkg/binding/binding.go b/go-sdk/pkg/binding/binding.go index 399d3ebb2b420..9f9f7771b956b 100644 --- a/go-sdk/pkg/binding/binding.go +++ b/go-sdk/pkg/binding/binding.go @@ -111,13 +111,13 @@ type Arg interface { } // XComArg sources the argument from an upstream task's return-value XCom. -// Kind is carried by the generated shape but unused here: the Go type itself -// is the discriminant. +// Kind carries the wire discriminant ("xcom") from the generated shape; the +// resolve path dispatches on the Go type itself and never reads it. type XComArg genmodels.XComArgBinding -// LiteralArg carries an inline value from the Dag file. Kind is carried by -// the generated shape but unused here: the Go type itself is the -// discriminant. +// LiteralArg carries an inline value from the Dag file. Kind carries the wire +// discriminant ("literal") from the generated shape; the resolve path +// dispatches on the Go type itself and never reads it. type LiteralArg genmodels.LiteralArgBinding func (a XComArg) ArgName() string { return a.Name } @@ -231,6 +231,10 @@ func Analyze(fnType reflect.Type, fnName string) (*Plan, error) { // struct's fields claim entries out of args by name; plain flat data // parameters consume args in declaration order (the two shapes are mutually // exclusive; see Analyze). An error fails the task before its body runs. +// +// client must be the full sdk.Client -- not just the sdk.XComClient the +// resolve helpers narrow to -- because a paramClient parameter receives the +// client itself, verbatim. func (p *Plan) Resolve( ctx context.Context, logger *slog.Logger, @@ -292,17 +296,19 @@ func (p *Plan) Resolve( } argCursor := 0 - nextUnclaimed := func() Arg { + nextUnclaimed := func() (Arg, bool) { for argCursor < len(args) { i := argCursor argCursor++ if !claimed[i] { - return args[i] + return args[i], true } } - // Unreachable: the remaining/numData check above guarantees enough - // unclaimed args exist for every flat parameter still to be filled. - panic("binding: exhausted unclaimed args despite a passing arity check") + // Unreachable in practice: the remaining/numData check above guarantees + // enough unclaimed args exist for every flat parameter still to be + // filled. Reported as an error rather than a panic so a bookkeeping bug + // fails one task instead of crashing the worker. + return nil, false } flatIdx := 0 @@ -326,7 +332,14 @@ func (p *Plan) Resolve( case paramTaskInput: // Already resolved above. case paramData: - v, err := p.resolveData(ctx, client, plan, nextUnclaimed(), flatIdx) + arg, ok := nextUnclaimed() + if !ok { + return nil, fmt.Errorf( + "task function %s: internal error: exhausted unclaimed args despite a passing arity check", + p.fnName, + ) + } + v, err := p.resolveData(ctx, client, plan, arg, flatIdx) if err != nil { return nil, err } @@ -360,7 +373,7 @@ func (p *Plan) resolveData( // rather than failing the task. func (p *Plan) resolveTaskInput( ctx context.Context, - client sdk.Client, + c sdk.XComClient, plan paramPlan, args []Arg, byName map[string]int, @@ -385,7 +398,7 @@ func (p *Plan) resolveTaskInput( claimed[idx] = true v, err := p.resolveOne( - ctx, client, tif.fieldType, args[idx], + ctx, c, tif.fieldType, args[idx], fmt.Sprintf("TaskInput field %s (parameter %d)", tif.goName, plan.index), fmt.Sprintf("TaskInput field %s", tif.goName), ) diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 1567c5f65adde..2eaa87c4f8e02 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -479,6 +479,57 @@ func TestTaskRunnerArgBindingsMalformedElement(t *testing.T) { assert.False(t, ran, "the task body must not run on a malformed binding element") } +// TestTaskRunnerArgBindingsMissingRequiredFields: a wire spec entry without a +// usable name, or an xcom entry without a task_id, fails the task before the +// body runs instead of silently binding empty strings. +func TestTaskRunnerArgBindingsMissingRequiredFields(t *testing.T) { + cases := []struct { + name string + spec map[string]any + }{ + {name: "missing name", spec: map[string]any{"kind": "literal", "value": "x"}}, + {name: "empty name", spec: map[string]any{"name": "", "kind": "literal", "value": "x"}}, + {name: "xcom missing task_id", spec: map[string]any{"name": "country", "kind": "xcom"}}, + { + name: "xcom empty task_id", + spec: map[string]any{"name": "country", "kind": "xcom", "task_id": ""}, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + ran := false + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTaskWithName("transform", + func(country string) error { + ran = true + return nil + }) + }) + + details := &genmodels.StartupDetails{ + TI: genmodels.TaskInstance{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "transform", + RunID: "run1", + MapIndex: ptr(-1), + }, + BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, + TIContext: genmodels.TIRunContext{ + ArgBindings: &genmodels.ArgBindings{tc.spec}, + }, + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(context.Background(), bundle, details, comm, logger) + assertTaskState(t, result, genmodels.TaskStateStateFailed) + assert.False(t, ran, "the task body must not run on an incomplete binding spec") + }) + } +} + 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/task_runner.go b/go-sdk/pkg/execution/task_runner.go index 3d18aae33157e..e4be65f061bd8 100644 --- a/go-sdk/pkg/execution/task_runner.go +++ b/go-sdk/pkg/execution/task_runner.go @@ -157,17 +157,32 @@ func convertArgBindings(specsPtr *genmodels.ArgBindings) ([]binding.Arg, error) if !ok { return nil, fmt.Errorf("arg_bindings[%d]: unexpected wire shape %T", i, raw) } - name, _ := m["name"].(string) + name, ok := m["name"].(string) + if !ok || name == "" { + return nil, fmt.Errorf("arg_bindings[%d]: missing or empty name", i) + } dataType := binding.DataTypeAny if s, ok := m["data_type"].(string); ok && s != "" { dataType = binding.DataType(s) } switch kind, _ := m["kind"].(string); kind { case "xcom": - taskID, _ := m["task_id"].(string) - args[i] = binding.XComArg{Name: name, TaskID: taskID, DataType: dataType} + taskID, ok := m["task_id"].(string) + if !ok || taskID == "" { + return nil, fmt.Errorf( + "arg_bindings[%d] (%q): missing or empty task_id for xcom kind", + i, + name, + ) + } + args[i] = binding.XComArg{Kind: kind, Name: name, TaskID: taskID, DataType: dataType} case "literal": - args[i] = binding.LiteralArg{Name: name, Value: m["value"], DataType: dataType} + args[i] = binding.LiteralArg{ + Kind: kind, + Name: name, + Value: m["value"], + DataType: dataType, + } default: return nil, fmt.Errorf("arg_bindings[%d]: unknown kind %q", i, kind) } diff --git a/providers/common/compat/src/airflow/providers/common/compat/sdk.py b/providers/common/compat/src/airflow/providers/common/compat/sdk.py index 93174df7b2a28..772650f5499e5 100644 --- a/providers/common/compat/src/airflow/providers/common/compat/sdk.py +++ b/providers/common/compat/src/airflow/providers/common/compat/sdk.py @@ -83,9 +83,13 @@ from airflow.sdk.bases.sensor import poke_mode_only as poke_mode_only from airflow.sdk.bases.skipmixin import SkipMixin as SkipMixin from airflow.sdk.configuration import conf as conf - from airflow.sdk.definitions.context import context_merge as context_merge + from airflow.sdk.definitions.context import ( + KNOWN_CONTEXT_KEYS as KNOWN_CONTEXT_KEYS, + context_merge as context_merge, + ) from airflow.sdk.definitions.mappedoperator import MappedOperator as MappedOperator from airflow.sdk.definitions.template import literal as literal + from airflow.sdk.definitions.xcom_arg import PlainXComArg as PlainXComArg from airflow.sdk.exceptions import ( AirflowConfigException as AirflowConfigException, AirflowException as AirflowException, @@ -192,6 +196,7 @@ "DAG": ("airflow.sdk", "airflow.models.dag"), "Param": ("airflow.sdk", "airflow.models.param"), "XComArg": ("airflow.sdk", "airflow.models.xcom_arg"), + "PlainXComArg": ("airflow.sdk.definitions.xcom_arg", "airflow.models.xcom_arg"), "DecoratedOperator": ("airflow.sdk.bases.decorator", "airflow.decorators.base"), "DecoratedMappedOperator": ("airflow.sdk.bases.decorator", "airflow.decorators.base"), "MappedOperator": ("airflow.sdk.definitions.mappedoperator", "airflow.models.mappedoperator"), @@ -246,6 +251,7 @@ # ============================================================================ "Context": ("airflow.sdk", "airflow.utils.context"), "context_merge": ("airflow.sdk.definitions.context", "airflow.utils.context"), + "KNOWN_CONTEXT_KEYS": ("airflow.sdk.definitions.context", "airflow.utils.context"), "context_to_airflow_vars": ("airflow.sdk.execution_time.context", "airflow.utils.operator_helpers"), "AIRFLOW_VAR_NAME_FORMAT_MAPPING": ( "airflow.sdk.execution_time.context", diff --git a/providers/standard/pyproject.toml b/providers/standard/pyproject.toml index b385a9e9eb7f6..373d16bba6c4e 100644 --- a/providers/standard/pyproject.toml +++ b/providers/standard/pyproject.toml @@ -60,7 +60,7 @@ requires-python = ">=3.10" # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build`` dependencies = [ "apache-airflow>=2.11.0", - "apache-airflow-providers-common-compat>=1.14.1", + "apache-airflow-providers-common-compat>=1.14.1", # use next version ] # The optional dependencies should be modified in place in the generated file diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index ddfd64fd31eb0..3834be9f08a8b 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -27,24 +27,17 @@ from typing import TYPE_CHECKING, Any, Union from airflow.providers.common.compat.sdk import ( + KNOWN_CONTEXT_KEYS, DecoratedOperator, + PlainXComArg, TaskDecorator, + XComArg, 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] - try: from airflow.sdk.api.datamodels._generated import ArgBindingDataType -except ImportError: # Airflow 2 -- no task-sdk execution-API generated models +except ImportError: # Airflow < 3.4 -- the generated models do not carry the enum yet class ArgBindingDataType(str, enum.Enum): # type: ignore[no-redef] """Language-neutral value type a stub-task argument binds to in the foreign runtime.""" @@ -73,7 +66,7 @@ def _data_type_from_annotation(annotation: Any) -> ArgBindingDataType: return ArgBindingDataType.ANY origin = typing.get_origin(annotation) if origin is not None: - if origin is Union or origin is getattr(types, "UnionType", None): + if origin is Union or origin is types.UnionType: 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]) @@ -138,7 +131,7 @@ def _build_arg_bindings( try: hints = typing.get_type_hints(python_callable) - except Exception: + except (NameError, TypeError): # Annotations that cannot be resolved at parse time (e.g. names behind # TYPE_CHECKING with ``from __future__ import annotations``) degrade to "any". hints = {} @@ -222,10 +215,10 @@ def __init__( module = ast.parse(self.get_python_source()) if len(module.body) != 1: - raise RuntimeError("Expected a single statement") + raise ValueError("Expected a single statement") fn = module.body[0] if not isinstance(fn, ast.FunctionDef): - raise RuntimeError("Expected a single sync function") + raise ValueError("Expected a single sync function") for stmt in fn.body: if isinstance(stmt, ast.Pass): continue @@ -266,8 +259,9 @@ def stub( 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. + outputs or JSON-serializable literals; the resulting argument-binding spec (parameter + names, declared types, and values, in declaration order) 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 02bf3c7c58e54..6babbbd75de1f 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -23,7 +23,7 @@ import pytest from airflow.providers.common.compat.sdk import DAG -from airflow.providers.standard.decorators.stub import _data_type_from_annotation, stub +from airflow.providers.standard.decorators.stub import ArgBindingDataType, _data_type_from_annotation, stub from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS, AIRFLOW_V_3_4_PLUS @@ -167,9 +167,27 @@ def test_non_json_literal_rejected(self): def test_mapped_xcom_arg_rejected(self): with DAG(dag_id="d"): extracted = stub(fn_extract)() - with pytest.raises(ValueError, match="MapXComArg"): + with pytest.raises(ValueError, match="only direct upstream task outputs"): stub(fn_transform)("uk", extracted.map(lambda v: v)) + def test_arg_bindings_survive_dag_serialization_round_trip(self): + """The captured spec must survive whichever core serializer the provider runs against.""" + try: + from airflow.serialization.serialized_objects import DagSerialization + except ImportError: # Airflow 2 exposes the round-trip API on SerializedDAG + from airflow.serialization.serialized_objects import SerializedDAG as DagSerialization + + with DAG(dag_id="d") as dag: + extracted = stub(fn_extract)() + stub(fn_transform)("uk", extracted) + + round_tripped = DagSerialization.from_dict(DagSerialization.to_dict(dag)) + assert round_tripped.task_dict["fn_transform"]._arg_bindings == [ + {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, + {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "fn_extract"}, + {"name": "retries_num", "kind": "literal", "data_type": "integer", "value": 3}, + ] + @pytest.mark.skipif( not AIRFLOW_V_3_4_PLUS, reason="task-sdk honors the supports_expand opt-out from Airflow 3.4" ) @@ -182,27 +200,35 @@ def test_expand_rejected_at_parse_time(self): @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"), + pytest.param(str, ArgBindingDataType.STRING, id="str"), + pytest.param(bool, ArgBindingDataType.BOOLEAN, id="bool"), + pytest.param(int, ArgBindingDataType.INTEGER, id="int"), + pytest.param(float, ArgBindingDataType.NUMBER, id="float"), + pytest.param(dict, ArgBindingDataType.OBJECT, id="dict"), + pytest.param(dict[str, int], ArgBindingDataType.OBJECT, id="dict-parameterized"), + pytest.param(typing.Mapping[str, int], ArgBindingDataType.OBJECT, id="mapping"), + pytest.param(list, ArgBindingDataType.ARRAY, id="list"), + pytest.param(list[int], ArgBindingDataType.ARRAY, id="list-parameterized"), + pytest.param(tuple, ArgBindingDataType.ARRAY, id="tuple"), + pytest.param(set, ArgBindingDataType.ARRAY, id="set"), + pytest.param(typing.Sequence[int], ArgBindingDataType.ARRAY, id="sequence"), + pytest.param(Any, ArgBindingDataType.ANY, id="any"), + pytest.param(None, ArgBindingDataType.ANY, id="none"), + pytest.param(bytes, ArgBindingDataType.ANY, id="bytes"), + pytest.param( + typing.Optional[str], # noqa: UP045 -- legacy form on purpose + ArgBindingDataType.STRING, + id="optional-str", + ), + pytest.param( + typing.Union[int, str], # noqa: UP007 -- legacy form on purpose + ArgBindingDataType.ANY, + id="union", + ), + pytest.param(str | None, ArgBindingDataType.STRING, id="pep604-optional"), + pytest.param(int | str, ArgBindingDataType.ANY, id="pep604-union"), + pytest.param(contextlib.AbstractContextManager, ArgBindingDataType.ANY, id="custom-class"), ], ) def test_data_type_from_annotation(annotation, expected): - assert _data_type_from_annotation(annotation) == expected + assert _data_type_from_annotation(annotation) is expected diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 48b598d90915e..696fdcc6c1365 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -27,7 +27,7 @@ from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, RootModel -API_VERSION: Final[str] = "2026-06-30" +API_VERSION: Final[str] = "2026-07-30" class ArgBindingDataType(str, Enum): 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 266aeb0c19736..59640c587473a 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,11 +37,13 @@ def get_bundle() -> VersionBundle: """ from cadwyn import HeadVersion, Version, VersionBundle - from airflow.sdk.execution_time.schema.versions.v2026_07_30 import AddArgBindingsToTIRunContext + from airflow.sdk.execution_time.schema.versions.v2026_07_30 import ( + AddArgBindingsToSupervisorTIRunContext, + ) return VersionBundle( HeadVersion(), - Version("2026-07-30", AddArgBindingsToTIRunContext), + Version("2026-07-30", AddArgBindingsToSupervisorTIRunContext), 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 index 119963de7e855..e6b93f5dea805 100644 --- 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 @@ -19,21 +19,18 @@ from cadwyn import VersionChange, schema -from airflow.sdk.api.datamodels._generated import LiteralArgBinding, TIRunContext, XComArgBinding +from airflow.sdk.api.datamodels._generated import TIRunContext -class AddArgBindingsToTIRunContext(VersionChange): +class AddArgBindingsToSupervisorTIRunContext(VersionChange): """ - Add the ``arg_bindings`` positional-argument binding spec for stub (foreign-runtime) tasks. + Add the ``arg_bindings`` argument-binding spec for stub (foreign-runtime) tasks. Each entry is a discriminated union of ``XComArgBinding`` and ``LiteralArgBinding`` - keyed on ``kind``. + keyed on ``kind``. The supervisor-schema mirror of the execution API's + ``AddArgBindingsToTIRunContext``, named apart so the two migrations are not confused. """ description = __doc__ - instructions_to_migrate_to_previous_version = ( - schema(TIRunContext).field("arg_bindings").didnt_exist, - schema(XComArgBinding).field("name").didnt_exist, - schema(LiteralArgBinding).field("name").didnt_exist, - ) + instructions_to_migrate_to_previous_version = (schema(TIRunContext).field("arg_bindings").didnt_exist,) diff --git a/task-sdk/tests/task_sdk/bases/test_decorator.py b/task-sdk/tests/task_sdk/bases/test_decorator.py index 860eb6e7b3312..0578467aa26c1 100644 --- a/task-sdk/tests/task_sdk/bases/test_decorator.py +++ b/task-sdk/tests/task_sdk/bases/test_decorator.py @@ -383,3 +383,29 @@ def sync_task_fn(): return 42 assert not is_async_callable(sync_task_fn) + + +class DummyNoExpandDecoratedOperator(DecoratedOperator): + custom_operator_name = "@task.dummy_no_expand" + + supports_expand = False + + +class TestSupportsExpandOptOut: + """An operator class that sets ``supports_expand = False`` rejects dynamic task mapping at parse time.""" + + @pytest.fixture + def no_expand_task(self): + from airflow.sdk.bases.decorator import task_decorator_factory + + def fn(a): ... + + return task_decorator_factory(fn, decorated_operator_class=DummyNoExpandDecoratedOperator) + + def test_expand_rejected_with_operator_name(self, no_expand_task): + with pytest.raises(TypeError, match="@task.dummy_no_expand tasks do not support dynamic task"): + no_expand_task.expand(a=[1, 2]) + + def test_expand_kwargs_rejected(self, no_expand_task): + with pytest.raises(TypeError, match="do not support dynamic task mapping"): + no_expand_task.expand_kwargs([{"a": 1}]) 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 44d22a255e218..ce26296519b2d 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 @@ -372,7 +372,7 @@ class TestRealBundleArgBindingsDowngrade: """ Drive the *real* supervisor bundle through the ``arg_bindings`` migration. - ``AddArgBindingsToTIRunContext`` is the bundle's first ``schema(...)`` + ``AddArgBindingsToSupervisorTIRunContext`` is the bundle's first ``schema(...)`` instruction on a model *nested* inside a registered body (``StartupDetails.ti_context``); this pins that the downgrade re-validation strips the nested field on the wire for a runtime From edc6b91027ec91372a7fb1af37f544acf1ce0705 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 21 Jul 2026 16:48:03 +0000 Subject: [PATCH 16/40] Harden the stub-task TaskFlow arg-binding contract after review A multi-angle review of the branch surfaced gaps at the edges of the new binding contract: - An unrecognized serialized spec escaped ti_run as an opaque 500 on provider/core version skew; it now returns a structured invalid_arg_bindings error, per the route-boundary convention. - The new parse-time signature checks broke previously importable argless stub Dags (e.g. a **kwargs or ti parameter); they now fire only when a TaskFlow call actually passes arguments. - Stubs called with arguments inside a mapped task group serialized a spec with no map-index dimension and failed (or mis-bound) at runtime; they are now rejected at parse time. - A Go TaskInput struct had to mirror every stub parameter -- including defaulted ones the author never passed -- or fail every run, while an empty spec silently zero-filled the whole struct, contradicting the documented fail-loud behavior. Literal entries captured from signature defaults now carry from_default on the wire (inside the still-in-progress 2026-07-30 schema) and may go unclaimed; a spec that never arrives fails loudly when the struct declares bindable fields. - NaN/Infinity literals passed the parse-time JSON check only to fail far away (or silently bind 0.0); json.dumps now rejects them. - A malformed spec bypassed ShouldRetry while an equally permanent arity error retried; both now share retry semantics. - ti_run no longer issues two queries and re-parses the serialized-Dag blob on every stub-task start: single joinedload query plus a per-(dag_version, task) cache of the immutable extracted spec, and XCom pulls in the Go binding path now run concurrently. --- .../datamodels/task_arg_binding.py | 3 + .../execution_api/routes/task_instances.py | 47 +- .../versions/head/test_task_instances.py | 40 +- .../v2026_07_30/test_task_instances.py | 3 +- .../test_go_sdk_taskflow_binding.py | 9 +- go-sdk/README.md | 22 +- go-sdk/bundle/bundlev1/task.go | 7 +- go-sdk/dags/go_examples.py | 11 +- .../bundle/taskflowbinding/taskflowbinding.go | 22 +- go-sdk/pkg/binding/binding.go | 408 ++++++++++-------- go-sdk/pkg/binding/binding_test.go | 80 +++- go-sdk/pkg/execution/genmodels/models.gen.go | 132 +++--- go-sdk/pkg/execution/integration_test.go | 82 ++++ go-sdk/pkg/execution/task_runner.go | 18 +- .../providers/standard/decorators/stub.py | 88 ++-- .../unit/standard/decorators/test_stub.py | 99 +++-- .../airflow/sdk/api/datamodels/_generated.py | 1 + .../sdk/execution_time/schema/schema.json | 5 + .../execution_time/schema/test_migrator.py | 12 +- 19 files changed, 726 insertions(+), 363 deletions(-) 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 index 1171a5fc6f1ce..58866a82fc98f 100644 --- 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 @@ -75,6 +75,9 @@ class LiteralArgBinding(BaseModel): value: JsonValue | None = None """The literal value from the Dag file.""" + from_default: bool = False + """True when the value was filled from the stub signature's default rather than passed in the call.""" + # A named alias (TypeAliasType, not a bare Annotated) so the union lands in every # schema as its own named definition instead of an anonymous field-title-derived one. 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 668f21520e171..43db192ad6212 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -32,7 +32,7 @@ from opentelemetry import trace from opentelemetry.trace import StatusCode from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator -from pydantic import JsonValue +from pydantic import JsonValue, ValidationError from sqlalchemy import and_, func, or_, tuple_, update from sqlalchemy.engine import CursorResult from sqlalchemy.exc import DataError, NoResultFound, SQLAlchemyError @@ -81,6 +81,7 @@ from airflow.models.asset import AssetActive from airflow.models.base import ID_LEN from airflow.models.dag import DagModel +from airflow.models.dag_version import DagVersion from airflow.models.dagrun import DagRun as DR from airflow.models.hitl import HITLDetail from airflow.models.log import Log @@ -90,6 +91,7 @@ from airflow.models.trigger import Trigger, handle_event_submit from airflow.models.xcom import XComModel from airflow.serialization.definitions.assets import SerializedAsset, SerializedAssetUniqueKey +from airflow.serialization.serialized_objects import LazyDeserializedDAG from airflow.state import get_state_backend from airflow.triggers.base import TriggerEvent from airflow.utils.sqlalchemy import get_dialect_name @@ -113,25 +115,33 @@ # 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. +# serialized-Dag lookup for ``arg_bindings`` so regular tasks never pay for it. +# The gate matches the exact class name; a subclass would need its own entry here. _STUB_TASK_TYPE = "_StubOperator" +# Specs are immutable per (Dag version, task): a re-serialized Dag gets a new +# version id, so entries never go stale and the cap only bounds memory. +_ARG_BINDINGS_CACHE: dict[tuple[UUID, str], list[dict] | None] = {} +_ARG_BINDINGS_CACHE_MAX_ENTRIES = 4096 + def _get_arg_bindings(dag_version_id: UUID | None, task_id: str, *, session) -> list[dict] | None: """Extract the stub task's serialized arg spec from its Dag version's serialized 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.serialized_objects import LazyDeserializedDAG - if dag_version_id is None: return None - dag_version = session.get(DagVersion, dag_version_id) + cache_key = (dag_version_id, task_id) + if cache_key in _ARG_BINDINGS_CACHE: + return _ARG_BINDINGS_CACHE[cache_key] + dag_version = session.get(DagVersion, dag_version_id, options=[joinedload(DagVersion.serialized_dag)]) if dag_version is None or dag_version.serialized_dag is None: return None if not (data := dag_version.serialized_dag.data): return None - return LazyDeserializedDAG(data=data).get_task_arg_bindings(task_id) + bindings = LazyDeserializedDAG(data=data).get_task_arg_bindings(task_id) + if len(_ARG_BINDINGS_CACHE) >= _ARG_BINDINGS_CACHE_MAX_ENTRIES: + _ARG_BINDINGS_CACHE.clear() + _ARG_BINDINGS_CACHE[cache_key] = bindings + return bindings @ti_id_router.patch( @@ -340,7 +350,24 @@ def ti_run( if ti.operator == _STUB_TASK_TYPE and ( arg_bindings := _get_arg_bindings(ti.dag_version_id, ti.task_id, session=session) ): - context.arg_bindings = get_arg_bindings_adapter().validate_python(arg_bindings) + try: + context.arg_bindings = get_arg_bindings_adapter().validate_python(arg_bindings) + except ValidationError: + log.exception( + "Serialized arg_bindings spec failed validation", + dag_id=ti.dag_id, + task_id=ti.task_id, + ) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={ + "reason": "invalid_arg_bindings", + "message": ( + "The serialized TaskFlow arg spec for this stub task is not valid on " + "this Airflow version; it may come from a newer providers release." + ), + }, + ) # Only set if they are non-null if ti.next_method: 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 85d25ea2a5a0b..22ff27f835e3f 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 @@ -54,6 +54,7 @@ from airflow.models.taskinstancehistory import TaskInstanceHistory from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk import Asset, TaskGroup, TriggerRule, task, task_group +from airflow.serialization.serialized_objects import LazyDeserializedDAG from airflow.state.metastore import MetastoreBackend from airflow.utils.state import DagRunState, State, TaskInstanceState, TerminalTIState @@ -372,14 +373,14 @@ async def workload_token(request: Request) -> TIToken: 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.""" + """A stub task's TaskFlow arg spec is extracted from the serialized Dag and returned.""" with dag_maker("test_arg_bindings_dag", serialized=True): @task.stub def extract(): ... @task.stub - def transform(country: str, extracted: dict): ... + def transform(country: str, extracted: dict, limit: int = 10): ... transform("uk", extract()) @@ -402,6 +403,7 @@ def transform(country: str, extracted: dict): ... assert response.json()["arg_bindings"] == [ {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract"}, + {"name": "limit", "kind": "literal", "data_type": "integer", "value": 10, "from_default": True}, ] # An argless stub has no captured spec, so the field stays unset. @@ -409,6 +411,40 @@ def transform(country: str, extracted: dict): ... assert response.status_code == 200 assert "arg_bindings" not in response.json() + @mock.patch.object( + LazyDeserializedDAG, + "get_task_arg_bindings", + autospec=True, + return_value=[{"name": "country", "kind": "hologram", "value": "uk"}], + ) + def test_ti_run_reports_invalid_arg_bindings_spec(self, _, client, dag_maker): + """A serialized spec this core version cannot validate fails with a structured error, not a bare 500.""" + with dag_maker("test_invalid_arg_bindings_dag", serialized=True): + + @task.stub + def transform(country: str): ... + + transform("uk") + + dr = dag_maker.create_dagrun() + (ti,) = dr.get_task_instances() + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch( + f"/execution/task-instances/{ti.id}/run", + json={ + "state": "running", + "hostname": "random-hostname", + "unixname": "random-unixname", + "pid": 100, + "start_date": "2024-09-30T12:00:00Z", + }, + ) + + assert response.status_code == 500 + assert response.json()["detail"]["reason"] == "invalid_arg_bindings" + def test_arg_bindings_adapter_rejects_unknown_kind(self): """The discriminated union refuses serialized specs with an unrecognised kind.""" from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import get_arg_bindings_adapter diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_07_30/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_07_30/test_task_instances.py index 9821baf162d90..cca4a78f9d859 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_07_30/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_07_30/test_task_instances.py @@ -63,7 +63,7 @@ def stub_ti(self, dag_maker): def extract(): ... @task.stub - def transform(country: str, extracted: dict): ... + def transform(country: str, extracted: dict, limit: int = 10): ... transform("uk", extract()) @@ -85,4 +85,5 @@ def test_head_version_includes_arg_bindings(self, client, stub_ti): assert response.json()["arg_bindings"] == [ {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract"}, + {"name": "limit", "kind": "literal", "data_type": "integer", "value": 10, "from_default": True}, ] 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 index f6f8033a23118..8e35177a33d3f 100644 --- 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 @@ -134,9 +134,12 @@ def test_via_struct_arg_tag_reflects_bound_arguments(completed_run: _CompletedRu def test_via_struct_unmatched_arg_reflects_zero_valued_field(completed_run: _CompletedRun): - """``via_struct_unmatched_arg`` demonstrates that a struct field whose name has - no corresponding TaskFlow call argument stays at its Go zero value instead of - failing the task -- kwarg-style, an unpassed name simply isn't bound.""" + """``via_struct_unmatched_arg`` demonstrates mismatch tolerance in both directions: + a struct field whose name has no corresponding TaskFlow call argument stays at its + Go zero value instead of failing the task (kwarg-style, an unpassed name simply + isn't bound), and the stub's defaulted ``sample_rate`` -- captured into the spec as + ``from_default`` -- needs no matching struct field. The task succeeding at all + proves the second half.""" assert completed_run.xcom("via_struct_unmatched_arg") == { "region": "eu-west-1", "missing_was_empty": True, diff --git a/go-sdk/README.md b/go-sdk/README.md index 34c9a9dcc2475..7b2c0e856a044 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -108,12 +108,13 @@ 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. +(`transform(..., extract())`) is pulled from that task's XCom in the current Dag run and decoded into +the parameter's type (independent XCom pulls run concurrently). The runtime fails the task loudly when +the argument count doesn't match the number of data parameters or a declared type can't bind to the Go +type. Data parameters must be JSON-decodable (no func/chan/unsafe-pointer, no non-empty interfaces) — +checked once at registration. TaskFlow argument binding arrives over the coordinator protocol, so it is +coordinator-mode only today; on the Edge Worker path a task with data parameters fails with the arity +error (and a `TaskInput` struct with bindable fields fails the same way, since nothing can fill them). ```go func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (any, error) { @@ -175,6 +176,15 @@ only binds an argument literally spelled `Threshold`, so a snake_case Python par explicit tag. If no TaskFlow call argument carries that name, the field is simply left at its Go zero value — it does not fail the task, kwarg-style (see `ViaStructUnmatchedArg` below). +The matching is checked in the other direction too: every argument the Dag author **explicitly +passed** in the TaskFlow call must be claimed by some field, so a typo'd field name fails the task +instead of silently dropping the value. Stub parameters the author left at their Python defaults +are the exception — the Python side captures them into the spec (marked `from_default` on the +wire), and the struct is free not to mirror them, the same way a Python callee never sees which +defaulted kwargs went unpassed. And when no argument spec arrives at all (an argless stub call, or +the Edge Worker path) a `TaskInput` struct with bindable fields fails loudly rather than running +fully zero-valued. + A plain custom struct type *without* the `sdk.TaskInput` embed is unaffected by any of this — it keeps working as a single flat data parameter, JSON-decoded whole from one TaskFlow argument (see `Config` in diff --git a/go-sdk/bundle/bundlev1/task.go b/go-sdk/bundle/bundlev1/task.go index c4d21083a1c0f..cb18a4e951519 100644 --- a/go-sdk/bundle/bundlev1/task.go +++ b/go-sdk/bundle/bundlev1/task.go @@ -56,7 +56,12 @@ func NewTaskFunction(fn any) (Task, error) { v := reflect.ValueOf(fn) fullName := runtime.FuncForPC(v.Pointer()).Name() f := &taskFunction{fn: v, fullName: fullName} - return f, f.validateFn(v.Type()) + if err := f.validateFn(v.Type()); err != nil { + // A half-built task (nil binding plan) would panic in Execute; a caller + // that mishandles the error must not be able to run it. + return nil, err + } + return f, nil } func (f *taskFunction) Execute(ctx context.Context, logger *slog.Logger) error { diff --git a/go-sdk/dags/go_examples.py b/go-sdk/dags/go_examples.py index 121a59f2cb9b3..01e8821f1cf62 100644 --- a/go-sdk/dags/go_examples.py +++ b/go-sdk/dags/go_examples.py @@ -157,7 +157,7 @@ def via_struct_arg_tag(region_code: str, threshold: float): ... @task.stub(queue="golang") -def via_struct_unmatched_arg(region_code: str): ... +def via_struct_unmatched_arg(region_code: str, sample_rate: float = 0.1): ... @dag(dag_id="taskflow_binding_dag") @@ -194,9 +194,12 @@ def taskflow_binding_dag(): ``arg:`` tag -- ``Region`` is genuinely renamed to ``region_code``, and ``Threshold`` is tagged ``threshold`` to pull the snake_case argument its verbatim field name would miss. - * ``via_struct_unmatched_arg``: the Go struct declares a field with no - corresponding argument in this TaskFlow call at all -- it stays at its Go - zero value rather than failing the task. + * ``via_struct_unmatched_arg``: the mismatch tolerance in both directions. + The Go struct declares a field with no corresponding argument in this + TaskFlow call at all -- it stays at its Go zero value rather than failing + the task. And the stub's defaulted ``sample_rate`` is never passed, so its + captured-from-default entry needs no matching struct field (an explicitly + passed argument no field claims would fail the task instead). """ via_flat_args( "summary", diff --git a/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go b/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go index 47c70fe365828..7d7fdcd6dc0a1 100644 --- a/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go +++ b/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go @@ -224,11 +224,15 @@ func ViaStructArgTag( }, nil } -// ViaStructUnmatchedArgInput demonstrates that a field whose name has no -// corresponding TaskFlow call argument at all is left at its Go zero value -// rather than failing the task: Region binds normally, but Missing's arg -// name is never among this call's arguments -- conceptually, an unpassed -// keyword argument falling back to its default in a kwargs-style call. +// ViaStructUnmatchedArgInput demonstrates the mismatch tolerance in both +// directions. A field whose name has no corresponding TaskFlow call argument +// at all is left at its Go zero value rather than failing the task: Region +// binds normally, but Missing's arg name is never among this call's +// arguments -- conceptually, an unpassed keyword argument falling back to +// its default in a kwargs-style call. The reverse also holds: the stub's +// defaulted sample_rate parameter arrives marked from_default, so this +// struct is free not to mirror it (an explicitly passed argument no field +// claims would fail the task instead). type ViaStructUnmatchedArgInput struct { sdk.TaskInput Region string `arg:"region_code"` @@ -239,9 +243,11 @@ type ViaStructUnmatchedArgInput struct { // // via_struct_unmatched_arg(region_code=make_region()) // -// -- the stub only declares region_code (bound from make_region's XCom), so -// Missing's arg name never appears among the call's arguments and stays at -// its Go zero value (""). +// -- the stub declares region_code (bound from make_region's XCom) plus a +// defaulted sample_rate this struct deliberately omits, so Missing's arg +// name never appears among the call's arguments and stays at its Go zero +// value (""), while sample_rate's from_default entry goes unclaimed without +// failing the task. func ViaStructUnmatchedArg( ctx sdk.TIRunContext, log *slog.Logger, input ViaStructUnmatchedArgInput, ) (any, error) { diff --git a/go-sdk/pkg/binding/binding.go b/go-sdk/pkg/binding/binding.go index 9f9f7771b956b..f37e2c1a7e10e 100644 --- a/go-sdk/pkg/binding/binding.go +++ b/go-sdk/pkg/binding/binding.go @@ -28,7 +28,8 @@ // 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. +// named upstream task in the current Dag run first (independent pulls run +// concurrently). // - TaskInput structs: a struct that anonymously embeds sdk.TaskInput opts // into per-field, name-based binding instead of consuming one positional // slot as a whole-value decode target. Each exported field binds by name @@ -48,7 +49,13 @@ // binding: fields match by name, and a field whose name has no corresponding // TaskFlow call argument is simply left at its Go zero value instead of // failing the task -- the same way an unpassed keyword argument falls back -// to its default in a kwargs-style call. +// to its default in a kwargs-style call. The check runs both ways: an +// explicitly passed argument that no field claims fails the task (catching +// typo'd field names), while spec entries the Python side captured from the +// stub signature's defaults (from_default on the wire) may go unclaimed. And +// a TaskInput struct with bindable fields fails loudly when no argument spec +// arrives at all (e.g. the Edge Worker path, where nothing could ever fill +// them), matching the flat-parameter arity check. // // Analyze inspects a function once at registration and returns a Plan; Resolve // builds the call arguments for each execution from that Plan and the @@ -65,10 +72,12 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" "log/slog" "reflect" "strings" + "sync" "github.com/apache/airflow/go-sdk/pkg/api" "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" @@ -181,9 +190,6 @@ type Plan struct { hasTaskInput bool } -// 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); @@ -241,77 +247,7 @@ func (p *Plan) Resolve( client sdk.Client, args []Arg, ) ([]reflect.Value, error) { - byName := make(map[string]int, len(args)) - for i, a := range args { - // A nil entry can claim no name; it stays unclaimed and fails loudly in - // resolveOne when the flat-parameter cursor reaches it. - if a != nil { - byName[a.ArgName()] = i - } - } - claimed := make([]bool, len(args)) - out := make([]reflect.Value, len(p.params)) - for i, plan := range p.params { - if plan.kind != paramTaskInput { - continue - } - v, err := p.resolveTaskInput(ctx, client, plan, args, byName, claimed) - if err != nil { - return nil, err - } - out[i] = v - } - - remaining := 0 - for _, c := range claimed { - if !c { - remaining++ - } - } - if remaining != p.numData { - if p.hasTaskInput { - names := make([]string, 0, remaining) - for i, c := range claimed { - if c { - continue - } - name := "" - if args[i] != nil { - name = fmt.Sprintf("%q", args[i].ArgName()) - } - names = append(names, name) - } - return nil, fmt.Errorf( - "task function %s: %d TaskFlow call argument(s) not claimed by any TaskInput "+ - "field: %s", - p.fnName, remaining, strings.Join(names, ", "), - ) - } - 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, remaining, p.numData, - ) - } - - argCursor := 0 - nextUnclaimed := func() (Arg, bool) { - for argCursor < len(args) { - i := argCursor - argCursor++ - if !claimed[i] { - return args[i], true - } - } - // Unreachable in practice: the remaining/numData check above guarantees - // enough unclaimed args exist for every flat parameter still to be - // filled. Reported as an error rather than a panic so a bookkeeping bug - // fails one task instead of crashing the worker. - return nil, false - } - - flatIdx := 0 for i, plan := range p.params { switch plan.kind { case paramTIRunContext: @@ -329,158 +265,272 @@ func (p *Plan) Resolve( out[i] = reflect.ValueOf(logger) case paramClient: out[i] = reflect.ValueOf(client) - case paramTaskInput: - // Already resolved above. - case paramData: - arg, ok := nextUnclaimed() - if !ok { - return nil, fmt.Errorf( - "task function %s: internal error: exhausted unclaimed args despite a passing arity check", - p.fnName, - ) - } - v, err := p.resolveData(ctx, client, plan, arg, flatIdx) - if err != nil { - return nil, err - } - out[i] = v - flatIdx++ + case paramData, paramTaskInput: + // Filled below, once the spec is matched and XComs are pulled. } } - return out, nil + if p.hasTaskInput { + return p.resolveTaskInputParam(ctx, client, args, out) + } + return p.resolveFlatParams(ctx, client, args, out) } -// resolveData produces the value for one flat data parameter from its -// argument spec. -func (p *Plan) resolveData( +// resolveFlatParams fills the plain data parameters positionally -- args[0] +// onto the first data parameter, and so on -- with strict arity (see the +// package doc comment). +func (p *Plan) resolveFlatParams( ctx context.Context, c sdk.XComClient, - plan paramPlan, - arg Arg, - argIdx int, -) (reflect.Value, error) { - return p.resolveOne( - ctx, c, plan.typ, arg, - fmt.Sprintf("argument %d (parameter %d)", argIdx, plan.index), - fmt.Sprintf("argument %d", argIdx), - ) + args []Arg, + out []reflect.Value, +) ([]reflect.Value, error) { + if len(args) != p.numData { + return nil, fmt.Errorf( + "task function %s: argument count mismatch: the Dag passes %d positional argument(s) "+ + "but the Go function declares %d data parameter(s)", + p.fnName, len(args), p.numData, + ) + } + raws, err := p.fetchArgValues(ctx, c, args, nil) + if err != nil { + return nil, err + } + flatIdx := 0 + for i, plan := range p.params { + if plan.kind != paramData { + continue + } + v, err := p.decodeArg( + args[flatIdx], raws[flatIdx], plan.typ, + fmt.Sprintf("argument %d (parameter %d)", flatIdx, plan.index), + ) + if err != nil { + return nil, err + } + out[i] = v + flatIdx++ + } + return out, nil } -// resolveTaskInput builds the struct value for one TaskInput parameter. A -// field claiming an argument spec entry by name marks it claimed so the -// later flat-parameter cursor skips it. A field whose name claims nothing -// (kwarg-style: it was never "passed") is left unset at its Go zero value -// rather than failing the task. -func (p *Plan) resolveTaskInput( +// resolveTaskInputParam fills the single TaskInput struct parameter by name +// (kwarg-style). Every explicitly passed spec entry must be claimed by some +// field; entries the Python side captured from the stub signature's defaults +// (FromDefault) may go unclaimed, the same way an unpassed keyword argument +// never reaches the callee. +func (p *Plan) resolveTaskInputParam( ctx context.Context, c sdk.XComClient, - plan paramPlan, args []Arg, - byName map[string]int, - claimed []bool, -) (reflect.Value, error) { - structType := plan.typ - isPtr := structType.Kind() == reflect.Pointer - if isPtr { - structType = structType.Elem() + out []reflect.Value, +) ([]reflect.Value, error) { + var paramIdx int + var plan paramPlan + for i, pl := range p.params { + if pl.kind == paramTaskInput { + paramIdx, plan = i, pl + break + } } - structVal := reflect.New(structType).Elem() + if len(args) == 0 && len(plan.fields) > 0 { + return nil, fmt.Errorf( + "task function %s: no TaskFlow arg bindings arrived but the TaskInput struct declares "+ + "%d bindable field(s); nothing can fill them on this execution path", + p.fnName, len(plan.fields), + ) + } + + byName := make(map[string]int, len(args)) + for i, a := range args { + if a != nil { + byName[a.ArgName()] = i + } + } + + claimed := make([]bool, len(args)) + type fieldBind struct { + field taskInputField + argIdx int + } + binds := make([]fieldBind, 0, len(plan.fields)) for _, tif := range plan.fields { idx, ok := byName[tif.argName] if !ok { // No TaskFlow call argument carries this name -- kwarg-style, an // unpassed name leaves the field at its Go zero value rather than - // failing the task (unlike a flat data parameter, where arity is - // checked strictly; see the package doc comment). + // failing the task (see the package doc comment). continue } claimed[idx] = true + binds = append(binds, fieldBind{field: tif, argIdx: idx}) + } - v, err := p.resolveOne( - ctx, c, tif.fieldType, args[idx], - fmt.Sprintf("TaskInput field %s (parameter %d)", tif.goName, plan.index), - fmt.Sprintf("TaskInput field %s", tif.goName), + var unclaimed []string + for i, c := range claimed { + if c { + continue + } + if lit, ok := args[i].(LiteralArg); ok && lit.FromDefault { + // The Dag author never passed this argument; the Python side filled + // it from the stub signature's default. A struct that does not + // mirror the defaulted parameter is fine. + continue + } + name := "" + if args[i] != nil { + name = fmt.Sprintf("%q", args[i].ArgName()) + } + unclaimed = append(unclaimed, name) + } + if len(unclaimed) > 0 { + return nil, fmt.Errorf( + "task function %s: %d TaskFlow call argument(s) not claimed by any TaskInput "+ + "field: %s", + p.fnName, len(unclaimed), strings.Join(unclaimed, ", "), + ) + } + + raws, err := p.fetchArgValues(ctx, c, args, claimed) + if err != nil { + return nil, err + } + + structType := plan.typ + isPtr := structType.Kind() == reflect.Pointer + if isPtr { + structType = structType.Elem() + } + structVal := reflect.New(structType).Elem() + for _, b := range binds { + v, err := p.decodeArg( + args[b.argIdx], raws[b.argIdx], b.field.fieldType, + fmt.Sprintf("TaskInput field %s (parameter %d)", b.field.goName, plan.index), ) if err != nil { - return reflect.Value{}, err + return nil, err } - structVal.Field(tif.structIndex).Set(v) + structVal.Field(b.field.structIndex).Set(v) } if isPtr { - return structVal.Addr(), nil + out[paramIdx] = structVal.Addr() + } else { + out[paramIdx] = structVal } - return structVal, nil + return out, nil } -// resolveOne decodes one argument-spec entry into a value assignable to -// targetType: type-check against the declared Dag type, then decode a -// literal or pull-and-decode an XCom. Shared by resolveData (one flat -// parameter) and resolveTaskInput (one TaskInput struct field) so a literal- -// or xcom-kind entry resolves identically regardless of which parameter -// shape it fills. typeCheckCtx/generalCtx are error-message prefixes: the -// former (used only for the type-check error) additionally names the -// parameter index, matching this package's existing error conventions. -func (p *Plan) resolveOne( +// fetchArgValues produces the raw (pre-decode) value for each argument the +// caller will consume: a literal's inline value, or the upstream task's +// return-value XCom pulled over the API -- independent pulls run +// concurrently. needed selects which entries to fetch; nil means all. +func (p *Plan) fetchArgValues( ctx context.Context, c sdk.XComClient, - targetType reflect.Type, + args []Arg, + needed []bool, +) ([]any, error) { + raws := make([]any, len(args)) + var xcomIdxs []int + for i, a := range args { + if needed != nil && !needed[i] { + continue + } + switch a := a.(type) { + case LiteralArg: + raws[i] = a.Value + case XComArg: + xcomIdxs = append(xcomIdxs, i) + } + // A nil or foreign Arg implementation fails in decodeArg, which names + // the destination parameter/field in the error. + } + if len(xcomIdxs) == 0 { + return raws, nil + } + + workload, ok := ctx.Value(sdkcontext.WorkloadContextKey).(api.ExecuteTaskWorkload) + if !ok { + return nil, fmt.Errorf( + "task function %s: no workload in context, cannot resolve xcom arguments", p.fnName, + ) + } + // Always the return-value XCom -- a stub Dag cannot reference any other + // key. Pull from the upstream's unmapped instance (map_index nil); mapped + // upstream fan-in is out of scope for now. + pull := func(i int) error { + a := args[i].(XComArg) + raw, err := c.GetXCom( + ctx, workload.TI.DagId, workload.TI.RunId, a.TaskID, nil, api.XComReturnValueKey, nil, + ) + if err != nil { + return fmt.Errorf( + "task function %s: argument %q: pulling xcom from task %q: %w", + p.fnName, a.Name, a.TaskID, err, + ) + } + raws[i] = raw + return nil + } + if len(xcomIdxs) == 1 { + if err := pull(xcomIdxs[0]); err != nil { + return nil, err + } + return raws, nil + } + var wg sync.WaitGroup + errs := make([]error, len(xcomIdxs)) + for j, i := range xcomIdxs { + wg.Add(1) + go func() { + defer wg.Done() + errs[j] = pull(i) + }() + } + wg.Wait() + if err := errors.Join(errs...); err != nil { + return nil, err + } + return raws, nil +} + +// decodeArg decodes one argument-spec entry's raw value into targetType: +// type-check against the declared Dag type, then a strict decode. errCtx +// names the destination parameter or TaskInput field for error messages. +func (p *Plan) decodeArg( arg Arg, - typeCheckCtx string, - generalCtx string, + raw any, + targetType reflect.Type, + errCtx string, ) (reflect.Value, error) { if arg == nil { return reflect.Value{}, fmt.Errorf( - "task function %s: %s: nil argument binding", p.fnName, generalCtx, + "task function %s: %s: nil argument binding", p.fnName, errCtx, ) } if err := checkDataType(arg.DeclaredType(), targetType); err != nil { - return reflect.Value{}, fmt.Errorf("task function %s: %s: %w", p.fnName, typeCheckCtx, err) + return reflect.Value{}, fmt.Errorf("task function %s: %s: %w", p.fnName, errCtx, err) } + var source string switch a := arg.(type) { case LiteralArg: - v, err := decodeValue(a.Value, targetType) - if err != nil { - return reflect.Value{}, fmt.Errorf( - "task function %s: %s: decoding literal value into %s: %w", - p.fnName, generalCtx, targetType, err, - ) - } - return v, nil + source = "literal value" case XComArg: - workload, ok := ctx.Value(sdkcontext.WorkloadContextKey).(api.ExecuteTaskWorkload) - if !ok { - return reflect.Value{}, fmt.Errorf( - "task function %s: %s: no workload in context, cannot resolve xcom argument", - p.fnName, generalCtx, - ) - } - // Always the return-value XCom -- a stub Dag cannot reference any other - // key. Pull from the upstream's unmapped instance (map_index nil); - // mapped upstream fan-in is out of scope for now. - raw, err := c.GetXCom( - ctx, workload.TI.DagId, workload.TI.RunId, a.TaskID, nil, api.XComReturnValueKey, nil, - ) - if err != nil { - return reflect.Value{}, fmt.Errorf( - "task function %s: %s: pulling xcom from task %q: %w", - p.fnName, generalCtx, a.TaskID, err, - ) - } - v, err := decodeValue(raw, targetType) - if err != nil { - return reflect.Value{}, fmt.Errorf( - "task function %s: %s: decoding xcom from task %q into %s: %w", - p.fnName, generalCtx, a.TaskID, targetType, err, - ) - } - return v, nil + source = fmt.Sprintf("xcom from task %q", a.TaskID) default: return reflect.Value{}, fmt.Errorf( - "task function %s: %s: unsupported argument binding %T", p.fnName, generalCtx, arg, + "task function %s: %s: unsupported argument binding %T", p.fnName, errCtx, arg, + ) + } + v, err := decodeValue(raw, targetType) + if err != nil { + return reflect.Value{}, fmt.Errorf( + "task function %s: %s: decoding %s into %s: %w", + p.fnName, errCtx, source, targetType, err, ) } + return v, nil } // classifyParam decides how a single parameter is filled. Injectable runtime diff --git a/go-sdk/pkg/binding/binding_test.go b/go-sdk/pkg/binding/binding_test.go index b82324f19f1a3..ee17046e3e675 100644 --- a/go-sdk/pkg/binding/binding_test.go +++ b/go-sdk/pkg/binding/binding_test.go @@ -21,6 +21,7 @@ import ( "context" "log/slog" "reflect" + "sync" "testing" "github.com/google/uuid" @@ -40,10 +41,12 @@ func TestBindingSuite(t *testing.T) { } // fakeXComClient records GetXCom calls and returns preconfigured values. +// Resolve pulls XComs concurrently, so recording is mutex-guarded. type fakeXComClient struct { sdk.Client values map[string]any // "/" -> raw value + mu sync.Mutex calls []fakeXComCall err error } @@ -60,7 +63,9 @@ func (f *fakeXComClient) GetXCom( key string, _ any, ) (any, error) { + f.mu.Lock() f.calls = append(f.calls, fakeXComCall{dagID, runID, taskID, key, mapIndex}) + f.mu.Unlock() if f.err != nil { return nil, f.err } @@ -68,7 +73,7 @@ func (f *fakeXComClient) GetXCom( } // workloadCtx returns a context carrying an ExecuteTaskWorkload the resolver -// reads the dag/run identifiers from. +// reads the Dag/run identifiers from. func workloadCtx() context.Context { return context.WithValue( context.Background(), @@ -102,12 +107,12 @@ func (s *BindingSuite) TestAnalyzeClassification() { return nil }, ) - s.Equal(2, plan.NumData()) + s.Equal(2, plan.numData) - s.Zero(analyze(s, func() error { return nil }).NumData()) + s.Zero(analyze(s, func() error { return nil }).numData) s.Equal( 1, - analyze(s, func(x any) error { return nil }).NumData(), + analyze(s, func(x any) error { return nil }).numData, "an `any` parameter is a data parameter", ) } @@ -350,16 +355,20 @@ func (s *BindingSuite) TestResolveXComArgs() { s.Equal("probe-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 XCom argument always pulls the return-value key", - ) - s.Nil(client.calls[0].mapIndex, "v1 always pulls the unmapped upstream instance") - s.Equal(api.XComReturnValueKey, client.calls[1].key) + taskIDs := make([]string, 0, 2) + for _, call := range client.calls { + // Pulls run concurrently, so assert per-call properties order-independently. + taskIDs = append(taskIDs, call.taskID) + s.Equal("dag1", call.dagID) + s.Equal("run1", call.runID) + s.Equal( + api.XComReturnValueKey, + call.key, + "an XCom argument always pulls the return-value key", + ) + s.Nil(call.mapIndex, "v1 always pulls the unmapped upstream instance") + } + s.ElementsMatch([]string{"extract", "probe"}, taskIDs) } func (s *BindingSuite) TestResolveXComStrictStructDecode() { @@ -458,14 +467,14 @@ func (s *BindingSuite) TestResolveTIRunContextRebuild() { func (s *BindingSuite) TestAnalyzeTaskInputClassification() { plan := analyze(s, func(input simpleTaskInput) error { return nil }) - s.Zero(plan.NumData(), "a TaskInput struct claims by name, not by position") + s.Zero(plan.numData, "a TaskInput struct claims by name, not by position") ptrPlan := analyze(s, func(input *simpleTaskInput) error { return nil }) - s.Zero(ptrPlan.NumData(), "a pointer to a TaskInput struct is detected the same way") + s.Zero(ptrPlan.numData, "a pointer to a TaskInput struct is detected the same way") plainPlan := analyze(s, func(cfg nonEmbeddingStruct) error { return nil }) s.Equal( - 1, plainPlan.NumData(), + 1, plainPlan.numData, "a plain struct without the TaskInput sentinel stays a whole-value data parameter", ) } @@ -574,6 +583,43 @@ func (s *BindingSuite) TestResolveTaskInputUnclaimedArgFailsLoudly() { } } +func (s *BindingSuite) TestResolveTaskInputUnclaimedFromDefaultAllowed() { + fn := func(input simpleTaskInput) error { return nil } + got, err := s.resolve(fn, []Arg{ + LiteralArg{Name: "Name", Value: "widget", DataType: DataTypeString}, + // The Dag author never passed "threshold"; Python captured it from the + // stub signature's default. The struct need not mirror it. + LiteralArg{Name: "threshold", Value: 0.75, DataType: DataTypeNumber, FromDefault: true}, + }, &fakeXComClient{}) + s.Require().NoError(err) + s.Equal("widget", got[0].Interface().(simpleTaskInput).Name) +} + +func (s *BindingSuite) TestResolveTaskInputEmptySpecFailsLoudly() { + fn := func(input simpleTaskInput) error { return nil } + for name, args := range map[string][]Arg{"nil-spec": nil, "empty-spec": {}} { + s.Run(name, func() { + _, err := s.resolve(fn, args, &fakeXComClient{}) + // The Edge Worker path delivers no arg bindings; a struct with + // bindable fields must fail rather than run fully zero-valued. + if s.Assert().Error(err) { + s.Contains(err.Error(), "no TaskFlow arg bindings arrived") + } + }) + } +} + +func (s *BindingSuite) TestResolveTaskInputOnlyDefaultsSpecZeroValuesUnmatchedFields() { + fn := func(input twoFieldTaskInput) error { return nil } + got, err := s.resolve(fn, []Arg{ + LiteralArg{Name: "threshold", Value: 0.75, DataType: DataTypeNumber, FromDefault: true}, + }, &fakeXComClient{}) + s.Require().NoError(err) + input := got[0].Interface().(twoFieldTaskInput) + s.Equal("", input.Name, "no explicit entry arrived; fields keep kwarg-style zero values") + s.Equal("", input.Missing) +} + func (s *BindingSuite) TestResolveTaskInputUnmatchedArgNameZeroValuedAlongsideMatch() { fn := func(input twoFieldTaskInput) error { return nil } got, err := s.resolve(fn, []Arg{ diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index 33909abeecf4a..127f00d32f421 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -835,6 +835,13 @@ type GetAssetEventByAsset struct { // Name corresponds to the JSON schema field "name". Name interface{} `msgpack:"name"` + // PartitionKey corresponds to the JSON schema field "partition_key". + PartitionKey interface{} `msgpack:"partition_key,omitempty"` + + // PartitionKeyRegexpPattern corresponds to the JSON schema field + // "partition_key_regexp_pattern". + PartitionKeyRegexpPattern interface{} `msgpack:"partition_key_regexp_pattern,omitempty"` + // Type corresponds to the JSON schema field "type". Type string `msgpack:"type,omitempty"` @@ -861,6 +868,13 @@ type GetAssetEventByAssetAlias struct { // Limit corresponds to the JSON schema field "limit". Limit interface{} `msgpack:"limit,omitempty"` + // PartitionKey corresponds to the JSON schema field "partition_key". + PartitionKey interface{} `msgpack:"partition_key,omitempty"` + + // PartitionKeyRegexpPattern corresponds to the JSON schema field + // "partition_key_regexp_pattern". + PartitionKeyRegexpPattern interface{} `msgpack:"partition_key_regexp_pattern,omitempty"` + // Type corresponds to the JSON schema field "type". Type string `msgpack:"type,omitempty"` } @@ -1268,6 +1282,9 @@ type LiteralArgBinding struct { // DataType corresponds to the JSON schema field "data_type". DataType ArgBindingDataType `msgpack:"data_type,omitempty"` + // FromDefault corresponds to the JSON schema field "from_default". + FromDefault bool `msgpack:"from_default,omitempty"` + // Kind corresponds to the JSON schema field "kind". Kind string `msgpack:"kind"` @@ -1650,41 +1667,19 @@ type TaskBreadcrumbsResult struct { type TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{} -// Task callback status information. -// -// A Class with information about the success/failure TI callback to be executed. -// Currently, only failure -// callbacks when tasks are externally killed or experience heartbeat timeouts are -// run via DagFileProcessorProcess. -type TaskCallbackRequest struct { - // BundleName corresponds to the JSON schema field "bundle_name". - BundleName string `msgpack:"bundle_name"` - - // BundleVersion corresponds to the JSON schema field "bundle_version". - BundleVersion interface{} `msgpack:"bundle_version"` - - // ContextFromServer corresponds to the JSON schema field "context_from_server". - ContextFromServer *TIRunContext `msgpack:"context_from_server,omitempty"` - - // Filepath corresponds to the JSON schema field "filepath". - Filepath string `msgpack:"filepath"` - - // Msg corresponds to the JSON schema field "msg". - Msg interface{} `msgpack:"msg,omitempty"` - - // TaskCallbackType corresponds to the JSON schema field "task_callback_type". - TaskCallbackType interface{} `msgpack:"task_callback_type,omitempty"` - - // TI corresponds to the JSON schema field "ti". - TI TaskInstance `msgpack:"ti"` +type TriggerKwargs map[string]JsonValue - // Type corresponds to the JSON schema field "type". - 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"` - // VersionData corresponds to the JSON schema field "version_data". - VersionData *VersionData `msgpack:"version_data,omitempty"` + // Value corresponds to the JSON schema field "value". + Value interface{} `msgpack:"value"` } +type Warnings []interface{} + type TaskIds []string // Schema for TaskInstance model with minimal required fields needed for Runtime. @@ -1720,24 +1715,60 @@ type TaskInstance struct { TryNumber int `msgpack:"try_number"` } -type TaskInstanceState string - const TaskInstanceStateAwaitingInput TaskInstanceState = "awaiting_input" const TaskInstanceStateDeferred TaskInstanceState = "deferred" const TaskInstanceStateFailed TaskInstanceState = "failed" -const TaskInstanceStateQueued TaskInstanceState = "queued" -const TaskInstanceStateRemoved TaskInstanceState = "removed" const TaskInstanceStateRestarting TaskInstanceState = "restarting" -const TaskInstanceStateRunning TaskInstanceState = "running" -const TaskInstanceStateScheduled TaskInstanceState = "scheduled" const TaskInstanceStateSkipped TaskInstanceState = "skipped" const TaskInstanceStateSuccess TaskInstanceState = "success" +const TaskInstanceStateRunning TaskInstanceState = "running" const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" type TaskOutlets []AssetProfile +const TaskInstanceStateQueued TaskInstanceState = "queued" +const TaskInstanceStateScheduled TaskInstanceState = "scheduled" +const TaskInstanceStateRemoved TaskInstanceState = "removed" + +type TaskInstanceState string + +// Task callback status information. +// +// A Class with information about the success/failure TI callback to be executed. +// Currently, only failure +// callbacks when tasks are externally killed or experience heartbeat timeouts are +// run via DagFileProcessorProcess. +type TaskCallbackRequest struct { + // BundleName corresponds to the JSON schema field "bundle_name". + BundleName string `msgpack:"bundle_name"` + + // BundleVersion corresponds to the JSON schema field "bundle_version". + BundleVersion interface{} `msgpack:"bundle_version"` + + // ContextFromServer corresponds to the JSON schema field "context_from_server". + ContextFromServer *TIRunContext `msgpack:"context_from_server,omitempty"` + + // Filepath corresponds to the JSON schema field "filepath". + Filepath string `msgpack:"filepath"` + + // Msg corresponds to the JSON schema field "msg". + Msg interface{} `msgpack:"msg,omitempty"` + + // TaskCallbackType corresponds to the JSON schema field "task_callback_type". + TaskCallbackType interface{} `msgpack:"task_callback_type,omitempty"` + + // TI corresponds to the JSON schema field "ti". + TI TaskInstance `msgpack:"ti"` + + // 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"` +} + // Response containing the first reschedule date for a task instance. type TaskRescheduleStartDate struct { // StartDate corresponds to the JSON schema field "start_date". @@ -1747,6 +1778,12 @@ type TaskRescheduleStartDate struct { Type string `msgpack:"type,omitempty"` } +type TaskStateState string + +const TaskStateStateFailed TaskStateState = "failed" +const TaskStateStateSkipped TaskStateState = "skipped" +const TaskStateStateRemoved TaskStateState = "removed" + // Update a task's state. // // If a process exits without sending one of these the state will be derived from @@ -1767,12 +1804,6 @@ type TaskState struct { Type string `msgpack:"type,omitempty"` } -type TaskStateState string - -const TaskStateStateFailed TaskStateState = "failed" -const TaskStateStateRemoved TaskStateState = "removed" -const TaskStateStateSkipped TaskStateState = "skipped" - // Response to GetTaskStateStore; wraps the generated API response for supervisor // to worker comms. type TaskStateStoreResult struct { @@ -1822,7 +1853,7 @@ type TriggerDagRun struct { Type string `msgpack:"type,omitempty"` } -type TriggerKwargs map[string]JsonValue +type VersionData map[string]interface{} // Update the response content part of an existing Human-in-the-loop response. type UpdateHITLDetail struct { @@ -1858,19 +1889,6 @@ type VariableKeysResult struct { Type string `msgpack:"type,omitempty"` } -// Variable schema for responses with fields that are needed for Runtime. -type VariableResponse struct { - // Key corresponds to the JSON schema field "key". - Key string `msgpack:"key"` - - // Value corresponds to the JSON schema field "value". - Value interface{} `msgpack:"value"` -} - -type VersionData map[string]interface{} - -type Warnings []interface{} - type VariableResult struct { // Key corresponds to the JSON schema field "key". Key string `msgpack:"key"` diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index 2eaa87c4f8e02..bbd7369e6d028 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -373,6 +373,55 @@ func TestTaskRunnerBindsTaskInputStructArgs(t *testing.T) { assert.Equal(t, "eu-west-1", got.Region) } +// TestTaskRunnerTaskInputIgnoresUnclaimedDefault: convertArgBindings must +// propagate from_default so a spec entry the Python side filled from the stub +// signature's default may go unclaimed by the TaskInput struct. +func TestTaskRunnerTaskInputIgnoresUnclaimedDefault(t *testing.T) { + var got combineInput + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTaskWithName("transform", + func(input combineInput) error { + got = input + return nil + }) + }) + + details := &genmodels.StartupDetails{ + TI: genmodels.TaskInstance{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "transform", + RunID: "run1", + MapIndex: ptr(-1), + }, + BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, + TIContext: genmodels.TIRunContext{ + ArgBindings: &genmodels.ArgBindings{ + map[string]any{ + "name": "region", + "kind": "literal", + "data_type": "string", + "value": "eu-west-1", + }, + map[string]any{ + "name": "threshold", + "kind": "literal", + "data_type": "number", + "value": 0.75, + "from_default": true, + }, + }, + }, + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(context.Background(), bundle, details, comm, logger) + assertSucceedTask(t, result) + assert.Equal(t, "eu-west-1", got.Region) +} + // TestTaskRunnerArgBindingsTypeMismatch: a declared Dag type that cannot bind to // the Go parameter type fails the task loudly before the body runs. func TestTaskRunnerArgBindingsTypeMismatch(t *testing.T) { @@ -530,6 +579,39 @@ func TestTaskRunnerArgBindingsMissingRequiredFields(t *testing.T) { } } +// TestTaskRunnerMalformedSpecHonorsShouldRetry: a spec that fails +// convertArgBindings terminates with the same retry semantics as a binding +// failure inside executeTask, not an unconditional FAILED. +func TestTaskRunnerMalformedSpecHonorsShouldRetry(t *testing.T) { + bundle := buildBundle(t, func(r bundlev1.Registry) { + r.AddDag("test_dag").AddTaskWithName("transform", + func(country string) error { return nil }) + }) + + details := &genmodels.StartupDetails{ + TI: genmodels.TaskInstance{ + ID: "550e8400-e29b-41d4-a716-446655440000", + DagID: "test_dag", + TaskID: "transform", + RunID: "run1", + MapIndex: ptr(-1), + }, + BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, + TIContext: genmodels.TIRunContext{ + ShouldRetry: true, + ArgBindings: &genmodels.ArgBindings{ + map[string]any{"name": "country", "kind": "template", "value": "x"}, + }, + }, + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) + + result := RunTask(context.Background(), bundle, details, comm, logger) + assertRetryTask(t, result, `unknown kind "template"`) +} + func TestRunTaskHonorsContextCancellation(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { r.AddDag("test_dag").AddTaskWithName("ctxcheck", diff --git a/go-sdk/pkg/execution/task_runner.go b/go-sdk/pkg/execution/task_runner.go index e4be65f061bd8..4577bd5a5bde3 100644 --- a/go-sdk/pkg/execution/task_runner.go +++ b/go-sdk/pkg/execution/task_runner.go @@ -132,6 +132,14 @@ func RunTask( "task_id", details.TI.TaskID, "error", err, ) + // Same retry semantics as a binding failure inside executeTask: an + // equally permanent spec error must not terminate differently. + if details.TIContext.ShouldRetry { + return genmodels.RetryTask{ + EndDate: time.Now().UTC(), + RetryReason: err.Error(), + } + } return genmodels.TaskState{ State: genmodels.TaskStateStateFailed, EndDate: time.Now().UTC(), @@ -177,11 +185,13 @@ func convertArgBindings(specsPtr *genmodels.ArgBindings) ([]binding.Arg, error) } args[i] = binding.XComArg{Kind: kind, Name: name, TaskID: taskID, DataType: dataType} case "literal": + fromDefault, _ := m["from_default"].(bool) args[i] = binding.LiteralArg{ - Kind: kind, - Name: name, - Value: m["value"], - DataType: dataType, + Kind: kind, + Name: name, + Value: m["value"], + DataType: dataType, + FromDefault: fromDefault, } default: return nil, fmt.Errorf("arg_bindings[%d]: unknown kind %q", i, kind) diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index 3834be9f08a8b..ce4d00e99d85d 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -18,7 +18,6 @@ from __future__ import annotations import ast -import enum import inspect import json import types @@ -35,61 +34,47 @@ task_decorator_factory, ) -try: - from airflow.sdk.api.datamodels._generated import ArgBindingDataType -except ImportError: # Airflow < 3.4 -- the generated models do not carry the enum yet - - class ArgBindingDataType(str, enum.Enum): # type: ignore[no-redef] - """Language-neutral value type a stub-task argument binds to in the foreign runtime.""" - - STRING = "string" - INTEGER = "integer" - NUMBER = "number" - BOOLEAN = "boolean" - OBJECT = "object" - ARRAY = "array" - ANY = "any" - - if TYPE_CHECKING: from airflow.providers.common.compat.sdk import Context -def _data_type_from_annotation(annotation: Any) -> ArgBindingDataType: +def _infer_data_type(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. + The returned name is one of the execution API's ``ArgBindingDataType`` values + (``string``/``integer``/``number``/``boolean``/``object``/``array``/``any``); the foreign + runtime type-checks the bound value against it. 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 ArgBindingDataType.ANY + return "any" origin = typing.get_origin(annotation) if origin is not None: if origin is Union or origin is types.UnionType: 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 ArgBindingDataType.ANY + return _infer_data_type(members[0]) + return "any" annotation = origin if not isinstance(annotation, type): - return ArgBindingDataType.ANY + return "any" # bool subclasses int, and str/bytes are Sequences -- order matters. if issubclass(annotation, bool): - return ArgBindingDataType.BOOLEAN + return "boolean" if issubclass(annotation, int): - return ArgBindingDataType.INTEGER + return "integer" if issubclass(annotation, float): - return ArgBindingDataType.NUMBER + return "number" if issubclass(annotation, str): - return ArgBindingDataType.STRING + return "string" if issubclass(annotation, bytes): - return ArgBindingDataType.ANY + return "any" if issubclass(annotation, (dict, Mapping)): - return ArgBindingDataType.OBJECT + return "object" if issubclass(annotation, (list, tuple, set, frozenset, Sequence)): - return ArgBindingDataType.ARRAY - return ArgBindingDataType.ANY + return "array" + return "any" def _build_arg_bindings( @@ -106,8 +91,13 @@ def _build_arg_bindings( outputs, or a ``LiteralArgBinding`` (``kind="literal"``) for everything else. ``name`` is always the stub function's parameter name, so a foreign runtime can bind by name (e.g. the Go SDK's ``sdk.TaskInput`` struct fields) in addition to the existing positional order. - Returns ``None`` for parameterless stubs. + Returns ``None`` for argless calls: the binding contract (including the signature checks + below) applies only once a TaskFlow call actually passes arguments, so pre-TaskFlow stub + Dags whose call arguments were always ignored keep parsing. """ + if not op_args and not op_kwargs: + return None + signature = inspect.signature(python_callable) for param in signature.parameters.values(): @@ -123,10 +113,8 @@ def _build_arg_bindings( "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) + explicitly_bound = set(bound.arguments) bound.apply_defaults() try: @@ -136,7 +124,7 @@ def _build_arg_bindings( # TYPE_CHECKING with ``from __future__ import annotations``) degrade to "any". hints = {} - def annotation_for(name: str, param: inspect.Parameter) -> Any: + def get_annotation_for(name: str, param: inspect.Parameter) -> Any: if name in hints: return hints[name] if isinstance(param.annotation, str): @@ -146,7 +134,7 @@ def annotation_for(name: str, param: inspect.Parameter) -> Any: 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)) + data_type = _infer_data_type(get_annotation_for(name, param)) if isinstance(value, PlainXComArg): if value.key != "return_value": raise ValueError( @@ -170,14 +158,17 @@ def annotation_for(name: str, param: inspect.Parameter) -> Any: "language boundary -- .map()/.zip()/.concat() results are not supported" ) try: - json.dumps(value) + json.dumps(value, allow_nan=False) except (TypeError, ValueError): raise ValueError( f"@task.stub task {task_id!r} parameter {name!r} received a literal of type " f"{type(value).__name__} that is not JSON-serializable, so it cannot be passed " "to the foreign runtime" ) - spec.append({"name": name, "kind": "literal", "data_type": data_type, "value": value}) + entry: dict[str, Any] = {"name": name, "kind": "literal", "data_type": data_type, "value": value} + if name not in explicitly_bound: + entry["from_default"] = True + spec.append(entry) return spec @@ -185,8 +176,10 @@ 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. + # receive yet. The task-sdk decorator machinery rejects direct .expand() at parse time + # for operator classes that opt out on Airflow >= 3.4 (older cores cannot enforce it, + # and never serialize a spec for the mapped stub); stubs called with arguments inside + # a mapped task group are rejected in __init__ below. supports_expand: bool = False def __init__( @@ -235,6 +228,17 @@ def __init__( # 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) + # supports_expand only blocks direct .expand() on the stub itself; a mapped task + # group still creates per-map-index instances of every task inside it, and the + # captured spec has no map-index dimension to bind against. + in_mapped_group = getattr(self, "get_closest_mapped_task_group", lambda: None)() is not None + if self._arg_bindings is not None and in_mapped_group: + raise ValueError( + f"@task.stub task {self.task_id!r} passes TaskFlow call arguments inside a mapped " + "task group; per-map-index arg specs cannot cross the language boundary yet, so " + "stub tasks with arguments are not supported under a task group's .expand()" + ) + @classmethod def get_serialized_fields(cls): return super().get_serialized_fields() | {"_arg_bindings"} diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index 6babbbd75de1f..007e1c68a12da 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -22,8 +22,8 @@ import pytest -from airflow.providers.common.compat.sdk import DAG -from airflow.providers.standard.decorators.stub import ArgBindingDataType, _data_type_from_annotation, stub +from airflow.providers.common.compat.sdk import DAG, task_group +from airflow.providers.standard.decorators.stub import _infer_data_type, stub from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS, AIRFLOW_V_3_4_PLUS @@ -104,7 +104,13 @@ def test_literal_and_xcom_spec(self): assert op._arg_bindings == [ {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "fn_extract"}, - {"name": "retries_num", "kind": "literal", "data_type": "integer", "value": 3}, + { + "name": "retries_num", + "kind": "literal", + "data_type": "integer", + "value": 3, + "from_default": True, + }, ] assert op.upstream_task_ids == {"fn_extract"} @@ -150,20 +156,34 @@ def fn(x): ... def test_varargs_rejected(self): with pytest.raises(ValueError, match="fixed number of parameters"): - stub(fn_varargs)() + stub(fn_varargs)(1, 2) def test_varkw_rejected(self): with pytest.raises(ValueError, match="fixed number of parameters"): - stub(fn_kwonly_varkw)() + stub(fn_kwonly_varkw)(x=1) def test_context_key_param_rejected(self): with pytest.raises(ValueError, match="is an Airflow context key"): stub(fn_context_key)(1) + @pytest.mark.parametrize("fn", [fn_varargs, fn_kwonly_varkw, fn_context_key], ids=lambda f: f.__name__) + def test_argless_call_skips_signature_checks(self, fn): + """Pre-TaskFlow stub Dags never passed arguments; their signatures must keep parsing.""" + assert stub(fn)().operator._arg_bindings is None + + def test_argless_call_captures_no_spec_for_defaulted_params(self): + def fn(limit: int = 10): ... + + assert stub(fn)().operator._arg_bindings is None + def test_non_json_literal_rejected(self): with DAG(dag_id="d"), pytest.raises(ValueError, match="not JSON-serializable"): stub(fn_transform)("uk", object()) + def test_nan_literal_rejected(self): + with DAG(dag_id="d"), pytest.raises(ValueError, match="not JSON-serializable"): + stub(fn_transform)("uk", {"ratio": float("nan")}) + def test_mapped_xcom_arg_rejected(self): with DAG(dag_id="d"): extracted = stub(fn_extract)() @@ -185,7 +205,13 @@ def test_arg_bindings_survive_dag_serialization_round_trip(self): assert round_tripped.task_dict["fn_transform"]._arg_bindings == [ {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "fn_extract"}, - {"name": "retries_num", "kind": "literal", "data_type": "integer", "value": 3}, + { + "name": "retries_num", + "kind": "literal", + "data_type": "integer", + "value": 3, + "from_default": True, + }, ] @pytest.mark.skipif( @@ -196,39 +222,56 @@ def test_expand_rejected_at_parse_time(self): with pytest.raises(TypeError, match="do not support dynamic task mapping"): stub(fn_transform).expand(country=["uk", "fr"], extracted=[{}, {}]) + def test_stub_with_args_inside_mapped_task_group_rejected(self): + @task_group + def group(n): + stub(fn_transform)("uk", {}) + + with DAG(dag_id="d"): + with pytest.raises(ValueError, match="mapped task group"): + group.expand(n=[1, 2]) + + def test_argless_stub_inside_mapped_task_group_allowed(self): + @task_group + def group(n): + stub(fn_extract)() + + with DAG(dag_id="d"): + group.expand(n=[1, 2]) + @pytest.mark.parametrize( ("annotation", "expected"), [ - pytest.param(str, ArgBindingDataType.STRING, id="str"), - pytest.param(bool, ArgBindingDataType.BOOLEAN, id="bool"), - pytest.param(int, ArgBindingDataType.INTEGER, id="int"), - pytest.param(float, ArgBindingDataType.NUMBER, id="float"), - pytest.param(dict, ArgBindingDataType.OBJECT, id="dict"), - pytest.param(dict[str, int], ArgBindingDataType.OBJECT, id="dict-parameterized"), - pytest.param(typing.Mapping[str, int], ArgBindingDataType.OBJECT, id="mapping"), - pytest.param(list, ArgBindingDataType.ARRAY, id="list"), - pytest.param(list[int], ArgBindingDataType.ARRAY, id="list-parameterized"), - pytest.param(tuple, ArgBindingDataType.ARRAY, id="tuple"), - pytest.param(set, ArgBindingDataType.ARRAY, id="set"), - pytest.param(typing.Sequence[int], ArgBindingDataType.ARRAY, id="sequence"), - pytest.param(Any, ArgBindingDataType.ANY, id="any"), - pytest.param(None, ArgBindingDataType.ANY, id="none"), - pytest.param(bytes, ArgBindingDataType.ANY, id="bytes"), + 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], # noqa: UP045 -- legacy form on purpose - ArgBindingDataType.STRING, + "string", id="optional-str", ), pytest.param( typing.Union[int, str], # noqa: UP007 -- legacy form on purpose - ArgBindingDataType.ANY, + "any", id="union", ), - pytest.param(str | None, ArgBindingDataType.STRING, id="pep604-optional"), - pytest.param(int | str, ArgBindingDataType.ANY, id="pep604-union"), - pytest.param(contextlib.AbstractContextManager, ArgBindingDataType.ANY, id="custom-class"), + 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) is expected +def test_infer_data_type(annotation, expected): + assert _infer_data_type(annotation) == expected diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 696fdcc6c1365..d0bc470c7ae96 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -241,6 +241,7 @@ class LiteralArgBinding(BaseModel): name: Annotated[str, Field(title="Name")] data_type: ArgBindingDataType | None = ArgBindingDataType.ANY value: JsonValue | None = None + from_default: Annotated[bool | None, Field(title="From Default")] = False class PrevSuccessfulDagRunResponse(BaseModel): 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 c1298dff6c743..43d100d7a743e 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -4603,6 +4603,11 @@ } ], "default": null + }, + "from_default": { + "default": false, + "title": "From Default", + "type": "boolean" } }, "required": [ 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 ce26296519b2d..6785fde63be23 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 @@ -422,6 +422,13 @@ def startup_details(self): arg_bindings=[ {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract"}, + { + "name": "limit", + "kind": "literal", + "data_type": "integer", + "value": 10, + "from_default": True, + }, ], ), sentry_integration="", @@ -440,10 +447,13 @@ 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 - literal, xcom = (a.root for a in out.ti_context.arg_bindings) + literal, xcom, defaulted = (a.root for a in out.ti_context.arg_bindings) assert isinstance(literal, LiteralArgBinding) assert literal.value == "uk" assert literal.name == "country" + assert literal.from_default is False assert isinstance(xcom, XComArgBinding) assert xcom.task_id == "extract" assert xcom.name == "extracted" + assert isinstance(defaulted, LiteralArgBinding) + assert defaulted.from_default is True From 0f900688607533ffc4da1534031f7d8c1d1fbd66 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 22 Jul 2026 02:58:53 +0000 Subject: [PATCH 17/40] Move the Go SDK arg-binding runtime to a stacked follow-up PR Reviewing the Python-side arg-binding contract and the Go runtime that consumes it in one PR ties the core/task-sdk review to Go SDK internals. Scoping this PR to the contract lets it merge on its own; the Go SDK consumption (pkg/binding, task-runner dispatch, example bundle, e2e test) lands stacked on top from feature/go-sdk/taskflow-arg-binding. --- .../test_go_sdk_taskflow_binding.py | 146 ---- go-sdk/README.md | 71 +- .../0003-coordinator-protocol-msgpack-ipc.md | 12 +- go-sdk/bundle/bundlev1/task.go | 154 +++- go-sdk/bundle/bundlev1/task_test.go | 86 +- .../airflow-go-pack/pack_integration_test.go | 12 +- go-sdk/dags/go_examples.py | 120 +-- go-sdk/example/bundle/main.go | 28 +- go-sdk/example/bundle/main_test.go | 6 +- .../bundle/taskflowbinding/taskflowbinding.go | 272 ------ .../taskflowbinding/taskflowbinding_test.go | 129 --- go-sdk/pkg/binding/binding.go | 782 ------------------ go-sdk/pkg/binding/binding_test.go | 632 -------------- go-sdk/pkg/execution/frames.go | 7 - .../pkg/execution/genmodels/defaults.gen.go | 22 - go-sdk/pkg/execution/genmodels/models.gen.go | 186 ++--- go-sdk/pkg/execution/integration_test.go | 382 --------- go-sdk/pkg/execution/messages.go | 2 +- go-sdk/pkg/execution/task_runner.go | 95 +-- go-sdk/sdk/context.go | 22 - 20 files changed, 196 insertions(+), 2970 deletions(-) delete mode 100644 airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_taskflow_binding.py delete mode 100644 go-sdk/example/bundle/taskflowbinding/taskflowbinding.go delete mode 100644 go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go delete mode 100644 go-sdk/pkg/binding/binding.go delete mode 100644 go-sdk/pkg/binding/binding_test.go 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 deleted file mode 100644 index 8e35177a33d3f..0000000000000 --- a/airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_taskflow_binding.py +++ /dev/null @@ -1,146 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -"""E2E test for the Go SDK ``taskflow_binding_dag`` example. - -The stub Dag's single mixed positional/keyword TaskFlow call carries literals -of every scalar type, an array literal, a defaulted ``None``, and XComs from -two upstream Go tasks (an object bound onto a strict Go struct and an array -bound onto ``[]int``). The Go ``via_flat_args`` task verifies every bound -value and errors on any mismatch, so a green run *is* the binding assertion; -the tests here check the run outcome and the summary XCom it pushes. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from datetime import datetime, timezone - -import pytest - -from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient - -# Three short Go tasks; allow room for coordinator startup. -_GO_TASK_TIMEOUT = 300 - -_DAG_ID = "taskflow_binding_dag" - - -@dataclass -class _CompletedRun: - """The single ``taskflow_binding_dag`` run shared across this module's tests.""" - - client: AirflowClient - run_id: str - state: str - ti_states: dict[str, str] - - def xcom(self, task_id: str, key: str = "return_value"): - return self.client.get_xcom_value(dag_id=_DAG_ID, task_id=task_id, run_id=self.run_id, key=key).get( - "value" - ) - - -@pytest.fixture(scope="module") -def completed_run() -> _CompletedRun: - """Trigger ``taskflow_binding_dag`` once and wait for it to finish.""" - client = AirflowClient() - resp = client.trigger_dag(_DAG_ID, json={"logical_date": datetime.now(timezone.utc).isoformat()}) - run_id = resp["dag_run_id"] - state = client.wait_for_dag_run(dag_id=_DAG_ID, run_id=run_id, timeout=_GO_TASK_TIMEOUT) - ti_resp = client.get_task_instances(dag_id=_DAG_ID, run_id=run_id) - ti_states = {ti["task_id"]: ti.get("state") for ti in ti_resp.get("task_instances", [])} - return _CompletedRun(client=client, run_id=run_id, state=state, ti_states=ti_states) - - -def test_all_tasks_succeeded(completed_run: _CompletedRun): - """The Go ``via_flat_args`` task errors on any mis-bound argument, so success here - proves every literal, XCom, keyword, and defaulted-None binding was correct.""" - assert completed_run.state == "success", ( - f"expected the run to succeed; got {completed_run.state!r}. task states: {completed_run.ti_states}" - ) - for task_id in ( - "make_config", - "make_numbers", - "make_region", - "via_flat_args", - "via_struct_no_tags", - "via_struct_arg_tag", - "via_struct_unmatched_arg", - ): - assert completed_run.ti_states.get(task_id) == "success", completed_run.ti_states - - -def test_upstream_xcoms_keep_their_shapes(completed_run: _CompletedRun): - """The Go struct arrives as an object XCom, the ``[]int`` as an array, the region as a string.""" - assert completed_run.xcom("make_config") == { - "environment": "production", - "region": "eu-west-1", - "debug": True, - } - assert completed_run.xcom("make_numbers") == [1, 1, 2, 3, 5, 8] - assert completed_run.xcom("make_region") == "eu-west-1" - - -def test_via_flat_args_summary_reflects_bound_arguments(completed_run: _CompletedRun): - """``via_flat_args`` re-emits every bound value, confirming types survived the - Python literal / XCom -> Go parameter -> XCom round trip.""" - assert completed_run.xcom("via_flat_args") == { - "name": "summary", - "count": 3, - "ratio": 2.5, - "enabled": True, - "tags": ["metrics", "hourly"], - "environment": "production", - "debug": True, - "sum": 20, - "note_was_null": True, - } - - -def test_via_struct_no_tags_reflects_bound_arguments(completed_run: _CompletedRun): - """``via_struct_no_tags`` demonstrates the Go SDK's ``sdk.TaskInput`` struct-field - injection mode with no field tags at all: each field binds the TaskFlow argument - spelled exactly like its Go field name (``RegionCode``, ``Threshold``). The region - is ``make_region``'s XCom, so a struct field binds an XCom-sourced value here.""" - assert completed_run.xcom("via_struct_no_tags") == { - "region_code": "eu-west-1", - "threshold": 0.75, - } - - -def test_via_struct_arg_tag_reflects_bound_arguments(completed_run: _CompletedRun): - """``via_struct_arg_tag`` demonstrates explicit ``arg:`` tags: ``Region`` is - genuinely renamed to ``region_code`` (bound from ``make_region``'s XCom), and - ``Threshold`` is tagged ``threshold`` to pull the snake_case literal its - verbatim field name would miss.""" - assert completed_run.xcom("via_struct_arg_tag") == { - "region": "eu-west-1", - "threshold": 0.75, - } - - -def test_via_struct_unmatched_arg_reflects_zero_valued_field(completed_run: _CompletedRun): - """``via_struct_unmatched_arg`` demonstrates mismatch tolerance in both directions: - a struct field whose name has no corresponding TaskFlow call argument stays at its - Go zero value instead of failing the task (kwarg-style, an unpassed name simply - isn't bound), and the stub's defaulted ``sample_rate`` -- captured into the spec as - ``from_default`` -- needs no matching struct field. The task succeeding at all - proves the second half.""" - assert completed_run.xcom("via_struct_unmatched_arg") == { - "region": "eu-west-1", - "missing_was_empty": True, - } diff --git a/go-sdk/README.md b/go-sdk/README.md index 7b2c0e856a044..7bd6b3d13c81a 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -105,17 +105,6 @@ A task is an ordinary Go function. The runtime inspects its signature and inject `sdk.VariableClient`). An optional `(any, error)` return becomes the task's XCom; an `error` return marks the task failed. -Any other parameter is a **data parameter**: in declaration order, data parameters receive the -positional arguments of the Python stub Dag's TaskFlow call. A JSON-serializable literal in the Dag -file (`transform("uk", ...)`) decodes straight into the parameter; an upstream task output -(`transform(..., extract())`) is pulled from that task's XCom in the current Dag run and decoded into -the parameter's type (independent XCom pulls run concurrently). The runtime fails the task loudly when -the argument count doesn't match the number of data parameters or a declared type can't bind to the Go -type. Data parameters must be JSON-decodable (no func/chan/unsafe-pointer, no non-empty interfaces) — -checked once at registration. TaskFlow argument binding arrives over the coordinator protocol, so it is -coordinator-mode only today; on the Edge Worker path a task with data parameters fails with the arity -error (and a `TaskInput` struct with bindable fields fails the same way, since nothing can fill them). - ```go func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (any, error) { conn, err := client.GetConnection(ctx, "test_http") @@ -123,17 +112,12 @@ func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (any, er return map[string]any{"go_version": runtime.Version()}, nil } -// 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 { +func transform(ctx sdk.TIRunContext, client sdk.VariableClient, log *slog.Logger) error { val, err := client.GetVariable(ctx, "my_variable") if err != nil { return err } - log.Info("Obtained variable", "my_variable", val, "country", country) + log.Info("Obtained variable", "my_variable", val) return nil } ``` @@ -142,57 +126,6 @@ Asking for the narrowest interface a task needs (e.g. `sdk.VariableClient` inste unit testing easier and documents which Airflow features the task touches. `RegisterDags` is the single source of truth for which `dag_id`s and `task_id`s a bundle can run. -### TaskInput structs - -A struct that anonymously embeds `sdk.TaskInput` opts into **per-field, name-based** binding instead -of a long flat parameter list. At most one such parameter is allowed per function, and it cannot be -combined with plain flat data parameters — a task function declares one shape or the other, and -registration fails on a signature that mixes them. - -Conceptually, a plain flat parameter list is **positional-argument** binding: order matters, and -every parameter must be filled or the task fails before its body runs. A `TaskInput` struct is -closer to **keyword-argument** binding: fields match by name instead of position, and (see the -`arg:` bullet below) a field whose name has no corresponding TaskFlow call argument is simply left -at its Go zero value rather than failing the task — the same way an unpassed keyword argument falls -back to a caller-side default in a kwargs-style call. - -```go -type CombineInput struct { - sdk.TaskInput // one-line opt-in, zero runtime cost - Region string `arg:"region_code"` // named lookup against the TaskFlow call argument "region_code" - Threshold float64 `arg:"threshold"` // tags also bridge Go's UpperCamelCase to a snake_case argument -} - -func Combine(ctx sdk.TIRunContext, log *slog.Logger, input CombineInput) (any, error) { - // input.Region and input.Threshold are both populated. - return nil, nil -} -``` - -Each exported field binds from the TaskFlow call argument named by its optional `arg:""` tag -(matched against the stub function's Python parameter name, independent of declaration order on -either side). With no tag, the field's own Go name is matched verbatim — an untagged `Threshold` -only binds an argument literally spelled `Threshold`, so a snake_case Python parameter needs an -explicit tag. If no TaskFlow call argument carries that name, the field is simply left at its Go -zero value — it does not fail the task, kwarg-style (see `ViaStructUnmatchedArg` below). - -The matching is checked in the other direction too: every argument the Dag author **explicitly -passed** in the TaskFlow call must be claimed by some field, so a typo'd field name fails the task -instead of silently dropping the value. Stub parameters the author left at their Python defaults -are the exception — the Python side captures them into the spec (marked `from_default` on the -wire), and the struct is free not to mirror them, the same way a Python callee never sees which -defaulted kwargs went unpassed. And when no argument spec arrives at all (an argless stub call, or -the Edge Worker path) a `TaskInput` struct with bindable fields fails loudly rather than running -fully zero-valued. - -A plain custom struct type *without* the `sdk.TaskInput` embed is unaffected by any of -this — it keeps working as a single flat data parameter, JSON-decoded whole from one TaskFlow -argument (see `Config` in -[`example/bundle/taskflowbinding/taskflowbinding.go`](./example/bundle/taskflowbinding/taskflowbinding.go)), -which is a different mechanism from per-field `TaskInput` binding. See -[`ViaStructNoTags`, `ViaStructArgTag`, and `ViaStructUnmatchedArg`](./example/bundle/taskflowbinding/taskflowbinding.go) -for a full worked example of each field-binding mode — and the unmatched-field case — in isolation. - ### Reading the task runtime context Declare an `sdk.TIRunContext` parameter on a task to read the identifiers and scheduling timestamps of the diff --git a/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md b/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md index e7535ab6bea75..82798bdeb10f4 100644 --- a/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md +++ b/go-sdk/adr/0003-coordinator-protocol-msgpack-ipc.md @@ -217,11 +217,7 @@ Supervisor Bundle binary (Go) │ │ ├── StartupDetails ────────────────────►│ │ (ti, dag_rel_path, bundle_info, │ - │ start_date, ti_context; the │ - │ ti_context carries arg_bindings, │ - │ the positional-argument spec │ - │ captured from the stub Dag's │ - │ TaskFlow call) │ + │ start_date, ti_context) │ │ │ │ ├── lookup task: │ │ bundle.dags[ti.dag_id] @@ -229,12 +225,6 @@ 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 cb18a4e951519..d31fea84b73f3 100644 --- a/go-sdk/bundle/bundlev1/task.go +++ b/go-sdk/bundle/bundlev1/task.go @@ -25,58 +25,30 @@ 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 _ TaskWithArgs = (*taskFunction)(nil) +var _ Task = (*taskFunction)(nil) // NewTaskFunction wraps a plain Go function as a Task, validating its signature -// (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. +// (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. func NewTaskFunction(fn any) (Task, error) { v := reflect.ValueOf(fn) fullName := runtime.FuncForPC(v.Pointer()).Name() - f := &taskFunction{fn: v, fullName: fullName} - if err := f.validateFn(v.Type()); err != nil { - // A half-built task (nil binding plan) would panic in Execute; a caller - // that mishandles the error must not be able to run it. - return nil, err - } - return f, nil + f := &taskFunction{v, fullName} + return f, f.validateFn(v.Type()) } func (f *taskFunction) Execute(ctx context.Context, logger *slog.Logger) error { - 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 { + fnType := f.fn.Type() var sdkClient sdk.Client if injected, ok := ctx.Value(sdkcontext.SdkClientContextKey).(sdk.Client); ok { sdkClient = injected @@ -84,14 +56,41 @@ func (f *taskFunction) ExecuteArgs( sdkClient = sdk.NewClient() } - reflectArgs, err := f.plan.Resolve(ctx, logger, sdkClient, args) - if err != nil { - return err + 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) + } } slog.Debug("Attempting to call fn", "fn", f.fn, "args", reflectArgs) retValues := f.fn.Call(reflectArgs) - err = nil + var err error if errResult := retValues[len(retValues)-1].Interface(); errResult != nil { var ok bool if err, ok = errResult.(error); !ok { @@ -151,11 +150,11 @@ func (f *taskFunction) validateFn(fnType reflect.Type) error { ) } - plan, err := binding.Analyze(fnType, f.fullName) - if err != nil { - return err + 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) + } } - f.plan = plan return nil } @@ -169,8 +168,75 @@ func isValidResultType(inType reflect.Type) bool { return true } -var errorType = reflect.TypeFor[error]() +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]() +) 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 6136529a0806d..3f58e930b8721 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,9 +141,21 @@ 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 -// parameters Execute can neither inject nor bind a task argument to. This -// replaces the historical silent zero-fill of unrecognized parameters. +// interface parameters Execute cannot inject. func (s *TaskSuite) TestNonInjectableParamsAreRejected() { cases := map[string]struct { fn any @@ -162,9 +174,9 @@ func (s *TaskSuite) TestNonInjectableParamsAreRejected() { }, "method GetVariable is func(context.Context, string) (string, error) on sdk.Client", }, - "func-param": { - func(cb func()) error { return nil }, - "cannot receive a task argument", + "empty-interface": { + func(x any) error { return nil }, + "empty interfaces cannot be injected", }, "context-with-extra-methods": { func(x interface { @@ -189,68 +201,6 @@ func (s *TaskSuite) TestNonInjectableParamsAreRejected() { } } -// TestExecuteArgsBindsDataParameters covers the TaskFlow path end to end at the -// task level: literals decode onto data parameters interleaved with -// injectables, and Execute (nil args) keeps working for argless functions. -func (s *TaskSuite) TestExecuteArgsBindsDataParameters() { - var gotCountry string - var gotMeta map[string]any - task, err := NewTaskFunction(func(log *slog.Logger, country string, meta map[string]any) error { - gotCountry = country - gotMeta = meta - return nil - }) - s.Require().NoError(err) - - tw, ok := task.(TaskWithArgs) - s.Require().True(ok, "taskFunction must implement TaskWithArgs") - - err = tw.ExecuteArgs(context.Background(), slog.New(logging.NewTeeLogger()), []binding.Arg{ - binding.LiteralArg{Value: "uk", DataType: binding.DataTypeString}, - binding.LiteralArg{ - 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{ - binding.LiteralArg{Value: "uk"}, - binding.LiteralArg{Value: "de"}, - }, - ) - if s.Assert().Error(err) { - s.Contains(err.Error(), "argument count mismatch") - s.Contains(err.Error(), "passes 2 positional argument(s)") - } -} - // probeKey is an unexported context key used to confirm the live task context // (not a freshly built one) backs the injected sdk.TIRunContext. type probeKeyType struct{} diff --git a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go index 3be151660a8a3..77725b0ac4efc 100644 --- a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go +++ b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go @@ -34,7 +34,6 @@ 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 @@ -143,7 +142,7 @@ func TestPack_CrossArchExecutableWithMetadataFile(t *testing.T) { sdk: language: "go" version: "` + sdkVersion + `" - supervisor_schema_version: "` + execution.SupervisorSchemaVersion + `" + supervisor_schema_version: "2026-06-16" source: "main.go" dags: concurrent_xcom_dag: @@ -154,15 +153,6 @@ dags: - "extract" - "transform" - "load" - taskflow_binding_dag: - tasks: - - "make_config" - - "make_numbers" - - "make_region" - - "via_flat_args" - - "via_struct_no_tags" - - "via_struct_arg_tag" - - "via_struct_unmatched_arg" ` assert.Equal(t, expectedManifest, string(metadata)) diff --git a/go-sdk/dags/go_examples.py b/go-sdk/dags/go_examples.py index 01e8821f1cf62..23e02dd5e49dc 100644 --- a/go-sdk/dags/go_examples.py +++ b/go-sdk/dags/go_examples.py @@ -17,14 +17,9 @@ """ Python stub Dags mirroring the Go SDK example bundle (``go-sdk/example/bundle``). -Three Dags, all backed by the same Go bundle: ``simple_dag`` (extract/transform/ -load, below), ``concurrent_xcom_dag`` (one ``pull_xcoms_concurrently`` task -timing sequential vs goroutine XCom pulls), and ``taskflow_binding_dag`` -(stressing the TaskFlow argument-binding surface -- the flat, positional -parameter list ``via_flat_args`` binds onto, plus three ``sdk.TaskInput`` -(keyword-style) struct examples, ``via_struct_no_tags``/``via_struct_arg_tag``/ -``via_struct_unmatched_arg``, each isolating one field-binding mode; see its -Dag function below). +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). ``simple_dag`` sandwiches the Go tasks between two native Python tasks so the run exercises XCom across the language boundary, the same way @@ -38,10 +33,6 @@ 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 @@ -74,7 +65,7 @@ def extract(): ... @task.stub(queue="golang") -def transform(country: str, extracted: dict): ... +def transform(): ... # ``load`` fails on its first attempt and succeeds on the retry, exercising the @@ -95,10 +86,7 @@ def python_task_2(extracted): @dag(dag_id="simple_dag") def simple_dag(): extracted = extract() - # 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) + transformed = transform() 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 @@ -119,101 +107,3 @@ def concurrent_xcom_dag(): concurrent_xcom_dag() - - -@task.stub(queue="golang") -def make_config(): ... - - -@task.stub(queue="golang") -def make_numbers(): ... - - -@task.stub(queue="golang") -def make_region(): ... - - -@task.stub(queue="golang") -def via_flat_args( - name: str, - count: int, - ratio: float, - enabled: bool, - tags: list, - config: dict, - numbers: list, - note: str | None = None, -): ... - - -# Capitalized parameters on purpose: with no ``arg:`` tags on the Go side, each -# struct field binds the argument spelled exactly like its Go field name. -@task.stub(queue="golang") -def via_struct_no_tags(RegionCode: str, Threshold: float): ... - - -@task.stub(queue="golang") -def via_struct_arg_tag(region_code: str, threshold: float): ... - - -@task.stub(queue="golang") -def via_struct_unmatched_arg(region_code: str, sample_rate: float = 0.1): ... - - -@dag(dag_id="taskflow_binding_dag") -def taskflow_binding_dag(): - """ - Stress the TaskFlow argument-binding surface beyond ``simple_dag``'s transform. - - Conceptually, the flat parameter list is *positional-argument* binding: order - matters, and every parameter must be filled or the task fails before it runs. - ``sdk.TaskInput`` structs are closer to *keyword-argument* binding: fields match - by name, and (see ``via_struct_unmatched_arg`` below) a field whose name has no - corresponding TaskFlow call argument simply stays at its zero value instead of - failing the task -- the same way an unpassed keyword argument falls back to a - default in kwargs-style calls. - - ``via_flat_args``'s one mixed positional/keyword call carries literals of every - scalar type plus an array literal, and fans in XComs from *two* upstream Go - tasks: ``make_config`` returns an object that binds onto a strictly-decoded Go - struct, ``make_numbers`` an array that binds onto ``[]int``. ``note`` is not - passed, so its ``None`` default is captured and arrives in Go as a nil - ``*string``. The Go ``via_flat_args`` (``go-sdk/example/bundle/taskflowbinding``) - verifies every bound value and fails the task on any mismatch. - - Three further tasks demonstrate the Go SDK's ``sdk.TaskInput`` struct injection - mode. Each call mixes a literal (``threshold``) with an XCom reference: the - ``region_code`` argument is ``make_region``'s output, so every struct example - also proves an XCom-sourced value binds onto a struct field. One field-binding - mode at a time: - - * ``via_struct_no_tags``: no ``arg:`` tags at all -- each struct field binds - the argument spelled exactly like its Go field name, hence this stub's - capitalized ``RegionCode``/``Threshold`` parameters. - * ``via_struct_arg_tag``: every field names its argument via an explicit - ``arg:`` tag -- ``Region`` is genuinely renamed to ``region_code``, and - ``Threshold`` is tagged ``threshold`` to pull the snake_case argument its - verbatim field name would miss. - * ``via_struct_unmatched_arg``: the mismatch tolerance in both directions. - The Go struct declares a field with no corresponding argument in this - TaskFlow call at all -- it stays at its Go zero value rather than failing - the task. And the stub's defaulted ``sample_rate`` is never passed, so its - captured-from-default entry needs no matching struct field (an explicitly - passed argument no field claims would fail the task instead). - """ - via_flat_args( - "summary", - 3, - 2.5, - True, - ["metrics", "hourly"], - config=make_config(), - numbers=make_numbers(), - ) - region = make_region() - via_struct_no_tags(RegionCode=region, Threshold=0.75) - via_struct_arg_tag(region_code=region, threshold=0.75) - via_struct_unmatched_arg(region_code=region) - - -taskflow_binding_dag() diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go index f4d4696c4ed42..23e60bd1dd46e 100644 --- a/go-sdk/example/bundle/main.go +++ b/go-sdk/example/bundle/main.go @@ -27,7 +27,6 @@ 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" ) @@ -56,15 +55,6 @@ func (m *myBundle) RegisterDags(dagbag v1.Registry) error { concurrentDag := dagbag.AddDag("concurrent_xcom_dag") concurrentDag.AddTaskWithName("pull_xcoms_concurrently", concurrentxcom.PullXComsConcurrently) - bindingDag := dagbag.AddDag("taskflow_binding_dag") - bindingDag.AddTaskWithName("make_config", taskflowbinding.MakeConfig) - bindingDag.AddTaskWithName("make_numbers", taskflowbinding.MakeNumbers) - bindingDag.AddTaskWithName("make_region", taskflowbinding.MakeRegion) - bindingDag.AddTaskWithName("via_flat_args", taskflowbinding.ViaFlatArgs) - bindingDag.AddTaskWithName("via_struct_no_tags", taskflowbinding.ViaStructNoTags) - bindingDag.AddTaskWithName("via_struct_arg_tag", taskflowbinding.ViaStructArgTag) - bindingDag.AddTaskWithName("via_struct_unmatched_arg", taskflowbinding.ViaStructUnmatchedArg) - return nil } @@ -137,28 +127,12 @@ func extract(ctx sdk.TIRunContext, client sdk.Client, log *slog.Logger) (any, er return ret, nil } -// 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 { +func transform(ctx sdk.TIRunContext, client sdk.VariableClient, log *slog.Logger) 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 84bb174533812..474a8406d6224 100644 --- a/go-sdk/example/bundle/main_test.go +++ b/go-sdk/example/bundle/main_test.go @@ -52,10 +52,8 @@ 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. The data parameters - // (country, extracted) are passed directly, exactly as the runtime would bind them from the - // stub Dag's TaskFlow call. + // sdk.NewTIRunContext wraps any context to build a TIRunContext in a test. ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) - err := transform(ctx, &mockVars{}, log, "uk", map[string]any{"go_version": "go1.24"}) + err := transform(ctx, &mockVars{}, log) assert.NoError(t, err) } diff --git a/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go b/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go deleted file mode 100644 index 7d7fdcd6dc0a1..0000000000000 --- a/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go +++ /dev/null @@ -1,272 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -// Package taskflowbinding holds the taskflow_binding_dag tasks. ViaFlatArgs is -// positional-argument binding pushed to its limit: literals of every scalar -// type, an array literal, keyword arguments, a defaulted null, and XCom fan-in -// from two upstream Go tasks decoded into a strict struct and a typed slice -- -// where simple_dag's transform shows the minimal case (one literal, one -// XCom), this shows the full argument surface. The ViaStruct* functions -// instead show the sdk.TaskInput struct-field injection mode -- conceptually -// keyword-argument binding, where fields match by name and an unmatched name -// is left at its zero value rather than failing the task -- one field-binding -// mode at a time: ViaStructNoTags (verbatim field-name fallback), -// ViaStructArgTag (explicit `arg:` naming), and ViaStructUnmatchedArg (a -// field whose name has no corresponding TaskFlow call argument at all). Each -// ViaStruct* call binds MakeRegion's XCom onto its region field alongside a -// literal, so struct fields are exercised with both argument sources. -package taskflowbinding - -import ( - "fmt" - "log/slog" - "reflect" - - "github.com/apache/airflow/go-sdk/sdk" -) - -// Config is the object make_config returns as its XCom; via_flat_args declares -// the same struct as a parameter, so the round trip exercises strict struct -// decoding (an unknown or renamed key fails the task rather than silently -// zeroing a field). -type Config struct { - Environment string `json:"environment"` - Region string `json:"region"` - Debug bool `json:"debug"` -} - -// MakeConfig pushes an object XCom that via_flat_args binds onto its Config parameter. -func MakeConfig(log *slog.Logger) (any, error) { - cfg := Config{Environment: "production", Region: "eu-west-1", Debug: true} - log.Info( - "Pushing config", - "environment", - cfg.Environment, - "region", - cfg.Region, - "debug", - cfg.Debug, - ) - return cfg, nil -} - -// MakeNumbers pushes an array XCom that via_flat_args binds onto its []int parameter. -func MakeNumbers(log *slog.Logger) (any, error) { - numbers := []int{1, 1, 2, 3, 5, 8} - log.Info("Pushing numbers", "numbers", fmt.Sprint(numbers)) - return numbers, nil -} - -// MakeRegion pushes a string XCom that every ViaStruct* task binds onto a -// struct field, so each field-binding mode is exercised with an XCom-sourced -// argument and not just literals. -func MakeRegion(log *slog.Logger) (any, error) { - region := "eu-west-1" - log.Info("Pushing region", "region", region) - return region, nil -} - -// ViaFlatArgs receives every argument shape the stub Dag can express as plain, -// positional data parameters. The Python side calls it as -// -// via_flat_args("summary", 3, 2.5, True, ["metrics", "hourly"], -// config=make_config(), numbers=make_numbers()) -// -// so the bound values are fixed; any mismatch below is a binding regression -// and fails the task loudly. note is never passed and falls back to the stub's -// None default, arriving as a nil *string. -func ViaFlatArgs( - ctx sdk.TIRunContext, - log *slog.Logger, - name string, - count int, - ratio float64, - enabled bool, - tags []string, - config Config, - numbers []int, - note *string, -) (any, error) { - if name != "summary" || count != 3 || ratio != 2.5 || !enabled { - return nil, fmt.Errorf( - "scalar literals bound incorrectly: name=%q count=%d ratio=%v enabled=%v", - name, count, ratio, enabled, - ) - } - if want := []string{"metrics", "hourly"}; !reflect.DeepEqual(tags, want) { - return nil, fmt.Errorf("array literal bound incorrectly: tags=%v, want %v", tags, want) - } - if want := (Config{Environment: "production", Region: "eu-west-1", Debug: true}); config != want { - return nil, fmt.Errorf("object XCom bound incorrectly: config=%+v, want %+v", config, want) - } - if want := []int{1, 1, 2, 3, 5, 8}; !reflect.DeepEqual(numbers, want) { - return nil, fmt.Errorf("array XCom bound incorrectly: numbers=%v, want %v", numbers, want) - } - if note != nil { - return nil, fmt.Errorf("defaulted None bound incorrectly: note=%q, want nil", *note) - } - - sum := 0 - for _, n := range numbers { - sum += n - } - log.InfoContext(ctx, "Bound TaskFlow arguments", - "name", name, - "count", count, - "ratio", ratio, - "enabled", enabled, - "tags", fmt.Sprint(tags), - "environment", config.Environment, - "sum", sum, - ) - return map[string]any{ - "name": name, - "count": count, - "ratio": ratio, - "enabled": enabled, - "tags": tags, - "environment": config.Environment, - "debug": config.Debug, - "sum": sum, - "note_was_null": note == nil, - }, nil -} - -// ViaStructNoTagsInput demonstrates the sdk.TaskInput struct-field injection -// mode with no field tags at all: each field binds the TaskFlow call argument -// spelled exactly like its Go field name ("RegionCode", "Threshold"), which -// is why the stub declares capitalized parameters. -type ViaStructNoTagsInput struct { - sdk.TaskInput - RegionCode string - Threshold float64 -} - -// ViaStructNoTags is called as -// -// via_struct_no_tags(RegionCode=make_region(), Threshold=0.75) -// -// so RegionCode arrives via make_region's XCom and Threshold as a literal. -func ViaStructNoTags( - ctx sdk.TIRunContext, - log *slog.Logger, - input ViaStructNoTagsInput, -) (any, error) { - if input.RegionCode != "eu-west-1" || input.Threshold != 0.75 { - return nil, fmt.Errorf( - "TaskInput fields bound incorrectly: region_code=%q threshold=%v", - input.RegionCode, - input.Threshold, - ) - } - - log.InfoContext(ctx, "Bound TaskInput struct (no tags)", - "region_code", input.RegionCode, - "threshold", input.Threshold, - ) - return map[string]any{ - "region_code": input.RegionCode, - "threshold": input.Threshold, - }, nil -} - -// ViaStructArgTagInput demonstrates the sdk.TaskInput struct-field injection -// mode with explicit arg: tags: Region binds to the "region_code" TaskFlow -// argument under a renamed Go field, proving the tag remaps the name rather -// than coincidentally matching it; Threshold is intentionally tagged -// "threshold" because an untagged field would only match an argument spelled -// exactly "Threshold". -type ViaStructArgTagInput struct { - sdk.TaskInput - Region string `arg:"region_code"` - Threshold float64 `arg:"threshold"` -} - -// ViaStructArgTag is called as -// -// via_struct_arg_tag(region_code=make_region(), threshold=0.75) -// -// so Region arrives via make_region's XCom and Threshold as a literal. -func ViaStructArgTag( - ctx sdk.TIRunContext, - log *slog.Logger, - input ViaStructArgTagInput, -) (any, error) { - if input.Region != "eu-west-1" || input.Threshold != 0.75 { - return nil, fmt.Errorf( - "TaskInput fields bound incorrectly: region=%q threshold=%v", - input.Region, - input.Threshold, - ) - } - - log.InfoContext(ctx, "Bound TaskInput struct (arg: tag)", - "region", input.Region, - "threshold", input.Threshold, - ) - return map[string]any{ - "region": input.Region, - "threshold": input.Threshold, - }, nil -} - -// ViaStructUnmatchedArgInput demonstrates the mismatch tolerance in both -// directions. A field whose name has no corresponding TaskFlow call argument -// at all is left at its Go zero value rather than failing the task: Region -// binds normally, but Missing's arg name is never among this call's -// arguments -- conceptually, an unpassed keyword argument falling back to -// its default in a kwargs-style call. The reverse also holds: the stub's -// defaulted sample_rate parameter arrives marked from_default, so this -// struct is free not to mirror it (an explicitly passed argument no field -// claims would fail the task instead). -type ViaStructUnmatchedArgInput struct { - sdk.TaskInput - Region string `arg:"region_code"` - Missing string `arg:"does_not_exist"` -} - -// ViaStructUnmatchedArg is called as -// -// via_struct_unmatched_arg(region_code=make_region()) -// -// -- the stub declares region_code (bound from make_region's XCom) plus a -// defaulted sample_rate this struct deliberately omits, so Missing's arg -// name never appears among the call's arguments and stays at its Go zero -// value (""), while sample_rate's from_default entry goes unclaimed without -// failing the task. -func ViaStructUnmatchedArg( - ctx sdk.TIRunContext, log *slog.Logger, input ViaStructUnmatchedArgInput, -) (any, error) { - if input.Region != "eu-west-1" { - return nil, fmt.Errorf("TaskInput field bound incorrectly: region=%q", input.Region) - } - if input.Missing != "" { - return nil, fmt.Errorf( - "expected the unmatched field to stay at its Go zero value, got missing=%q", - input.Missing, - ) - } - - log.InfoContext(ctx, "Bound TaskInput struct (unmatched arg)", - "region", input.Region, - "missing_was_empty", input.Missing == "", - ) - return map[string]any{ - "region": input.Region, - "missing_was_empty": input.Missing == "", - }, nil -} diff --git a/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go b/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go deleted file mode 100644 index 648dec1473fde..0000000000000 --- a/go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go +++ /dev/null @@ -1,129 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -package taskflowbinding - -import ( - "context" - "log/slog" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - - "github.com/apache/airflow/go-sdk/sdk" -) - -// Like example/bundle/main_test.go, this shows a task fn is unit-testable by -// passing the data parameters directly, exactly as the runtime binds them. -func TestViaFlatArgs(t *testing.T) { - ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) - got, err := ViaFlatArgs(ctx, slog.Default(), - "summary", 3, 2.5, true, - []string{"metrics", "hourly"}, - Config{Environment: "production", Region: "eu-west-1", Debug: true}, - []int{1, 1, 2, 3, 5, 8}, - nil, - ) - require.NoError(t, err) - - summary, ok := got.(map[string]any) - require.True(t, ok, "ViaFlatArgs should return a map summary, got %T", got) - assert.Equal(t, 20, summary["sum"]) - assert.Equal(t, true, summary["note_was_null"]) -} - -func TestViaFlatArgsRejectsWrongBinding(t *testing.T) { - ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) - _, err := ViaFlatArgs(ctx, slog.Default(), - "summary", 3, 2.5, true, - []string{"metrics", "hourly"}, - Config{}, - []int{1, 1, 2, 3, 5, 8}, - nil, - ) - assert.ErrorContains(t, err, "object XCom bound incorrectly") -} - -func TestViaStructNoTags(t *testing.T) { - ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) - got, err := ViaStructNoTags(ctx, slog.Default(), ViaStructNoTagsInput{ - RegionCode: "eu-west-1", - Threshold: 0.75, - }) - require.NoError(t, err) - - summary, ok := got.(map[string]any) - require.True(t, ok, "ViaStructNoTags should return a map summary, got %T", got) - assert.Equal(t, "eu-west-1", summary["region_code"]) -} - -func TestViaStructNoTagsRejectsWrongBinding(t *testing.T) { - ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) - _, err := ViaStructNoTags(ctx, slog.Default(), ViaStructNoTagsInput{ - RegionCode: "wrong-region", - Threshold: 0.75, - }) - assert.ErrorContains(t, err, "TaskInput fields bound incorrectly") -} - -func TestViaStructArgTag(t *testing.T) { - ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) - got, err := ViaStructArgTag(ctx, slog.Default(), ViaStructArgTagInput{ - Region: "eu-west-1", - Threshold: 0.75, - }) - require.NoError(t, err) - - summary, ok := got.(map[string]any) - require.True(t, ok, "ViaStructArgTag should return a map summary, got %T", got) - assert.Equal(t, "eu-west-1", summary["region"]) -} - -func TestViaStructArgTagRejectsWrongBinding(t *testing.T) { - ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) - _, err := ViaStructArgTag(ctx, slog.Default(), ViaStructArgTagInput{ - Region: "wrong-region", - Threshold: 0.75, - }) - assert.ErrorContains(t, err, "TaskInput fields bound incorrectly") -} - -func TestViaStructUnmatchedArg(t *testing.T) { - ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) - // Missing is left at its Go zero value, exactly as binding.Resolve leaves an - // unmatched TaskInput field -- this task fn is unit-testable independent of - // the binding package precisely because it declares that expectation itself. - got, err := ViaStructUnmatchedArg(ctx, slog.Default(), ViaStructUnmatchedArgInput{ - Region: "eu-west-1", - Missing: "", - }) - require.NoError(t, err) - - summary, ok := got.(map[string]any) - require.True(t, ok, "ViaStructUnmatchedArg should return a map summary, got %T", got) - assert.Equal(t, true, summary["missing_was_empty"]) -} - -func TestViaStructUnmatchedArgRejectsNonZeroMissingField(t *testing.T) { - ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{}) - _, err := ViaStructUnmatchedArg(ctx, slog.Default(), ViaStructUnmatchedArgInput{ - Region: "eu-west-1", - Missing: "unexpected", - }) - assert.ErrorContains(t, err, "expected the unmatched field to stay at its Go zero value") -} diff --git a/go-sdk/pkg/binding/binding.go b/go-sdk/pkg/binding/binding.go deleted file mode 100644 index f37e2c1a7e10e..0000000000000 --- a/go-sdk/pkg/binding/binding.go +++ /dev/null @@ -1,782 +0,0 @@ -// 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. -// -// Three 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 (except TaskInput structs, below), in -// declaration order. They receive the positional arguments the Python -// stub Dag captured at parse time from the TaskFlow call -// (“transform("uk", extract())“) and delivered in StartupDetails. A -// literal argument decodes directly; an XCom argument is pulled from the -// named upstream task in the current Dag run first (independent pulls run -// concurrently). -// - TaskInput structs: a struct that anonymously embeds sdk.TaskInput opts -// into per-field, name-based binding instead of consuming one positional -// slot as a whole-value decode target. Each exported field binds by name -// against the Dag's TaskFlow call arguments: an `arg:""` tag names -// the argument to claim, and a field with no tag claims its own Go field -// name, verbatim. At most one such parameter is allowed per function. -// -// The two data-parameter shapes are mutually exclusive: a function declares -// either plain flat data parameters or one TaskInput struct, never both. -// Analyze rejects a signature that mixes them -- splitting one TaskFlow -// call's arguments between by-name claiming and positional order is too -// ambiguous to reason about. -// -// Conceptually, flat data parameters are positional-argument binding: order -// matters, and every parameter must be filled or Resolve fails the task -// before its body runs. A TaskInput struct is closer to keyword-argument -// binding: fields match by name, and a field whose name has no corresponding -// TaskFlow call argument is simply left at its Go zero value instead of -// failing the task -- the same way an unpassed keyword argument falls back -// to its default in a kwargs-style call. The check runs both ways: an -// explicitly passed argument that no field claims fails the task (catching -// typo'd field names), while spec entries the Python side captured from the -// stub signature's defaults (from_default on the wire) may go unclaimed. And -// a TaskInput struct with bindable fields fails loudly when no argument spec -// arrives at all (e.g. the Edge Worker path, where nothing could ever fill -// them), matching the flat-parameter arity check. -// -// 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 of -// flat data parameters. 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" - "errors" - "fmt" - "log/slog" - "reflect" - "strings" - "sync" - - "github.com/apache/airflow/go-sdk/pkg/api" - "github.com/apache/airflow/go-sdk/pkg/execution/genmodels" - "github.com/apache/airflow/go-sdk/pkg/sdkcontext" - "github.com/apache/airflow/go-sdk/sdk" -) - -// DataType is the language-neutral value type the Dag declared for an -// argument (from the stub function's annotation on the Python side). It -// aliases the enum generated from the supervisor schema so the vocabulary -// cannot drift from the wire model. -type DataType = genmodels.ArgBindingDataType - -const ( - DataTypeString = genmodels.ArgBindingDataTypeString - DataTypeInteger = genmodels.ArgBindingDataTypeInteger - DataTypeNumber = genmodels.ArgBindingDataTypeNumber - DataTypeBoolean = genmodels.ArgBindingDataTypeBoolean - DataTypeObject = genmodels.ArgBindingDataTypeObject - DataTypeArray = genmodels.ArgBindingDataTypeArray - DataTypeAny = genmodels.ArgBindingDataTypeAny -) - -// Arg is one positional argument for a task function's data parameters, in -// declaration order: an XComArg or a LiteralArg. A sealed sum type over the -// wire model's XComArgBinding/LiteralArgBinding split; each variant is -// defined in terms of its generated schema struct so the fields cannot drift -// from the coordinator protocol. -type Arg interface { - // ArgName is the stub function's parameter name this binding fills; used - // to match a TaskInput struct field's `arg:` tag (or its verbatim - // field-name fallback). - ArgName() string - // DeclaredType is the Dag-declared language-neutral type for the argument; - // empty is treated as DataTypeAny. - DeclaredType() DataType - // sealedArg restricts implementations to this package, keeping - // resolveOne's type switch the single exhaustive consumer. - sealedArg() -} - -// XComArg sources the argument from an upstream task's return-value XCom. -// Kind carries the wire discriminant ("xcom") from the generated shape; the -// resolve path dispatches on the Go type itself and never reads it. -type XComArg genmodels.XComArgBinding - -// LiteralArg carries an inline value from the Dag file. Kind carries the wire -// discriminant ("literal") from the generated shape; the resolve path -// dispatches on the Go type itself and never reads it. -type LiteralArg genmodels.LiteralArgBinding - -func (a XComArg) ArgName() string { return a.Name } -func (a LiteralArg) ArgName() string { return a.Name } - -func (a XComArg) DeclaredType() DataType { return a.DataType } -func (a LiteralArg) DeclaredType() DataType { return a.DataType } - -func (XComArg) sealedArg() {} -func (LiteralArg) sealedArg() {} - -// paramKind classifies how a task-function parameter is filled at execution. -type paramKind int - -const ( - paramTIRunContext paramKind = iota - paramContext - paramLogger - paramClient - paramData - // paramTaskInput is a struct that anonymously embeds sdk.TaskInput, - // opting into per-field, name-based binding instead of consuming one - // positional slot as a whole-value decode target. - paramTaskInput -) - -// taskInputField describes how Resolve fills one exported field of a -// TaskInput-embedding struct. Precomputed once by Analyze. -type taskInputField struct { - // structIndex is the field's index within the struct, for - // reflect.Value.Field. - structIndex int - // goName is the Go field name, for error messages. - goName string - fieldType reflect.Type - // argName is the name to claim from the argument spec: the field's `arg:` - // tag, or its Go field name verbatim when the tag is omitted. - argName string -} - -// paramPlan describes how Resolve fills a single task-function parameter. -type paramPlan struct { - kind paramKind - // typ is the declared Go type of a data parameter (kind == paramData) or - // the struct type (kind == paramTaskInput). - typ reflect.Type - // index is the parameter's position in the function signature, for error - // messages. - index int - // fields describes each exported field's binding. Set only for - // kind == paramTaskInput. - fields []taskInputField -} - -// 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 - hasTaskInput bool -} - -// Analyze inspects the parameters of a task function type and builds a Plan. -// fnName appears in error messages only. Every parameter must be an injectable -// runtime type or a type that can receive a task argument (JSON-decodable); -// anything else is a registration error. -func Analyze(fnType reflect.Type, fnName string) (*Plan, error) { - p := &Plan{fnName: fnName, params: make([]paramPlan, fnType.NumIn())} - seenTaskInput := -1 - for i := range fnType.NumIn() { - plan, err := classifyParam(fnName, fnType.In(i), i) - if err != nil { - return nil, err - } - if plan.kind == paramTaskInput { - if seenTaskInput >= 0 { - return nil, fmt.Errorf( - "task function %s: parameter %d: only one TaskInput struct parameter is allowed "+ - "per function (parameter %d already is one)", - fnName, - i, - seenTaskInput, - ) - } - seenTaskInput = i - } - if plan.kind == paramData { - p.numData++ - } - p.params[i] = plan - } - if seenTaskInput >= 0 { - p.hasTaskInput = true - if p.numData > 0 { - return nil, fmt.Errorf( - "task function %s: cannot mix a TaskInput struct parameter (parameter %d) with "+ - "plain data parameters; declare either flat data parameters or one TaskInput struct", - fnName, seenTaskInput, - ) - } - } - return p, nil -} - -// Resolve builds the ordered argument values for one call. Injectable -// parameters receive values derived from ctx, logger, or client. A TaskInput -// struct's fields claim entries out of args by name; plain flat data -// parameters consume args in declaration order (the two shapes are mutually -// exclusive; see Analyze). An error fails the task before its body runs. -// -// client must be the full sdk.Client -- not just the sdk.XComClient the -// resolve helpers narrow to -- because a paramClient parameter receives the -// client itself, verbatim. -func (p *Plan) Resolve( - ctx context.Context, - logger *slog.Logger, - client sdk.Client, - args []Arg, -) ([]reflect.Value, error) { - out := make([]reflect.Value, len(p.params)) - for i, plan := range p.params { - switch plan.kind { - case paramTIRunContext: - // The runtime stores the identifiers/timestamps under - // RuntimeContextKey; rebuild the value around the live task context. - var ti sdk.TaskInstance - var dagRun sdk.DagRun - if stored, ok := ctx.Value(sdkcontext.RuntimeContextKey).(sdk.TIRunContext); ok { - ti, dagRun = stored.TaskInstance(), stored.DagRun() - } - out[i] = reflect.ValueOf(sdk.NewTIRunContext(ctx, ti, dagRun)) - case paramContext: - out[i] = reflect.ValueOf(ctx) - case paramLogger: - out[i] = reflect.ValueOf(logger) - case paramClient: - out[i] = reflect.ValueOf(client) - case paramData, paramTaskInput: - // Filled below, once the spec is matched and XComs are pulled. - } - } - if p.hasTaskInput { - return p.resolveTaskInputParam(ctx, client, args, out) - } - return p.resolveFlatParams(ctx, client, args, out) -} - -// resolveFlatParams fills the plain data parameters positionally -- args[0] -// onto the first data parameter, and so on -- with strict arity (see the -// package doc comment). -func (p *Plan) resolveFlatParams( - ctx context.Context, - c sdk.XComClient, - args []Arg, - out []reflect.Value, -) ([]reflect.Value, error) { - if len(args) != p.numData { - return nil, fmt.Errorf( - "task function %s: argument count mismatch: the Dag passes %d positional argument(s) "+ - "but the Go function declares %d data parameter(s)", - p.fnName, len(args), p.numData, - ) - } - raws, err := p.fetchArgValues(ctx, c, args, nil) - if err != nil { - return nil, err - } - flatIdx := 0 - for i, plan := range p.params { - if plan.kind != paramData { - continue - } - v, err := p.decodeArg( - args[flatIdx], raws[flatIdx], plan.typ, - fmt.Sprintf("argument %d (parameter %d)", flatIdx, plan.index), - ) - if err != nil { - return nil, err - } - out[i] = v - flatIdx++ - } - return out, nil -} - -// resolveTaskInputParam fills the single TaskInput struct parameter by name -// (kwarg-style). Every explicitly passed spec entry must be claimed by some -// field; entries the Python side captured from the stub signature's defaults -// (FromDefault) may go unclaimed, the same way an unpassed keyword argument -// never reaches the callee. -func (p *Plan) resolveTaskInputParam( - ctx context.Context, - c sdk.XComClient, - args []Arg, - out []reflect.Value, -) ([]reflect.Value, error) { - var paramIdx int - var plan paramPlan - for i, pl := range p.params { - if pl.kind == paramTaskInput { - paramIdx, plan = i, pl - break - } - } - - if len(args) == 0 && len(plan.fields) > 0 { - return nil, fmt.Errorf( - "task function %s: no TaskFlow arg bindings arrived but the TaskInput struct declares "+ - "%d bindable field(s); nothing can fill them on this execution path", - p.fnName, len(plan.fields), - ) - } - - byName := make(map[string]int, len(args)) - for i, a := range args { - if a != nil { - byName[a.ArgName()] = i - } - } - - claimed := make([]bool, len(args)) - type fieldBind struct { - field taskInputField - argIdx int - } - binds := make([]fieldBind, 0, len(plan.fields)) - for _, tif := range plan.fields { - idx, ok := byName[tif.argName] - if !ok { - // No TaskFlow call argument carries this name -- kwarg-style, an - // unpassed name leaves the field at its Go zero value rather than - // failing the task (see the package doc comment). - continue - } - claimed[idx] = true - binds = append(binds, fieldBind{field: tif, argIdx: idx}) - } - - var unclaimed []string - for i, c := range claimed { - if c { - continue - } - if lit, ok := args[i].(LiteralArg); ok && lit.FromDefault { - // The Dag author never passed this argument; the Python side filled - // it from the stub signature's default. A struct that does not - // mirror the defaulted parameter is fine. - continue - } - name := "" - if args[i] != nil { - name = fmt.Sprintf("%q", args[i].ArgName()) - } - unclaimed = append(unclaimed, name) - } - if len(unclaimed) > 0 { - return nil, fmt.Errorf( - "task function %s: %d TaskFlow call argument(s) not claimed by any TaskInput "+ - "field: %s", - p.fnName, len(unclaimed), strings.Join(unclaimed, ", "), - ) - } - - raws, err := p.fetchArgValues(ctx, c, args, claimed) - if err != nil { - return nil, err - } - - structType := plan.typ - isPtr := structType.Kind() == reflect.Pointer - if isPtr { - structType = structType.Elem() - } - structVal := reflect.New(structType).Elem() - for _, b := range binds { - v, err := p.decodeArg( - args[b.argIdx], raws[b.argIdx], b.field.fieldType, - fmt.Sprintf("TaskInput field %s (parameter %d)", b.field.goName, plan.index), - ) - if err != nil { - return nil, err - } - structVal.Field(b.field.structIndex).Set(v) - } - - if isPtr { - out[paramIdx] = structVal.Addr() - } else { - out[paramIdx] = structVal - } - return out, nil -} - -// fetchArgValues produces the raw (pre-decode) value for each argument the -// caller will consume: a literal's inline value, or the upstream task's -// return-value XCom pulled over the API -- independent pulls run -// concurrently. needed selects which entries to fetch; nil means all. -func (p *Plan) fetchArgValues( - ctx context.Context, - c sdk.XComClient, - args []Arg, - needed []bool, -) ([]any, error) { - raws := make([]any, len(args)) - var xcomIdxs []int - for i, a := range args { - if needed != nil && !needed[i] { - continue - } - switch a := a.(type) { - case LiteralArg: - raws[i] = a.Value - case XComArg: - xcomIdxs = append(xcomIdxs, i) - } - // A nil or foreign Arg implementation fails in decodeArg, which names - // the destination parameter/field in the error. - } - if len(xcomIdxs) == 0 { - return raws, nil - } - - workload, ok := ctx.Value(sdkcontext.WorkloadContextKey).(api.ExecuteTaskWorkload) - if !ok { - return nil, fmt.Errorf( - "task function %s: no workload in context, cannot resolve xcom arguments", p.fnName, - ) - } - // Always the return-value XCom -- a stub Dag cannot reference any other - // key. Pull from the upstream's unmapped instance (map_index nil); mapped - // upstream fan-in is out of scope for now. - pull := func(i int) error { - a := args[i].(XComArg) - raw, err := c.GetXCom( - ctx, workload.TI.DagId, workload.TI.RunId, a.TaskID, nil, api.XComReturnValueKey, nil, - ) - if err != nil { - return fmt.Errorf( - "task function %s: argument %q: pulling xcom from task %q: %w", - p.fnName, a.Name, a.TaskID, err, - ) - } - raws[i] = raw - return nil - } - if len(xcomIdxs) == 1 { - if err := pull(xcomIdxs[0]); err != nil { - return nil, err - } - return raws, nil - } - var wg sync.WaitGroup - errs := make([]error, len(xcomIdxs)) - for j, i := range xcomIdxs { - wg.Add(1) - go func() { - defer wg.Done() - errs[j] = pull(i) - }() - } - wg.Wait() - if err := errors.Join(errs...); err != nil { - return nil, err - } - return raws, nil -} - -// decodeArg decodes one argument-spec entry's raw value into targetType: -// type-check against the declared Dag type, then a strict decode. errCtx -// names the destination parameter or TaskInput field for error messages. -func (p *Plan) decodeArg( - arg Arg, - raw any, - targetType reflect.Type, - errCtx string, -) (reflect.Value, error) { - if arg == nil { - return reflect.Value{}, fmt.Errorf( - "task function %s: %s: nil argument binding", p.fnName, errCtx, - ) - } - if err := checkDataType(arg.DeclaredType(), targetType); err != nil { - return reflect.Value{}, fmt.Errorf("task function %s: %s: %w", p.fnName, errCtx, err) - } - var source string - switch a := arg.(type) { - case LiteralArg: - source = "literal value" - case XComArg: - source = fmt.Sprintf("xcom from task %q", a.TaskID) - default: - return reflect.Value{}, fmt.Errorf( - "task function %s: %s: unsupported argument binding %T", p.fnName, errCtx, arg, - ) - } - v, err := decodeValue(raw, targetType) - if err != nil { - return reflect.Value{}, fmt.Errorf( - "task function %s: %s: decoding %s into %s: %w", - p.fnName, errCtx, source, targetType, err, - ) - } - return v, nil -} - -// classifyParam decides how a single parameter is filled. Injectable runtime -// types map to their paramKind; anything else is a data parameter and must be -// a type a task argument can decode into. -func classifyParam(fnName string, in reflect.Type, index int) (paramPlan, error) { - switch { - case isTIRunContext(in): - // sdk.TIRunContext embeds context.Context, so it also satisfies - // isContext - this case must come first. - return paramPlan{kind: paramTIRunContext, index: index}, nil - case isContext(in): - // The plain task context injected here cannot satisfy extra methods. - if !contextType.Implements(in) { - return paramPlan{}, fmt.Errorf( - "task function %s: parameter %d: interface %s adds methods on top of "+ - "context.Context; declare sdk.TIRunContext or a separate parameter instead", - fnName, index, in, - ) - } - return paramPlan{kind: paramContext, index: index}, nil - case isLogger(in): - return paramPlan{kind: paramLogger, index: index}, nil - case isClient(in): - return paramPlan{kind: paramClient, index: index}, nil - } - if in.Kind() == reflect.Interface && in.NumMethod() > 0 { - return paramPlan{}, fmt.Errorf( - "task function %s: parameter %d: interface %s is not injectable "+ - "(want context.Context, sdk.TIRunContext, or a subset of sdk.Client): %s", - fnName, index, in, explainClientMismatch(in), - ) - } - if structType := taskInputStructType(in); structType != nil { - fields, err := buildTaskInputFields(fnName, structType, index) - if err != nil { - return paramPlan{}, err - } - return paramPlan{kind: paramTaskInput, typ: in, index: index, fields: fields}, nil - } - 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 -} - -// taskInputStructType reports whether in (after dereferencing one pointer -// level, matching checkDataType's convention) is a struct that anonymously -// embeds sdk.TaskInput, returning that struct type or nil. -func taskInputStructType(in reflect.Type) reflect.Type { - t := in - if t.Kind() == reflect.Pointer { - t = t.Elem() - } - if t.Kind() != reflect.Struct { - return nil - } - for i := range t.NumField() { - f := t.Field(i) - if f.Anonymous && f.Type == taskInputType { - return t - } - } - return nil -} - -// buildTaskInputFields validates and precomputes the field-binding plan for a -// TaskInput-embedding struct parameter. It runs once at registration time so -// a misconfigured struct fails loudly before any task ever executes. -func buildTaskInputFields( - fnName string, - structType reflect.Type, - paramIndex int, -) ([]taskInputField, error) { - var fields []taskInputField - seenArgNames := make(map[string]string) // resolved arg name -> Go field name that claims it - - for i := range structType.NumField() { - f := structType.Field(i) - if f.Anonymous && f.Type == taskInputType { - continue - } - if !f.IsExported() { - continue - } - - if !isDecodableType(f.Type) { - return nil, fmt.Errorf( - "task function %s: parameter %d: TaskInput field %s: type %s cannot receive a task "+ - "argument (func/chan/unsafe-pointer values cannot be decoded)", - fnName, - paramIndex, - f.Name, - f.Type, - ) - } - - tif := taskInputField{structIndex: i, goName: f.Name, fieldType: f.Type} - tif.argName = f.Tag.Get("arg") - if tif.argName == "" { - tif.argName = f.Name - } - if existing, ok := seenArgNames[tif.argName]; ok { - return nil, fmt.Errorf( - "task function %s: parameter %d: TaskInput fields %s and %s both bind arg name %q", - fnName, paramIndex, existing, f.Name, tif.argName, - ) - } - seenArgNames[tif.argName] = f.Name - fields = append(fields, tif) - } - return fields, 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]() - taskInputType = reflect.TypeFor[sdk.TaskInput]() -) - -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 deleted file mode 100644 index ee17046e3e675..0000000000000 --- a/go-sdk/pkg/binding/binding_test.go +++ /dev/null @@ -1,632 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -package binding - -import ( - "context" - "log/slog" - "reflect" - "sync" - "testing" - - "github.com/google/uuid" - "github.com/stretchr/testify/suite" - - "github.com/apache/airflow/go-sdk/pkg/api" - "github.com/apache/airflow/go-sdk/pkg/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. -// Resolve pulls XComs concurrently, so recording is mutex-guarded. -type fakeXComClient struct { - sdk.Client - - values map[string]any // "/" -> raw value - mu sync.Mutex - calls []fakeXComCall - err error -} - -type fakeXComCall struct { - dagID, runID, taskID, key string - mapIndex *int -} - -func (f *fakeXComClient) GetXCom( - ctx context.Context, - dagID, runID, taskID string, - mapIndex *int, - key string, - _ any, -) (any, error) { - f.mu.Lock() - f.calls = append(f.calls, fakeXComCall{dagID, runID, taskID, key, mapIndex}) - f.mu.Unlock() - if f.err != nil { - return nil, f.err - } - return f.values[taskID+"/"+key], nil -} - -// workloadCtx returns a context carrying an ExecuteTaskWorkload the resolver -// reads the Dag/run identifiers from. -func workloadCtx() context.Context { - return context.WithValue( - context.Background(), - sdkcontext.WorkloadContextKey, - api.ExecuteTaskWorkload{ - TI: api.TaskInstance{ - Id: uuid.New(), - DagId: "dag1", - RunId: "run1", - TaskId: "transform", - }, - }, - ) -} - -func analyze(s *BindingSuite, fn any) *Plan { - plan, err := Analyze(reflect.TypeOf(fn), "testFn") - s.Require().NoError(err) - return plan -} - -func (s *BindingSuite) resolve(fn any, args []Arg, client sdk.Client) ([]reflect.Value, error) { - plan := analyze(s, fn) - return plan.Resolve(workloadCtx(), slog.Default(), client, args) -} - -func (s *BindingSuite) TestAnalyzeClassification() { - plan := analyze( - s, - func(ctx sdk.TIRunContext, log *slog.Logger, c sdk.VariableClient, country string, extracted map[string]any) error { - return nil - }, - ) - s.Equal(2, plan.numData) - - s.Zero(analyze(s, func() error { return nil }).numData) - s.Equal( - 1, - analyze(s, func(x any) error { return nil }).numData, - "an `any` parameter is a data parameter", - ) -} - -func (s *BindingSuite) TestAnalyzeRejections() { - cases := map[string]struct { - fn any - errContains string - }{ - "func-param": { - func(cb func()) error { return nil }, - "cannot receive a task argument", - }, - "chan-param": { - func(ch chan int) error { return nil }, - "cannot receive a task argument", - }, - "non-client-interface": { - func(x interface{ NotAClientMethod() }) error { return nil }, - "sdk.Client has no method NotAClientMethod", - }, - "context-with-extra-methods": { - func(x interface { - context.Context - TaskInstance() sdk.TaskInstance - }, - ) error { - return nil - }, - "adds methods on top of context.Context", - }, - } - for name, tt := range cases { - s.Run(name, func() { - _, err := Analyze(reflect.TypeOf(tt.fn), "testFn") - if s.Assert().Error(err) { - s.Assert().Contains(err.Error(), tt.errContains) - } - }) - } -} - -// TestNamedClientInterfacesAreInjectable guards against sdk.Client dropping an -// embedded interface, which would break tasks declaring it. -func (s *BindingSuite) TestNamedClientInterfacesAreInjectable() { - for name, typ := range map[string]reflect.Type{ - "Client": reflect.TypeFor[sdk.Client](), - "VariableClient": reflect.TypeFor[sdk.VariableClient](), - "ConnectionClient": reflect.TypeFor[sdk.ConnectionClient](), - "XComClient": reflect.TypeFor[sdk.XComClient](), - } { - s.True(isClient(typ), "sdk.%s must stay injectable", name) - } -} - -func (s *BindingSuite) TestResolveArityMismatch() { - fn := func(country string) error { return nil } - _, err := s.resolve(fn, nil, &fakeXComClient{}) - if s.Assert().Error(err) { - s.Contains(err.Error(), "argument count mismatch") - s.Contains(err.Error(), "passes 0 positional argument(s)") - s.Contains(err.Error(), "declares 1 data parameter(s)") - } - - _, err = s.resolve( - func() error { return nil }, - []Arg{LiteralArg{Value: "uk"}}, - &fakeXComClient{}, - ) - if s.Assert().Error(err) { - s.Contains(err.Error(), "argument count mismatch") - } -} - -func (s *BindingSuite) TestResolveLiterals() { - fn := func(country string, count int, ratio float64, on bool, tags []string, meta map[string]any) error { - return nil - } - got, err := s.resolve(fn, []Arg{ - LiteralArg{Value: "uk", DataType: DataTypeString}, - LiteralArg{Value: 3, DataType: DataTypeInteger}, - LiteralArg{Value: 1.5, DataType: DataTypeNumber}, - LiteralArg{Value: true, DataType: DataTypeBoolean}, - LiteralArg{Value: []any{"a", "b"}, DataType: DataTypeArray}, - LiteralArg{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{ - LiteralArg{Value: "uk", DataType: DataTypeString}, - LiteralArg{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{LiteralArg{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{LiteralArg{Value: "uk"}}, &fakeXComClient{}) - if s.Assert().Error(err) { - s.Contains(err.Error(), "decoding literal value into int") - } -} - -type extractResult struct { - GoVersion string `json:"go_version"` - Timestamp int64 `json:"timestamp"` -} - -// simpleTaskInput is the minimal TaskInput struct: one field, no tags, so it -// falls back to matching its own field name, verbatim. -type simpleTaskInput struct { - sdk.TaskInput - Name string -} - -// twoFieldTaskInput has one field the args always match (Name) and one whose -// arg name is never present in the tests that use it (Missing), to prove an -// unmatched field is left at its Go zero value instead of failing the task. -type twoFieldTaskInput struct { - sdk.TaskInput - Name string - Missing string `arg:"missing"` -} - -// nonEmbeddingStruct has no sdk.TaskInput sentinel, so it must keep resolving -// as today's whole-value decode target, not per-field TaskInput binding. -type nonEmbeddingStruct struct { - Name string -} - -// combineInput exercises both TaskInput field-binding modes side by side: -// Name falls back to its verbatim field name, Count is explicitly named -// via its `arg:` tag. -type combineInput struct { - sdk.TaskInput - Name string - Count int `arg:"count"` -} - -// reportInput deliberately declares Ratio before Region, the reverse of the -// wire order those names appear in, to prove field declaration order is -// irrelevant to by-name claiming. -type reportInput struct { - sdk.TaskInput - Ratio float64 - Region string `arg:"region"` -} - -// mixedInput pairs a TaskInput struct with a plain data parameter in the -// functions that assert Analyze rejects that combination. -type mixedInput struct { - sdk.TaskInput - Name string -} - -func (s *BindingSuite) TestResolveXComArgs() { - client := &fakeXComClient{values: map[string]any{ - "extract/return_value": map[string]any{"go_version": "go1.24", "timestamp": int64(42)}, - "probe/return_value": "probe-value", - }} - - fn := func(res extractResult, probe string) error { return nil } - got, err := s.resolve(fn, []Arg{ - XComArg{TaskID: "extract", DataType: DataTypeObject}, - XComArg{TaskID: "probe", DataType: DataTypeString}, - }, client) - s.Require().NoError(err) - s.Equal(extractResult{GoVersion: "go1.24", Timestamp: 42}, got[0].Interface()) - s.Equal("probe-value", got[1].Interface()) - - s.Require().Len(client.calls, 2) - taskIDs := make([]string, 0, 2) - for _, call := range client.calls { - // Pulls run concurrently, so assert per-call properties order-independently. - taskIDs = append(taskIDs, call.taskID) - s.Equal("dag1", call.dagID) - s.Equal("run1", call.runID) - s.Equal( - api.XComReturnValueKey, - call.key, - "an XCom argument always pulls the return-value key", - ) - s.Nil(call.mapIndex, "v1 always pulls the unmapped upstream instance") - } - s.ElementsMatch([]string{"extract", "probe"}, taskIDs) -} - -func (s *BindingSuite) TestResolveXComStrictStructDecode() { - client := &fakeXComClient{values: map[string]any{ - "extract/return_value": map[string]any{"go_version": "go1.24", "renamed_field": 1}, - }} - fn := func(res extractResult) error { return nil } - _, err := s.resolve(fn, []Arg{XComArg{TaskID: "extract"}}, client) - if s.Assert().Error(err) { - s.Contains(err.Error(), `decoding xcom from task "extract"`) - s.Contains(err.Error(), "unknown field") - } -} - -func (s *BindingSuite) TestResolveXComPullFailure() { - client := &fakeXComClient{err: sdk.XComNotFound} - fn := func(res map[string]any) error { return nil } - _, err := s.resolve(fn, []Arg{XComArg{TaskID: "extract"}}, client) - if s.Assert().Error(err) { - s.Contains(err.Error(), `pulling xcom from task "extract"`) - } -} - -func (s *BindingSuite) TestResolveXComWithoutWorkload() { - plan := analyze(s, func(res map[string]any) error { return nil }) - _, err := plan.Resolve( - context.Background(), slog.Default(), &fakeXComClient{}, - []Arg{XComArg{TaskID: "extract"}}, - ) - if s.Assert().Error(err) { - s.Contains(err.Error(), "no workload in context") - } -} - -func (s *BindingSuite) TestResolveNullHandling() { - fn := func(meta map[string]any) error { return nil } - got, err := s.resolve( - fn, - []Arg{LiteralArg{Value: nil, DataType: DataTypeObject}}, - &fakeXComClient{}, - ) - s.Require().NoError(err) - s.Nil(got[0].Interface()) - - fnStr := func(country string) error { return nil } - _, err = s.resolve(fnStr, []Arg{LiteralArg{Value: nil}}, &fakeXComClient{}) - if s.Assert().Error(err) { - s.Contains(err.Error(), "not nilable") - } -} - -// fakeArg is an out-of-catalogue Arg variant: the compiler seals the sum type -// to this package, so the defensive default branch can only be reached from -// inside it. -type fakeArg struct{} - -func (fakeArg) ArgName() string { return "fake" } -func (fakeArg) DeclaredType() DataType { return DataTypeAny } -func (fakeArg) sealedArg() {} - -func (s *BindingSuite) TestResolveUnsupportedVariant() { - fn := func(country string) error { return nil } - _, err := s.resolve(fn, []Arg{fakeArg{}}, &fakeXComClient{}) - if s.Assert().Error(err) { - s.Contains(err.Error(), "unsupported argument binding binding.fakeArg") - } -} - -func (s *BindingSuite) TestResolveNilArg() { - fn := func(country string) error { return nil } - _, err := s.resolve(fn, []Arg{nil}, &fakeXComClient{}) - if s.Assert().Error(err) { - s.Contains(err.Error(), "nil argument binding") - } -} - -func (s *BindingSuite) TestResolveTIRunContextRebuild() { - ti := sdk.TaskInstance{DagID: "dag1", RunID: "run1", TaskID: "transform"} - dagRun := sdk.DagRun{DagID: "dag1", RunID: "run1"} - ctx := context.WithValue( - workloadCtx(), - sdkcontext.RuntimeContextKey, - sdk.NewTIRunContext(context.Background(), ti, dagRun), - ) - - plan := analyze(s, func(rc sdk.TIRunContext, country string) error { return nil }) - got, err := plan.Resolve(ctx, slog.Default(), &fakeXComClient{}, []Arg{ - LiteralArg{Value: "uk", 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()) -} - -func (s *BindingSuite) TestAnalyzeTaskInputClassification() { - plan := analyze(s, func(input simpleTaskInput) error { return nil }) - s.Zero(plan.numData, "a TaskInput struct claims by name, not by position") - - ptrPlan := analyze(s, func(input *simpleTaskInput) error { return nil }) - s.Zero(ptrPlan.numData, "a pointer to a TaskInput struct is detected the same way") - - plainPlan := analyze(s, func(cfg nonEmbeddingStruct) error { return nil }) - s.Equal( - 1, plainPlan.numData, - "a plain struct without the TaskInput sentinel stays a whole-value data parameter", - ) -} - -func (s *BindingSuite) TestAnalyzeTaskInputValidation() { - type duplicateArgNames struct { - sdk.TaskInput - A string - B string `arg:"A"` - } - type nonDecodableField struct { - sdk.TaskInput - Bad chan int - } - - cases := map[string]struct { - fn any - errContains string - }{ - "duplicate-arg-names": { - func(input duplicateArgNames) error { return nil }, - `fields A and B both bind arg name "A"`, - }, - "non-decodable-field": { - func(input nonDecodableField) error { return nil }, - "cannot receive a task argument", - }, - "two-taskinput-params": { - func(a simpleTaskInput, b simpleTaskInput) error { return nil }, - "only one TaskInput struct parameter is allowed", - }, - "mixed-with-flat-data-param": { - func(prefix string, input mixedInput) error { return nil }, - "cannot mix a TaskInput struct parameter", - }, - "mixed-with-trailing-flat-data-param": { - func(input mixedInput, suffix string) error { return nil }, - "cannot mix a TaskInput struct parameter", - }, - } - for name, tt := range cases { - s.Run(name, func() { - _, err := Analyze(reflect.TypeOf(tt.fn), "testFn") - if s.Assert().Error(err) { - s.Assert().Contains(err.Error(), tt.errContains) - } - }) - } -} - -func (s *BindingSuite) TestResolveTaskInputAllStruct() { - fn := func(input combineInput) error { return nil } - got, err := s.resolve(fn, []Arg{ - LiteralArg{Name: "Name", Value: "widget", DataType: DataTypeString}, - LiteralArg{Name: "count", Value: 7, DataType: DataTypeInteger}, - }, &fakeXComClient{}) - s.Require().NoError(err) - - input := got[0].Interface().(combineInput) - s.Equal("widget", input.Name, "the untagged field claims its verbatim field name") - s.Equal(7, input.Count, "the `arg:` tag claims its named entry") -} - -func (s *BindingSuite) TestResolveTaskInputXComArg() { - fn := func(log *slog.Logger, input reportInput) error { return nil } - got, err := s.resolve(fn, []Arg{ - XComArg{Name: "region", TaskID: "make_region", DataType: DataTypeString}, - LiteralArg{Name: "Ratio", Value: 0.5, DataType: DataTypeNumber}, - }, &fakeXComClient{values: map[string]any{"make_region/return_value": "east"}}) - s.Require().NoError(err) - - input := got[1].Interface().(reportInput) - s.Equal("east", input.Region, "Region resolves by name despite being declared after Ratio") - s.Equal(0.5, input.Ratio) -} - -func (s *BindingSuite) TestResolveTaskInputLiteralThroughArgName() { - fn := func(input simpleTaskInput) error { return nil } - got, err := s.resolve(fn, []Arg{ - LiteralArg{Name: "Name", Value: "widget", DataType: DataTypeString}, - }, &fakeXComClient{}) - s.Require().NoError(err) - s.Equal("widget", got[0].Interface().(simpleTaskInput).Name) -} - -func (s *BindingSuite) TestResolveTaskInputPointerStruct() { - fn := func(input *simpleTaskInput) error { return nil } - got, err := s.resolve(fn, []Arg{ - LiteralArg{Name: "Name", Value: "widget", DataType: DataTypeString}, - }, &fakeXComClient{}) - s.Require().NoError(err) - input := got[0].Interface().(*simpleTaskInput) - s.Require().NotNil(input) - s.Equal("widget", input.Name) -} - -func (s *BindingSuite) TestResolveTaskInputUnclaimedArgFailsLoudly() { - fn := func(input simpleTaskInput) error { return nil } - _, err := s.resolve(fn, []Arg{ - LiteralArg{Name: "different_name", Value: "x", DataType: DataTypeString}, - }, &fakeXComClient{}) - // No TaskInput field claims "different_name"; the leftover argument fails - // the task rather than being dropped silently. - if s.Assert().Error(err) { - s.Contains(err.Error(), `not claimed by any TaskInput field: "different_name"`) - } -} - -func (s *BindingSuite) TestResolveTaskInputUnclaimedFromDefaultAllowed() { - fn := func(input simpleTaskInput) error { return nil } - got, err := s.resolve(fn, []Arg{ - LiteralArg{Name: "Name", Value: "widget", DataType: DataTypeString}, - // The Dag author never passed "threshold"; Python captured it from the - // stub signature's default. The struct need not mirror it. - LiteralArg{Name: "threshold", Value: 0.75, DataType: DataTypeNumber, FromDefault: true}, - }, &fakeXComClient{}) - s.Require().NoError(err) - s.Equal("widget", got[0].Interface().(simpleTaskInput).Name) -} - -func (s *BindingSuite) TestResolveTaskInputEmptySpecFailsLoudly() { - fn := func(input simpleTaskInput) error { return nil } - for name, args := range map[string][]Arg{"nil-spec": nil, "empty-spec": {}} { - s.Run(name, func() { - _, err := s.resolve(fn, args, &fakeXComClient{}) - // The Edge Worker path delivers no arg bindings; a struct with - // bindable fields must fail rather than run fully zero-valued. - if s.Assert().Error(err) { - s.Contains(err.Error(), "no TaskFlow arg bindings arrived") - } - }) - } -} - -func (s *BindingSuite) TestResolveTaskInputOnlyDefaultsSpecZeroValuesUnmatchedFields() { - fn := func(input twoFieldTaskInput) error { return nil } - got, err := s.resolve(fn, []Arg{ - LiteralArg{Name: "threshold", Value: 0.75, DataType: DataTypeNumber, FromDefault: true}, - }, &fakeXComClient{}) - s.Require().NoError(err) - input := got[0].Interface().(twoFieldTaskInput) - s.Equal("", input.Name, "no explicit entry arrived; fields keep kwarg-style zero values") - s.Equal("", input.Missing) -} - -func (s *BindingSuite) TestResolveTaskInputUnmatchedArgNameZeroValuedAlongsideMatch() { - fn := func(input twoFieldTaskInput) error { return nil } - got, err := s.resolve(fn, []Arg{ - LiteralArg{Name: "Name", Value: "widget", DataType: DataTypeString}, - }, &fakeXComClient{}) - s.Require().NoError(err) - input := got[0].Interface().(twoFieldTaskInput) - s.Equal("widget", input.Name, "the matched field binds normally") - s.Equal("", input.Missing, "the unmatched field is left at its Go zero value, not an error") -} diff --git a/go-sdk/pkg/execution/frames.go b/go-sdk/pkg/execution/frames.go index 947346316c57c..f9a246286efce 100644 --- a/go-sdk/pkg/execution/frames.go +++ b/go-sdk/pkg/execution/frames.go @@ -62,13 +62,6 @@ 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 0b885e9d6e446..a4e9b1212823d 100644 --- a/go-sdk/pkg/execution/genmodels/defaults.gen.go +++ b/go-sdk/pkg/execution/genmodels/defaults.gen.go @@ -186,17 +186,6 @@ func (m *HITLDetailRequestResult) DecodeMsgpack(dec *msgpack.Decoder) error { return nil } -// DecodeMsgpack applies LiteralArgBinding's schema defaults that msgpack would otherwise skip. -func (m *LiteralArgBinding) DecodeMsgpack(dec *msgpack.Decoder) error { - type alias LiteralArgBinding - v := alias{DataType: ArgBindingDataType("any")} - if err := dec.Decode(&v); err != nil { - return err - } - *m = LiteralArgBinding(v) - return nil -} - // DecodeMsgpack applies PreviousTIResponse's schema defaults that msgpack would otherwise skip. func (m *PreviousTIResponse) DecodeMsgpack(dec *msgpack.Decoder) error { type alias PreviousTIResponse @@ -327,14 +316,3 @@ func (m *TriggerDagRun) DecodeMsgpack(dec *msgpack.Decoder) error { *m = TriggerDagRun(v) return nil } - -// DecodeMsgpack applies XComArgBinding's schema defaults that msgpack would otherwise skip. -func (m *XComArgBinding) DecodeMsgpack(dec *msgpack.Decoder) error { - type alias XComArgBinding - v := alias{DataType: ArgBindingDataType("any")} - if err := dec.Decode(&v); err != nil { - return err - } - *m = XComArgBinding(v) - return nil -} diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index 127f00d32f421..e6861d8c8add5 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -20,18 +20,6 @@ package genmodels import "time" -type ArgBindingDataType string - -const ArgBindingDataTypeAny ArgBindingDataType = "any" -const ArgBindingDataTypeArray ArgBindingDataType = "array" -const ArgBindingDataTypeBoolean ArgBindingDataType = "boolean" -const ArgBindingDataTypeInteger ArgBindingDataType = "integer" -const ArgBindingDataTypeNumber ArgBindingDataType = "number" -const ArgBindingDataTypeObject ArgBindingDataType = "object" -const ArgBindingDataTypeString ArgBindingDataType = "string" - -type ArgBindings []TaskArgBinding - // Schema for AssetAliasModel used in AssetEventDagRunReference. type AssetAliasReferenceAssetEventDagRun struct { // Name corresponds to the JSON schema field "name". @@ -382,9 +370,6 @@ 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. @@ -764,9 +749,6 @@ 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 @@ -826,22 +808,12 @@ type GetAssetEventByAsset struct { // Before corresponds to the JSON schema field "before". Before interface{} `msgpack:"before,omitempty"` - // Extra corresponds to the JSON schema field "extra". - Extra *Extra `msgpack:"extra,omitempty"` - // Limit corresponds to the JSON schema field "limit". Limit interface{} `msgpack:"limit,omitempty"` // Name corresponds to the JSON schema field "name". Name interface{} `msgpack:"name"` - // PartitionKey corresponds to the JSON schema field "partition_key". - PartitionKey interface{} `msgpack:"partition_key,omitempty"` - - // PartitionKeyRegexpPattern corresponds to the JSON schema field - // "partition_key_regexp_pattern". - PartitionKeyRegexpPattern interface{} `msgpack:"partition_key_regexp_pattern,omitempty"` - // Type corresponds to the JSON schema field "type". Type string `msgpack:"type,omitempty"` @@ -862,19 +834,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"` - // PartitionKey corresponds to the JSON schema field "partition_key". - PartitionKey interface{} `msgpack:"partition_key,omitempty"` - - // PartitionKeyRegexpPattern corresponds to the JSON schema field - // "partition_key_regexp_pattern". - PartitionKeyRegexpPattern interface{} `msgpack:"partition_key_regexp_pattern,omitempty"` - // Type corresponds to the JSON schema field "type". Type string `msgpack:"type,omitempty"` } @@ -1277,24 +1239,6 @@ type LazyDeserializedDAG struct { LastLoaded interface{} `msgpack:"last_loaded,omitempty"` } -// One positional stub-task argument carrying an inline literal from the Dag file. -type LiteralArgBinding struct { - // DataType corresponds to the JSON schema field "data_type". - DataType ArgBindingDataType `msgpack:"data_type,omitempty"` - - // FromDefault corresponds to the JSON schema field "from_default". - FromDefault bool `msgpack:"from_default,omitempty"` - - // Kind corresponds to the JSON schema field "kind". - Kind string `msgpack:"kind"` - - // Name corresponds to the JSON schema field "name". - Name string `msgpack:"name"` - - // Value corresponds to the JSON schema field "value". - Value interface{} `msgpack:"value,omitempty"` -} - type LogicalDates []time.Time // Add a new value to be redacted in task logs. @@ -1620,9 +1564,6 @@ 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"` @@ -1655,8 +1596,6 @@ type TIRunContext struct { XcomKeysToClear []string `msgpack:"xcom_keys_to_clear,omitempty"` } -type TaskArgBinding interface{} - type TaskBreadcrumbsResult struct { // Breadcrumbs corresponds to the JSON schema field "breadcrumbs". Breadcrumbs []TaskBreadcrumbsResultBreadcrumbsElem `msgpack:"breadcrumbs"` @@ -1667,18 +1606,37 @@ type TaskBreadcrumbsResult struct { type TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{} -type TriggerKwargs map[string]JsonValue +// Task callback status information. +// +// A Class with information about the success/failure TI callback to be executed. +// Currently, only failure +// callbacks when tasks are externally killed or experience heartbeat timeouts are +// run via DagFileProcessorProcess. +type TaskCallbackRequest struct { + // BundleName corresponds to the JSON schema field "bundle_name". + BundleName string `msgpack:"bundle_name"` -// 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"` + // BundleVersion corresponds to the JSON schema field "bundle_version". + BundleVersion interface{} `msgpack:"bundle_version"` - // Value corresponds to the JSON schema field "value". - Value interface{} `msgpack:"value"` -} + // ContextFromServer corresponds to the JSON schema field "context_from_server". + ContextFromServer *TIRunContext `msgpack:"context_from_server,omitempty"` -type Warnings []interface{} + // Filepath corresponds to the JSON schema field "filepath". + Filepath string `msgpack:"filepath"` + + // Msg corresponds to the JSON schema field "msg". + Msg interface{} `msgpack:"msg,omitempty"` + + // TaskCallbackType corresponds to the JSON schema field "task_callback_type". + TaskCallbackType interface{} `msgpack:"task_callback_type,omitempty"` + + // TI corresponds to the JSON schema field "ti". + TI TaskInstance `msgpack:"ti"` + + // Type corresponds to the JSON schema field "type". + Type string `msgpack:"type,omitempty"` +} type TaskIds []string @@ -1715,60 +1673,24 @@ type TaskInstance struct { TryNumber int `msgpack:"try_number"` } +type TaskInstanceState string + const TaskInstanceStateAwaitingInput TaskInstanceState = "awaiting_input" const TaskInstanceStateDeferred TaskInstanceState = "deferred" const TaskInstanceStateFailed TaskInstanceState = "failed" +const TaskInstanceStateQueued TaskInstanceState = "queued" +const TaskInstanceStateRemoved TaskInstanceState = "removed" const TaskInstanceStateRestarting TaskInstanceState = "restarting" +const TaskInstanceStateRunning TaskInstanceState = "running" +const TaskInstanceStateScheduled TaskInstanceState = "scheduled" const TaskInstanceStateSkipped TaskInstanceState = "skipped" const TaskInstanceStateSuccess TaskInstanceState = "success" -const TaskInstanceStateRunning TaskInstanceState = "running" const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" type TaskOutlets []AssetProfile -const TaskInstanceStateQueued TaskInstanceState = "queued" -const TaskInstanceStateScheduled TaskInstanceState = "scheduled" -const TaskInstanceStateRemoved TaskInstanceState = "removed" - -type TaskInstanceState string - -// Task callback status information. -// -// A Class with information about the success/failure TI callback to be executed. -// Currently, only failure -// callbacks when tasks are externally killed or experience heartbeat timeouts are -// run via DagFileProcessorProcess. -type TaskCallbackRequest struct { - // BundleName corresponds to the JSON schema field "bundle_name". - BundleName string `msgpack:"bundle_name"` - - // BundleVersion corresponds to the JSON schema field "bundle_version". - BundleVersion interface{} `msgpack:"bundle_version"` - - // ContextFromServer corresponds to the JSON schema field "context_from_server". - ContextFromServer *TIRunContext `msgpack:"context_from_server,omitempty"` - - // Filepath corresponds to the JSON schema field "filepath". - Filepath string `msgpack:"filepath"` - - // Msg corresponds to the JSON schema field "msg". - Msg interface{} `msgpack:"msg,omitempty"` - - // TaskCallbackType corresponds to the JSON schema field "task_callback_type". - TaskCallbackType interface{} `msgpack:"task_callback_type,omitempty"` - - // TI corresponds to the JSON schema field "ti". - TI TaskInstance `msgpack:"ti"` - - // 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"` -} - // Response containing the first reschedule date for a task instance. type TaskRescheduleStartDate struct { // StartDate corresponds to the JSON schema field "start_date". @@ -1778,12 +1700,6 @@ type TaskRescheduleStartDate struct { Type string `msgpack:"type,omitempty"` } -type TaskStateState string - -const TaskStateStateFailed TaskStateState = "failed" -const TaskStateStateSkipped TaskStateState = "skipped" -const TaskStateStateRemoved TaskStateState = "removed" - // Update a task's state. // // If a process exits without sending one of these the state will be derived from @@ -1804,6 +1720,12 @@ type TaskState struct { Type string `msgpack:"type,omitempty"` } +type TaskStateState string + +const TaskStateStateFailed TaskStateState = "failed" +const TaskStateStateRemoved TaskStateState = "removed" +const TaskStateStateSkipped TaskStateState = "skipped" + // Response to GetTaskStateStore; wraps the generated API response for supervisor // to worker comms. type TaskStateStoreResult struct { @@ -1853,6 +1775,19 @@ type TriggerDagRun struct { Type string `msgpack:"type,omitempty"` } +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. @@ -1900,21 +1835,6 @@ type VariableResult struct { Value interface{} `msgpack:"value,omitempty"` } -// One positional stub-task argument pulled from an upstream task's XCom. -type XComArgBinding struct { - // DataType corresponds to the JSON schema field "data_type". - DataType ArgBindingDataType `msgpack:"data_type,omitempty"` - - // Kind corresponds to the JSON schema field "kind". - Kind string `msgpack:"kind"` - - // Name corresponds to the JSON schema field "name". - Name string `msgpack:"name"` - - // TaskID corresponds to the JSON schema field "task_id". - TaskID string `msgpack:"task_id"` -} - type XComCountResponse struct { // Len corresponds to the JSON schema field "len". Len int `msgpack:"len"` diff --git a/go-sdk/pkg/execution/integration_test.go b/go-sdk/pkg/execution/integration_test.go index bbd7369e6d028..2559359453866 100644 --- a/go-sdk/pkg/execution/integration_test.go +++ b/go-sdk/pkg/execution/integration_test.go @@ -230,388 +230,6 @@ func TestTaskRunnerPanicRetry(t *testing.T) { assertRetryTask(t, result, "panic: something went wrong") } -// TestTaskRunnerBindsArgs covers the TaskFlow path through RunTask: the -// positional-argument spec in ti_context.arg_bindings binds literals onto the -// task function's data parameters. -func TestTaskRunnerBindsArgs(t *testing.T) { - var gotCountry string - var gotMeta map[string]any - bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTaskWithName("transform", - func(log *slog.Logger, country string, meta map[string]any) error { - gotCountry = country - gotMeta = meta - return nil - }) - }) - - details := &genmodels.StartupDetails{ - TI: genmodels.TaskInstance{ - ID: "550e8400-e29b-41d4-a716-446655440000", - DagID: "test_dag", - TaskID: "transform", - RunID: "run1", - MapIndex: ptr(-1), - }, - BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, - TIContext: genmodels.TIRunContext{ - ArgBindings: &genmodels.ArgBindings{ - map[string]any{ - "name": "country", - "kind": "literal", - "data_type": "string", - "value": "uk", - }, - map[string]any{ - "name": "meta", - "kind": "literal", - "data_type": "object", - "value": map[string]any{"k": "v"}, - }, - }, - }, - } - - logger := slog.New(slog.NewTextHandler(io.Discard, nil)) - comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) - - result := RunTask(context.Background(), bundle, details, comm, logger) - assertSucceedTask(t, result) - assert.Equal(t, "uk", gotCountry) - assert.Equal(t, map[string]any{"k": "v"}, gotMeta) -} - -// TestTaskRunnerArgBindingsArityMismatch: an argument spec that does not match -// the function's data parameters fails the task loudly instead of running it -// with zero values. -func TestTaskRunnerArgBindingsArityMismatch(t *testing.T) { - ran := false - bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTaskWithName("transform", - func(country string, meta map[string]any) error { - ran = true - return nil - }) - }) - - details := &genmodels.StartupDetails{ - TI: genmodels.TaskInstance{ - ID: "550e8400-e29b-41d4-a716-446655440000", - DagID: "test_dag", - TaskID: "transform", - RunID: "run1", - MapIndex: ptr(-1), - }, - BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, - TIContext: genmodels.TIRunContext{ - ArgBindings: &genmodels.ArgBindings{ - map[string]any{ - "name": "country", - "kind": "literal", - "data_type": "string", - "value": "uk", - }, - }, - }, - } - - logger := slog.New(slog.NewTextHandler(io.Discard, nil)) - comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) - - result := RunTask(context.Background(), bundle, details, comm, logger) - assertTaskState(t, result, genmodels.TaskStateStateFailed) - assert.False(t, ran, "the task body must not run on an arity mismatch") -} - -// combineInput is a TaskInput struct whose sole field claims a named entry -// out of ti_context.arg_bindings. -type combineInput struct { - sdk.TaskInput - Region string `arg:"region"` -} - -// TestTaskRunnerBindsTaskInputStructArgs covers the TaskFlow path through -// RunTask for a TaskInput struct parameter: convertArgBindings must propagate -// each spec's Name through to binding.Arg so the struct's `arg:"region"` field -// can claim it by name. -func TestTaskRunnerBindsTaskInputStructArgs(t *testing.T) { - var got combineInput - bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTaskWithName("transform", - func(input combineInput) error { - got = input - return nil - }) - }) - - details := &genmodels.StartupDetails{ - TI: genmodels.TaskInstance{ - ID: "550e8400-e29b-41d4-a716-446655440000", - DagID: "test_dag", - TaskID: "transform", - RunID: "run1", - MapIndex: ptr(-1), - }, - BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, - TIContext: genmodels.TIRunContext{ - ArgBindings: &genmodels.ArgBindings{ - map[string]any{ - "name": "region", - "kind": "literal", - "data_type": "string", - "value": "eu-west-1", - }, - }, - }, - } - - logger := slog.New(slog.NewTextHandler(io.Discard, nil)) - comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) - - result := RunTask(context.Background(), bundle, details, comm, logger) - assertSucceedTask(t, result) - assert.Equal(t, "eu-west-1", got.Region) -} - -// TestTaskRunnerTaskInputIgnoresUnclaimedDefault: convertArgBindings must -// propagate from_default so a spec entry the Python side filled from the stub -// signature's default may go unclaimed by the TaskInput struct. -func TestTaskRunnerTaskInputIgnoresUnclaimedDefault(t *testing.T) { - var got combineInput - bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTaskWithName("transform", - func(input combineInput) error { - got = input - return nil - }) - }) - - details := &genmodels.StartupDetails{ - TI: genmodels.TaskInstance{ - ID: "550e8400-e29b-41d4-a716-446655440000", - DagID: "test_dag", - TaskID: "transform", - RunID: "run1", - MapIndex: ptr(-1), - }, - BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, - TIContext: genmodels.TIRunContext{ - ArgBindings: &genmodels.ArgBindings{ - map[string]any{ - "name": "region", - "kind": "literal", - "data_type": "string", - "value": "eu-west-1", - }, - map[string]any{ - "name": "threshold", - "kind": "literal", - "data_type": "number", - "value": 0.75, - "from_default": true, - }, - }, - }, - } - - logger := slog.New(slog.NewTextHandler(io.Discard, nil)) - comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) - - result := RunTask(context.Background(), bundle, details, comm, logger) - assertSucceedTask(t, result) - assert.Equal(t, "eu-west-1", got.Region) -} - -// TestTaskRunnerArgBindingsTypeMismatch: a declared Dag type that cannot bind to -// the Go parameter type fails the task loudly before the body runs. -func TestTaskRunnerArgBindingsTypeMismatch(t *testing.T) { - bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTaskWithName("transform", - func(count int) error { return nil }) - }) - - details := &genmodels.StartupDetails{ - TI: genmodels.TaskInstance{ - ID: "550e8400-e29b-41d4-a716-446655440000", - DagID: "test_dag", - TaskID: "transform", - RunID: "run1", - MapIndex: ptr(-1), - }, - BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, - TIContext: genmodels.TIRunContext{ - ArgBindings: &genmodels.ArgBindings{ - map[string]any{ - "name": "count", - "kind": "literal", - "data_type": "string", - "value": "uk", - }, - }, - }, - } - - logger := slog.New(slog.NewTextHandler(io.Discard, nil)) - comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) - - result := RunTask(context.Background(), bundle, details, comm, logger) - assertTaskState(t, result, genmodels.TaskStateStateFailed) -} - -// TestTaskRunnerArgBindingsUnknownKind: a wire spec whose kind is neither xcom -// nor literal fails the task before the body runs. -func TestTaskRunnerArgBindingsUnknownKind(t *testing.T) { - ran := false - bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTaskWithName("transform", - func(country string) error { - ran = true - return nil - }) - }) - - details := &genmodels.StartupDetails{ - TI: genmodels.TaskInstance{ - ID: "550e8400-e29b-41d4-a716-446655440000", - DagID: "test_dag", - TaskID: "transform", - RunID: "run1", - MapIndex: ptr(-1), - }, - BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, - TIContext: genmodels.TIRunContext{ - ArgBindings: &genmodels.ArgBindings{ - map[string]any{"name": "country", "kind": "template", "value": "x"}, - }, - }, - } - - logger := slog.New(slog.NewTextHandler(io.Discard, nil)) - comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) - - result := RunTask(context.Background(), bundle, details, comm, logger) - assertTaskState(t, result, genmodels.TaskStateStateFailed) - assert.False(t, ran, "the task body must not run on an unknown binding kind") -} - -// TestTaskRunnerArgBindingsMalformedElement: a wire spec element that is not a -// map at all fails the task before the body runs. -func TestTaskRunnerArgBindingsMalformedElement(t *testing.T) { - ran := false - bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTaskWithName("transform", - func(country string) error { - ran = true - return nil - }) - }) - - details := &genmodels.StartupDetails{ - TI: genmodels.TaskInstance{ - ID: "550e8400-e29b-41d4-a716-446655440000", - DagID: "test_dag", - TaskID: "transform", - RunID: "run1", - MapIndex: ptr(-1), - }, - BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, - TIContext: genmodels.TIRunContext{ - ArgBindings: &genmodels.ArgBindings{"bogus"}, - }, - } - - logger := slog.New(slog.NewTextHandler(io.Discard, nil)) - comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) - - result := RunTask(context.Background(), bundle, details, comm, logger) - assertTaskState(t, result, genmodels.TaskStateStateFailed) - assert.False(t, ran, "the task body must not run on a malformed binding element") -} - -// TestTaskRunnerArgBindingsMissingRequiredFields: a wire spec entry without a -// usable name, or an xcom entry without a task_id, fails the task before the -// body runs instead of silently binding empty strings. -func TestTaskRunnerArgBindingsMissingRequiredFields(t *testing.T) { - cases := []struct { - name string - spec map[string]any - }{ - {name: "missing name", spec: map[string]any{"kind": "literal", "value": "x"}}, - {name: "empty name", spec: map[string]any{"name": "", "kind": "literal", "value": "x"}}, - {name: "xcom missing task_id", spec: map[string]any{"name": "country", "kind": "xcom"}}, - { - name: "xcom empty task_id", - spec: map[string]any{"name": "country", "kind": "xcom", "task_id": ""}, - }, - } - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - ran := false - bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTaskWithName("transform", - func(country string) error { - ran = true - return nil - }) - }) - - details := &genmodels.StartupDetails{ - TI: genmodels.TaskInstance{ - ID: "550e8400-e29b-41d4-a716-446655440000", - DagID: "test_dag", - TaskID: "transform", - RunID: "run1", - MapIndex: ptr(-1), - }, - BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, - TIContext: genmodels.TIRunContext{ - ArgBindings: &genmodels.ArgBindings{tc.spec}, - }, - } - - logger := slog.New(slog.NewTextHandler(io.Discard, nil)) - comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) - - result := RunTask(context.Background(), bundle, details, comm, logger) - assertTaskState(t, result, genmodels.TaskStateStateFailed) - assert.False(t, ran, "the task body must not run on an incomplete binding spec") - }) - } -} - -// TestTaskRunnerMalformedSpecHonorsShouldRetry: a spec that fails -// convertArgBindings terminates with the same retry semantics as a binding -// failure inside executeTask, not an unconditional FAILED. -func TestTaskRunnerMalformedSpecHonorsShouldRetry(t *testing.T) { - bundle := buildBundle(t, func(r bundlev1.Registry) { - r.AddDag("test_dag").AddTaskWithName("transform", - func(country string) error { return nil }) - }) - - details := &genmodels.StartupDetails{ - TI: genmodels.TaskInstance{ - ID: "550e8400-e29b-41d4-a716-446655440000", - DagID: "test_dag", - TaskID: "transform", - RunID: "run1", - MapIndex: ptr(-1), - }, - BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"}, - TIContext: genmodels.TIRunContext{ - ShouldRetry: true, - ArgBindings: &genmodels.ArgBindings{ - map[string]any{"name": "country", "kind": "template", "value": "x"}, - }, - }, - } - - logger := slog.New(slog.NewTextHandler(io.Discard, nil)) - comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger) - - result := RunTask(context.Background(), bundle, details, comm, logger) - assertRetryTask(t, result, `unknown kind "template"`) -} - func TestRunTaskHonorsContextCancellation(t *testing.T) { bundle := buildBundle(t, func(r bundlev1.Registry) { r.AddDag("test_dag").AddTaskWithName("ctxcheck", diff --git a/go-sdk/pkg/execution/messages.go b/go-sdk/pkg/execution/messages.go index 378cc931dda50..72d451a866632 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-07-30" +const SupervisorSchemaVersion = "2026-06-16" // 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 4577bd5a5bde3..c656b4cfd71f2 100644 --- a/go-sdk/pkg/execution/task_runner.go +++ b/go-sdk/pkg/execution/task_runner.go @@ -28,7 +28,6 @@ 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" @@ -125,79 +124,7 @@ func RunTask( ctx = context.WithValue(ctx, sdkcontext.SdkClientContextKey, sdk.Client(client)) ctx = context.WithValue(ctx, sdkcontext.RuntimeContextKey, runtimeContext) - args, err := convertArgBindings(details.TIContext.ArgBindings) - if err != nil { - logger.Error("Invalid arg_bindings spec from supervisor", - "dag_id", details.TI.DagID, - "task_id", details.TI.TaskID, - "error", err, - ) - // Same retry semantics as a binding failure inside executeTask: an - // equally permanent spec error must not terminate differently. - if details.TIContext.ShouldRetry { - return genmodels.RetryTask{ - EndDate: time.Now().UTC(), - RetryReason: err.Error(), - } - } - return genmodels.TaskState{ - State: genmodels.TaskStateStateFailed, - EndDate: time.Now().UTC(), - } - } - - return executeTask(ctx, task, args, details.TIContext.ShouldRetry, logger) -} - -// convertArgBindings maps the wire-model positional-argument spec (captured from -// the Python stub Dag's TaskFlow call) onto the runtime binding sum type. The -// wire union generates untyped items (msgpack delivers each XComArgBinding / -// LiteralArgBinding as a plain map), so the kind dispatch and the schema -// default (data_type "any") are applied here. -func convertArgBindings(specsPtr *genmodels.ArgBindings) ([]binding.Arg, error) { - if specsPtr == nil || len(*specsPtr) == 0 { - return nil, nil - } - specs := *specsPtr - args := make([]binding.Arg, len(specs)) - for i, raw := range specs { - m, ok := raw.(map[string]any) - if !ok { - return nil, fmt.Errorf("arg_bindings[%d]: unexpected wire shape %T", i, raw) - } - name, ok := m["name"].(string) - if !ok || name == "" { - return nil, fmt.Errorf("arg_bindings[%d]: missing or empty name", i) - } - dataType := binding.DataTypeAny - if s, ok := m["data_type"].(string); ok && s != "" { - dataType = binding.DataType(s) - } - switch kind, _ := m["kind"].(string); kind { - case "xcom": - taskID, ok := m["task_id"].(string) - if !ok || taskID == "" { - return nil, fmt.Errorf( - "arg_bindings[%d] (%q): missing or empty task_id for xcom kind", - i, - name, - ) - } - args[i] = binding.XComArg{Kind: kind, Name: name, TaskID: taskID, DataType: dataType} - case "literal": - fromDefault, _ := m["from_default"].(bool) - args[i] = binding.LiteralArg{ - Kind: kind, - Name: name, - Value: m["value"], - DataType: dataType, - FromDefault: fromDefault, - } - default: - return nil, fmt.Errorf("arg_bindings[%d]: unknown kind %q", i, kind) - } - } - return args, nil + return executeTask(ctx, task, details.TIContext.ShouldRetry, logger) } // mapIndexPtr normalizes the supervisor's map_index into the optional form @@ -215,15 +142,9 @@ 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) { @@ -247,19 +168,7 @@ func executeTask( } }() - 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 { + if err := task.Execute(ctx, logger); 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/go-sdk/sdk/context.go b/go-sdk/sdk/context.go index 0db323b15b602..afd0dd3705051 100644 --- a/go-sdk/sdk/context.go +++ b/go-sdk/sdk/context.go @@ -108,25 +108,3 @@ type DagRun struct { DataIntervalStart *time.Time DataIntervalEnd *time.Time } - -// TaskInput is a zero-size marker embedded anonymously in a struct to opt -// that struct into per-field, name-based TaskFlow argument binding -- an -// ergonomic alternative to a long flat parameter list. Each exported field of -// such a struct may carry an `arg:""` tag naming the stub's TaskFlow -// argument to bind, falling back to the field's own name, verbatim, when -// omitted -- so an untagged Name binds the argument "Name", and a snake_case -// argument like "count" needs an explicit tag: -// -// type CombineInput struct { -// sdk.TaskInput -// Name string -// Count int `arg:"count"` -// } -// -// func Combine(ctx sdk.TIRunContext, log *slog.Logger, input CombineInput) (any, error) -// -// Embedding costs nothing at runtime: the marker occupies zero bytes and is -// never read; its only purpose is for the binding package to detect it -// reflectively at task-registration time. See the go-sdk README's "TaskInput -// structs" section for the full field-tag reference. -type TaskInput struct{} From b4e41cc194d5e3ebd9d5227b1e02f67de653f5f8 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 22 Jul 2026 03:51:03 +0000 Subject: [PATCH 18/40] Resolve stub arg bindings through the shared Dag cache in ti_run The API server already holds a DBDagBag with configurable LRU+TTL caching of deserialized Dags; a route-local raw-blob accessor plus a hand-rolled module-level cache duplicated that machinery with its own eviction story. The serialized _arg_bindings field survives full deserialization onto the task object, so ti_run can read it off the dag_bag-resolved Dag version directly, and the LazyDeserializedDAG.get_task_arg_bindings accessor goes away. --- .../execution_api/routes/task_instances.py | 29 +++++-------------- .../serialization/serialized_objects.py | 15 ---------- .../versions/head/test_task_instances.py | 6 ++-- 3 files changed, 10 insertions(+), 40 deletions(-) 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 43db192ad6212..cd7ed440b391f 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 @@ -81,7 +81,6 @@ from airflow.models.asset import AssetActive from airflow.models.base import ID_LEN from airflow.models.dag import DagModel -from airflow.models.dag_version import DagVersion from airflow.models.dagrun import DagRun as DR from airflow.models.hitl import HITLDetail from airflow.models.log import Log @@ -91,7 +90,6 @@ from airflow.models.trigger import Trigger, handle_event_submit from airflow.models.xcom import XComModel from airflow.serialization.definitions.assets import SerializedAsset, SerializedAssetUniqueKey -from airflow.serialization.serialized_objects import LazyDeserializedDAG from airflow.state import get_state_backend from airflow.triggers.base import TriggerEvent from airflow.utils.sqlalchemy import get_dialect_name @@ -119,29 +117,18 @@ # The gate matches the exact class name; a subclass would need its own entry here. _STUB_TASK_TYPE = "_StubOperator" -# Specs are immutable per (Dag version, task): a re-serialized Dag gets a new -# version id, so entries never go stale and the cap only bounds memory. -_ARG_BINDINGS_CACHE: dict[tuple[UUID, str], list[dict] | None] = {} -_ARG_BINDINGS_CACHE_MAX_ENTRIES = 4096 - -def _get_arg_bindings(dag_version_id: UUID | None, task_id: str, *, session) -> list[dict] | None: - """Extract the stub task's serialized arg spec from its Dag version's serialized blob.""" +def _get_arg_bindings( + dag_bag: DagBagDep, dag_version_id: UUID | None, task_id: str, *, session +) -> list | None: + """Extract the stub task's captured TaskFlow arg spec from its Dag version.""" if dag_version_id is None: return None - cache_key = (dag_version_id, task_id) - if cache_key in _ARG_BINDINGS_CACHE: - return _ARG_BINDINGS_CACHE[cache_key] - dag_version = session.get(DagVersion, dag_version_id, options=[joinedload(DagVersion.serialized_dag)]) - if dag_version is None or dag_version.serialized_dag is None: + if (dag := dag_bag.get_dag(dag_version_id, session=session)) is None: return None - if not (data := dag_version.serialized_dag.data): + if (task := dag.task_dict.get(task_id)) is None: return None - bindings = LazyDeserializedDAG(data=data).get_task_arg_bindings(task_id) - if len(_ARG_BINDINGS_CACHE) >= _ARG_BINDINGS_CACHE_MAX_ENTRIES: - _ARG_BINDINGS_CACHE.clear() - _ARG_BINDINGS_CACHE[cache_key] = bindings - return bindings + return getattr(task, "_arg_bindings", None) @ti_id_router.patch( @@ -348,7 +335,7 @@ def ti_run( # 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) + arg_bindings := _get_arg_bindings(dag_bag, ti.dag_version_id, ti.task_id, session=session) ): try: context.arg_bindings = get_arg_bindings_adapter().validate_python(arg_bindings) diff --git a/airflow-core/src/airflow/serialization/serialized_objects.py b/airflow-core/src/airflow/serialization/serialized_objects.py index fd154f00b96ee..54bc3389c64ce 100644 --- a/airflow-core/src/airflow/serialization/serialized_objects.py +++ b/airflow-core/src/airflow/serialization/serialized_objects.py @@ -2277,21 +2277,6 @@ def __getattr__(self, name: str, /) -> Any: def timetable(self) -> Timetable: return decode_timetable(self.data["dag"]["timetable"]) - def get_task_arg_bindings(self, task_id: str) -> list | None: - """ - Extract one task's serialized ``_arg_bindings`` spec without deserializing the Dag. - - The spec is captured at parse time from a stub task's TaskFlow call; ``None`` for - tasks that carry no spec (regular tasks, and stub tasks with no parameters). - """ - for task in self.data["dag"]["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 - @property def has_task_concurrency_limits(self) -> bool: return any( 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 22ff27f835e3f..a47fd974e5274 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 @@ -54,7 +54,6 @@ from airflow.models.taskinstancehistory import TaskInstanceHistory from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk import Asset, TaskGroup, TriggerRule, task, task_group -from airflow.serialization.serialized_objects import LazyDeserializedDAG from airflow.state.metastore import MetastoreBackend from airflow.utils.state import DagRunState, State, TaskInstanceState, TerminalTIState @@ -411,9 +410,8 @@ def transform(country: str, extracted: dict, limit: int = 10): ... assert response.status_code == 200 assert "arg_bindings" not in response.json() - @mock.patch.object( - LazyDeserializedDAG, - "get_task_arg_bindings", + @mock.patch( + "airflow.api_fastapi.execution_api.routes.task_instances._get_arg_bindings", autospec=True, return_value=[{"name": "country", "kind": "hologram", "value": "uk"}], ) From 0e3ad780ba1c1e9a9c9eaf7513acb4de1e989fd3 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 22 Jul 2026 06:38:48 +0000 Subject: [PATCH 19/40] Test that stub args explicitly passed at their default stay unflagged The from_default flag records provenance, not value equality: the Go TaskInput struct mode fails unclaimed explicit arguments but tolerates unclaimed defaults, so an author-passed value that happens to equal the signature default must not be flagged. Pin that boundary at the capture site. --- .../tests/unit/standard/decorators/test_stub.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index 007e1c68a12da..d853ab8ca4eae 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -125,6 +125,20 @@ def test_kwargs_normalize_to_declaration_order(self): {"name": "retries_num", "kind": "literal", "data_type": "integer", "value": 7}, ] + def test_explicitly_passing_the_default_value_is_not_from_default(self): + """The flag tracks provenance, not value equality: an author-passed argument is explicit + even when it equals the signature default, so keyword-style consumers must still claim it.""" + with DAG(dag_id="d"): + extracted = stub(fn_extract)() + result = stub(fn_transform)("uk", extracted, retries_num=3) + + assert result.operator._arg_bindings[2] == { + "name": "retries_num", + "kind": "literal", + "data_type": "integer", + "value": 3, + } + def test_custom_xcom_key_rejected(self): with DAG(dag_id="d"): extracted = stub(fn_extract)() From eed8bb14d3216f5551bdf3c3192e5cb466916bf9 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 22 Jul 2026 07:39:39 +0000 Subject: [PATCH 20/40] Pin Go SDK and ts-sdk to the 2026-07-30 supervisor schema The supervisor schema bump ships in this PR, and the version-pin guards (TestSupervisorSchemaVersionMatchesSnapshot, the ts-sdk supervisor schema hook) rightly insist the pinned constant and the generated TypeScript models follow the schema in the same change. The Go runtime consumption of the new arg_bindings field stays in the stacked feature/go-sdk/taskflow-arg-binding branch. --- go-sdk/pkg/execution/messages.go | 2 +- ts-sdk/src/generated/supervisor.ts | 132 ++++++++++++++++++++++------- 2 files changed, 104 insertions(+), 30 deletions(-) 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/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index ab2632831ab95..ef58ba4589737 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -22,6 +22,20 @@ // // Re-run with: pnpm run generate:supervisor +/** + * Language-neutral value type a stub-task argument binds to in the foreign runtime. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "ArgBindingDataType". + */ +export type ArgBindingDataType = + | "string" + | "integer" + | "number" + | "boolean" + | "object" + | "array" + | "any"; export type Name = string; export type Id = number; export type Timestamp = string; @@ -245,6 +259,40 @@ export type NextKwargs1 = export type XcomKeysToClear = string[]; export type ShouldRetry = boolean; export type StartDate2 = string | null; +export type ArgBindings = TaskArgBinding[] | null; +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskArgBinding". + */ +export type TaskArgBinding = XComArgBinding | LiteralArgBinding; +export type Kind = "xcom"; +export type Name8 = string; +/** + * Language-neutral value type a stub-task argument binds to in the foreign runtime. + */ +export type ArgBindingDataType1 = + | "string" + | "integer" + | "number" + | "boolean" + | "object" + | "array" + | "any"; +export type TaskId1 = string; +export type Kind1 = "literal"; +export type Name9 = string; +/** + * Language-neutral value type a stub-task argument binds to in the foreign runtime. + */ +export type ArgBindingDataType2 = + | "string" + | "integer" + | "number" + | "boolean" + | "object" + | "array" + | "any"; +export type FromDefault = boolean; export type Type13 = "TaskCallbackRequest"; export type Filepath2 = string; export type BundleName3 = string; @@ -310,7 +358,7 @@ export type NextKwargs2 = { } | null; export type RenderedMapIndex1 = string | null; export type Type20 = "DeferTask"; -export type Name8 = string; +export type Name10 = string; export type Key1 = string; export type Type21 = "DeleteAssetStateStoreByName"; export type Uri5 = string; @@ -324,7 +372,7 @@ export type Type24 = "DeleteVariable"; export type Key5 = string; export type DagId6 = string; export type RunId5 = string; -export type TaskId1 = string; +export type TaskId2 = string; export type MapIndex1 = number | null; export type Type25 = "DeleteXCom"; /** @@ -362,11 +410,11 @@ export type ErrorType1 = | "PERMISSION_DENIED" | "GENERIC_ERROR" | "API_SERVER_ERROR"; -export type Name9 = string; +export type Name11 = string; export type Type27 = "GetAssetByName"; export type Uri6 = string; export type Type28 = "GetAssetByUri"; -export type Name10 = string | null; +export type Name12 = string | null; export type Uri7 = string | null; export type After = string | null; export type Before = string | null; @@ -389,7 +437,7 @@ export type Extra8 = { [k: string]: string; } | null; export type Type30 = "GetAssetEventByAssetAlias"; -export type Name11 = string; +export type Name13 = string; export type Key6 = string; export type Type31 = "GetAssetStateStoreByName"; export type Uri8 = string; @@ -421,7 +469,7 @@ export type LogicalDate3 = string; export type State3 = string | null; export type Type41 = "GetPreviousDagRun"; export type DagId12 = string; -export type TaskId2 = string; +export type TaskId3 = string; export type LogicalDate4 = string | null; export type MapIndex2 = number; export type Type42 = "GetPreviousTI"; @@ -458,25 +506,25 @@ export type Type49 = "GetVariableKeys"; export type Key10 = string; export type DagId16 = string; export type RunId9 = string; -export type TaskId3 = string; +export type TaskId4 = string; export type MapIndex5 = number | null; export type IncludePriorDates = boolean; export type Type50 = "GetXCom"; export type Key11 = string; export type DagId17 = string; export type RunId10 = string; -export type TaskId4 = string; +export type TaskId5 = string; export type Type51 = "GetXComCount"; export type Key12 = string; export type DagId18 = string; export type RunId11 = string; -export type TaskId5 = string; +export type TaskId6 = string; export type Offset1 = number; export type Type52 = "GetXComSequenceItem"; export type Key13 = string; export type DagId19 = string; export type RunId12 = string; -export type TaskId6 = string; +export type TaskId7 = string; export type Start = number | null; export type Stop = number | null; export type Step = number | null; @@ -498,7 +546,7 @@ export type AssignedUsers1 = HITLUser[] | null; export type Type54 = "HITLDetailRequestResult"; export type InactiveAssets = AssetProfile[] | null; export type Type55 = "InactiveAssetsResult"; -export type Name12 = string | null; +export type Name14 = string | null; export type Type56 = "MaskSecret"; export type Ok = boolean; export type Type57 = "OKResponse"; @@ -508,7 +556,7 @@ export type StartDate4 = string | null; export type EndDate3 = string | null; export type Type58 = "PrevSuccessfulDagRunResult"; export type Type59 = "PreviousDagRunResult"; -export type TaskId7 = string; +export type TaskId8 = string; export type DagId20 = string; export type RunId13 = string; export type LogicalDate5 = string | null; @@ -536,7 +584,7 @@ export type RetryReason = string | null; export type Type64 = "RetryTask"; export type Type65 = "SentFDs"; export type Fds = number[]; -export type Name13 = string; +export type Name15 = string; export type Key15 = string; export type Type66 = "SetAssetStateStoreByName"; export type Uri9 = string; @@ -552,7 +600,7 @@ export type Type70 = "SetTaskStateStore"; export type Key18 = string; export type DagId21 = string; export type RunId14 = string; -export type TaskId8 = string; +export type TaskId9 = string; export type MapIndex7 = number | null; export type DagResult1 = boolean; export type MappedLength = number | null; @@ -1003,6 +1051,7 @@ export interface TIRunContext { xcom_keys_to_clear?: XcomKeysToClear; should_retry?: ShouldRetry; start_date?: StartDate2; + arg_bindings?: ArgBindings; } /** * Variable schema for responses with fields that are needed for Runtime. @@ -1030,6 +1079,31 @@ export interface ConnectionResponse { port: Port1; extra: Extra6; } +/** + * One positional stub-task argument pulled from an upstream task's XCom. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "XComArgBinding". + */ +export interface XComArgBinding { + kind: Kind; + name: Name8; + data_type?: ArgBindingDataType1; + task_id: TaskId1; +} +/** + * One positional stub-task argument carrying an inline literal from the Dag file. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "LiteralArgBinding". + */ +export interface LiteralArgBinding { + kind: Kind1; + name: Name9; + data_type?: ArgBindingDataType2; + value?: unknown; + from_default?: FromDefault; +} /** * Email notification request for task failures/retries. * @@ -1149,7 +1223,7 @@ export interface DeferTask { * via the `definition` "DeleteAssetStateStoreByName". */ export interface DeleteAssetStateStoreByName { - name: Name8; + name: Name10; key: Key1; type?: Type21; } @@ -1187,7 +1261,7 @@ export interface DeleteXCom { key: Key5; dag_id: DagId6; run_id: RunId5; - task_id: TaskId1; + task_id: TaskId2; map_index?: MapIndex1; type?: Type25; } @@ -1205,7 +1279,7 @@ export interface ErrorResponse { * via the `definition` "GetAssetByName". */ export interface GetAssetByName { - name: Name9; + name: Name11; type?: Type27; } /** @@ -1221,7 +1295,7 @@ export interface GetAssetByUri { * via the `definition` "GetAssetEventByAsset". */ export interface GetAssetEventByAsset { - name: Name10; + name: Name12; uri: Uri7; after?: After; before?: Before; @@ -1252,7 +1326,7 @@ export interface GetAssetEventByAssetAlias { * via the `definition` "GetAssetStateStoreByName". */ export interface GetAssetStateStoreByName { - name: Name11; + name: Name13; key: Key6; type?: Type31; } @@ -1354,7 +1428,7 @@ export interface GetPreviousDagRun { */ export interface GetPreviousTI { dag_id: DagId12; - task_id: TaskId2; + task_id: TaskId3; logical_date?: LogicalDate4; map_index?: MapIndex2; state?: TaskInstanceState | null; @@ -1440,7 +1514,7 @@ export interface GetXCom { key: Key10; dag_id: DagId16; run_id: RunId9; - task_id: TaskId3; + task_id: TaskId4; map_index?: MapIndex5; include_prior_dates?: IncludePriorDates; type?: Type50; @@ -1455,7 +1529,7 @@ export interface GetXComCount { key: Key11; dag_id: DagId17; run_id: RunId10; - task_id: TaskId4; + task_id: TaskId5; type?: Type51; } /** @@ -1466,7 +1540,7 @@ export interface GetXComSequenceItem { key: Key12; dag_id: DagId18; run_id: RunId11; - task_id: TaskId5; + task_id: TaskId6; offset: Offset1; type?: Type52; } @@ -1478,7 +1552,7 @@ export interface GetXComSequenceSlice { key: Key13; dag_id: DagId19; run_id: RunId12; - task_id: TaskId6; + task_id: TaskId7; start: Start; stop: Stop; step: Step; @@ -1520,7 +1594,7 @@ export interface InactiveAssetsResult { */ export interface MaskSecret { value: JsonValue; - name?: Name12; + name?: Name14; type?: Type56; } /** @@ -1559,7 +1633,7 @@ export interface PreviousDagRunResult { * via the `definition` "PreviousTIResponse". */ export interface PreviousTIResponse { - task_id: TaskId7; + task_id: TaskId8; dag_id: DagId20; run_id: RunId13; logical_date?: LogicalDate5; @@ -1636,7 +1710,7 @@ export interface SentFDs { * via the `definition` "SetAssetStateStoreByName". */ export interface SetAssetStateStoreByName { - name: Name13; + name: Name15; key: Key15; value: JsonValue; type?: Type66; @@ -1694,7 +1768,7 @@ export interface SetXCom { value: JsonValue; dag_id: DagId21; run_id: RunId14; - task_id: TaskId8; + task_id: TaskId9; map_index?: MapIndex7; dag_result?: DagResult1; mapped_length?: MappedLength; @@ -1896,4 +1970,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; From 49a7a35466b933000105e48f263b347a6001d708 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 24 Jul 2026 02:44:58 +0000 Subject: [PATCH 21/40] Carry stub arg types as JSON-schema fragments instead of a custom enum Review feedback on the arg-binding contract asked to reuse JSON Schema rather than invent an ArgBindingDataType vocabulary, so foreign runtimes can validate bound values with plain JSON-schema semantics. Each binding now ships an optional value_schema fragment carrying the standard type and format keywords: unions and Optionals map to type lists instead of degrading to "any", and int/float/datetime/date/time/timedelta annotations gain int64/double/date-time/date/time/duration formats the bare type name cannot convey. An unconstrained argument omits the field entirely, and unknown keywords from newer providers are tolerated. The kind discriminator stays required: giving it a server-side default drops it from the OpenAPI required list and datamodel-codegen then emits Literal | None, which pydantic rejects as a tagged-union discriminator. Also retargets the execution API and supervisor schema version to 2026-10-30 to match the Airflow 3.4 release train, documents the optional _arg_bindings task property in the serialized-Dag schema (no serializer version bump: optional field, unchanged logic), adds the Dag version id to the ti_run spec-validation failure log, and rewords the stub .expand() rejection rationale: arg types are uniform across map indexes, the blocker is per-index value resolution at runtime. --- .../datamodels/task_arg_binding.py | 52 ++- .../execution_api/routes/task_instances.py | 1 + .../execution_api/versions/__init__.py | 4 +- .../{v2026_07_30.py => v2026_10_30.py} | 0 .../src/airflow/serialization/schema.json | 49 ++- .../versions/head/test_task_instances.py | 29 +- .../{v2026_07_30 => v2026_10_30}/__init__.py | 0 .../test_task_instances.py | 12 +- .../serialization/test_dag_serialization.py | 18 +- go-sdk/pkg/execution/messages.go | 2 +- .../providers/standard/decorators/stub.py | 138 ++++-- .../unit/standard/decorators/test_stub.py | 123 ++++-- task-sdk/pyproject.toml | 1 - .../airflow/sdk/api/datamodels/_generated.py | 87 ++-- .../sdk/execution_time/schema/schema.json | 142 ++++-- .../schema/versions/__init__.py | 4 +- .../{v2026_07_30.py => v2026_10_30.py} | 0 .../execution_time/schema/test_migrator.py | 19 +- ts-sdk/src/generated/supervisor.ts | 414 +++++++++--------- 19 files changed, 679 insertions(+), 416 deletions(-) rename airflow-core/src/airflow/api_fastapi/execution_api/versions/{v2026_07_30.py => v2026_10_30.py} (100%) rename airflow-core/tests/unit/api_fastapi/execution_api/versions/{v2026_07_30 => v2026_10_30}/__init__.py (100%) rename airflow-core/tests/unit/api_fastapi/execution_api/versions/{v2026_07_30 => v2026_10_30}/test_task_instances.py (86%) rename task-sdk/src/airflow/sdk/execution_time/schema/versions/{v2026_07_30.py => v2026_10_30.py} (100%) 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 index 58866a82fc98f..e26559b1d090c 100644 --- 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 @@ -24,7 +24,6 @@ from __future__ import annotations -from enum import Enum from functools import cache from typing import Annotated, Literal @@ -33,29 +32,52 @@ from airflow.api_fastapi.core_api.base import BaseModel +# A named alias (like TaskArgBinding below) so both union branches of ArgValueSchema.type +# reference one shared schema definition instead of two inlined enum copies; the explicit +# title lets the supervisor-schema dump merge this def with the task-sdk-generated twin. +JsonSchemaType = TypeAliasType( + "JsonSchemaType", + Annotated[ + Literal["string", "integer", "number", "boolean", "object", "array", "null"], + Field(title="JsonSchemaType"), + ], +) +"""JSON-schema primitive type names a stub-task argument annotation can map to.""" + + +class ArgValueSchema(BaseModel): + """ + JSON-schema fragment constraining the value a stub-task argument binds to. + + Only the ``type`` and ``format`` keywords are carried today, with their standard + JSON-schema semantics: ``type`` is asserted by the runtime, ``format`` is an + annotation a runtime may additionally check. Unknown keywords from newer providers + are ignored rather than rejected (as JSON-schema consumers do), so a core on this + version keeps serving specs written by a newer provider. + """ -class ArgBindingDataType(str, Enum): - """Language-neutral value type a stub-task argument binds to in the foreign runtime.""" + type: JsonSchemaType | list[JsonSchemaType] | None = None + """A single type name, a union of type names, or ``None`` when unconstrained.""" - STRING = "string" - INTEGER = "integer" - NUMBER = "number" - BOOLEAN = "boolean" - OBJECT = "object" - ARRAY = "array" - ANY = "any" + format: str | None = None + """Wire representation the type name alone cannot convey (``int64``, ``date-time``, + ``duration``, ...); open vocabulary, per JSON schema.""" class XComArgBinding(BaseModel): """One positional stub-task argument pulled from an upstream task's XCom.""" + # No default on the discriminator: a default drops ``kind`` from ``required`` in the + # OpenAPI schema, and the generated task-sdk client then types it ``Literal | None``, + # which pydantic rejects as a tagged-union discriminator. kind: Literal["xcom"] name: str """The stub function's parameter name this binding fills, in declaration order.""" - data_type: ArgBindingDataType = ArgBindingDataType.ANY - """Declared type from the stub function's annotation; runtimes type-check against it.""" + value_schema: ArgValueSchema | None = None + """JSON-schema fragment from the stub function's annotation; runtimes validate the + bound value against it. Omitted when the annotation gives no constraint.""" task_id: str """Upstream task id to pull the XCom from; the ``return_value`` XCom is always the one pulled.""" @@ -65,12 +87,14 @@ class LiteralArgBinding(BaseModel): """One positional stub-task argument carrying an inline literal from the Dag file.""" kind: Literal["literal"] + """No default, for the same generated-client reason as ``XComArgBinding.kind``.""" name: str """The stub function's parameter name this binding fills, in declaration order.""" - data_type: ArgBindingDataType = ArgBindingDataType.ANY - """Declared type from the stub function's annotation; runtimes type-check against it.""" + value_schema: ArgValueSchema | None = None + """JSON-schema fragment from the stub function's annotation; runtimes validate the + bound value against it. Omitted when the annotation gives no constraint.""" value: JsonValue | None = None """The literal value from the Dag file.""" 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 cd7ed440b391f..b1fca0aaaa51d 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 @@ -344,6 +344,7 @@ def ti_run( "Serialized arg_bindings spec failed validation", dag_id=ti.dag_id, task_id=ti.task_id, + dag_version_id=ti.dag_version_id, ) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, 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 03d50957e8a84..d56ec735c8f13 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py @@ -51,11 +51,11 @@ AddTeamNameField, AddVariableKeysEndpoint, ) -from airflow.api_fastapi.execution_api.versions.v2026_07_30 import AddArgBindingsToTIRunContext +from airflow.api_fastapi.execution_api.versions.v2026_10_30 import AddArgBindingsToTIRunContext bundle = VersionBundle( HeadVersion(), - Version("2026-07-30", AddArgBindingsToTIRunContext), + Version("2026-10-30", AddArgBindingsToTIRunContext), Version( "2026-06-30", AddVariableKeysEndpoint, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_07_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py similarity index 100% rename from airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_07_30.py rename to airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py diff --git a/airflow-core/src/airflow/serialization/schema.json b/airflow-core/src/airflow/serialization/schema.json index 872c3a1331ee3..fb954f0615b9d 100644 --- a/airflow-core/src/airflow/serialization/schema.json +++ b/airflow-core/src/airflow/serialization/schema.json @@ -142,6 +142,48 @@ "description": "A python dictionary containing values of any type", "type": "object" }, + "typed_dict": { + "type": "object", + "properties": { + "__type": { + "type": "string", + "const": "dict" + }, + "__var": { "$ref": "#/definitions/dict" } + }, + "required": [ + "__type", + "__var" + ], + "additionalProperties": false + }, + "arg_binding": { + "$comment": "One captured TaskFlow call argument of a @task.stub task, in dict-encoded form. The inner object stays open so future binding fields keep validating on older cores", + "type": "object", + "properties": { + "__type": { + "type": "string", + "const": "dict" + }, + "__var": { + "type": "object", + "properties": { + "name": { "type": "string" }, + "kind": { "type": "string" }, + "value_schema": { "$ref": "#/definitions/typed_dict" }, + "task_id": { "type": "string" }, + "value": {}, + "from_default": { "type": "boolean" } + }, + "required": [ "name", "kind" ] + } + }, + "required": [ + "__type", + "__var" + ], + "additionalProperties": false + }, "color": { "type": "string", "pattern": "^#[a-fA-F0-9]{3,6}$" @@ -345,7 +387,12 @@ "is_teardown": {"type": "boolean", "default": false}, "on_failure_fail_dagrun": {"type": "boolean", "default": false}, "max_active_tis_per_dag": {"type": "integer"}, - "max_active_tis_per_dagrun": {"type": "integer"} + "max_active_tis_per_dagrun": {"type": "integer"}, + "_arg_bindings": { + "$comment": "Only present on @task.stub tasks called with TaskFlow arguments", + "type": "array", + "items": { "$ref": "#/definitions/arg_binding" } + } }, "dependencies": { "expand_input": ["partial_kwargs", "_is_mapped"], diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index a47fd974e5274..288301c8a63f0 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 @@ -400,9 +400,15 @@ def transform(country: str, extracted: dict, limit: int = 10): ... response = client.patch(f"/execution/task-instances/{tis['transform'].id}/run", json=payload) assert response.status_code == 200 assert response.json()["arg_bindings"] == [ - {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, - {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract"}, - {"name": "limit", "kind": "literal", "data_type": "integer", "value": 10, "from_default": True}, + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, + {"name": "extracted", "kind": "xcom", "value_schema": {"type": "object"}, "task_id": "extract"}, + { + "name": "limit", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 10, + "from_default": True, + }, ] # An argless stub has no captured spec, so the field stays unset. @@ -452,6 +458,23 @@ def test_arg_bindings_adapter_rejects_unknown_kind(self): [{"name": "country", "kind": "template", "value": "x"}] ) + def test_arg_bindings_adapter_tolerates_unknown_value_schema_keyword(self): + """A JSON-schema keyword this core version does not know (e.g. from a newer provider) + must not fail the spec -- JSON-schema consumers ignore unknown keywords by design.""" + from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import get_arg_bindings_adapter + + (binding,) = get_arg_bindings_adapter().validate_python( + [ + { + "name": "tags", + "kind": "literal", + "value_schema": {"type": "array", "items": {"type": "string"}}, + "value": ["a"], + } + ] + ) + assert binding.value_schema.type == "array" + 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_07_30/__init__.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/__init__.py similarity index 100% rename from airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_07_30/__init__.py rename to airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/__init__.py diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_07_30/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py similarity index 86% rename from airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_07_30/test_task_instances.py rename to airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py index cca4a78f9d859..0e0972afed8da 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_07_30/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py @@ -83,7 +83,13 @@ def test_head_version_includes_arg_bindings(self, client, stub_ti): response = client.patch(f"/execution/task-instances/{stub_ti.id}/run", json=RUN_PATCH_BODY) assert response.status_code == 200 assert response.json()["arg_bindings"] == [ - {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, - {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract"}, - {"name": "limit", "kind": "literal", "data_type": "integer", "value": 10, "from_default": True}, + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, + {"name": "extracted", "kind": "xcom", "value_schema": {"type": "object"}, "task_id": "extract"}, + { + "name": "limit", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 10, + "from_default": True, + }, ] diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index 15b98987be7a4..76e937bd8fa15 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -3425,18 +3425,28 @@ def transform(country: str, extracted: dict): ... assert encoded_tasks["transform"]["_arg_bindings"] == [ { Encoding.TYPE: DAT.DICT, - Encoding.VAR: {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, + Encoding.VAR: { + "name": "country", + "kind": "literal", + "value_schema": {Encoding.TYPE: DAT.DICT, Encoding.VAR: {"type": "string"}}, + "value": "uk", + }, }, { Encoding.TYPE: DAT.DICT, - Encoding.VAR: {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract"}, + Encoding.VAR: { + "name": "extracted", + "kind": "xcom", + "value_schema": {Encoding.TYPE: DAT.DICT, Encoding.VAR: {"type": "object"}}, + "task_id": "extract", + }, }, ] round_tripped = DagSerialization.from_dict(ser_dag) assert round_tripped.task_dict["transform"]._arg_bindings == [ - {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, - {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract"}, + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, + {"name": "extracted", "kind": "xcom", "value_schema": {"type": "object"}, "task_id": "extract"}, ] assert not hasattr(round_tripped.task_dict["extract"], "_arg_bindings") or ( round_tripped.task_dict["extract"]._arg_bindings is None diff --git a/go-sdk/pkg/execution/messages.go b/go-sdk/pkg/execution/messages.go index 378cc931dda50..bb81d60c0a4ff 100644 --- a/go-sdk/pkg/execution/messages.go +++ b/go-sdk/pkg/execution/messages.go @@ -32,7 +32,7 @@ import ( // reported in a bundle's airflow-metadata manifest as // sdk.supervisor_schema_version so the supervisor can down/upgrade messages to // a shape the bundle understands. -const SupervisorSchemaVersion = "2026-07-30" +const SupervisorSchemaVersion = "2026-10-30" // The message-type discriminator strings (genmodels.Type*) are generated from the // schema's "type" consts in discriminators.gen.go; outbound messages stamp the diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index ce4d00e99d85d..0ed17b355c868 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -18,6 +18,7 @@ from __future__ import annotations import ast +import datetime import inspect import json import types @@ -38,43 +39,84 @@ from airflow.providers.common.compat.sdk import Context -def _infer_data_type(annotation: Any) -> str: - """ - Map a stub function parameter annotation to the language-neutral arg-type vocabulary. - - The returned name is one of the execution API's ``ArgBindingDataType`` values - (``string``/``integer``/``number``/``boolean``/``object``/``array``/``any``); the foreign - runtime type-checks the bound value against it. 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" +def _json_schema_fragment(annotation: Any) -> dict[str, Any] | None: + """Map one non-union annotation to a JSON-schema fragment; ``None`` = unclassifiable.""" + if annotation is type(None): + return {"type": "null"} origin = typing.get_origin(annotation) if origin is not None: - if origin is Union or origin is types.UnionType: - members = [a for a in typing.get_args(annotation) if a is not type(None)] - if len(members) == 1: - return _infer_data_type(members[0]) - return "any" annotation = origin if not isinstance(annotation, type): - return "any" - # bool subclasses int, and str/bytes are Sequences -- order matters. + return None + # bool subclasses int, str/bytes are Sequences, and datetime subclasses date -- order matters. if issubclass(annotation, bool): - return "boolean" + return {"type": "boolean"} if issubclass(annotation, int): - return "integer" + return {"type": "integer", "format": "int64"} if issubclass(annotation, float): - return "number" + return {"type": "number", "format": "double"} if issubclass(annotation, str): - return "string" + return {"type": "string"} if issubclass(annotation, bytes): - return "any" + return None + if issubclass(annotation, datetime.datetime): + return {"type": "string", "format": "date-time"} + if issubclass(annotation, datetime.date): + return {"type": "string", "format": "date"} + if issubclass(annotation, datetime.time): + return {"type": "string", "format": "time"} + if issubclass(annotation, datetime.timedelta): + return {"type": "string", "format": "duration"} if issubclass(annotation, (dict, Mapping)): - return "object" + return {"type": "object"} if issubclass(annotation, (list, tuple, set, frozenset, Sequence)): - return "array" - return "any" + return {"type": "array"} + return None + + +def _infer_value_schema(annotation: Any) -> dict[str, Any] | None: + """ + Map a stub function parameter annotation to a JSON-schema fragment. + + Fragments carry the standard ``type`` keyword -- a single name, or a list for union + annotations (set semantics; member order follows the annotation) -- plus a ``format`` + annotation where the Python type implies a wire representation the type name alone + cannot (``int`` -> ``int64``, ``datetime`` -> ``date-time``, ``timedelta`` -> + ``duration``, ...). Returns ``None`` when the annotation gives no constraint; the + binding then omits ``value_schema`` and the foreign runtime falls back to a + decode-only check. + """ + if annotation is inspect.Parameter.empty or annotation is None or annotation is Any: + return None + if annotation is type(None): + # get_type_hints normalizes a bare ``None`` annotation to NoneType; a parameter + # that can only ever be None constrains nothing worth shipping. + return None + origin = typing.get_origin(annotation) + if origin is Union or origin is types.UnionType: + fragments: list[dict[str, Any]] = [] + for member in typing.get_args(annotation): + fragment = _json_schema_fragment(member) + if fragment is None: + # One unclassifiable member makes the whole union unconstrained: a + # partial schema would wrongly reject that member's values. + return None + if fragment not in fragments: + fragments.append(fragment) + data_fragments = [f for f in fragments if f["type"] != "null"] + if len(data_fragments) == 1: + # A single data type (+ optional null) keeps its format: format only + # constrains values of its own type, so null passes it untouched. + schema = dict(data_fragments[0]) + if len(fragments) > len(data_fragments): + schema["type"] = [schema["type"], "null"] + return schema + # Mixed-type union: formats are per-member and inexpressible in one flat + # fragment, so carry the type names only. + type_names = [f["type"] for f in fragments] + deduped = list(dict.fromkeys(type_names)) + return {"type": deduped if len(deduped) > 1 else deduped[0]} + return _json_schema_fragment(annotation) def _build_arg_bindings( @@ -134,7 +176,7 @@ def get_annotation_for(name: str, param: inspect.Parameter) -> Any: spec: list[dict[str, Any]] = [] for name, param in signature.parameters.items(): value = bound.arguments[name] - data_type = _infer_data_type(get_annotation_for(name, param)) + value_schema = _infer_value_schema(get_annotation_for(name, param)) if isinstance(value, PlainXComArg): if value.key != "return_value": raise ValueError( @@ -142,14 +184,10 @@ def get_annotation_for(name: str, param: inspect.Parameter) -> Any: f"{value.key!r}; only an upstream task's return value can cross the language " "boundary -- indexing an output by a custom key is not supported" ) - spec.append( - { - "name": name, - "kind": "xcom", - "data_type": data_type, - "task_id": value.operator.task_id, - } - ) + xcom_entry: dict[str, Any] = {"name": name, "kind": "xcom", "task_id": value.operator.task_id} + if value_schema is not None: + xcom_entry["value_schema"] = value_schema + spec.append(xcom_entry) continue if isinstance(value, XComArg): raise ValueError( @@ -165,7 +203,11 @@ def get_annotation_for(name: str, param: inspect.Parameter) -> Any: f"{type(value).__name__} that is not JSON-serializable, so it cannot be passed " "to the foreign runtime" ) - entry: dict[str, Any] = {"name": name, "kind": "literal", "data_type": data_type, "value": value} + entry: dict[str, Any] = {"name": name, "kind": "literal", "value": value} + if value_schema is not None: + # Key omission (never ``None``) is the wire contract for "unconstrained": + # ti_run responds with ``exclude_unset``, so an absent key stays absent. + entry["value_schema"] = value_schema if name not in explicitly_bound: entry["from_default"] = True spec.append(entry) @@ -175,11 +217,14 @@ def get_annotation_for(name: str, param: inspect.Parameter) -> Any: 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 direct .expand() at parse time - # for operator classes that opt out on Airflow >= 3.4 (older cores cannot enforce it, - # and never serialize a spec for the mapped stub); stubs called with arguments inside - # a mapped task group are rejected in __init__ below. + # A mapped stub's arg *types* are uniform across map indexes (same function), but its + # arg *values* only resolve per map index at runtime, while the spec below is captured + # at parse time and the wire contract has no mapped-binding kind yet -- so .expand() + # is rejected rather than shipped with a wrong or empty spec. The task-sdk decorator + # machinery rejects direct .expand() at parse time for operator classes that opt out + # on Airflow >= 3.4 (older cores cannot enforce it, and never serialize a spec for the + # mapped stub); stubs called with arguments inside a mapped task group are rejected in + # __init__ below. supports_expand: bool = False def __init__( @@ -229,14 +274,15 @@ def __init__( self._arg_bindings = _build_arg_bindings(python_callable, self.op_args, self.op_kwargs, self.task_id) # supports_expand only blocks direct .expand() on the stub itself; a mapped task - # group still creates per-map-index instances of every task inside it, and the - # captured spec has no map-index dimension to bind against. + # group still creates per-map-index instances of every task inside it, whose arg + # values resolve per map index at runtime -- after this parse-time capture. in_mapped_group = getattr(self, "get_closest_mapped_task_group", lambda: None)() is not None if self._arg_bindings is not None and in_mapped_group: raise ValueError( f"@task.stub task {self.task_id!r} passes TaskFlow call arguments inside a mapped " - "task group; per-map-index arg specs cannot cross the language boundary yet, so " - "stub tasks with arguments are not supported under a task group's .expand()" + "task group; the captured spec cannot carry values that resolve per map index at " + "runtime, so stub tasks with arguments are not supported under a task group's " + ".expand()" ) @classmethod @@ -264,7 +310,7 @@ def stub( Stub functions may declare parameters and be called TaskFlow-style with upstream task outputs or JSON-serializable literals; the resulting argument-binding spec (parameter - names, declared types, and values, in declaration order) is delivered to the foreign + names, value schemas, and values, in declaration order) is delivered to the foreign runtime, which binds the values onto the native task function. """ return task_decorator_factory( diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index d853ab8ca4eae..a8c3ae8634a48 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -17,13 +17,15 @@ from __future__ import annotations import contextlib +import datetime import typing from typing import Any +import pendulum import pytest from airflow.providers.common.compat.sdk import DAG, task_group -from airflow.providers.standard.decorators.stub import _infer_data_type, stub +from airflow.providers.standard.decorators.stub import _infer_value_schema, stub from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS, AIRFLOW_V_3_4_PLUS @@ -102,12 +104,17 @@ def test_literal_and_xcom_spec(self): op = result.operator assert op._arg_bindings == [ - {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, - {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "fn_extract"}, + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object"}, + "task_id": "fn_extract", + }, { "name": "retries_num", "kind": "literal", - "data_type": "integer", + "value_schema": {"type": "integer", "format": "int64"}, "value": 3, "from_default": True, }, @@ -120,9 +127,19 @@ def test_kwargs_normalize_to_declaration_order(self): result = stub(fn_transform)(extracted=extracted, country="fr", retries_num=7) assert result.operator._arg_bindings == [ - {"name": "country", "kind": "literal", "data_type": "string", "value": "fr"}, - {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "fn_extract"}, - {"name": "retries_num", "kind": "literal", "data_type": "integer", "value": 7}, + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "fr"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object"}, + "task_id": "fn_extract", + }, + { + "name": "retries_num", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 7, + }, ] def test_explicitly_passing_the_default_value_is_not_from_default(self): @@ -135,7 +152,7 @@ def test_explicitly_passing_the_default_value_is_not_from_default(self): assert result.operator._arg_bindings[2] == { "name": "retries_num", "kind": "literal", - "data_type": "integer", + "value_schema": {"type": "integer", "format": "int64"}, "value": 3, } @@ -148,25 +165,24 @@ def test_custom_xcom_key_rejected(self): 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): + def test_untyped_params_omit_value_schema(self): + """Key absence (never ``None``) is the wire contract for an unconstrained argument.""" with DAG(dag_id="d"): result = stub(fn_untyped)(1, "x") assert result.operator._arg_bindings == [ - {"name": "a", "kind": "literal", "data_type": "any", "value": 1}, - {"name": "b", "kind": "literal", "data_type": "any", "value": "x"}, + {"name": "a", "kind": "literal", "value": 1}, + {"name": "b", "kind": "literal", "value": "x"}, ] - def test_unresolvable_annotation_degrades_to_any(self): + def test_unresolvable_annotation_omits_value_schema(self): def fn(x): ... fn.__annotations__ = {"x": "NotARealType"} with DAG(dag_id="d"): result = stub(fn)("v") - assert result.operator._arg_bindings == [ - {"name": "x", "kind": "literal", "data_type": "any", "value": "v"} - ] + assert result.operator._arg_bindings == [{"name": "x", "kind": "literal", "value": "v"}] def test_varargs_rejected(self): with pytest.raises(ValueError, match="fixed number of parameters"): @@ -217,12 +233,17 @@ def test_arg_bindings_survive_dag_serialization_round_trip(self): round_tripped = DagSerialization.from_dict(DagSerialization.to_dict(dag)) assert round_tripped.task_dict["fn_transform"]._arg_bindings == [ - {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, - {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "fn_extract"}, + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object"}, + "task_id": "fn_extract", + }, { "name": "retries_num", "kind": "literal", - "data_type": "integer", + "value_schema": {"type": "integer", "format": "int64"}, "value": 3, "from_default": True, }, @@ -257,35 +278,57 @@ def group(n): @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(str, {"type": "string"}, id="str"), + pytest.param(bool, {"type": "boolean"}, id="bool"), + pytest.param(int, {"type": "integer", "format": "int64"}, id="int"), + pytest.param(float, {"type": "number", "format": "double"}, id="float"), + pytest.param(dict, {"type": "object"}, id="dict"), + pytest.param(dict[str, int], {"type": "object"}, id="dict-parameterized"), + pytest.param(typing.Mapping[str, int], {"type": "object"}, id="mapping"), + pytest.param(list, {"type": "array"}, id="list"), + pytest.param(list[int], {"type": "array"}, id="list-parameterized"), + pytest.param(tuple, {"type": "array"}, id="tuple"), + pytest.param(set, {"type": "array"}, id="set"), + pytest.param(typing.Sequence[int], {"type": "array"}, id="sequence"), + pytest.param(datetime.datetime, {"type": "string", "format": "date-time"}, id="datetime"), + pytest.param(pendulum.DateTime, {"type": "string", "format": "date-time"}, id="pendulum-datetime"), + pytest.param(datetime.date, {"type": "string", "format": "date"}, id="date"), + pytest.param(datetime.time, {"type": "string", "format": "time"}, id="time"), + pytest.param(datetime.timedelta, {"type": "string", "format": "duration"}, id="timedelta"), + pytest.param(Any, None, id="any"), + pytest.param(None, None, id="none"), + pytest.param(type(None), None, id="nonetype"), + pytest.param(bytes, None, id="bytes"), pytest.param( typing.Optional[str], # noqa: UP045 -- legacy form on purpose - "string", + {"type": ["string", "null"]}, id="optional-str", ), pytest.param( typing.Union[int, str], # noqa: UP007 -- legacy form on purpose - "any", + {"type": ["integer", "string"]}, id="union", ), - pytest.param(str | None, "string", id="pep604-optional"), - pytest.param(int | str, "any", id="pep604-union"), - pytest.param(contextlib.AbstractContextManager, "any", id="custom-class"), + pytest.param(str | None, {"type": ["string", "null"]}, id="pep604-optional"), + pytest.param(int | None, {"type": ["integer", "null"], "format": "int64"}, id="optional-int"), + pytest.param( + datetime.datetime | None, + {"type": ["string", "null"], "format": "date-time"}, + id="optional-datetime", + ), + pytest.param(int | str, {"type": ["integer", "string"]}, id="pep604-union"), + pytest.param(dict | bool, {"type": ["object", "boolean"]}, id="union-dict-bool"), + pytest.param(bool | int, {"type": ["boolean", "integer"]}, id="union-bool-int-order"), + pytest.param(str | int | None, {"type": ["string", "integer", "null"]}, id="union-with-null"), + pytest.param(list | tuple, {"type": "array"}, id="union-collapses-to-one-type"), + pytest.param( + datetime.datetime | str, + {"type": "string"}, + id="mixed-format-union-drops-format", + ), + pytest.param(str | bytes, None, id="union-unclassifiable-member"), + pytest.param(contextlib.AbstractContextManager, None, id="custom-class"), ], ) -def test_infer_data_type(annotation, expected): - assert _infer_data_type(annotation) == expected +def test_infer_value_schema(annotation, expected): + assert _infer_value_schema(annotation) == expected diff --git a/task-sdk/pyproject.toml b/task-sdk/pyproject.toml index ea394d57b12c1..75eac99a93888 100644 --- a/task-sdk/pyproject.toml +++ b/task-sdk/pyproject.toml @@ -253,7 +253,6 @@ 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 d0bc470c7ae96..1de6d574278e2 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -27,21 +27,7 @@ from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, RootModel -API_VERSION: Final[str] = "2026-07-30" - - -class ArgBindingDataType(str, Enum): - """ - Language-neutral value type a stub-task argument binds to in the foreign runtime. - """ - - STRING = "string" - INTEGER = "integer" - NUMBER = "number" - BOOLEAN = "boolean" - OBJECT = "object" - ARRAY = "array" - ANY = "any" +API_VERSION: Final[str] = "2026-10-30" class AssetAliasReferenceAssetEventDagRun(BaseModel): @@ -232,16 +218,14 @@ class IntermediateTIState(str, Enum): AWAITING_INPUT = "awaiting_input" -class LiteralArgBinding(BaseModel): - """ - One positional stub-task argument carrying an inline literal from the Dag file. - """ - - kind: Annotated[Literal["literal"], Field(title="Kind")] - name: Annotated[str, Field(title="Name")] - data_type: ArgBindingDataType | None = ArgBindingDataType.ANY - value: JsonValue | None = None - from_default: Annotated[bool | None, Field(title="From Default")] = False +class JsonSchemaType(str, Enum): + STRING = "string" + INTEGER = "integer" + NUMBER = "number" + BOOLEAN = "boolean" + OBJECT = "object" + ARRAY = "array" + NULL = "null" class PrevSuccessfulDagRunResponse(BaseModel): @@ -539,17 +523,6 @@ class VariableResponse(BaseModel): value: Annotated[str | None, Field(title="Value")] = None -class XComArgBinding(BaseModel): - """ - One positional stub-task argument pulled from an upstream task's XCom. - """ - - kind: Annotated[Literal["xcom"], Field(title="Kind")] - name: Annotated[str, Field(title="Name")] - data_type: ArgBindingDataType | None = ArgBindingDataType.ANY - task_id: Annotated[str, Field(title="Task Id")] - - class XComResponse(BaseModel): """ XCom schema for responses with fields that are needed for Runtime. @@ -645,6 +618,21 @@ class DagAttributeTypes(str, Enum): TASK_GROUP = "taskgroup" +class ArgValueSchema(BaseModel): + """ + JSON-schema fragment constraining the value a stub-task argument binds to. + + Only the ``type`` and ``format`` keywords are carried today, with their standard + JSON-schema semantics: ``type`` is asserted by the runtime, ``format`` is an + annotation a runtime may additionally check. Unknown keywords from newer providers + are ignored rather than rejected (as JSON-schema consumers do), so a core on this + version keeps serving specs written by a newer provider. + """ + + type: Annotated[JsonSchemaType | list[JsonSchemaType] | None, Field(title="Type")] = None + format: Annotated[str | None, Field(title="Format")] = None + + class AssetReferenceAssetEventDagRun(BaseModel): """ Schema for AssetModel used in AssetEventDagRunReference. @@ -734,6 +722,18 @@ class HTTPValidationError(BaseModel): detail: Annotated[list[ValidationError] | None, Field(title="Detail")] = None +class LiteralArgBinding(BaseModel): + """ + One positional stub-task argument carrying an inline literal from the Dag file. + """ + + kind: Annotated[Literal["literal"], Field(title="Kind")] + name: Annotated[str, Field(title="Name")] + value_schema: ArgValueSchema | None = None + value: JsonValue | None = None + from_default: Annotated[bool | None, Field(title="From Default")] = False + + class TITerminalStatePayload(BaseModel): """ Schema for updating TaskInstance to a terminal state except SUCCESS state. @@ -747,8 +747,15 @@ class TITerminalStatePayload(BaseModel): rendered_map_index: Annotated[str | None, Field(title="Rendered Map Index")] = None -class TaskArgBinding(RootModel[XComArgBinding | LiteralArgBinding]): - root: Annotated[XComArgBinding | LiteralArgBinding, Field(discriminator="kind", title="TaskArgBinding")] +class XComArgBinding(BaseModel): + """ + One positional stub-task argument pulled from an upstream task's XCom. + """ + + kind: Annotated[Literal["xcom"], Field(title="Kind")] + name: Annotated[str, Field(title="Name")] + value_schema: ArgValueSchema | None = None + task_id: Annotated[str, Field(title="Task Id")] class AssetEventDagRunReference(BaseModel): @@ -823,6 +830,10 @@ class DagRun(BaseModel): team_name: Annotated[str | None, Field(title="Team Name")] = None +class TaskArgBinding(RootModel[XComArgBinding | LiteralArgBinding]): + root: Annotated[XComArgBinding | LiteralArgBinding, Field(discriminator="kind", title="TaskArgBinding")] + + class TIRunContext(BaseModel): """ Response schema for TaskInstance run context. 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 43d100d7a743e..3c3f112e0866b 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -1,21 +1,44 @@ { "$schema": "https://json-schema.org/draft/2020-12/schema", - "api_version": "2026-07-30", + "api_version": "2026-10-30", "description": "Apache Airflow SDK Supervisor Schema", "$defs": { - "ArgBindingDataType": { - "description": "Language-neutral value type a stub-task argument binds to in the foreign runtime.", - "enum": [ - "string", - "integer", - "number", - "boolean", - "object", - "array", - "any" - ], - "title": "ArgBindingDataType", - "type": "string" + "ArgValueSchema": { + "description": "JSON-schema fragment constraining the value a stub-task argument binds to.\n\nOnly the ``type`` and ``format`` keywords are carried today, with their standard\nJSON-schema semantics: ``type`` is asserted by the runtime, ``format`` is an\nannotation a runtime may additionally check. Unknown keywords from newer providers\nare ignored rather than rejected (as JSON-schema consumers do), so a core on this\nversion keeps serving specs written by a newer provider.", + "properties": { + "type": { + "anyOf": [ + { + "$ref": "#/$defs/JsonSchemaType" + }, + { + "items": { + "$ref": "#/$defs/JsonSchemaType" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Type" + }, + "format": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Format" + } + }, + "title": "ArgValueSchema", + "type": "object" }, "AssetAliasReferenceAssetEventDagRun": { "additionalProperties": false, @@ -3059,6 +3082,19 @@ "title": "InactiveAssetsResult", "type": "object" }, + "JsonSchemaType": { + "enum": [ + "string", + "integer", + "number", + "boolean", + "object", + "array", + "null" + ], + "title": "JsonSchemaType", + "type": "string" + }, "JsonValue": {}, "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.", @@ -4403,6 +4439,42 @@ "title": "VariableResult", "type": "object" }, + "XComArgBinding": { + "description": "One positional stub-task argument pulled from an upstream task's XCom.", + "properties": { + "kind": { + "const": "xcom", + "title": "Kind", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "value_schema": { + "anyOf": [ + { + "$ref": "#/$defs/ArgValueSchema" + }, + { + "type": "null" + } + ], + "default": null + }, + "task_id": { + "title": "Task Id", + "type": "string" + } + }, + "required": [ + "kind", + "name", + "task_id" + ], + "title": "XComArgBinding", + "type": "object" + }, "XComCountResponse": { "properties": { "len": { @@ -4589,9 +4661,16 @@ "title": "Name", "type": "string" }, - "data_type": { - "$ref": "#/$defs/ArgBindingDataType", - "default": "any" + "value_schema": { + "anyOf": [ + { + "$ref": "#/$defs/ArgValueSchema" + }, + { + "type": "null" + } + ], + "default": null }, "value": { "anyOf": [ @@ -4635,35 +4714,6 @@ ], "title": "TaskArgBinding" }, - "XComArgBinding": { - "description": "One positional stub-task argument pulled from an upstream task's XCom.", - "properties": { - "kind": { - "const": "xcom", - "title": "Kind", - "type": "string" - }, - "name": { - "title": "Name", - "type": "string" - }, - "data_type": { - "$ref": "#/$defs/ArgBindingDataType", - "default": "any" - }, - "task_id": { - "title": "Task Id", - "type": "string" - } - }, - "required": [ - "kind", - "name", - "task_id" - ], - "title": "XComArgBinding", - "type": "object" - }, "AssetEventDagRunReference": { "additionalProperties": false, "description": "Schema for AssetEvent model used in DagRun.", 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 59640c587473a..7e5ce93f86bdc 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py @@ -37,13 +37,13 @@ def get_bundle() -> VersionBundle: """ from cadwyn import HeadVersion, Version, VersionBundle - from airflow.sdk.execution_time.schema.versions.v2026_07_30 import ( + from airflow.sdk.execution_time.schema.versions.v2026_10_30 import ( AddArgBindingsToSupervisorTIRunContext, ) return VersionBundle( HeadVersion(), - Version("2026-07-30", AddArgBindingsToSupervisorTIRunContext), + Version("2026-10-30", AddArgBindingsToSupervisorTIRunContext), 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_10_30.py similarity index 100% rename from task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_07_30.py rename to task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py 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 6785fde63be23..518426ad9f010 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 @@ -420,12 +420,18 @@ def startup_details(self): ), max_tries=1, arg_bindings=[ - {"name": "country", "kind": "literal", "data_type": "string", "value": "uk"}, - {"name": "extracted", "kind": "xcom", "data_type": "object", "task_id": "extract"}, + # No value_schema: the unconstrained ("any") case rides through the migrator too. + {"name": "country", "kind": "literal", "value": "uk"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object"}, + "task_id": "extract", + }, { "name": "limit", "kind": "literal", - "data_type": "integer", + "value_schema": {"type": "integer", "format": "int64"}, "value": 10, "from_default": True, }, @@ -443,17 +449,20 @@ def test_downgrade_strips_arg_bindings_for_previous_version(self, real_migrator, assert "arg_bindings" not in out["ti_context"] def test_head_version_keeps_arg_bindings(self, real_migrator, startup_details): - from airflow.sdk.api.datamodels._generated import LiteralArgBinding, XComArgBinding + from airflow.sdk.api.datamodels._generated import JsonSchemaType, LiteralArgBinding, XComArgBinding - out = real_migrator.downgrade(startup_details, "2026-07-30") + out = real_migrator.downgrade(startup_details, "2026-10-30") assert out.ti_context.arg_bindings is not None literal, xcom, defaulted = (a.root for a in out.ti_context.arg_bindings) assert isinstance(literal, LiteralArgBinding) assert literal.value == "uk" assert literal.name == "country" assert literal.from_default is False + assert literal.value_schema is None assert isinstance(xcom, XComArgBinding) assert xcom.task_id == "extract" assert xcom.name == "extracted" + assert xcom.value_schema.type == JsonSchemaType.OBJECT assert isinstance(defaulted, LiteralArgBinding) assert defaulted.from_default is True + assert defaulted.value_schema.format == "int64" diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index ef58ba4589737..9596a42274bf9 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -22,20 +22,20 @@ // // Re-run with: pnpm run generate:supervisor +export type Type = JsonSchemaType | JsonSchemaType[] | null; /** - * Language-neutral value type a stub-task argument binds to in the foreign runtime. - * * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema - * via the `definition` "ArgBindingDataType". + * via the `definition` "JsonSchemaType". */ -export type ArgBindingDataType = +export type JsonSchemaType = | "string" | "integer" | "number" | "boolean" | "object" | "array" - | "any"; + | "null"; +export type Format = string | null; export type Name = string; export type Id = number; export type Timestamp = string; @@ -69,10 +69,10 @@ export type SourceRunId = string | null; export type SourceMapIndex = number | null; export type PartitionKey1 = string | null; export type AssetEvents = AssetEventResponse[]; -export type Type = "AssetEventsResult"; +export type Type1 = "AssetEventsResult"; export type Name2 = string | null; export type Uri1 = string | null; -export type Type1 = string; +export type Type2 = string; export type Name3 = string; export type Uri2 = string; export type Name4 = string; @@ -81,10 +81,10 @@ export type Group1 = string; export type Extra3 = { [k: string]: JsonValue; } | null; -export type Type2 = "AssetResult"; -export type Type3 = "AssetStateStoreResult"; +export type Type3 = "AssetResult"; +export type Type4 = "AssetStateStoreResult"; export type Assets = AssetResult[]; -export type Type4 = "AssetsByAliasResult"; +export type Type5 = "AssetsByAliasResult"; export type State1 = "awaiting_input" | null; export type Timeout = string | null; export type NextMethod = string; @@ -92,18 +92,18 @@ export type NextKwargs = { [k: string]: JsonValue; } | null; export type RenderedMapIndex = string | null; -export type Type5 = "AwaitInputTask"; +export type Type6 = "AwaitInputTask"; export type Name5 = string; export type Version = string | null; export type VersionData = { [k: string]: unknown; } | null; export type Name6 = string; -export type Type6 = "ClearAssetStateStoreByName"; +export type Type7 = "ClearAssetStateStoreByName"; export type Uri4 = string; -export type Type7 = "ClearAssetStateStoreByUri"; +export type Type8 = "ClearAssetStateStoreByUri"; export type TiId = string; -export type Type8 = "ClearTaskStateStore"; +export type Type9 = "ClearTaskStateStore"; export type ConnId = string; export type ConnType = string; export type Host = string | null; @@ -112,7 +112,7 @@ export type Login = string | null; export type Password = string | null; export type Port = number | null; export type Extra4 = string | null; -export type Type9 = "ConnectionResult"; +export type Type10 = "ConnectionResult"; export type TiId1 = string; /** * @minItems 1 @@ -128,9 +128,9 @@ export type Params = { export type AssignedUsers = HITLUser[] | null; export type Id1 = string; export type Name7 = string; -export type Type10 = "CreateHITLDetailPayload"; +export type Type11 = "CreateHITLDetailPayload"; export type Count = number; -export type Type11 = "DRCount"; +export type Type12 = "DRCount"; export type Filepath = string; export type BundleName = string; export type BundleVersion = string | null; @@ -202,7 +202,7 @@ export type ContextCarrier = { } | null; export type Queue = string; export type IsFailureCallback = boolean | null; -export type Type12 = "DagCallbackRequest"; +export type Type13 = "DagCallbackRequest"; export type File = string; export type BundlePath = string; export type BundleName1 = string; @@ -267,33 +267,11 @@ export type ArgBindings = TaskArgBinding[] | null; export type TaskArgBinding = XComArgBinding | LiteralArgBinding; export type Kind = "xcom"; export type Name8 = string; -/** - * Language-neutral value type a stub-task argument binds to in the foreign runtime. - */ -export type ArgBindingDataType1 = - | "string" - | "integer" - | "number" - | "boolean" - | "object" - | "array" - | "any"; export type TaskId1 = string; export type Kind1 = "literal"; export type Name9 = string; -/** - * Language-neutral value type a stub-task argument binds to in the foreign runtime. - */ -export type ArgBindingDataType2 = - | "string" - | "integer" - | "number" - | "boolean" - | "object" - | "array" - | "any"; export type FromDefault = boolean; -export type Type13 = "TaskCallbackRequest"; +export type Type14 = "TaskCallbackRequest"; export type Filepath2 = string; export type BundleName3 = string; export type BundleVersion2 = string | null; @@ -302,9 +280,9 @@ export type VersionData3 = { } | null; export type Msg2 = string | null; export type EmailType = "failure" | "retry"; -export type Type14 = "EmailRequest"; +export type Type15 = "EmailRequest"; export type CallbackRequests = (DagCallbackRequest | TaskCallbackRequest | EmailRequest)[]; -export type Type15 = "DagFileParseRequest"; +export type Type16 = "DagFileParseRequest"; export type Fileloc = string; export type LastLoaded = string | null; export type SerializedDags = LazyDeserializedDAG[]; @@ -312,7 +290,7 @@ export type Warnings = unknown[] | null; export type ImportErrors = { [k: string]: string; } | null; -export type Type16 = "DagFileParsingResult"; +export type Type17 = "DagFileParsingResult"; export type DagId4 = string; export type IsPaused = boolean; export type BundleName4 = string | null; @@ -321,7 +299,7 @@ export type RelativeFileloc = string | null; export type Owners = string | null; export type Tags = string[]; export type NextDagrun = string | null; -export type Type17 = "DagResult"; +export type Type18 = "DagResult"; export type DagId5 = string; export type RunId4 = string; export type LogicalDate2 = string | null; @@ -340,8 +318,8 @@ export type PartitionKey4 = string | null; export type PartitionDate1 = string | null; export type Note1 = string | null; export type TeamName1 = string | null; -export type Type18 = "DagRunResult"; -export type Type19 = "DagRunStateResult"; +export type Type19 = "DagRunResult"; +export type Type20 = "DagRunStateResult"; export type State2 = "deferred" | null; export type Classpath = string; export type TriggerKwargs = @@ -357,24 +335,24 @@ export type NextKwargs2 = { [k: string]: JsonValue; } | null; export type RenderedMapIndex1 = string | null; -export type Type20 = "DeferTask"; +export type Type21 = "DeferTask"; export type Name10 = string; export type Key1 = string; -export type Type21 = "DeleteAssetStateStoreByName"; +export type Type22 = "DeleteAssetStateStoreByName"; export type Uri5 = string; export type Key2 = string; -export type Type22 = "DeleteAssetStateStoreByUri"; +export type Type23 = "DeleteAssetStateStoreByUri"; export type TiId2 = string; export type Key3 = string; -export type Type23 = "DeleteTaskStateStore"; +export type Type24 = "DeleteTaskStateStore"; export type Key4 = string; -export type Type24 = "DeleteVariable"; +export type Type25 = "DeleteVariable"; export type Key5 = string; export type DagId6 = string; export type RunId5 = string; export type TaskId2 = string; export type MapIndex1 = number | null; -export type Type25 = "DeleteXCom"; +export type Type26 = "DeleteXCom"; /** * Error types used in the API client. */ @@ -392,7 +370,7 @@ export type ErrorType = export type Detail = { [k: string]: unknown; } | null; -export type Type26 = "ErrorResponse"; +export type Type27 = "ErrorResponse"; /** * Error types used in the API client. * @@ -411,9 +389,9 @@ export type ErrorType1 = | "GENERIC_ERROR" | "API_SERVER_ERROR"; export type Name11 = string; -export type Type27 = "GetAssetByName"; +export type Type28 = "GetAssetByName"; export type Uri6 = string; -export type Type28 = "GetAssetByUri"; +export type Type29 = "GetAssetByUri"; export type Name12 = string | null; export type Uri7 = string | null; export type After = string | null; @@ -425,7 +403,7 @@ export type PartitionKeyRegexpPattern = string | null; export type Extra7 = { [k: string]: string; } | null; -export type Type29 = "GetAssetEventByAsset"; +export type Type30 = "GetAssetEventByAsset"; export type AliasName = string; export type After1 = string | null; export type Before1 = string | null; @@ -436,43 +414,43 @@ export type PartitionKeyRegexpPattern1 = string | null; export type Extra8 = { [k: string]: string; } | null; -export type Type30 = "GetAssetEventByAssetAlias"; +export type Type31 = "GetAssetEventByAssetAlias"; export type Name13 = string; export type Key6 = string; -export type Type31 = "GetAssetStateStoreByName"; +export type Type32 = "GetAssetStateStoreByName"; export type Uri8 = string; export type Key7 = string; -export type Type32 = "GetAssetStateStoreByUri"; +export type Type33 = "GetAssetStateStoreByUri"; export type AliasName1 = string; -export type Type33 = "GetAssetsByAlias"; +export type Type34 = "GetAssetsByAlias"; export type ConnId2 = string; -export type Type34 = "GetConnection"; +export type Type35 = "GetConnection"; export type DagId7 = string; export type LogicalDates = string[] | null; export type RunIds = string[] | null; export type States = string[] | null; -export type Type35 = "GetDRCount"; +export type Type36 = "GetDRCount"; export type DagId8 = string; -export type Type36 = "GetDag"; +export type Type37 = "GetDag"; export type DagId9 = string; export type RunId6 = string; -export type Type37 = "GetDagRun"; +export type Type38 = "GetDagRun"; export type DagId10 = string; export type RunId7 = string; -export type Type38 = "GetDagRunState"; +export type Type39 = "GetDagRunState"; export type TiId3 = string; -export type Type39 = "GetHITLDetailResponse"; +export type Type40 = "GetHITLDetailResponse"; export type TiId4 = string; -export type Type40 = "GetPrevSuccessfulDagRun"; +export type Type41 = "GetPrevSuccessfulDagRun"; export type DagId11 = string; export type LogicalDate3 = string; export type State3 = string | null; -export type Type41 = "GetPreviousDagRun"; +export type Type42 = "GetPreviousDagRun"; export type DagId12 = string; export type TaskId3 = string; export type LogicalDate4 = string | null; export type MapIndex2 = number; -export type Type42 = "GetPreviousTI"; +export type Type43 = "GetPreviousTI"; export type DagId13 = string; export type MapIndex3 = number | null; export type TaskIds = string[] | null; @@ -480,47 +458,47 @@ export type TaskGroupId = string | null; export type LogicalDates1 = string[] | null; export type RunIds1 = string[] | null; export type States1 = string[] | null; -export type Type43 = "GetTICount"; +export type Type44 = "GetTICount"; export type DagId14 = string; export type RunId8 = string; -export type Type44 = "GetTaskBreadcrumbs"; +export type Type45 = "GetTaskBreadcrumbs"; export type TiId5 = string; export type TryNumber1 = number; -export type Type45 = "GetTaskRescheduleStartDate"; +export type Type46 = "GetTaskRescheduleStartDate"; export type TiId6 = string; export type Key8 = string; -export type Type46 = "GetTaskStateStore"; +export type Type47 = "GetTaskStateStore"; export type DagId15 = string; export type MapIndex4 = number | null; export type TaskIds1 = string[] | null; export type TaskGroupId1 = string | null; export type LogicalDates2 = string[] | null; export type RunIds2 = string[] | null; -export type Type47 = "GetTaskStates"; +export type Type48 = "GetTaskStates"; export type Key9 = string; -export type Type48 = "GetVariable"; +export type Type49 = "GetVariable"; export type Prefix = string | null; export type Limit2 = number; export type Offset = number; -export type Type49 = "GetVariableKeys"; +export type Type50 = "GetVariableKeys"; export type Key10 = string; export type DagId16 = string; export type RunId9 = string; export type TaskId4 = string; export type MapIndex5 = number | null; export type IncludePriorDates = boolean; -export type Type50 = "GetXCom"; +export type Type51 = "GetXCom"; export type Key11 = string; export type DagId17 = string; export type RunId10 = string; export type TaskId5 = string; -export type Type51 = "GetXComCount"; +export type Type52 = "GetXComCount"; export type Key12 = string; export type DagId18 = string; export type RunId11 = string; export type TaskId6 = string; export type Offset1 = number; -export type Type52 = "GetXComSequenceItem"; +export type Type53 = "GetXComSequenceItem"; export type Key13 = string; export type DagId19 = string; export type RunId12 = string; @@ -529,7 +507,7 @@ export type Start = number | null; export type Stop = number | null; export type Step = number | null; export type IncludePriorDates1 = boolean; -export type Type53 = "GetXComSequenceSlice"; +export type Type54 = "GetXComSequenceSlice"; export type TiId7 = string; /** * @minItems 1 @@ -543,19 +521,19 @@ export type Params1 = { [k: string]: unknown; } | null; export type AssignedUsers1 = HITLUser[] | null; -export type Type54 = "HITLDetailRequestResult"; +export type Type55 = "HITLDetailRequestResult"; export type InactiveAssets = AssetProfile[] | null; -export type Type55 = "InactiveAssetsResult"; +export type Type56 = "InactiveAssetsResult"; export type Name14 = string | null; -export type Type56 = "MaskSecret"; +export type Type57 = "MaskSecret"; export type Ok = boolean; -export type Type57 = "OKResponse"; +export type Type58 = "OKResponse"; export type DataIntervalStart3 = string | null; export type DataIntervalEnd3 = string | null; export type StartDate4 = string | null; export type EndDate3 = string | null; -export type Type58 = "PrevSuccessfulDagRunResult"; -export type Type59 = "PreviousDagRunResult"; +export type Type59 = "PrevSuccessfulDagRunResult"; +export type Type60 = "PreviousDagRunResult"; export type TaskId8 = string; export type DagId20 = string; export type RunId13 = string; @@ -566,37 +544,37 @@ export type State4 = string | null; export type TryNumber2 = number; export type MapIndex6 = number | null; export type Duration = number | null; -export type Type60 = "PreviousTIResult"; +export type Type61 = "PreviousTIResult"; export type Key14 = string; export type Value1 = string | null; export type Description = string | null; -export type Type61 = "PutVariable"; +export type Type62 = "PutVariable"; export type State5 = "up_for_reschedule" | null; export type RescheduleDate = string; export type EndDate5 = string; -export type Type62 = "RescheduleTask"; -export type Type63 = "ResendLoggingFD"; +export type Type63 = "RescheduleTask"; +export type Type64 = "ResendLoggingFD"; export type State6 = "up_for_retry" | null; export type EndDate6 = string; export type RenderedMapIndex2 = string | null; export type RetryDelaySeconds = number | null; export type RetryReason = string | null; -export type Type64 = "RetryTask"; -export type Type65 = "SentFDs"; +export type Type65 = "RetryTask"; +export type Type66 = "SentFDs"; export type Fds = number[]; export type Name15 = string; export type Key15 = string; -export type Type66 = "SetAssetStateStoreByName"; +export type Type67 = "SetAssetStateStoreByName"; export type Uri9 = string; export type Key16 = string; -export type Type67 = "SetAssetStateStoreByUri"; -export type Type68 = "SetRenderedFields"; +export type Type68 = "SetAssetStateStoreByUri"; +export type Type69 = "SetRenderedFields"; export type RenderedMapIndex3 = string; -export type Type69 = "SetRenderedMapIndex"; +export type Type70 = "SetRenderedMapIndex"; export type TiId8 = string; export type Key17 = string; export type ExpiresAt = string | null; -export type Type70 = "SetTaskStateStore"; +export type Type71 = "SetTaskStateStore"; export type Key18 = string; export type DagId21 = string; export type RunId14 = string; @@ -604,13 +582,13 @@ export type TaskId9 = string; export type MapIndex7 = number | null; export type DagResult1 = boolean; export type MappedLength = number | null; -export type Type71 = "SetXCom"; +export type Type72 = "SetXCom"; export type Tasks = (string | [unknown, unknown])[]; -export type Type72 = "SkipDownstreamTasks"; +export type Type73 = "SkipDownstreamTasks"; export type DagRelPath = string; export type StartDate6 = string; export type SentryIntegration = string; -export type Type73 = "StartupDetails"; +export type Type74 = "StartupDetails"; export type State7 = "success" | null; export type EndDate7 = string; export type TaskOutlets = AssetProfile[] | null; @@ -620,21 +598,21 @@ export type OutletEvents = }[] | null; export type RenderedMapIndex4 = string | null; -export type Type74 = "SucceedTask"; +export type Type75 = "SucceedTask"; export type Count1 = number; -export type Type75 = "TICount"; +export type Type76 = "TICount"; export type Breadcrumbs = { [k: string]: unknown; }[]; -export type Type76 = "TaskBreadcrumbsResult"; +export type Type77 = "TaskBreadcrumbsResult"; export type StartDate7 = string | null; -export type Type77 = "TaskRescheduleStartDate"; +export type Type78 = "TaskRescheduleStartDate"; export type State8 = "failed" | "skipped" | "removed"; export type EndDate8 = string | null; -export type Type78 = "TaskState"; +export type Type79 = "TaskState"; export type RenderedMapIndex5 = string | null; -export type Type79 = "TaskStateStoreResult"; -export type Type80 = "TaskStatesResult"; +export type Type80 = "TaskStateStoreResult"; +export type Type81 = "TaskStatesResult"; export type LogicalDate6 = string | null; export type RunAfter2 = string | null; export type Conf2 = { @@ -645,7 +623,7 @@ export type PartitionKey7 = string | null; export type Note2 = string | null; export type DagId22 = string; export type DagRunId = string; -export type Type81 = "TriggerDagRun"; +export type Type82 = "TriggerDagRun"; export type TiId9 = string; /** * @minItems 1 @@ -654,24 +632,40 @@ export type ChosenOptions = [string, ...string[]]; export type ParamsInput = { [k: string]: unknown; } | null; -export type Type82 = "UpdateHITLDetail"; +export type Type83 = "UpdateHITLDetail"; export type TiId10 = string; -export type Type83 = "ValidateInletsAndOutlets"; +export type Type84 = "ValidateInletsAndOutlets"; export type Keys = string[]; export type TotalEntries = number; -export type Type84 = "VariableKeysResult"; +export type Type85 = "VariableKeysResult"; export type Key19 = string; export type Value2 = string | null; -export type Type85 = "VariableResult"; +export type Type86 = "VariableResult"; export type Len = number; -export type Type86 = "XComCountResponse"; +export type Type87 = "XComCountResponse"; export type Key20 = string; -export type Type87 = "XComResult"; -export type Type88 = "XComSequenceIndexResult"; +export type Type88 = "XComResult"; +export type Type89 = "XComSequenceIndexResult"; export type Root = JsonValue[]; -export type Type89 = "XComSequenceSliceResult"; +export type Type90 = "XComSequenceSliceResult"; export interface SupervisorWireSchema {} +/** + * JSON-schema fragment constraining the value a stub-task argument binds to. + * + * Only the ``type`` and ``format`` keywords are carried today, with their standard + * JSON-schema semantics: ``type`` is asserted by the runtime, ``format`` is an + * annotation a runtime may additionally check. Unknown keywords from newer providers + * are ignored rather than rejected (as JSON-schema consumers do), so a core on this + * version keeps serving specs written by a newer provider. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "ArgValueSchema". + */ +export interface ArgValueSchema { + type?: Type; + format?: Format; +} /** * Schema for AssetAliasModel used in AssetEventDagRunReference. * @@ -736,7 +730,7 @@ export interface DagRunAssetReference { */ export interface AssetEventsResult { asset_events: AssetEvents; - type?: Type; + type?: Type1; } /** * Profile of an asset-like object. @@ -756,7 +750,7 @@ export interface AssetEventsResult { export interface AssetProfile { name?: Name2; uri?: Uri1; - type: Type1; + type: Type2; } /** * Schema for AssetModel used in AssetEventDagRunReference. @@ -783,7 +777,7 @@ export interface AssetResult { uri: Uri3; group: Group1; extra?: Extra3; - type?: Type2; + type?: Type3; } /** * Response to GetAssetStateStore; wraps the generated API response for supervisor to worker comms. @@ -793,7 +787,7 @@ export interface AssetResult { */ export interface AssetStateStoreResult { value: JsonValue; - type?: Type3; + type?: Type4; } /** * Response to GetAssetsByAlias; list of concrete assets resolved from an alias. @@ -803,7 +797,7 @@ export interface AssetStateStoreResult { */ export interface AssetsByAliasResult { assets: Assets; - type?: Type4; + type?: Type5; } /** * Park a task instance awaiting human input (Human-in-the-loop), without a trigger. @@ -817,7 +811,7 @@ export interface AwaitInputTask { next_method: NextMethod; next_kwargs?: NextKwargs; rendered_map_index?: RenderedMapIndex; - type?: Type5; + type?: Type6; } /** * Schema for telling task which bundle to run with. @@ -836,7 +830,7 @@ export interface BundleInfo { */ export interface ClearAssetStateStoreByName { name: Name6; - type?: Type6; + type?: Type7; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -844,7 +838,7 @@ export interface ClearAssetStateStoreByName { */ export interface ClearAssetStateStoreByUri { uri: Uri4; - type?: Type7; + type?: Type8; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -852,7 +846,7 @@ export interface ClearAssetStateStoreByUri { */ export interface ClearTaskStateStore { ti_id: TiId; - type?: Type8; + type?: Type9; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -867,7 +861,7 @@ export interface ConnectionResult { password?: Password; port?: Port; extra?: Extra4; - type?: Type9; + type?: Type10; } /** * Add the input request part of a Human-in-the-loop response. @@ -884,7 +878,7 @@ export interface CreateHITLDetailPayload { multiple?: Multiple; params?: Params; assigned_users?: AssignedUsers; - type?: Type10; + type?: Type11; } /** * Schema for a Human-in-the-loop users. @@ -904,7 +898,7 @@ export interface HITLUser { */ export interface DRCount { count: Count; - type?: Type11; + type?: Type12; } /** * A Class with information about the success/failure DAG callback to be executed. @@ -922,7 +916,7 @@ export interface DagCallbackRequest { run_id: RunId1; context_from_server?: DagRunContext | null; is_failure_callback?: IsFailureCallback; - type?: Type12; + type?: Type13; } /** * Class to pass context info from the server to build a Execution context object. @@ -1012,7 +1006,7 @@ export interface DagFileParseRequest { bundle_path: BundlePath; bundle_name: BundleName1; callback_requests?: CallbackRequests; - type?: Type15; + type?: Type16; } /** * Task callback status information. @@ -1032,7 +1026,7 @@ export interface TaskCallbackRequest { ti: TaskInstance; task_callback_type?: TaskInstanceState | null; context_from_server?: TIRunContext | null; - type?: Type13; + type?: Type14; } /** * Response schema for TaskInstance run context. @@ -1088,7 +1082,7 @@ export interface ConnectionResponse { export interface XComArgBinding { kind: Kind; name: Name8; - data_type?: ArgBindingDataType1; + value_schema?: ArgValueSchema | null; task_id: TaskId1; } /** @@ -1100,7 +1094,7 @@ export interface XComArgBinding { export interface LiteralArgBinding { kind: Kind1; name: Name9; - data_type?: ArgBindingDataType2; + value_schema?: ArgValueSchema | null; value?: unknown; from_default?: FromDefault; } @@ -1119,7 +1113,7 @@ export interface EmailRequest { ti: TaskInstance; email_type?: EmailType; context_from_server: TIRunContext; - type?: Type14; + type?: Type15; } /** * Result of DAG File Parsing. @@ -1135,7 +1129,7 @@ export interface DagFileParsingResult { serialized_dags: SerializedDags; warnings?: Warnings; import_errors?: ImportErrors; - type?: Type16; + type?: Type17; } /** * Lazily build information from the serialized DAG structure. @@ -1166,7 +1160,7 @@ export interface DagResult { owners?: Owners; tags: Tags; next_dagrun?: NextDagrun; - type?: Type17; + type?: Type18; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1191,7 +1185,7 @@ export interface DagRunResult { partition_date?: PartitionDate1; note?: Note1; team_name?: TeamName1; - type?: Type18; + type?: Type19; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1199,7 +1193,7 @@ export interface DagRunResult { */ export interface DagRunStateResult { state: DagRunState; - type?: Type19; + type?: Type20; } /** * Update a task instance state to deferred. @@ -1216,7 +1210,7 @@ export interface DeferTask { next_method: NextMethod2; next_kwargs?: NextKwargs2; rendered_map_index?: RenderedMapIndex1; - type?: Type20; + type?: Type21; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1225,7 +1219,7 @@ export interface DeferTask { export interface DeleteAssetStateStoreByName { name: Name10; key: Key1; - type?: Type21; + type?: Type22; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1234,7 +1228,7 @@ export interface DeleteAssetStateStoreByName { export interface DeleteAssetStateStoreByUri { uri: Uri5; key: Key2; - type?: Type22; + type?: Type23; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1243,7 +1237,7 @@ export interface DeleteAssetStateStoreByUri { export interface DeleteTaskStateStore { ti_id: TiId2; key: Key3; - type?: Type23; + type?: Type24; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1251,7 +1245,7 @@ export interface DeleteTaskStateStore { */ export interface DeleteVariable { key: Key4; - type?: Type24; + type?: Type25; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1263,7 +1257,7 @@ export interface DeleteXCom { run_id: RunId5; task_id: TaskId2; map_index?: MapIndex1; - type?: Type25; + type?: Type26; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1272,7 +1266,7 @@ export interface DeleteXCom { export interface ErrorResponse { error?: ErrorType; detail?: Detail; - type?: Type26; + type?: Type27; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1280,7 +1274,7 @@ export interface ErrorResponse { */ export interface GetAssetByName { name: Name11; - type?: Type27; + type?: Type28; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1288,7 +1282,7 @@ export interface GetAssetByName { */ export interface GetAssetByUri { uri: Uri6; - type?: Type28; + type?: Type29; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1304,7 +1298,7 @@ export interface GetAssetEventByAsset { partition_key?: PartitionKey5; partition_key_regexp_pattern?: PartitionKeyRegexpPattern; extra?: Extra7; - type?: Type29; + type?: Type30; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1319,7 +1313,7 @@ export interface GetAssetEventByAssetAlias { partition_key?: PartitionKey6; partition_key_regexp_pattern?: PartitionKeyRegexpPattern1; extra?: Extra8; - type?: Type30; + type?: Type31; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1328,7 +1322,7 @@ export interface GetAssetEventByAssetAlias { export interface GetAssetStateStoreByName { name: Name13; key: Key6; - type?: Type31; + type?: Type32; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1337,7 +1331,7 @@ export interface GetAssetStateStoreByName { export interface GetAssetStateStoreByUri { uri: Uri8; key: Key7; - type?: Type32; + type?: Type33; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1345,7 +1339,7 @@ export interface GetAssetStateStoreByUri { */ export interface GetAssetsByAlias { alias_name: AliasName1; - type?: Type33; + type?: Type34; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1353,7 +1347,7 @@ export interface GetAssetsByAlias { */ export interface GetConnection { conn_id: ConnId2; - type?: Type34; + type?: Type35; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1364,7 +1358,7 @@ export interface GetDRCount { logical_dates?: LogicalDates; run_ids?: RunIds; states?: States; - type?: Type35; + type?: Type36; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1372,7 +1366,7 @@ export interface GetDRCount { */ export interface GetDag { dag_id: DagId8; - type?: Type36; + type?: Type37; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1381,7 +1375,7 @@ export interface GetDag { export interface GetDagRun { dag_id: DagId9; run_id: RunId6; - type?: Type37; + type?: Type38; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1390,7 +1384,7 @@ export interface GetDagRun { export interface GetDagRunState { dag_id: DagId10; run_id: RunId7; - type?: Type38; + type?: Type39; } /** * Get the response content part of a Human-in-the-loop response. @@ -1400,7 +1394,7 @@ export interface GetDagRunState { */ export interface GetHITLDetailResponse { ti_id: TiId3; - type?: Type39; + type?: Type40; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1408,7 +1402,7 @@ export interface GetHITLDetailResponse { */ export interface GetPrevSuccessfulDagRun { ti_id: TiId4; - type?: Type40; + type?: Type41; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1418,7 +1412,7 @@ export interface GetPreviousDagRun { dag_id: DagId11; logical_date: LogicalDate3; state?: State3; - type?: Type41; + type?: Type42; } /** * Request to get previous task instance. @@ -1432,7 +1426,7 @@ export interface GetPreviousTI { logical_date?: LogicalDate4; map_index?: MapIndex2; state?: TaskInstanceState | null; - type?: Type42; + type?: Type43; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1446,7 +1440,7 @@ export interface GetTICount { logical_dates?: LogicalDates1; run_ids?: RunIds1; states?: States1; - type?: Type43; + type?: Type44; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1455,7 +1449,7 @@ export interface GetTICount { export interface GetTaskBreadcrumbs { dag_id: DagId14; run_id: RunId8; - type?: Type44; + type?: Type45; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1464,7 +1458,7 @@ export interface GetTaskBreadcrumbs { export interface GetTaskRescheduleStartDate { ti_id: TiId5; try_number?: TryNumber1; - type?: Type45; + type?: Type46; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1473,7 +1467,7 @@ export interface GetTaskRescheduleStartDate { export interface GetTaskStateStore { ti_id: TiId6; key: Key8; - type?: Type46; + type?: Type47; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1486,7 +1480,7 @@ export interface GetTaskStates { task_group_id?: TaskGroupId1; logical_dates?: LogicalDates2; run_ids?: RunIds2; - type?: Type47; + type?: Type48; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1494,7 +1488,7 @@ export interface GetTaskStates { */ export interface GetVariable { key: Key9; - type?: Type48; + type?: Type49; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1504,7 +1498,7 @@ export interface GetVariableKeys { prefix?: Prefix; limit?: Limit2; offset?: Offset; - type?: Type49; + type?: Type50; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1517,7 +1511,7 @@ export interface GetXCom { task_id: TaskId4; map_index?: MapIndex5; include_prior_dates?: IncludePriorDates; - type?: Type50; + type?: Type51; } /** * Get the number of (mapped) XCom values available. @@ -1530,7 +1524,7 @@ export interface GetXComCount { dag_id: DagId17; run_id: RunId10; task_id: TaskId5; - type?: Type51; + type?: Type52; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1542,7 +1536,7 @@ export interface GetXComSequenceItem { run_id: RunId11; task_id: TaskId6; offset: Offset1; - type?: Type52; + type?: Type53; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1557,7 +1551,7 @@ export interface GetXComSequenceSlice { stop: Stop; step: Step; include_prior_dates?: IncludePriorDates1; - type?: Type53; + type?: Type54; } /** * Response to CreateHITLDetailPayload request. @@ -1574,7 +1568,7 @@ export interface HITLDetailRequestResult { multiple?: Multiple1; params?: Params1; assigned_users?: AssignedUsers1; - type?: Type54; + type?: Type55; } /** * Response of InactiveAssets requests. @@ -1584,7 +1578,7 @@ export interface HITLDetailRequestResult { */ export interface InactiveAssetsResult { inactive_assets?: InactiveAssets; - type?: Type55; + type?: Type56; } /** * Add a new value to be redacted in task logs. @@ -1595,7 +1589,7 @@ export interface InactiveAssetsResult { export interface MaskSecret { value: JsonValue; name?: Name14; - type?: Type56; + type?: Type57; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1603,7 +1597,7 @@ export interface MaskSecret { */ export interface OKResponse { ok: Ok; - type?: Type57; + type?: Type58; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1614,7 +1608,7 @@ export interface PrevSuccessfulDagRunResult { data_interval_end?: DataIntervalEnd3; start_date?: StartDate4; end_date?: EndDate3; - type?: Type58; + type?: Type59; } /** * Response containing previous Dag run information. @@ -1624,7 +1618,7 @@ export interface PrevSuccessfulDagRunResult { */ export interface PreviousDagRunResult { dag_run?: DagRun | null; - type?: Type59; + type?: Type60; } /** * Schema for response with previous TaskInstance information. @@ -1652,7 +1646,7 @@ export interface PreviousTIResponse { */ export interface PreviousTIResult { task_instance?: PreviousTIResponse | null; - type?: Type60; + type?: Type61; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1662,7 +1656,7 @@ export interface PutVariable { key: Key14; value: Value1; description: Description; - type?: Type61; + type?: Type62; } /** * Update a task instance state to reschedule/up_for_reschedule. @@ -1674,14 +1668,14 @@ export interface RescheduleTask { state?: State5; reschedule_date: RescheduleDate; end_date: EndDate5; - type?: Type62; + type?: Type63; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema * via the `definition` "ResendLoggingFD". */ export interface ResendLoggingFD { - type?: Type63; + type?: Type64; } /** * Update a task instance state to up_for_retry. @@ -1695,14 +1689,14 @@ export interface RetryTask { rendered_map_index?: RenderedMapIndex2; retry_delay_seconds?: RetryDelaySeconds; retry_reason?: RetryReason; - type?: Type64; + type?: Type65; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema * via the `definition` "SentFDs". */ export interface SentFDs { - type?: Type65; + type?: Type66; fds: Fds; } /** @@ -1713,7 +1707,7 @@ export interface SetAssetStateStoreByName { name: Name15; key: Key15; value: JsonValue; - type?: Type66; + type?: Type67; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1723,7 +1717,7 @@ export interface SetAssetStateStoreByUri { uri: Uri9; key: Key16; value: JsonValue; - type?: Type67; + type?: Type68; } /** * Payload for setting RTIF for a task instance. @@ -1733,7 +1727,7 @@ export interface SetAssetStateStoreByUri { */ export interface SetRenderedFields { rendered_fields: RenderedFields; - type?: Type68; + type?: Type69; } export interface RenderedFields { [k: string]: JsonValue; @@ -1746,7 +1740,7 @@ export interface RenderedFields { */ export interface SetRenderedMapIndex { rendered_map_index: RenderedMapIndex3; - type?: Type69; + type?: Type70; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1757,7 +1751,7 @@ export interface SetTaskStateStore { key: Key17; value: JsonValue; expires_at: ExpiresAt; - type?: Type70; + type?: Type71; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1772,7 +1766,7 @@ export interface SetXCom { map_index?: MapIndex7; dag_result?: DagResult1; mapped_length?: MappedLength; - type?: Type71; + type?: Type72; } /** * Update state of downstream tasks within a task instance to 'skipped', while updating current task to success state. @@ -1782,7 +1776,7 @@ export interface SetXCom { */ export interface SkipDownstreamTasks { tasks: Tasks; - type?: Type72; + type?: Type73; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1795,7 +1789,7 @@ export interface StartupDetails { start_date: StartDate6; ti_context: TIRunContext; sentry_integration: SentryIntegration; - type?: Type73; + type?: Type74; } /** * Update a task's state to success. Includes task_outlets and outlet_events for registering asset events. @@ -1809,7 +1803,7 @@ export interface SucceedTask { task_outlets?: TaskOutlets; outlet_events?: OutletEvents; rendered_map_index?: RenderedMapIndex4; - type?: Type74; + type?: Type75; } /** * Response containing count of Task Instances matching certain filters. @@ -1819,7 +1813,7 @@ export interface SucceedTask { */ export interface TICount { count: Count1; - type?: Type75; + type?: Type76; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1827,7 +1821,7 @@ export interface TICount { */ export interface TaskBreadcrumbsResult { breadcrumbs: Breadcrumbs; - type?: Type76; + type?: Type77; } /** * Response containing the first reschedule date for a task instance. @@ -1837,7 +1831,7 @@ export interface TaskBreadcrumbsResult { */ export interface TaskRescheduleStartDate { start_date: StartDate7; - type?: Type77; + type?: Type78; } /** * Update a task's state. @@ -1852,7 +1846,7 @@ export interface TaskRescheduleStartDate { export interface TaskState { state: State8; end_date?: EndDate8; - type?: Type78; + type?: Type79; rendered_map_index?: RenderedMapIndex5; } /** @@ -1863,7 +1857,7 @@ export interface TaskState { */ export interface TaskStateStoreResult { value: JsonValue; - type?: Type79; + type?: Type80; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1871,7 +1865,7 @@ export interface TaskStateStoreResult { */ export interface TaskStatesResult { task_states: TaskStates; - type?: Type80; + type?: Type81; } export interface TaskStates { [k: string]: unknown; @@ -1889,7 +1883,7 @@ export interface TriggerDagRun { note?: Note2; dag_id: DagId22; run_id: DagRunId; - type?: Type81; + type?: Type82; } /** * Update the response content part of an existing Human-in-the-loop response. @@ -1901,7 +1895,7 @@ export interface UpdateHITLDetail { ti_id: TiId9; chosen_options: ChosenOptions; params_input?: ParamsInput; - type?: Type82; + type?: Type83; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1909,7 +1903,7 @@ export interface UpdateHITLDetail { */ export interface ValidateInletsAndOutlets { ti_id: TiId10; - type?: Type83; + type?: Type84; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1918,7 +1912,7 @@ export interface ValidateInletsAndOutlets { export interface VariableKeysResult { keys: Keys; total_entries: TotalEntries; - type?: Type84; + type?: Type85; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1927,7 +1921,7 @@ export interface VariableKeysResult { export interface VariableResult { key: Key19; value?: Value2; - type?: Type85; + type?: Type86; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1935,7 +1929,7 @@ export interface VariableResult { */ export interface XComCountResponse { len: Len; - type?: Type86; + type?: Type87; } /** * Response to ReadXCom request. @@ -1946,7 +1940,7 @@ export interface XComCountResponse { export interface XComResult { key: Key20; value: JsonValue; - type?: Type87; + type?: Type88; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1954,7 +1948,7 @@ export interface XComResult { */ export interface XComSequenceIndexResult { root: JsonValue; - type?: Type88; + type?: Type89; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1962,7 +1956,7 @@ export interface XComSequenceIndexResult { */ export interface XComSequenceSliceResult { root: Root; - type?: Type89; + type?: Type90; } /** Cadwyn schema version this SDK was generated against. @@ -1970,4 +1964,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-07-30" as const; +export const SUPERVISOR_API_VERSION = "2026-10-30" as const; From 258ac8d8744548e8b03aba072a505d5f99e2f9ac Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 24 Jul 2026 04:44:30 +0000 Subject: [PATCH 22/40] Generate stub arg value schemas with pydantic instead of a hand-rolled mapper The hand-written annotation-to-fragment mapper duplicated what pydantic already exposes publicly: TypeAdapter(annotation).json_schema(), with a GenerateJsonSchema subclass layering the int64/double numeric formats a foreign runtime needs. Delegating to pydantic removes the bespoke union logic and buys richer, standard fragments for free -- anyOf for unions (per-member formats survive mixed unions now), items/additionalProperties for parameterized containers, and enum for Literal annotations -- while anything pydantic cannot schema (arbitrary classes anywhere in the annotation) still degrades to a decode-only binding. value_schema becomes a free-form JSON object on the wire instead of a typed model: a typed model silently strips every keyword it does not know when the spec is re-serialized along the server-to-supervisor delivery path, which would corrupt exactly the open-vocabulary fragments this contract promises to carry verbatim. The provider declares its now-direct pydantic dependency (edge3/http precedent). Trade-offs accepted: fragment shapes follow the pydantic version active at parse time, and pendulum.DateTime annotations degrade to decode-only since pydantic has no schema for arbitrary datetime subclasses. No new execution-API version: the field's shape changes inside the still-unreleased 2026-10-30 version this PR introduces. --- .../datamodels/task_arg_binding.py | 41 +- .../versions/head/test_task_instances.py | 25 +- .../v2026_10_30/test_task_instances.py | 7 +- .../serialization/test_dag_serialization.py | 12 +- providers/standard/README.rst | 1 + providers/standard/docs/index.rst | 1 + providers/standard/pyproject.toml | 2 + .../providers/standard/decorators/stub.py | 105 ++--- .../unit/standard/decorators/test_stub.py | 92 ++-- .../airflow/sdk/api/datamodels/_generated.py | 25 +- .../sdk/execution_time/schema/schema.json | 129 ++---- .../execution_time/schema/test_migrator.py | 6 +- ts-sdk/src/generated/supervisor.ts | 397 +++++++++--------- uv.lock | 4 +- 14 files changed, 383 insertions(+), 464 deletions(-) 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 index e26559b1d090c..157612f3edb5f 100644 --- 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 @@ -32,36 +32,21 @@ from airflow.api_fastapi.core_api.base import BaseModel -# A named alias (like TaskArgBinding below) so both union branches of ArgValueSchema.type -# reference one shared schema definition instead of two inlined enum copies; the explicit -# title lets the supervisor-schema dump merge this def with the task-sdk-generated twin. -JsonSchemaType = TypeAliasType( - "JsonSchemaType", - Annotated[ - Literal["string", "integer", "number", "boolean", "object", "array", "null"], - Field(title="JsonSchemaType"), - ], +# A named alias (like TaskArgBinding below) so every schema keeps one shared, titled +# ArgValueSchema definition; a free-form JSON object rather than a typed model because +# the fragment is whatever pydantic generates from the stub annotation at parse time +# (``anyOf``, ``items``, ``enum``, ``$defs``, ...) and a typed model would silently strip +# any keyword it does not know when the spec is re-serialized along the delivery path. +ArgValueSchema = TypeAliasType( + "ArgValueSchema", Annotated[dict[str, JsonValue], Field(title="ArgValueSchema")] ) -"""JSON-schema primitive type names a stub-task argument annotation can map to.""" - - -class ArgValueSchema(BaseModel): - """ - JSON-schema fragment constraining the value a stub-task argument binds to. - - Only the ``type`` and ``format`` keywords are carried today, with their standard - JSON-schema semantics: ``type`` is asserted by the runtime, ``format`` is an - annotation a runtime may additionally check. Unknown keywords from newer providers - are ignored rather than rejected (as JSON-schema consumers do), so a core on this - version keeps serving specs written by a newer provider. - """ - - type: JsonSchemaType | list[JsonSchemaType] | None = None - """A single type name, a union of type names, or ``None`` when unconstrained.""" +""" +JSON-schema fragment constraining the value a stub-task argument binds to. - format: str | None = None - """Wire representation the type name alone cannot convey (``int64``, ``date-time``, - ``duration``, ...); open vocabulary, per JSON schema.""" +Generated by pydantic from the stub function's parameter annotation and carried verbatim; +consumers validate the keywords they understand and ignore the rest, per JSON-schema +semantics, so newer producers never break an older core or runtime. +""" class XComArgBinding(BaseModel): 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 288301c8a63f0..a281291bf84b7 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 @@ -401,7 +401,12 @@ def transform(country: str, extracted: dict, limit: int = 10): ... assert response.status_code == 200 assert response.json()["arg_bindings"] == [ {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, - {"name": "extracted", "kind": "xcom", "value_schema": {"type": "object"}, "task_id": "extract"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "extract", + }, { "name": "limit", "kind": "literal", @@ -458,22 +463,16 @@ def test_arg_bindings_adapter_rejects_unknown_kind(self): [{"name": "country", "kind": "template", "value": "x"}] ) - def test_arg_bindings_adapter_tolerates_unknown_value_schema_keyword(self): - """A JSON-schema keyword this core version does not know (e.g. from a newer provider) - must not fail the spec -- JSON-schema consumers ignore unknown keywords by design.""" + def test_arg_bindings_adapter_carries_value_schema_fragments_verbatim(self): + """The fragment is free-form JSON schema: every keyword the provider generated must + survive validation untouched -- a typed model would silently strip what it doesn't know.""" from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import get_arg_bindings_adapter + fragment = {"anyOf": [{"type": "array", "items": {"type": "string"}}, {"type": "null"}]} (binding,) = get_arg_bindings_adapter().validate_python( - [ - { - "name": "tags", - "kind": "literal", - "value_schema": {"type": "array", "items": {"type": "string"}}, - "value": ["a"], - } - ] + [{"name": "tags", "kind": "literal", "value_schema": fragment, "value": ["a"]}] ) - assert binding.value_schema.type == "array" + assert binding.value_schema == fragment def test_dynamic_task_mapping_with_parse_time_value(self, client, dag_maker): """Test that dynamic task mapping works correctly with parse-time values.""" diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py index 0e0972afed8da..a4b98bd10206e 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py @@ -84,7 +84,12 @@ def test_head_version_includes_arg_bindings(self, client, stub_ti): assert response.status_code == 200 assert response.json()["arg_bindings"] == [ {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, - {"name": "extracted", "kind": "xcom", "value_schema": {"type": "object"}, "task_id": "extract"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "extract", + }, { "name": "limit", "kind": "literal", diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index 76e937bd8fa15..e88445d75e5f8 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -3437,7 +3437,10 @@ def transform(country: str, extracted: dict): ... Encoding.VAR: { "name": "extracted", "kind": "xcom", - "value_schema": {Encoding.TYPE: DAT.DICT, Encoding.VAR: {"type": "object"}}, + "value_schema": { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: {"type": "object", "additionalProperties": True}, + }, "task_id": "extract", }, }, @@ -3446,7 +3449,12 @@ def transform(country: str, extracted: dict): ... round_tripped = DagSerialization.from_dict(ser_dag) assert round_tripped.task_dict["transform"]._arg_bindings == [ {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, - {"name": "extracted", "kind": "xcom", "value_schema": {"type": "object"}, "task_id": "extract"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "extract", + }, ] assert not hasattr(round_tripped.task_dict["extract"], "_arg_bindings") or ( round_tripped.task_dict["extract"]._arg_bindings is None diff --git a/providers/standard/README.rst b/providers/standard/README.rst index 839b70f10786e..ce275c510f528 100644 --- a/providers/standard/README.rst +++ b/providers/standard/README.rst @@ -55,6 +55,7 @@ PIP package Version required ========================================== ================== ``apache-airflow`` ``>=2.11.0`` ``apache-airflow-providers-common-compat`` ``>=1.14.1`` +``pydantic`` ``>=2.11.0`` ========================================== ================== Optional cross provider package dependencies diff --git a/providers/standard/docs/index.rst b/providers/standard/docs/index.rst index 1dc9dcdf2d7dd..0df86ce95290c 100644 --- a/providers/standard/docs/index.rst +++ b/providers/standard/docs/index.rst @@ -91,6 +91,7 @@ PIP package Version required ========================================== ================== ``apache-airflow`` ``>=2.11.0`` ``apache-airflow-providers-common-compat`` ``>=1.14.1`` +``pydantic`` ``>=2.11.0`` ========================================== ================== Optional cross provider package dependencies diff --git a/providers/standard/pyproject.toml b/providers/standard/pyproject.toml index 373d16bba6c4e..fdc8379b1c9ba 100644 --- a/providers/standard/pyproject.toml +++ b/providers/standard/pyproject.toml @@ -61,6 +61,8 @@ requires-python = ">=3.10" dependencies = [ "apache-airflow>=2.11.0", "apache-airflow-providers-common-compat>=1.14.1", # use next version + # The stub decorator generates arg-binding value schemas from parameter annotations + "pydantic>=2.11.0", ] # The optional dependencies should be modified in place in the generated file diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index 0ed17b355c868..19f8e83563d4d 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -18,13 +18,14 @@ from __future__ import annotations import ast -import datetime import inspect import json -import types import typing -from collections.abc import Callable, Collection, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Union +from collections.abc import Callable, Collection, Mapping +from typing import TYPE_CHECKING, Any + +from pydantic import PydanticSchemaGenerationError, TypeAdapter +from pydantic.json_schema import GenerateJsonSchema from airflow.providers.common.compat.sdk import ( KNOWN_CONTEXT_KEYS, @@ -39,52 +40,34 @@ from airflow.providers.common.compat.sdk import Context -def _json_schema_fragment(annotation: Any) -> dict[str, Any] | None: - """Map one non-union annotation to a JSON-schema fragment; ``None`` = unclassifiable.""" - if annotation is type(None): - return {"type": "null"} - origin = typing.get_origin(annotation) - if origin is not None: - annotation = origin - if not isinstance(annotation, type): - return None - # bool subclasses int, str/bytes are Sequences, and datetime subclasses date -- order matters. - if issubclass(annotation, bool): - return {"type": "boolean"} - if issubclass(annotation, int): - return {"type": "integer", "format": "int64"} - if issubclass(annotation, float): - return {"type": "number", "format": "double"} - if issubclass(annotation, str): - return {"type": "string"} - if issubclass(annotation, bytes): - return None - if issubclass(annotation, datetime.datetime): - return {"type": "string", "format": "date-time"} - if issubclass(annotation, datetime.date): - return {"type": "string", "format": "date"} - if issubclass(annotation, datetime.time): - return {"type": "string", "format": "time"} - if issubclass(annotation, datetime.timedelta): - return {"type": "string", "format": "duration"} - if issubclass(annotation, (dict, Mapping)): - return {"type": "object"} - if issubclass(annotation, (list, tuple, set, frozenset, Sequence)): - return {"type": "array"} - return None +class _ValueSchemaGenerator(GenerateJsonSchema): + """ + Pydantic's stock JSON-schema generation plus OpenAPI's fixed-width numeric formats. + + A foreign runtime decodes numbers into machine types, which the bare + ``integer``/``number`` type names cannot convey; ``format`` is an annotation per + JSON schema, so runtimes that don't know these names simply skip them. + """ + + def int_schema(self, schema): + return {**super().int_schema(schema), "format": "int64"} + + def float_schema(self, schema): + return {**super().float_schema(schema), "format": "double"} def _infer_value_schema(annotation: Any) -> dict[str, Any] | None: """ - Map a stub function parameter annotation to a JSON-schema fragment. - - Fragments carry the standard ``type`` keyword -- a single name, or a list for union - annotations (set semantics; member order follows the annotation) -- plus a ``format`` - annotation where the Python type implies a wire representation the type name alone - cannot (``int`` -> ``int64``, ``datetime`` -> ``date-time``, ``timedelta`` -> - ``duration``, ...). Returns ``None`` when the annotation gives no constraint; the - binding then omits ``value_schema`` and the foreign runtime falls back to a - decode-only check. + Build the JSON-schema fragment for one stub parameter annotation, via pydantic. + + Whatever ``pydantic.TypeAdapter(annotation).json_schema()`` produces is shipped + verbatim (``anyOf`` for unions, ``items``/``additionalProperties`` for parameterized + containers, ``enum`` for Literals, ...), so the fragment's exact shape follows the + pydantic version active at parse time and runtimes must treat it as open-vocabulary + JSON schema. Returns ``None`` when the annotation constrains nothing (missing, + ``Any``, bare ``None``) or pydantic cannot generate a schema for it (arbitrary + classes, including anywhere inside a union); the binding then omits ``value_schema`` + and the foreign runtime falls back to a decode-only check. """ if annotation is inspect.Parameter.empty or annotation is None or annotation is Any: return None @@ -92,31 +75,11 @@ def _infer_value_schema(annotation: Any) -> dict[str, Any] | None: # get_type_hints normalizes a bare ``None`` annotation to NoneType; a parameter # that can only ever be None constrains nothing worth shipping. return None - origin = typing.get_origin(annotation) - if origin is Union or origin is types.UnionType: - fragments: list[dict[str, Any]] = [] - for member in typing.get_args(annotation): - fragment = _json_schema_fragment(member) - if fragment is None: - # One unclassifiable member makes the whole union unconstrained: a - # partial schema would wrongly reject that member's values. - return None - if fragment not in fragments: - fragments.append(fragment) - data_fragments = [f for f in fragments if f["type"] != "null"] - if len(data_fragments) == 1: - # A single data type (+ optional null) keeps its format: format only - # constrains values of its own type, so null passes it untouched. - schema = dict(data_fragments[0]) - if len(fragments) > len(data_fragments): - schema["type"] = [schema["type"], "null"] - return schema - # Mixed-type union: formats are per-member and inexpressible in one flat - # fragment, so carry the type names only. - type_names = [f["type"] for f in fragments] - deduped = list(dict.fromkeys(type_names)) - return {"type": deduped if len(deduped) > 1 else deduped[0]} - return _json_schema_fragment(annotation) + try: + schema = TypeAdapter(annotation).json_schema(schema_generator=_ValueSchemaGenerator) + except PydanticSchemaGenerationError: + return None + return schema or None def _build_arg_bindings( diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index a8c3ae8634a48..cfd85d1f31f48 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -108,7 +108,7 @@ def test_literal_and_xcom_spec(self): { "name": "extracted", "kind": "xcom", - "value_schema": {"type": "object"}, + "value_schema": {"type": "object", "additionalProperties": True}, "task_id": "fn_extract", }, { @@ -131,7 +131,7 @@ def test_kwargs_normalize_to_declaration_order(self): { "name": "extracted", "kind": "xcom", - "value_schema": {"type": "object"}, + "value_schema": {"type": "object", "additionalProperties": True}, "task_id": "fn_extract", }, { @@ -237,7 +237,7 @@ def test_arg_bindings_survive_dag_serialization_round_trip(self): { "name": "extracted", "kind": "xcom", - "value_schema": {"type": "object"}, + "value_schema": {"type": "object", "additionalProperties": True}, "task_id": "fn_extract", }, { @@ -282,51 +282,91 @@ def group(n): pytest.param(bool, {"type": "boolean"}, id="bool"), pytest.param(int, {"type": "integer", "format": "int64"}, id="int"), pytest.param(float, {"type": "number", "format": "double"}, id="float"), - pytest.param(dict, {"type": "object"}, id="dict"), - pytest.param(dict[str, int], {"type": "object"}, id="dict-parameterized"), - pytest.param(typing.Mapping[str, int], {"type": "object"}, id="mapping"), - pytest.param(list, {"type": "array"}, id="list"), - pytest.param(list[int], {"type": "array"}, id="list-parameterized"), - pytest.param(tuple, {"type": "array"}, id="tuple"), - pytest.param(set, {"type": "array"}, id="set"), - pytest.param(typing.Sequence[int], {"type": "array"}, id="sequence"), + pytest.param(dict, {"type": "object", "additionalProperties": True}, id="dict"), + pytest.param( + dict[str, int], + {"type": "object", "additionalProperties": {"type": "integer", "format": "int64"}}, + id="dict-parameterized", + ), + pytest.param( + typing.Mapping[str, int], + {"type": "object", "additionalProperties": {"type": "integer", "format": "int64"}}, + id="mapping", + ), + pytest.param(list, {"type": "array", "items": {}}, id="list"), + pytest.param( + list[int], + {"type": "array", "items": {"type": "integer", "format": "int64"}}, + id="list-parameterized", + ), + pytest.param(tuple, {"type": "array", "items": {}}, id="tuple"), + pytest.param(set, {"type": "array", "items": {}, "uniqueItems": True}, id="set"), + pytest.param( + typing.Sequence[int], + {"type": "array", "items": {"type": "integer", "format": "int64"}}, + id="sequence", + ), pytest.param(datetime.datetime, {"type": "string", "format": "date-time"}, id="datetime"), - pytest.param(pendulum.DateTime, {"type": "string", "format": "date-time"}, id="pendulum-datetime"), pytest.param(datetime.date, {"type": "string", "format": "date"}, id="date"), pytest.param(datetime.time, {"type": "string", "format": "time"}, id="time"), pytest.param(datetime.timedelta, {"type": "string", "format": "duration"}, id="timedelta"), + pytest.param(bytes, {"type": "string", "format": "binary"}, id="bytes"), + pytest.param( + typing.Literal["a", "b"], + {"type": "string", "enum": ["a", "b"]}, + id="literal", + ), pytest.param(Any, None, id="any"), pytest.param(None, None, id="none"), pytest.param(type(None), None, id="nonetype"), - pytest.param(bytes, None, id="bytes"), + pytest.param( + pendulum.DateTime, + None, + id="pendulum-datetime", + # pydantic has no schema for arbitrary datetime subclasses; decode-only fallback. + ), pytest.param( typing.Optional[str], # noqa: UP045 -- legacy form on purpose - {"type": ["string", "null"]}, + {"anyOf": [{"type": "string"}, {"type": "null"}]}, id="optional-str", ), pytest.param( typing.Union[int, str], # noqa: UP007 -- legacy form on purpose - {"type": ["integer", "string"]}, + {"anyOf": [{"type": "integer", "format": "int64"}, {"type": "string"}]}, id="union", ), - pytest.param(str | None, {"type": ["string", "null"]}, id="pep604-optional"), - pytest.param(int | None, {"type": ["integer", "null"], "format": "int64"}, id="optional-int"), + pytest.param(str | None, {"anyOf": [{"type": "string"}, {"type": "null"}]}, id="pep604-optional"), + pytest.param( + int | None, + {"anyOf": [{"type": "integer", "format": "int64"}, {"type": "null"}]}, + id="optional-int", + ), pytest.param( datetime.datetime | None, - {"type": ["string", "null"], "format": "date-time"}, + {"anyOf": [{"type": "string", "format": "date-time"}, {"type": "null"}]}, id="optional-datetime", ), - pytest.param(int | str, {"type": ["integer", "string"]}, id="pep604-union"), - pytest.param(dict | bool, {"type": ["object", "boolean"]}, id="union-dict-bool"), - pytest.param(bool | int, {"type": ["boolean", "integer"]}, id="union-bool-int-order"), - pytest.param(str | int | None, {"type": ["string", "integer", "null"]}, id="union-with-null"), - pytest.param(list | tuple, {"type": "array"}, id="union-collapses-to-one-type"), + pytest.param( + dict | bool, + {"anyOf": [{"type": "object", "additionalProperties": True}, {"type": "boolean"}]}, + id="union-dict-bool", + ), + pytest.param( + str | int | None, + {"anyOf": [{"type": "string"}, {"type": "integer", "format": "int64"}, {"type": "null"}]}, + id="union-with-null", + ), + pytest.param(list | tuple, {"type": "array", "items": {}}, id="union-dedupes-equal-members"), pytest.param( datetime.datetime | str, - {"type": "string"}, - id="mixed-format-union-drops-format", + {"anyOf": [{"type": "string", "format": "date-time"}, {"type": "string"}]}, + id="mixed-format-union-keeps-both", + ), + pytest.param( + str | contextlib.AbstractContextManager, + None, + id="union-unclassifiable-member", ), - pytest.param(str | bytes, None, id="union-unclassifiable-member"), pytest.param(contextlib.AbstractContextManager, None, id="custom-class"), ], ) diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 1de6d574278e2..94a56701d5294 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -218,16 +218,6 @@ class IntermediateTIState(str, Enum): AWAITING_INPUT = "awaiting_input" -class JsonSchemaType(str, Enum): - STRING = "string" - INTEGER = "integer" - NUMBER = "number" - BOOLEAN = "boolean" - OBJECT = "object" - ARRAY = "array" - NULL = "null" - - class PrevSuccessfulDagRunResponse(BaseModel): """ Schema for response with previous successful DagRun information for Task Template Context. @@ -618,19 +608,8 @@ class DagAttributeTypes(str, Enum): TASK_GROUP = "taskgroup" -class ArgValueSchema(BaseModel): - """ - JSON-schema fragment constraining the value a stub-task argument binds to. - - Only the ``type`` and ``format`` keywords are carried today, with their standard - JSON-schema semantics: ``type`` is asserted by the runtime, ``format`` is an - annotation a runtime may additionally check. Unknown keywords from newer providers - are ignored rather than rejected (as JSON-schema consumers do), so a core on this - version keeps serving specs written by a newer provider. - """ - - type: Annotated[JsonSchemaType | list[JsonSchemaType] | None, Field(title="Type")] = None - format: Annotated[str | None, Field(title="Format")] = None +class ArgValueSchema(RootModel[dict[str, JsonValue] | None]): + root: dict[str, JsonValue] | None = None class AssetReferenceAssetEventDagRun(BaseModel): 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 3c3f112e0866b..2467b5f7ea990 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -3,43 +3,6 @@ "api_version": "2026-10-30", "description": "Apache Airflow SDK Supervisor Schema", "$defs": { - "ArgValueSchema": { - "description": "JSON-schema fragment constraining the value a stub-task argument binds to.\n\nOnly the ``type`` and ``format`` keywords are carried today, with their standard\nJSON-schema semantics: ``type`` is asserted by the runtime, ``format`` is an\nannotation a runtime may additionally check. Unknown keywords from newer providers\nare ignored rather than rejected (as JSON-schema consumers do), so a core on this\nversion keeps serving specs written by a newer provider.", - "properties": { - "type": { - "anyOf": [ - { - "$ref": "#/$defs/JsonSchemaType" - }, - { - "items": { - "$ref": "#/$defs/JsonSchemaType" - }, - "type": "array" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Type" - }, - "format": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Format" - } - }, - "title": "ArgValueSchema", - "type": "object" - }, "AssetAliasReferenceAssetEventDagRun": { "additionalProperties": false, "description": "Schema for AssetAliasModel used in AssetEventDagRunReference.", @@ -3082,19 +3045,6 @@ "title": "InactiveAssetsResult", "type": "object" }, - "JsonSchemaType": { - "enum": [ - "string", - "integer", - "number", - "boolean", - "object", - "array", - "null" - ], - "title": "JsonSchemaType", - "type": "string" - }, "JsonValue": {}, "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.", @@ -4439,42 +4389,6 @@ "title": "VariableResult", "type": "object" }, - "XComArgBinding": { - "description": "One positional stub-task argument pulled from an upstream task's XCom.", - "properties": { - "kind": { - "const": "xcom", - "title": "Kind", - "type": "string" - }, - "name": { - "title": "Name", - "type": "string" - }, - "value_schema": { - "anyOf": [ - { - "$ref": "#/$defs/ArgValueSchema" - }, - { - "type": "null" - } - ], - "default": null - }, - "task_id": { - "title": "Task Id", - "type": "string" - } - }, - "required": [ - "kind", - "name", - "task_id" - ], - "title": "XComArgBinding", - "type": "object" - }, "XComCountResponse": { "properties": { "len": { @@ -4649,6 +4563,13 @@ "title": "ConnectionResponse", "type": "object" }, + "ArgValueSchema": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "ArgValueSchema", + "type": "object" + }, "LiteralArgBinding": { "description": "One positional stub-task argument carrying an inline literal from the Dag file.", "properties": { @@ -4714,6 +4635,42 @@ ], "title": "TaskArgBinding" }, + "XComArgBinding": { + "description": "One positional stub-task argument pulled from an upstream task's XCom.", + "properties": { + "kind": { + "const": "xcom", + "title": "Kind", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "value_schema": { + "anyOf": [ + { + "$ref": "#/$defs/ArgValueSchema" + }, + { + "type": "null" + } + ], + "default": null + }, + "task_id": { + "title": "Task Id", + "type": "string" + } + }, + "required": [ + "kind", + "name", + "task_id" + ], + "title": "XComArgBinding", + "type": "object" + }, "AssetEventDagRunReference": { "additionalProperties": false, "description": "Schema for AssetEvent model used in DagRun.", 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 518426ad9f010..8e2379e2fdbba 100644 --- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py @@ -449,7 +449,7 @@ def test_downgrade_strips_arg_bindings_for_previous_version(self, real_migrator, assert "arg_bindings" not in out["ti_context"] def test_head_version_keeps_arg_bindings(self, real_migrator, startup_details): - from airflow.sdk.api.datamodels._generated import JsonSchemaType, LiteralArgBinding, XComArgBinding + from airflow.sdk.api.datamodels._generated import LiteralArgBinding, XComArgBinding out = real_migrator.downgrade(startup_details, "2026-10-30") assert out.ti_context.arg_bindings is not None @@ -462,7 +462,7 @@ def test_head_version_keeps_arg_bindings(self, real_migrator, startup_details): assert isinstance(xcom, XComArgBinding) assert xcom.task_id == "extract" assert xcom.name == "extracted" - assert xcom.value_schema.type == JsonSchemaType.OBJECT + assert xcom.value_schema.root == {"type": "object"} assert isinstance(defaulted, LiteralArgBinding) assert defaulted.from_default is True - assert defaulted.value_schema.format == "int64" + assert defaulted.value_schema.root == {"type": "integer", "format": "int64"} diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 9596a42274bf9..757429e7f29fd 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -22,20 +22,6 @@ // // Re-run with: pnpm run generate:supervisor -export type Type = JsonSchemaType | JsonSchemaType[] | null; -/** - * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema - * via the `definition` "JsonSchemaType". - */ -export type JsonSchemaType = - | "string" - | "integer" - | "number" - | "boolean" - | "object" - | "array" - | "null"; -export type Format = string | null; export type Name = string; export type Id = number; export type Timestamp = string; @@ -69,10 +55,10 @@ export type SourceRunId = string | null; export type SourceMapIndex = number | null; export type PartitionKey1 = string | null; export type AssetEvents = AssetEventResponse[]; -export type Type1 = "AssetEventsResult"; +export type Type = "AssetEventsResult"; export type Name2 = string | null; export type Uri1 = string | null; -export type Type2 = string; +export type Type1 = string; export type Name3 = string; export type Uri2 = string; export type Name4 = string; @@ -81,10 +67,10 @@ export type Group1 = string; export type Extra3 = { [k: string]: JsonValue; } | null; -export type Type3 = "AssetResult"; -export type Type4 = "AssetStateStoreResult"; +export type Type2 = "AssetResult"; +export type Type3 = "AssetStateStoreResult"; export type Assets = AssetResult[]; -export type Type5 = "AssetsByAliasResult"; +export type Type4 = "AssetsByAliasResult"; export type State1 = "awaiting_input" | null; export type Timeout = string | null; export type NextMethod = string; @@ -92,18 +78,18 @@ export type NextKwargs = { [k: string]: JsonValue; } | null; export type RenderedMapIndex = string | null; -export type Type6 = "AwaitInputTask"; +export type Type5 = "AwaitInputTask"; export type Name5 = string; export type Version = string | null; export type VersionData = { [k: string]: unknown; } | null; export type Name6 = string; -export type Type7 = "ClearAssetStateStoreByName"; +export type Type6 = "ClearAssetStateStoreByName"; export type Uri4 = string; -export type Type8 = "ClearAssetStateStoreByUri"; +export type Type7 = "ClearAssetStateStoreByUri"; export type TiId = string; -export type Type9 = "ClearTaskStateStore"; +export type Type8 = "ClearTaskStateStore"; export type ConnId = string; export type ConnType = string; export type Host = string | null; @@ -112,7 +98,7 @@ export type Login = string | null; export type Password = string | null; export type Port = number | null; export type Extra4 = string | null; -export type Type10 = "ConnectionResult"; +export type Type9 = "ConnectionResult"; export type TiId1 = string; /** * @minItems 1 @@ -128,9 +114,9 @@ export type Params = { export type AssignedUsers = HITLUser[] | null; export type Id1 = string; export type Name7 = string; -export type Type11 = "CreateHITLDetailPayload"; +export type Type10 = "CreateHITLDetailPayload"; export type Count = number; -export type Type12 = "DRCount"; +export type Type11 = "DRCount"; export type Filepath = string; export type BundleName = string; export type BundleVersion = string | null; @@ -202,7 +188,7 @@ export type ContextCarrier = { } | null; export type Queue = string; export type IsFailureCallback = boolean | null; -export type Type13 = "DagCallbackRequest"; +export type Type12 = "DagCallbackRequest"; export type File = string; export type BundlePath = string; export type BundleName1 = string; @@ -271,7 +257,7 @@ export type TaskId1 = string; export type Kind1 = "literal"; export type Name9 = string; export type FromDefault = boolean; -export type Type14 = "TaskCallbackRequest"; +export type Type13 = "TaskCallbackRequest"; export type Filepath2 = string; export type BundleName3 = string; export type BundleVersion2 = string | null; @@ -280,9 +266,9 @@ export type VersionData3 = { } | null; export type Msg2 = string | null; export type EmailType = "failure" | "retry"; -export type Type15 = "EmailRequest"; +export type Type14 = "EmailRequest"; export type CallbackRequests = (DagCallbackRequest | TaskCallbackRequest | EmailRequest)[]; -export type Type16 = "DagFileParseRequest"; +export type Type15 = "DagFileParseRequest"; export type Fileloc = string; export type LastLoaded = string | null; export type SerializedDags = LazyDeserializedDAG[]; @@ -290,7 +276,7 @@ export type Warnings = unknown[] | null; export type ImportErrors = { [k: string]: string; } | null; -export type Type17 = "DagFileParsingResult"; +export type Type16 = "DagFileParsingResult"; export type DagId4 = string; export type IsPaused = boolean; export type BundleName4 = string | null; @@ -299,7 +285,7 @@ export type RelativeFileloc = string | null; export type Owners = string | null; export type Tags = string[]; export type NextDagrun = string | null; -export type Type18 = "DagResult"; +export type Type17 = "DagResult"; export type DagId5 = string; export type RunId4 = string; export type LogicalDate2 = string | null; @@ -318,8 +304,8 @@ export type PartitionKey4 = string | null; export type PartitionDate1 = string | null; export type Note1 = string | null; export type TeamName1 = string | null; -export type Type19 = "DagRunResult"; -export type Type20 = "DagRunStateResult"; +export type Type18 = "DagRunResult"; +export type Type19 = "DagRunStateResult"; export type State2 = "deferred" | null; export type Classpath = string; export type TriggerKwargs = @@ -335,24 +321,24 @@ export type NextKwargs2 = { [k: string]: JsonValue; } | null; export type RenderedMapIndex1 = string | null; -export type Type21 = "DeferTask"; +export type Type20 = "DeferTask"; export type Name10 = string; export type Key1 = string; -export type Type22 = "DeleteAssetStateStoreByName"; +export type Type21 = "DeleteAssetStateStoreByName"; export type Uri5 = string; export type Key2 = string; -export type Type23 = "DeleteAssetStateStoreByUri"; +export type Type22 = "DeleteAssetStateStoreByUri"; export type TiId2 = string; export type Key3 = string; -export type Type24 = "DeleteTaskStateStore"; +export type Type23 = "DeleteTaskStateStore"; export type Key4 = string; -export type Type25 = "DeleteVariable"; +export type Type24 = "DeleteVariable"; export type Key5 = string; export type DagId6 = string; export type RunId5 = string; export type TaskId2 = string; export type MapIndex1 = number | null; -export type Type26 = "DeleteXCom"; +export type Type25 = "DeleteXCom"; /** * Error types used in the API client. */ @@ -370,7 +356,7 @@ export type ErrorType = export type Detail = { [k: string]: unknown; } | null; -export type Type27 = "ErrorResponse"; +export type Type26 = "ErrorResponse"; /** * Error types used in the API client. * @@ -389,9 +375,9 @@ export type ErrorType1 = | "GENERIC_ERROR" | "API_SERVER_ERROR"; export type Name11 = string; -export type Type28 = "GetAssetByName"; +export type Type27 = "GetAssetByName"; export type Uri6 = string; -export type Type29 = "GetAssetByUri"; +export type Type28 = "GetAssetByUri"; export type Name12 = string | null; export type Uri7 = string | null; export type After = string | null; @@ -403,7 +389,7 @@ export type PartitionKeyRegexpPattern = string | null; export type Extra7 = { [k: string]: string; } | null; -export type Type30 = "GetAssetEventByAsset"; +export type Type29 = "GetAssetEventByAsset"; export type AliasName = string; export type After1 = string | null; export type Before1 = string | null; @@ -414,43 +400,43 @@ export type PartitionKeyRegexpPattern1 = string | null; export type Extra8 = { [k: string]: string; } | null; -export type Type31 = "GetAssetEventByAssetAlias"; +export type Type30 = "GetAssetEventByAssetAlias"; export type Name13 = string; export type Key6 = string; -export type Type32 = "GetAssetStateStoreByName"; +export type Type31 = "GetAssetStateStoreByName"; export type Uri8 = string; export type Key7 = string; -export type Type33 = "GetAssetStateStoreByUri"; +export type Type32 = "GetAssetStateStoreByUri"; export type AliasName1 = string; -export type Type34 = "GetAssetsByAlias"; +export type Type33 = "GetAssetsByAlias"; export type ConnId2 = string; -export type Type35 = "GetConnection"; +export type Type34 = "GetConnection"; export type DagId7 = string; export type LogicalDates = string[] | null; export type RunIds = string[] | null; export type States = string[] | null; -export type Type36 = "GetDRCount"; +export type Type35 = "GetDRCount"; export type DagId8 = string; -export type Type37 = "GetDag"; +export type Type36 = "GetDag"; export type DagId9 = string; export type RunId6 = string; -export type Type38 = "GetDagRun"; +export type Type37 = "GetDagRun"; export type DagId10 = string; export type RunId7 = string; -export type Type39 = "GetDagRunState"; +export type Type38 = "GetDagRunState"; export type TiId3 = string; -export type Type40 = "GetHITLDetailResponse"; +export type Type39 = "GetHITLDetailResponse"; export type TiId4 = string; -export type Type41 = "GetPrevSuccessfulDagRun"; +export type Type40 = "GetPrevSuccessfulDagRun"; export type DagId11 = string; export type LogicalDate3 = string; export type State3 = string | null; -export type Type42 = "GetPreviousDagRun"; +export type Type41 = "GetPreviousDagRun"; export type DagId12 = string; export type TaskId3 = string; export type LogicalDate4 = string | null; export type MapIndex2 = number; -export type Type43 = "GetPreviousTI"; +export type Type42 = "GetPreviousTI"; export type DagId13 = string; export type MapIndex3 = number | null; export type TaskIds = string[] | null; @@ -458,47 +444,47 @@ export type TaskGroupId = string | null; export type LogicalDates1 = string[] | null; export type RunIds1 = string[] | null; export type States1 = string[] | null; -export type Type44 = "GetTICount"; +export type Type43 = "GetTICount"; export type DagId14 = string; export type RunId8 = string; -export type Type45 = "GetTaskBreadcrumbs"; +export type Type44 = "GetTaskBreadcrumbs"; export type TiId5 = string; export type TryNumber1 = number; -export type Type46 = "GetTaskRescheduleStartDate"; +export type Type45 = "GetTaskRescheduleStartDate"; export type TiId6 = string; export type Key8 = string; -export type Type47 = "GetTaskStateStore"; +export type Type46 = "GetTaskStateStore"; export type DagId15 = string; export type MapIndex4 = number | null; export type TaskIds1 = string[] | null; export type TaskGroupId1 = string | null; export type LogicalDates2 = string[] | null; export type RunIds2 = string[] | null; -export type Type48 = "GetTaskStates"; +export type Type47 = "GetTaskStates"; export type Key9 = string; -export type Type49 = "GetVariable"; +export type Type48 = "GetVariable"; export type Prefix = string | null; export type Limit2 = number; export type Offset = number; -export type Type50 = "GetVariableKeys"; +export type Type49 = "GetVariableKeys"; export type Key10 = string; export type DagId16 = string; export type RunId9 = string; export type TaskId4 = string; export type MapIndex5 = number | null; export type IncludePriorDates = boolean; -export type Type51 = "GetXCom"; +export type Type50 = "GetXCom"; export type Key11 = string; export type DagId17 = string; export type RunId10 = string; export type TaskId5 = string; -export type Type52 = "GetXComCount"; +export type Type51 = "GetXComCount"; export type Key12 = string; export type DagId18 = string; export type RunId11 = string; export type TaskId6 = string; export type Offset1 = number; -export type Type53 = "GetXComSequenceItem"; +export type Type52 = "GetXComSequenceItem"; export type Key13 = string; export type DagId19 = string; export type RunId12 = string; @@ -507,7 +493,7 @@ export type Start = number | null; export type Stop = number | null; export type Step = number | null; export type IncludePriorDates1 = boolean; -export type Type54 = "GetXComSequenceSlice"; +export type Type53 = "GetXComSequenceSlice"; export type TiId7 = string; /** * @minItems 1 @@ -521,19 +507,19 @@ export type Params1 = { [k: string]: unknown; } | null; export type AssignedUsers1 = HITLUser[] | null; -export type Type55 = "HITLDetailRequestResult"; +export type Type54 = "HITLDetailRequestResult"; export type InactiveAssets = AssetProfile[] | null; -export type Type56 = "InactiveAssetsResult"; +export type Type55 = "InactiveAssetsResult"; export type Name14 = string | null; -export type Type57 = "MaskSecret"; +export type Type56 = "MaskSecret"; export type Ok = boolean; -export type Type58 = "OKResponse"; +export type Type57 = "OKResponse"; export type DataIntervalStart3 = string | null; export type DataIntervalEnd3 = string | null; export type StartDate4 = string | null; export type EndDate3 = string | null; -export type Type59 = "PrevSuccessfulDagRunResult"; -export type Type60 = "PreviousDagRunResult"; +export type Type58 = "PrevSuccessfulDagRunResult"; +export type Type59 = "PreviousDagRunResult"; export type TaskId8 = string; export type DagId20 = string; export type RunId13 = string; @@ -544,37 +530,37 @@ export type State4 = string | null; export type TryNumber2 = number; export type MapIndex6 = number | null; export type Duration = number | null; -export type Type61 = "PreviousTIResult"; +export type Type60 = "PreviousTIResult"; export type Key14 = string; export type Value1 = string | null; export type Description = string | null; -export type Type62 = "PutVariable"; +export type Type61 = "PutVariable"; export type State5 = "up_for_reschedule" | null; export type RescheduleDate = string; export type EndDate5 = string; -export type Type63 = "RescheduleTask"; -export type Type64 = "ResendLoggingFD"; +export type Type62 = "RescheduleTask"; +export type Type63 = "ResendLoggingFD"; export type State6 = "up_for_retry" | null; export type EndDate6 = string; export type RenderedMapIndex2 = string | null; export type RetryDelaySeconds = number | null; export type RetryReason = string | null; -export type Type65 = "RetryTask"; -export type Type66 = "SentFDs"; +export type Type64 = "RetryTask"; +export type Type65 = "SentFDs"; export type Fds = number[]; export type Name15 = string; export type Key15 = string; -export type Type67 = "SetAssetStateStoreByName"; +export type Type66 = "SetAssetStateStoreByName"; export type Uri9 = string; export type Key16 = string; -export type Type68 = "SetAssetStateStoreByUri"; -export type Type69 = "SetRenderedFields"; +export type Type67 = "SetAssetStateStoreByUri"; +export type Type68 = "SetRenderedFields"; export type RenderedMapIndex3 = string; -export type Type70 = "SetRenderedMapIndex"; +export type Type69 = "SetRenderedMapIndex"; export type TiId8 = string; export type Key17 = string; export type ExpiresAt = string | null; -export type Type71 = "SetTaskStateStore"; +export type Type70 = "SetTaskStateStore"; export type Key18 = string; export type DagId21 = string; export type RunId14 = string; @@ -582,13 +568,13 @@ export type TaskId9 = string; export type MapIndex7 = number | null; export type DagResult1 = boolean; export type MappedLength = number | null; -export type Type72 = "SetXCom"; +export type Type71 = "SetXCom"; export type Tasks = (string | [unknown, unknown])[]; -export type Type73 = "SkipDownstreamTasks"; +export type Type72 = "SkipDownstreamTasks"; export type DagRelPath = string; export type StartDate6 = string; export type SentryIntegration = string; -export type Type74 = "StartupDetails"; +export type Type73 = "StartupDetails"; export type State7 = "success" | null; export type EndDate7 = string; export type TaskOutlets = AssetProfile[] | null; @@ -598,21 +584,21 @@ export type OutletEvents = }[] | null; export type RenderedMapIndex4 = string | null; -export type Type75 = "SucceedTask"; +export type Type74 = "SucceedTask"; export type Count1 = number; -export type Type76 = "TICount"; +export type Type75 = "TICount"; export type Breadcrumbs = { [k: string]: unknown; }[]; -export type Type77 = "TaskBreadcrumbsResult"; +export type Type76 = "TaskBreadcrumbsResult"; export type StartDate7 = string | null; -export type Type78 = "TaskRescheduleStartDate"; +export type Type77 = "TaskRescheduleStartDate"; export type State8 = "failed" | "skipped" | "removed"; export type EndDate8 = string | null; -export type Type79 = "TaskState"; +export type Type78 = "TaskState"; export type RenderedMapIndex5 = string | null; -export type Type80 = "TaskStateStoreResult"; -export type Type81 = "TaskStatesResult"; +export type Type79 = "TaskStateStoreResult"; +export type Type80 = "TaskStatesResult"; export type LogicalDate6 = string | null; export type RunAfter2 = string | null; export type Conf2 = { @@ -623,7 +609,7 @@ export type PartitionKey7 = string | null; export type Note2 = string | null; export type DagId22 = string; export type DagRunId = string; -export type Type82 = "TriggerDagRun"; +export type Type81 = "TriggerDagRun"; export type TiId9 = string; /** * @minItems 1 @@ -632,40 +618,24 @@ export type ChosenOptions = [string, ...string[]]; export type ParamsInput = { [k: string]: unknown; } | null; -export type Type83 = "UpdateHITLDetail"; +export type Type82 = "UpdateHITLDetail"; export type TiId10 = string; -export type Type84 = "ValidateInletsAndOutlets"; +export type Type83 = "ValidateInletsAndOutlets"; export type Keys = string[]; export type TotalEntries = number; -export type Type85 = "VariableKeysResult"; +export type Type84 = "VariableKeysResult"; export type Key19 = string; export type Value2 = string | null; -export type Type86 = "VariableResult"; +export type Type85 = "VariableResult"; export type Len = number; -export type Type87 = "XComCountResponse"; +export type Type86 = "XComCountResponse"; export type Key20 = string; -export type Type88 = "XComResult"; -export type Type89 = "XComSequenceIndexResult"; +export type Type87 = "XComResult"; +export type Type88 = "XComSequenceIndexResult"; export type Root = JsonValue[]; -export type Type90 = "XComSequenceSliceResult"; +export type Type89 = "XComSequenceSliceResult"; export interface SupervisorWireSchema {} -/** - * JSON-schema fragment constraining the value a stub-task argument binds to. - * - * Only the ``type`` and ``format`` keywords are carried today, with their standard - * JSON-schema semantics: ``type`` is asserted by the runtime, ``format`` is an - * annotation a runtime may additionally check. Unknown keywords from newer providers - * are ignored rather than rejected (as JSON-schema consumers do), so a core on this - * version keeps serving specs written by a newer provider. - * - * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema - * via the `definition` "ArgValueSchema". - */ -export interface ArgValueSchema { - type?: Type; - format?: Format; -} /** * Schema for AssetAliasModel used in AssetEventDagRunReference. * @@ -730,7 +700,7 @@ export interface DagRunAssetReference { */ export interface AssetEventsResult { asset_events: AssetEvents; - type?: Type1; + type?: Type; } /** * Profile of an asset-like object. @@ -750,7 +720,7 @@ export interface AssetEventsResult { export interface AssetProfile { name?: Name2; uri?: Uri1; - type: Type2; + type: Type1; } /** * Schema for AssetModel used in AssetEventDagRunReference. @@ -777,7 +747,7 @@ export interface AssetResult { uri: Uri3; group: Group1; extra?: Extra3; - type?: Type3; + type?: Type2; } /** * Response to GetAssetStateStore; wraps the generated API response for supervisor to worker comms. @@ -787,7 +757,7 @@ export interface AssetResult { */ export interface AssetStateStoreResult { value: JsonValue; - type?: Type4; + type?: Type3; } /** * Response to GetAssetsByAlias; list of concrete assets resolved from an alias. @@ -797,7 +767,7 @@ export interface AssetStateStoreResult { */ export interface AssetsByAliasResult { assets: Assets; - type?: Type5; + type?: Type4; } /** * Park a task instance awaiting human input (Human-in-the-loop), without a trigger. @@ -811,7 +781,7 @@ export interface AwaitInputTask { next_method: NextMethod; next_kwargs?: NextKwargs; rendered_map_index?: RenderedMapIndex; - type?: Type6; + type?: Type5; } /** * Schema for telling task which bundle to run with. @@ -830,7 +800,7 @@ export interface BundleInfo { */ export interface ClearAssetStateStoreByName { name: Name6; - type?: Type7; + type?: Type6; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -838,7 +808,7 @@ export interface ClearAssetStateStoreByName { */ export interface ClearAssetStateStoreByUri { uri: Uri4; - type?: Type8; + type?: Type7; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -846,7 +816,7 @@ export interface ClearAssetStateStoreByUri { */ export interface ClearTaskStateStore { ti_id: TiId; - type?: Type9; + type?: Type8; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -861,7 +831,7 @@ export interface ConnectionResult { password?: Password; port?: Port; extra?: Extra4; - type?: Type10; + type?: Type9; } /** * Add the input request part of a Human-in-the-loop response. @@ -878,7 +848,7 @@ export interface CreateHITLDetailPayload { multiple?: Multiple; params?: Params; assigned_users?: AssignedUsers; - type?: Type11; + type?: Type10; } /** * Schema for a Human-in-the-loop users. @@ -898,7 +868,7 @@ export interface HITLUser { */ export interface DRCount { count: Count; - type?: Type12; + type?: Type11; } /** * A Class with information about the success/failure DAG callback to be executed. @@ -916,7 +886,7 @@ export interface DagCallbackRequest { run_id: RunId1; context_from_server?: DagRunContext | null; is_failure_callback?: IsFailureCallback; - type?: Type13; + type?: Type12; } /** * Class to pass context info from the server to build a Execution context object. @@ -1006,7 +976,7 @@ export interface DagFileParseRequest { bundle_path: BundlePath; bundle_name: BundleName1; callback_requests?: CallbackRequests; - type?: Type16; + type?: Type15; } /** * Task callback status information. @@ -1026,7 +996,7 @@ export interface TaskCallbackRequest { ti: TaskInstance; task_callback_type?: TaskInstanceState | null; context_from_server?: TIRunContext | null; - type?: Type14; + type?: Type13; } /** * Response schema for TaskInstance run context. @@ -1085,6 +1055,13 @@ export interface XComArgBinding { value_schema?: ArgValueSchema | null; task_id: TaskId1; } +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "ArgValueSchema". + */ +export interface ArgValueSchema { + [k: string]: JsonValue; +} /** * One positional stub-task argument carrying an inline literal from the Dag file. * @@ -1113,7 +1090,7 @@ export interface EmailRequest { ti: TaskInstance; email_type?: EmailType; context_from_server: TIRunContext; - type?: Type15; + type?: Type14; } /** * Result of DAG File Parsing. @@ -1129,7 +1106,7 @@ export interface DagFileParsingResult { serialized_dags: SerializedDags; warnings?: Warnings; import_errors?: ImportErrors; - type?: Type17; + type?: Type16; } /** * Lazily build information from the serialized DAG structure. @@ -1160,7 +1137,7 @@ export interface DagResult { owners?: Owners; tags: Tags; next_dagrun?: NextDagrun; - type?: Type18; + type?: Type17; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1185,7 +1162,7 @@ export interface DagRunResult { partition_date?: PartitionDate1; note?: Note1; team_name?: TeamName1; - type?: Type19; + type?: Type18; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1193,7 +1170,7 @@ export interface DagRunResult { */ export interface DagRunStateResult { state: DagRunState; - type?: Type20; + type?: Type19; } /** * Update a task instance state to deferred. @@ -1210,7 +1187,7 @@ export interface DeferTask { next_method: NextMethod2; next_kwargs?: NextKwargs2; rendered_map_index?: RenderedMapIndex1; - type?: Type21; + type?: Type20; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1219,7 +1196,7 @@ export interface DeferTask { export interface DeleteAssetStateStoreByName { name: Name10; key: Key1; - type?: Type22; + type?: Type21; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1228,7 +1205,7 @@ export interface DeleteAssetStateStoreByName { export interface DeleteAssetStateStoreByUri { uri: Uri5; key: Key2; - type?: Type23; + type?: Type22; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1237,7 +1214,7 @@ export interface DeleteAssetStateStoreByUri { export interface DeleteTaskStateStore { ti_id: TiId2; key: Key3; - type?: Type24; + type?: Type23; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1245,7 +1222,7 @@ export interface DeleteTaskStateStore { */ export interface DeleteVariable { key: Key4; - type?: Type25; + type?: Type24; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1257,7 +1234,7 @@ export interface DeleteXCom { run_id: RunId5; task_id: TaskId2; map_index?: MapIndex1; - type?: Type26; + type?: Type25; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1266,7 +1243,7 @@ export interface DeleteXCom { export interface ErrorResponse { error?: ErrorType; detail?: Detail; - type?: Type27; + type?: Type26; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1274,7 +1251,7 @@ export interface ErrorResponse { */ export interface GetAssetByName { name: Name11; - type?: Type28; + type?: Type27; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1282,7 +1259,7 @@ export interface GetAssetByName { */ export interface GetAssetByUri { uri: Uri6; - type?: Type29; + type?: Type28; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1298,7 +1275,7 @@ export interface GetAssetEventByAsset { partition_key?: PartitionKey5; partition_key_regexp_pattern?: PartitionKeyRegexpPattern; extra?: Extra7; - type?: Type30; + type?: Type29; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1313,7 +1290,7 @@ export interface GetAssetEventByAssetAlias { partition_key?: PartitionKey6; partition_key_regexp_pattern?: PartitionKeyRegexpPattern1; extra?: Extra8; - type?: Type31; + type?: Type30; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1322,7 +1299,7 @@ export interface GetAssetEventByAssetAlias { export interface GetAssetStateStoreByName { name: Name13; key: Key6; - type?: Type32; + type?: Type31; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1331,7 +1308,7 @@ export interface GetAssetStateStoreByName { export interface GetAssetStateStoreByUri { uri: Uri8; key: Key7; - type?: Type33; + type?: Type32; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1339,7 +1316,7 @@ export interface GetAssetStateStoreByUri { */ export interface GetAssetsByAlias { alias_name: AliasName1; - type?: Type34; + type?: Type33; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1347,7 +1324,7 @@ export interface GetAssetsByAlias { */ export interface GetConnection { conn_id: ConnId2; - type?: Type35; + type?: Type34; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1358,7 +1335,7 @@ export interface GetDRCount { logical_dates?: LogicalDates; run_ids?: RunIds; states?: States; - type?: Type36; + type?: Type35; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1366,7 +1343,7 @@ export interface GetDRCount { */ export interface GetDag { dag_id: DagId8; - type?: Type37; + type?: Type36; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1375,7 +1352,7 @@ export interface GetDag { export interface GetDagRun { dag_id: DagId9; run_id: RunId6; - type?: Type38; + type?: Type37; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1384,7 +1361,7 @@ export interface GetDagRun { export interface GetDagRunState { dag_id: DagId10; run_id: RunId7; - type?: Type39; + type?: Type38; } /** * Get the response content part of a Human-in-the-loop response. @@ -1394,7 +1371,7 @@ export interface GetDagRunState { */ export interface GetHITLDetailResponse { ti_id: TiId3; - type?: Type40; + type?: Type39; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1402,7 +1379,7 @@ export interface GetHITLDetailResponse { */ export interface GetPrevSuccessfulDagRun { ti_id: TiId4; - type?: Type41; + type?: Type40; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1412,7 +1389,7 @@ export interface GetPreviousDagRun { dag_id: DagId11; logical_date: LogicalDate3; state?: State3; - type?: Type42; + type?: Type41; } /** * Request to get previous task instance. @@ -1426,7 +1403,7 @@ export interface GetPreviousTI { logical_date?: LogicalDate4; map_index?: MapIndex2; state?: TaskInstanceState | null; - type?: Type43; + type?: Type42; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1440,7 +1417,7 @@ export interface GetTICount { logical_dates?: LogicalDates1; run_ids?: RunIds1; states?: States1; - type?: Type44; + type?: Type43; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1449,7 +1426,7 @@ export interface GetTICount { export interface GetTaskBreadcrumbs { dag_id: DagId14; run_id: RunId8; - type?: Type45; + type?: Type44; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1458,7 +1435,7 @@ export interface GetTaskBreadcrumbs { export interface GetTaskRescheduleStartDate { ti_id: TiId5; try_number?: TryNumber1; - type?: Type46; + type?: Type45; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1467,7 +1444,7 @@ export interface GetTaskRescheduleStartDate { export interface GetTaskStateStore { ti_id: TiId6; key: Key8; - type?: Type47; + type?: Type46; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1480,7 +1457,7 @@ export interface GetTaskStates { task_group_id?: TaskGroupId1; logical_dates?: LogicalDates2; run_ids?: RunIds2; - type?: Type48; + type?: Type47; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1488,7 +1465,7 @@ export interface GetTaskStates { */ export interface GetVariable { key: Key9; - type?: Type49; + type?: Type48; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1498,7 +1475,7 @@ export interface GetVariableKeys { prefix?: Prefix; limit?: Limit2; offset?: Offset; - type?: Type50; + type?: Type49; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1511,7 +1488,7 @@ export interface GetXCom { task_id: TaskId4; map_index?: MapIndex5; include_prior_dates?: IncludePriorDates; - type?: Type51; + type?: Type50; } /** * Get the number of (mapped) XCom values available. @@ -1524,7 +1501,7 @@ export interface GetXComCount { dag_id: DagId17; run_id: RunId10; task_id: TaskId5; - type?: Type52; + type?: Type51; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1536,7 +1513,7 @@ export interface GetXComSequenceItem { run_id: RunId11; task_id: TaskId6; offset: Offset1; - type?: Type53; + type?: Type52; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1551,7 +1528,7 @@ export interface GetXComSequenceSlice { stop: Stop; step: Step; include_prior_dates?: IncludePriorDates1; - type?: Type54; + type?: Type53; } /** * Response to CreateHITLDetailPayload request. @@ -1568,7 +1545,7 @@ export interface HITLDetailRequestResult { multiple?: Multiple1; params?: Params1; assigned_users?: AssignedUsers1; - type?: Type55; + type?: Type54; } /** * Response of InactiveAssets requests. @@ -1578,7 +1555,7 @@ export interface HITLDetailRequestResult { */ export interface InactiveAssetsResult { inactive_assets?: InactiveAssets; - type?: Type56; + type?: Type55; } /** * Add a new value to be redacted in task logs. @@ -1589,7 +1566,7 @@ export interface InactiveAssetsResult { export interface MaskSecret { value: JsonValue; name?: Name14; - type?: Type57; + type?: Type56; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1597,7 +1574,7 @@ export interface MaskSecret { */ export interface OKResponse { ok: Ok; - type?: Type58; + type?: Type57; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1608,7 +1585,7 @@ export interface PrevSuccessfulDagRunResult { data_interval_end?: DataIntervalEnd3; start_date?: StartDate4; end_date?: EndDate3; - type?: Type59; + type?: Type58; } /** * Response containing previous Dag run information. @@ -1618,7 +1595,7 @@ export interface PrevSuccessfulDagRunResult { */ export interface PreviousDagRunResult { dag_run?: DagRun | null; - type?: Type60; + type?: Type59; } /** * Schema for response with previous TaskInstance information. @@ -1646,7 +1623,7 @@ export interface PreviousTIResponse { */ export interface PreviousTIResult { task_instance?: PreviousTIResponse | null; - type?: Type61; + type?: Type60; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1656,7 +1633,7 @@ export interface PutVariable { key: Key14; value: Value1; description: Description; - type?: Type62; + type?: Type61; } /** * Update a task instance state to reschedule/up_for_reschedule. @@ -1668,14 +1645,14 @@ export interface RescheduleTask { state?: State5; reschedule_date: RescheduleDate; end_date: EndDate5; - type?: Type63; + type?: Type62; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema * via the `definition` "ResendLoggingFD". */ export interface ResendLoggingFD { - type?: Type64; + type?: Type63; } /** * Update a task instance state to up_for_retry. @@ -1689,14 +1666,14 @@ export interface RetryTask { rendered_map_index?: RenderedMapIndex2; retry_delay_seconds?: RetryDelaySeconds; retry_reason?: RetryReason; - type?: Type65; + type?: Type64; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema * via the `definition` "SentFDs". */ export interface SentFDs { - type?: Type66; + type?: Type65; fds: Fds; } /** @@ -1707,7 +1684,7 @@ export interface SetAssetStateStoreByName { name: Name15; key: Key15; value: JsonValue; - type?: Type67; + type?: Type66; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1717,7 +1694,7 @@ export interface SetAssetStateStoreByUri { uri: Uri9; key: Key16; value: JsonValue; - type?: Type68; + type?: Type67; } /** * Payload for setting RTIF for a task instance. @@ -1727,7 +1704,7 @@ export interface SetAssetStateStoreByUri { */ export interface SetRenderedFields { rendered_fields: RenderedFields; - type?: Type69; + type?: Type68; } export interface RenderedFields { [k: string]: JsonValue; @@ -1740,7 +1717,7 @@ export interface RenderedFields { */ export interface SetRenderedMapIndex { rendered_map_index: RenderedMapIndex3; - type?: Type70; + type?: Type69; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1751,7 +1728,7 @@ export interface SetTaskStateStore { key: Key17; value: JsonValue; expires_at: ExpiresAt; - type?: Type71; + type?: Type70; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1766,7 +1743,7 @@ export interface SetXCom { map_index?: MapIndex7; dag_result?: DagResult1; mapped_length?: MappedLength; - type?: Type72; + type?: Type71; } /** * Update state of downstream tasks within a task instance to 'skipped', while updating current task to success state. @@ -1776,7 +1753,7 @@ export interface SetXCom { */ export interface SkipDownstreamTasks { tasks: Tasks; - type?: Type73; + type?: Type72; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1789,7 +1766,7 @@ export interface StartupDetails { start_date: StartDate6; ti_context: TIRunContext; sentry_integration: SentryIntegration; - type?: Type74; + type?: Type73; } /** * Update a task's state to success. Includes task_outlets and outlet_events for registering asset events. @@ -1803,7 +1780,7 @@ export interface SucceedTask { task_outlets?: TaskOutlets; outlet_events?: OutletEvents; rendered_map_index?: RenderedMapIndex4; - type?: Type75; + type?: Type74; } /** * Response containing count of Task Instances matching certain filters. @@ -1813,7 +1790,7 @@ export interface SucceedTask { */ export interface TICount { count: Count1; - type?: Type76; + type?: Type75; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1821,7 +1798,7 @@ export interface TICount { */ export interface TaskBreadcrumbsResult { breadcrumbs: Breadcrumbs; - type?: Type77; + type?: Type76; } /** * Response containing the first reschedule date for a task instance. @@ -1831,7 +1808,7 @@ export interface TaskBreadcrumbsResult { */ export interface TaskRescheduleStartDate { start_date: StartDate7; - type?: Type78; + type?: Type77; } /** * Update a task's state. @@ -1846,7 +1823,7 @@ export interface TaskRescheduleStartDate { export interface TaskState { state: State8; end_date?: EndDate8; - type?: Type79; + type?: Type78; rendered_map_index?: RenderedMapIndex5; } /** @@ -1857,7 +1834,7 @@ export interface TaskState { */ export interface TaskStateStoreResult { value: JsonValue; - type?: Type80; + type?: Type79; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1865,7 +1842,7 @@ export interface TaskStateStoreResult { */ export interface TaskStatesResult { task_states: TaskStates; - type?: Type81; + type?: Type80; } export interface TaskStates { [k: string]: unknown; @@ -1883,7 +1860,7 @@ export interface TriggerDagRun { note?: Note2; dag_id: DagId22; run_id: DagRunId; - type?: Type82; + type?: Type81; } /** * Update the response content part of an existing Human-in-the-loop response. @@ -1895,7 +1872,7 @@ export interface UpdateHITLDetail { ti_id: TiId9; chosen_options: ChosenOptions; params_input?: ParamsInput; - type?: Type83; + type?: Type82; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1903,7 +1880,7 @@ export interface UpdateHITLDetail { */ export interface ValidateInletsAndOutlets { ti_id: TiId10; - type?: Type84; + type?: Type83; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1912,7 +1889,7 @@ export interface ValidateInletsAndOutlets { export interface VariableKeysResult { keys: Keys; total_entries: TotalEntries; - type?: Type85; + type?: Type84; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1921,7 +1898,7 @@ export interface VariableKeysResult { export interface VariableResult { key: Key19; value?: Value2; - type?: Type86; + type?: Type85; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1929,7 +1906,7 @@ export interface VariableResult { */ export interface XComCountResponse { len: Len; - type?: Type87; + type?: Type86; } /** * Response to ReadXCom request. @@ -1940,7 +1917,7 @@ export interface XComCountResponse { export interface XComResult { key: Key20; value: JsonValue; - type?: Type88; + type?: Type87; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1948,7 +1925,7 @@ export interface XComResult { */ export interface XComSequenceIndexResult { root: JsonValue; - type?: Type89; + type?: Type88; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1956,7 +1933,7 @@ export interface XComSequenceIndexResult { */ export interface XComSequenceSliceResult { root: Root; - type?: Type90; + type?: Type89; } /** Cadwyn schema version this SDK was generated against. diff --git a/uv.lock b/uv.lock index ea07228988704..e411e55b60485 100644 --- a/uv.lock +++ b/uv.lock @@ -64,9 +64,9 @@ apache-airflow-providers-apache-cassandra = false apache-airflow-providers-asana = false apache-airflow-providers-oracle = false apache-airflow-providers-mysql = false +apache-airflow-providers-teradata = false apache-airflow-providers-alibaba = false apache-airflow-providers-microsoft-mssql = false -apache-airflow-providers-teradata = false apache-airflow-providers-jdbc = false apache-airflow-helm-chart = false apache-airflow-providers-anthropic = false @@ -7909,6 +7909,7 @@ source = { editable = "providers/standard" } dependencies = [ { name = "apache-airflow" }, { name = "apache-airflow-providers-common-compat" }, + { name = "pydantic" }, ] [package.optional-dependencies] @@ -7934,6 +7935,7 @@ requires-dist = [ { name = "apache-airflow", editable = "." }, { name = "apache-airflow-providers-common-compat", editable = "providers/common/compat" }, { name = "apache-airflow-providers-openlineage", marker = "extra == 'openlineage'", editable = "providers/openlineage" }, + { name = "pydantic", specifier = ">=2.11.0" }, ] provides-extras = ["openlineage"] From d5be0f44b0337baf232375e06a885c727baa3e81 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 24 Jul 2026 08:06:07 +0000 Subject: [PATCH 23/40] Support pendulum date/time annotations on stub task arguments pendulum.DateTime is the most common temporal annotation in Airflow Dag code, but pydantic cannot generate schemas for datetime subclasses, so stub arg specs silently degraded to decode-only checks. Temporal subclasses now normalize to their stdlib bases (recursively through unions and containers) before schema generation. Also drop the standard provider's explicit pydantic dependency (apache-airflow already provides it), trim the verbose comments introduced with the arg-binding spec, and shorten the invalid-arg-bindings error message. --- .../datamodels/task_arg_binding.py | 46 ++++++------------- .../execution_api/routes/task_instances.py | 5 +- providers/standard/README.rst | 1 - providers/standard/docs/index.rst | 1 - providers/standard/pyproject.toml | 2 - .../providers/standard/decorators/stub.py | 43 +++++++++++++---- .../unit/standard/decorators/test_stub.py | 14 +++++- uv.lock | 2 - 8 files changed, 62 insertions(+), 52 deletions(-) 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 index 157612f3edb5f..69b6453bb1c8e 100644 --- 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 @@ -17,9 +17,8 @@ """ 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. +Captured at parse time from the ``@task.stub`` TaskFlow call, stored in the serialized +Dag, and delivered to the lang-SDK runtime via ``TIRunContext.arg_bindings``. """ from __future__ import annotations @@ -32,40 +31,30 @@ from airflow.api_fastapi.core_api.base import BaseModel -# A named alias (like TaskArgBinding below) so every schema keeps one shared, titled -# ArgValueSchema definition; a free-form JSON object rather than a typed model because -# the fragment is whatever pydantic generates from the stub annotation at parse time -# (``anyOf``, ``items``, ``enum``, ``$defs``, ...) and a typed model would silently strip -# any keyword it does not know when the spec is re-serialized along the delivery path. +# A named, titled alias (like TaskArgBinding below) kept as free-form JSON rather than a +# typed model, so unknown JSON-schema keywords survive re-serialization along the way. ArgValueSchema = TypeAliasType( "ArgValueSchema", Annotated[dict[str, JsonValue], Field(title="ArgValueSchema")] ) -""" -JSON-schema fragment constraining the value a stub-task argument binds to. - -Generated by pydantic from the stub function's parameter annotation and carried verbatim; -consumers validate the keywords they understand and ignore the rest, per JSON-schema -semantics, so newer producers never break an older core or runtime. -""" +"""JSON-schema fragment constraining the value a stub-task argument binds to; generated +by pydantic from the stub annotation, carried verbatim, unknown keywords ignored.""" class XComArgBinding(BaseModel): """One positional stub-task argument pulled from an upstream task's XCom.""" - # No default on the discriminator: a default drops ``kind`` from ``required`` in the - # OpenAPI schema, and the generated task-sdk client then types it ``Literal | None``, - # which pydantic rejects as a tagged-union discriminator. + # No default: it would drop ``kind`` from ``required``, and the generated task-sdk + # client then types it ``Literal | None``, invalid as a tagged-union discriminator. kind: Literal["xcom"] name: str """The stub function's parameter name this binding fills, in declaration order.""" value_schema: ArgValueSchema | None = None - """JSON-schema fragment from the stub function's annotation; runtimes validate the - bound value against it. Omitted when the annotation gives no constraint.""" + """Schema fragment from the stub function's annotation; omitted when unconstrained.""" task_id: str - """Upstream task id to pull the XCom from; the ``return_value`` XCom is always the one pulled.""" + """Upstream task id whose ``return_value`` XCom is pulled.""" class LiteralArgBinding(BaseModel): @@ -78,8 +67,7 @@ class LiteralArgBinding(BaseModel): """The stub function's parameter name this binding fills, in declaration order.""" value_schema: ArgValueSchema | None = None - """JSON-schema fragment from the stub function's annotation; runtimes validate the - bound value against it. Omitted when the annotation gives no constraint.""" + """Schema fragment from the stub function's annotation; omitted when unconstrained.""" value: JsonValue | None = None """The literal value from the Dag file.""" @@ -88,10 +76,8 @@ class LiteralArgBinding(BaseModel): """True when the value was filled from the stub signature's default rather than passed in the call.""" -# A named alias (TypeAliasType, not a bare Annotated) so the union lands in every -# schema as its own named definition instead of an anonymous field-title-derived one. -# The explicit title lets the supervisor-schema dump merge this def with the -# task-sdk-generated twin (its core/SDK dedup keys on titles). +# A named alias with an explicit title so the union lands in every schema as its own +# named definition, which the supervisor-schema dump dedups with its task-sdk twin by title. TaskArgBinding = TypeAliasType( "TaskArgBinding", Annotated[XComArgBinding | LiteralArgBinding, Field(discriminator="kind", title="TaskArgBinding")], @@ -102,10 +88,8 @@ class LiteralArgBinding(BaseModel): @cache def get_arg_bindings_adapter() -> TypeAdapter[list[TaskArgBinding]]: """ - Validate serialized arg-binding dicts into the kind-discriminated ``TaskArgBinding`` union. + Build (lazily, then cache) the adapter validating serialized dicts into ``TaskArgBinding``. - Constructed lazily on first use (then cached): only the stub-task path in the - execution API ever needs the adapter, and most workloads are not stub operators, - so regular task runs never pay for building it. + Only the stub-task path in the execution API needs it, so regular runs never pay for it. """ return TypeAdapter(list[TaskArgBinding]) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index b1fca0aaaa51d..02055dcba6c42 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 @@ -350,10 +350,7 @@ def ti_run( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail={ "reason": "invalid_arg_bindings", - "message": ( - "The serialized TaskFlow arg spec for this stub task is not valid on " - "this Airflow version; it may come from a newer providers release." - ), + "message": "The serialized TaskFlow arg spec for this stub task is not valid.", }, ) diff --git a/providers/standard/README.rst b/providers/standard/README.rst index ce275c510f528..839b70f10786e 100644 --- a/providers/standard/README.rst +++ b/providers/standard/README.rst @@ -55,7 +55,6 @@ PIP package Version required ========================================== ================== ``apache-airflow`` ``>=2.11.0`` ``apache-airflow-providers-common-compat`` ``>=1.14.1`` -``pydantic`` ``>=2.11.0`` ========================================== ================== Optional cross provider package dependencies diff --git a/providers/standard/docs/index.rst b/providers/standard/docs/index.rst index 0df86ce95290c..1dc9dcdf2d7dd 100644 --- a/providers/standard/docs/index.rst +++ b/providers/standard/docs/index.rst @@ -91,7 +91,6 @@ PIP package Version required ========================================== ================== ``apache-airflow`` ``>=2.11.0`` ``apache-airflow-providers-common-compat`` ``>=1.14.1`` -``pydantic`` ``>=2.11.0`` ========================================== ================== Optional cross provider package dependencies diff --git a/providers/standard/pyproject.toml b/providers/standard/pyproject.toml index fdc8379b1c9ba..373d16bba6c4e 100644 --- a/providers/standard/pyproject.toml +++ b/providers/standard/pyproject.toml @@ -61,8 +61,6 @@ requires-python = ">=3.10" dependencies = [ "apache-airflow>=2.11.0", "apache-airflow-providers-common-compat>=1.14.1", # use next version - # The stub decorator generates arg-binding value schemas from parameter annotations - "pydantic>=2.11.0", ] # The optional dependencies should be modified in place in the generated file diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index 19f8e83563d4d..9efdd8c9add82 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -18,8 +18,10 @@ from __future__ import annotations import ast +import datetime import inspect import json +import types import typing from collections.abc import Callable, Collection, Mapping from typing import TYPE_CHECKING, Any @@ -56,18 +58,39 @@ def float_schema(self, schema): return {**super().float_schema(schema), "format": "double"} +_TEMPORAL_BASES = (datetime.datetime, datetime.date, datetime.time, datetime.timedelta) + + +def _normalize_temporal_annotation(annotation: Any) -> Any: + """ + Map temporal subclasses (e.g. ``pendulum.DateTime``) to their stdlib base. + + Applied recursively through unions and containers, since pydantic only generates + schemas for the stdlib temporal types themselves. + """ + if isinstance(annotation, type): + return next((base for base in _TEMPORAL_BASES if issubclass(annotation, base)), annotation) + origin = typing.get_origin(annotation) + args = typing.get_args(annotation) + if origin is None or not args: + return annotation + normalized = tuple(_normalize_temporal_annotation(arg) for arg in args) + if normalized == args: + return annotation + if origin in (typing.Union, types.UnionType): + return typing.Union[normalized] # noqa: UP007 -- runtime construction from a tuple + return origin[normalized] + + def _infer_value_schema(annotation: Any) -> dict[str, Any] | None: """ Build the JSON-schema fragment for one stub parameter annotation, via pydantic. - Whatever ``pydantic.TypeAdapter(annotation).json_schema()`` produces is shipped - verbatim (``anyOf`` for unions, ``items``/``additionalProperties`` for parameterized - containers, ``enum`` for Literals, ...), so the fragment's exact shape follows the - pydantic version active at parse time and runtimes must treat it as open-vocabulary - JSON schema. Returns ``None`` when the annotation constrains nothing (missing, - ``Any``, bare ``None``) or pydantic cannot generate a schema for it (arbitrary - classes, including anywhere inside a union); the binding then omits ``value_schema`` - and the foreign runtime falls back to a decode-only check. + The pydantic-generated schema ships verbatim, so runtimes must treat it as + open-vocabulary JSON schema. Returns ``None`` when the annotation constrains nothing + (missing, ``Any``, bare ``None``) or pydantic cannot generate a schema for it; the + binding then omits ``value_schema`` and the foreign runtime falls back to a + decode-only check. """ if annotation is inspect.Parameter.empty or annotation is None or annotation is Any: return None @@ -76,7 +99,9 @@ def _infer_value_schema(annotation: Any) -> dict[str, Any] | None: # that can only ever be None constrains nothing worth shipping. return None try: - schema = TypeAdapter(annotation).json_schema(schema_generator=_ValueSchemaGenerator) + schema = TypeAdapter(_normalize_temporal_annotation(annotation)).json_schema( + schema_generator=_ValueSchemaGenerator + ) except PydanticSchemaGenerationError: return None return schema or None diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index cfd85d1f31f48..0add1d0208c1d 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -321,10 +321,20 @@ def group(n): pytest.param(type(None), None, id="nonetype"), pytest.param( pendulum.DateTime, - None, + {"type": "string", "format": "date-time"}, id="pendulum-datetime", - # pydantic has no schema for arbitrary datetime subclasses; decode-only fallback. ), + pytest.param( + pendulum.DateTime | None, + {"anyOf": [{"type": "string", "format": "date-time"}, {"type": "null"}]}, + id="optional-pendulum-datetime", + ), + pytest.param( + list[pendulum.DateTime], + {"type": "array", "items": {"type": "string", "format": "date-time"}}, + id="list-pendulum-datetime", + ), + pytest.param(pendulum.Duration, {"type": "string", "format": "duration"}, id="pendulum-duration"), pytest.param( typing.Optional[str], # noqa: UP045 -- legacy form on purpose {"anyOf": [{"type": "string"}, {"type": "null"}]}, diff --git a/uv.lock b/uv.lock index e411e55b60485..7fbb52a5907a5 100644 --- a/uv.lock +++ b/uv.lock @@ -7909,7 +7909,6 @@ source = { editable = "providers/standard" } dependencies = [ { name = "apache-airflow" }, { name = "apache-airflow-providers-common-compat" }, - { name = "pydantic" }, ] [package.optional-dependencies] @@ -7935,7 +7934,6 @@ requires-dist = [ { name = "apache-airflow", editable = "." }, { name = "apache-airflow-providers-common-compat", editable = "providers/common/compat" }, { name = "apache-airflow-providers-openlineage", marker = "extra == 'openlineage'", editable = "providers/openlineage" }, - { name = "pydantic", specifier = ">=2.11.0" }, ] provides-extras = ["openlineage"] From b519d20ae974cad41c3063149fc78d2c0cc962db Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 24 Jul 2026 08:34:49 +0000 Subject: [PATCH 24/40] Keep @task.stub working when pydantic is missing or cannot schema an annotation Airflow 2.x base installs do not ship pydantic (it is an optional extra there), so after dropping the standard provider's explicit pydantic dependency, importing the stub decorator would crash Dag parsing on such installs. Value schemas now degrade to the decode-only fallback the wire contract already supports when pydantic is absent. Annotations pydantic can validate but not schema-ify (e.g. Callable) raise PydanticInvalidForJsonSchema, which escaped the existing handler and crashed Dag parsing; it is now caught alongside PydanticSchemaGenerationError. Temporal normalization now runs only as a retry after direct schema generation fails, so temporal subclasses that carry their own pydantic schema keep it. --- .../providers/standard/decorators/stub.py | 32 +++++++++++++------ .../unit/standard/decorators/test_stub.py | 12 +++++++ 2 files changed, 35 insertions(+), 9 deletions(-) diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index 9efdd8c9add82..374e7ffb48859 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -26,8 +26,15 @@ from collections.abc import Callable, Collection, Mapping from typing import TYPE_CHECKING, Any -from pydantic import PydanticSchemaGenerationError, TypeAdapter -from pydantic.json_schema import GenerateJsonSchema +try: + from pydantic import PydanticInvalidForJsonSchema, PydanticSchemaGenerationError, TypeAdapter + from pydantic.json_schema import GenerateJsonSchema +except ImportError: + # Airflow 3 always ships pydantic but Airflow 2.x base installs do not; without it, + # stub args carry no value schemas and runtimes keep their decode-only fallback. + GenerateJsonSchema = object # type: ignore[assignment,misc] + TypeAdapter = None # type: ignore[assignment,misc] + PydanticInvalidForJsonSchema = PydanticSchemaGenerationError = None # type: ignore[assignment,misc] from airflow.providers.common.compat.sdk import ( KNOWN_CONTEXT_KEYS, @@ -58,6 +65,7 @@ def float_schema(self, schema): return {**super().float_schema(schema), "format": "double"} +# Most-derived first: datetime subclasses date, so it must be matched before date. _TEMPORAL_BASES = (datetime.datetime, datetime.date, datetime.time, datetime.timedelta) @@ -65,8 +73,8 @@ def _normalize_temporal_annotation(annotation: Any) -> Any: """ Map temporal subclasses (e.g. ``pendulum.DateTime``) to their stdlib base. - Applied recursively through unions and containers, since pydantic only generates - schemas for the stdlib temporal types themselves. + Applied recursively through unions and containers, and only as a retry when direct + schema generation fails, so temporal types carrying their own pydantic schema keep it. """ if isinstance(annotation, type): return next((base for base in _TEMPORAL_BASES if issubclass(annotation, base)), annotation) @@ -92,6 +100,8 @@ def _infer_value_schema(annotation: Any) -> dict[str, Any] | None: binding then omits ``value_schema`` and the foreign runtime falls back to a decode-only check. """ + if TypeAdapter is None: + return None if annotation is inspect.Parameter.empty or annotation is None or annotation is Any: return None if annotation is type(None): @@ -99,11 +109,15 @@ def _infer_value_schema(annotation: Any) -> dict[str, Any] | None: # that can only ever be None constrains nothing worth shipping. return None try: - schema = TypeAdapter(_normalize_temporal_annotation(annotation)).json_schema( - schema_generator=_ValueSchemaGenerator - ) - except PydanticSchemaGenerationError: - return None + schema = TypeAdapter(annotation).json_schema(schema_generator=_ValueSchemaGenerator) + except (PydanticSchemaGenerationError, PydanticInvalidForJsonSchema): + normalized = _normalize_temporal_annotation(annotation) + if normalized is annotation: + return None + try: + schema = TypeAdapter(normalized).json_schema(schema_generator=_ValueSchemaGenerator) + except (PydanticSchemaGenerationError, PydanticInvalidForJsonSchema): + return None return schema or None diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index 0add1d0208c1d..f45474a994bfe 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -20,6 +20,7 @@ import datetime import typing from typing import Any +from unittest import mock import pendulum import pytest @@ -378,7 +379,18 @@ def group(n): id="union-unclassifiable-member", ), pytest.param(contextlib.AbstractContextManager, None, id="custom-class"), + pytest.param(typing.Callable[[int], str], None, id="callable-invalid-for-json-schema"), + pytest.param( + pendulum.DateTime | contextlib.AbstractContextManager, + None, + id="union-temporal-and-unclassifiable", + ), ], ) def test_infer_value_schema(annotation, expected): assert _infer_value_schema(annotation) == expected + + +@mock.patch("airflow.providers.standard.decorators.stub.TypeAdapter", None) +def test_infer_value_schema_without_pydantic(): + assert _infer_value_schema(str) is None From 501dc02fbf5c0e46b584a351c2edc42ae64909df Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 24 Jul 2026 09:27:23 +0000 Subject: [PATCH 25/40] Derive go-pack manifest expectation from the supervisor schema constant The cross-arch pack integration test asserted a hardcoded schema date, so bumping the Go SDK's supervisor schema version (2026-06-16 -> 2026-10-30 on this branch) broke it in CI while the unit tests, which inject the version explicitly, kept passing. The packed binary and the test compile from the same module tree, so referencing execution.SupervisorSchemaVersion asserts the same value the binary embeds and future bumps cannot drift. --- go-sdk/cmd/airflow-go-pack/pack_integration_test.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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..84e9a1045f4e8 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: From b9718e7d9ad1cf3643bff6ab17f8827477ebad8a Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 24 Jul 2026 10:59:39 +0000 Subject: [PATCH 26/40] Support dynamic task mapping on stub tasks Banning .expand() on @task.stub blocked a core dynamic-mapping pattern for foreign-runtime Dags with no workaround. A mapped stub never instantiates at parse time, so instead of a parse-time capture, ti_run now derives the per-map-index arg spec from the serialized expand input, mirroring the task-sdk's DictOfListsExpandInput index decomposition: literal expands resolve to their element server-side, expands over a mapped upstream bind that upstream's XCom row via the new map_index field, and expands over an unmapped upstream's output carry the new element_index field so the runtime picks the right element of the pulled list. Value schemas stay parse-time-only and are omitted for mapped stubs, falling back to the decode-only contract. The derivation grew into enough business logic that it lives in a new execution_api services package (mirroring core_api's services layout) rather than the routes module, and the generic map-index decomposition sits on SchedulerDictOfListsExpandInput beside its map-length helpers, mirroring where the task-sdk twin keeps the same arithmetic. The supports_expand opt-out this branch added to the task-sdk decorator machinery existed only for the stub ban, so it is reverted. Stubs with arguments inside a mapped task group stay rejected: those instances have no expand input of their own to derive bindings from. --- .../datamodels/task_arg_binding.py | 7 + .../execution_api/routes/task_instances.py | 24 +-- .../execution_api/services/__init__.py | 16 ++ .../execution_api/services/task_instances.py | 128 +++++++++++++++ .../src/airflow/models/expandinput.py | 22 +++ .../versions/head/test_task_instances.py | 154 +++++++++++++++++- .../providers/standard/decorators/stub.py | 17 +- .../unit/standard/decorators/test_stub.py | 11 +- .../airflow/sdk/api/datamodels/_generated.py | 2 + task-sdk/src/airflow/sdk/bases/decorator.py | 5 - .../sdk/execution_time/schema/schema.json | 17 ++ .../tests/task_sdk/bases/test_decorator.py | 26 --- 12 files changed, 355 insertions(+), 74 deletions(-) create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/services/__init__.py create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py 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 index 69b6453bb1c8e..f34548e53a589 100644 --- 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 @@ -56,6 +56,13 @@ class XComArgBinding(BaseModel): task_id: str """Upstream task id whose ``return_value`` XCom is pulled.""" + map_index: int = -1 + """Map index of the upstream XCom row to pull; -1 is the unmapped row.""" + + element_index: int | None = None + """When set, the pulled value is a sequence and this binding takes the element at + this index (the stub was expanded over an unmapped upstream's output).""" + class LiteralArgBinding(BaseModel): """One positional stub-task argument carrying an inline literal from the Dag file.""" 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 02055dcba6c42..481c158b69c96 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 @@ -76,6 +76,7 @@ get_team_name_for_ti, require_auth, ) +from airflow.api_fastapi.execution_api.services.task_instances import STUB_TASK_TYPE, get_arg_bindings from airflow.configuration import conf from airflow.exceptions import InvalidPartitionKeyError, TaskNotFound from airflow.models.asset import AssetActive @@ -111,25 +112,6 @@ 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. -# The gate matches the exact class name; a subclass would need its own entry here. -_STUB_TASK_TYPE = "_StubOperator" - - -def _get_arg_bindings( - dag_bag: DagBagDep, dag_version_id: UUID | None, task_id: str, *, session -) -> list | None: - """Extract the stub task's captured TaskFlow arg spec from its Dag version.""" - if dag_version_id is None: - return None - if (dag := dag_bag.get_dag(dag_version_id, session=session)) is None: - return None - if (task := dag.task_dict.get(task_id)) is None: - return None - return getattr(task, "_arg_bindings", None) - @ti_id_router.patch( "/{task_instance_id}/run", @@ -334,9 +316,7 @@ def ti_run( # 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(dag_bag, ti.dag_version_id, ti.task_id, session=session) - ): + if ti.operator == STUB_TASK_TYPE and (arg_bindings := get_arg_bindings(dag_bag, ti, session=session)): try: context.arg_bindings = get_arg_bindings_adapter().validate_python(arg_bindings) except ValidationError: diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/services/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/services/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/services/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py new file mode 100644 index 0000000000000..26b4e192a4e6c --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py @@ -0,0 +1,128 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Business logic backing the task-instance execution routes.""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Any, NoReturn + +from fastapi import HTTPException, status + +from airflow.models.expandinput import NotFullyPopulated, SchedulerDictOfListsExpandInput +from airflow.serialization.definitions.xcom_arg import SchedulerPlainXComArg, SchedulerXComArg + +if TYPE_CHECKING: + from sqlalchemy.orm import Session + + from airflow.models.dagbag import DBDagBag + +# Task type recorded on the TI row (``TaskInstance.operator``) for +# ``airflow.providers.standard.decorators.stub._StubOperator``. Used to gate the +# serialized-Dag lookup for ``arg_bindings`` so regular tasks never pay for it. +# The gate matches the exact class name; a subclass would need its own entry here. +STUB_TASK_TYPE = "_StubOperator" + + +def get_arg_bindings(dag_bag: DBDagBag, ti: Any, *, session: Session) -> list | None: + """Extract or derive the stub task's TaskFlow arg spec from its Dag version.""" + if ti.dag_version_id is None: + return None + if (dag := dag_bag.get_dag(ti.dag_version_id, session=session)) is None: + return None + if (task := dag.task_dict.get(ti.task_id)) is None: + return None + if task.is_mapped: + return _resolve_mapped_stub_arg_bindings(task, ti, session=session) + return getattr(task, "_arg_bindings", None) + + +def _unsupported_arg_bindings(detail: str) -> NoReturn: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={ + "reason": "invalid_arg_bindings", + "message": f"The stub task's TaskFlow arguments cannot be delivered: {detail}.", + }, + ) + + +def _resolve_mapped_stub_arg_bindings(task: Any, ti: Any, *, session: Session) -> list[dict[str, Any]]: + """ + Build the per-map-index arg spec for a mapped (``.expand()``) stub task. + + A mapped stub never instantiates at parse time, so no spec is captured in the + serialized Dag; it is derived here from the serialized expand input instead, with + the map-index decomposition delegated to + ``SchedulerDictOfListsExpandInput.resolve_expansion_sub_indexes``. + Value schemas come from the stub function's annotations, which are not available + server-side, so mapped bindings omit them (runtimes fall back to decode-only checks). + """ + expand_input = task._get_specified_expand_input() + if not isinstance(expand_input, SchedulerDictOfListsExpandInput): + _unsupported_arg_bindings("expand_kwargs() is not supported on stub tasks") + if ti.map_index < 0: + _unsupported_arg_bindings("the task instance has not been expanded to a map index") + + expand_value = expand_input.value + try: + sub_indexes = expand_input.resolve_expansion_sub_indexes(ti.map_index, ti.run_id, session=session) + except NotFullyPopulated as e: + _unsupported_arg_bindings(f"upstream map lengths are not yet known for {sorted(e.missing)}") + + spec = [ + _bind_mapped_stub_arg(name, value, sub_index=None) + for name, value in (task.partial_kwargs.get("op_kwargs") or {}).items() + ] + spec += [ + _bind_mapped_stub_arg(name, value, sub_index=sub_indexes[name]) + for name, value in expand_value.items() + ] + return spec + + +def _bind_mapped_stub_arg(name: str, value: Any, *, sub_index: int | None) -> dict[str, Any]: + """Build one arg-binding dict; ``sub_index`` is set for expanded kwargs, None for partial ones.""" + if isinstance(value, SchedulerPlainXComArg): + if value.key != "return_value": + _unsupported_arg_bindings(f"parameter {name!r} references the XCom key {value.key!r}") + entry: dict[str, Any] = {"name": name, "kind": "xcom", "task_id": value.operator.task_id} + if sub_index is not None: + if value.operator.is_mapped: + entry["map_index"] = sub_index + else: + entry["element_index"] = sub_index + return entry + if isinstance(value, SchedulerXComArg): + _unsupported_arg_bindings( + f"parameter {name!r} received a {type(value).__name__}; only direct upstream" + " task outputs and literals are supported" + ) + if sub_index is not None: + items = list(value.items()) if isinstance(value, dict) else value + try: + value = items[sub_index] + except (IndexError, KeyError, TypeError): + _unsupported_arg_bindings(f"parameter {name!r} has no element at expansion index {sub_index}") + try: + json.dumps(value, allow_nan=False) + except (TypeError, ValueError): + _unsupported_arg_bindings( + f"parameter {name!r} carries a {type(value).__name__} value, which cannot cross" + " the language boundary" + ) + return {"name": name, "kind": "literal", "value": value} diff --git a/airflow-core/src/airflow/models/expandinput.py b/airflow-core/src/airflow/models/expandinput.py index 0363bae92620a..4b2db92757e59 100644 --- a/airflow-core/src/airflow/models/expandinput.py +++ b/airflow-core/src/airflow/models/expandinput.py @@ -151,6 +151,28 @@ def get_total_map_length(self, run_id: str, *, session: Session) -> int: lengths = self._get_map_lengths(run_id, session=session) return functools.reduce(operator.mul, (lengths[name] for name in self.value), 1) + def resolve_expansion_sub_indexes( + self, map_index: int, run_id: str, *, session: Session + ) -> dict[str, int]: + """ + Decompose a task instance's map index into one index per expanded kwarg. + + Server-side counterpart of the index decomposition in the SDK's + ``DictOfListsExpandInput._expand_mapped_field``: the cross-product of the + expanded kwargs is ordered with the last kwarg varying fastest. A single + expanded kwarg maps one-to-one, skipping the upstream length lookups. + + :raises NotFullyPopulated: if upstream map lengths are not all known yet. + """ + if len(self.value) == 1: + return dict.fromkeys(self.value, map_index) + lengths = self._get_map_lengths(run_id, session=session) + sub_indexes = {} + for key in reversed(self.value): + sub_indexes[key] = map_index % lengths[key] + map_index //= lengths[key] + return sub_indexes + def iter_references(self) -> Iterable[tuple[Operator, str]]: from airflow.models.referencemixin import ReferenceMixin 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 a281291bf84b7..068048b5ab9fd 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 @@ -17,6 +17,7 @@ from __future__ import annotations +import itertools from datetime import datetime from typing import TYPE_CHECKING from unittest import mock @@ -422,7 +423,7 @@ def transform(country: str, extracted: dict, limit: int = 10): ... assert "arg_bindings" not in response.json() @mock.patch( - "airflow.api_fastapi.execution_api.routes.task_instances._get_arg_bindings", + "airflow.api_fastapi.execution_api.routes.task_instances.get_arg_bindings", autospec=True, return_value=[{"name": "country", "kind": "hologram", "value": "uk"}], ) @@ -454,6 +455,157 @@ def transform(country: str): ... assert response.status_code == 500 assert response.json()["detail"]["reason"] == "invalid_arg_bindings" + RUN_PAYLOAD = { + "state": "running", + "hostname": "random-hostname", + "unixname": "random-unixname", + "pid": 100, + "start_date": "2024-09-30T12:00:00Z", + } + + def test_ti_run_resolves_mapped_stub_literal_expand(self, client, dag_maker): + """Expanding a stub over a literal list resolves each map index to its element server-side.""" + with dag_maker("test_mapped_stub_literal", serialized=True): + + @task.stub + def transform(country: str): ... + + transform.expand(country=["uk", "fr", "de"]) + + dr = dag_maker.create_dagrun() + tis = {ti.map_index: ti for ti in dr.get_task_instances()} + assert set(tis) == {0, 1, 2} + for ti in tis.values(): + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + for map_index, country in enumerate(["uk", "fr", "de"]): + response = client.patch( + f"/execution/task-instances/{tis[map_index].id}/run", json=self.RUN_PAYLOAD + ) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "country", "kind": "literal", "value": country} + ] + + def test_ti_run_resolves_mapped_stub_over_unmapped_upstream(self, client, dag_maker): + """Expanding over an unmapped upstream's output binds the whole XCom plus an element index.""" + with dag_maker("test_mapped_stub_unmapped_upstream", serialized=True): + + @task.stub + def extract(): ... + + @task.stub + def transform(extracted: dict): ... + + transform.expand(extracted=extract()) + + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance("transform") + ti.map_index = 1 + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "extracted", "kind": "xcom", "task_id": "extract", "element_index": 1} + ] + + def test_ti_run_resolves_mapped_stub_over_mapped_upstream(self, client, dag_maker): + """Expanding over a mapped upstream binds the upstream XCom row at the same map index.""" + with dag_maker("test_mapped_stub_mapped_upstream", serialized=True): + + @task.stub + def seed(n: int): ... + + @task.stub + def transform(extracted: dict): ... + + transform.expand(extracted=seed.expand(n=[1, 2])) + + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance("transform") + ti.map_index = 1 + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "extracted", "kind": "xcom", "task_id": "seed", "map_index": 1} + ] + + def test_ti_run_decomposes_multi_kwarg_mapped_stub(self, client, dag_maker): + """Cross-product expansion decomposes the map index per kwarg like the task-sdk does.""" + with dag_maker("test_mapped_stub_multi_kwarg", serialized=True): + + @task.stub + def combine(a: str, b: int): ... + + combine.expand(a=["x", "y"], b=[1, 2, 3]) + + dr = dag_maker.create_dagrun() + tis = {ti.map_index: ti for ti in dr.get_task_instances()} + assert set(tis) == set(range(6)) + for ti in tis.values(): + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + for map_index, (a, b) in enumerate(itertools.product(["x", "y"], [1, 2, 3])): + response = client.patch( + f"/execution/task-instances/{tis[map_index].id}/run", json=self.RUN_PAYLOAD + ) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "a", "kind": "literal", "value": a}, + {"name": "b", "kind": "literal", "value": b}, + ] + + def test_ti_run_binds_partial_kwargs_of_mapped_stub(self, client, dag_maker): + """partial() kwargs bind like an unmapped TaskFlow call alongside the expanded ones.""" + with dag_maker("test_mapped_stub_partial", serialized=True): + + @task.stub + def transform(country: str, extracted: dict): ... + + transform.partial(country="uk").expand(extracted=[{"a": 1}, {"b": 2}]) + + dr = dag_maker.create_dagrun() + tis = {ti.map_index: ti for ti in dr.get_task_instances()} + for ti in tis.values(): + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + for map_index, extracted in enumerate([{"a": 1}, {"b": 2}]): + response = client.patch( + f"/execution/task-instances/{tis[map_index].id}/run", json=self.RUN_PAYLOAD + ) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "country", "kind": "literal", "value": "uk"}, + {"name": "extracted", "kind": "literal", "value": extracted}, + ] + + def test_ti_run_rejects_expand_kwargs_on_stub(self, client, dag_maker): + """expand_kwargs() has no per-parameter spec to derive, so delivery fails structurally.""" + with dag_maker("test_mapped_stub_expand_kwargs", serialized=True): + + @task.stub + def transform(country: str): ... + + transform.expand_kwargs([{"country": "uk"}]) + + dr = dag_maker.create_dagrun() + (ti,) = dr.get_task_instances() + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 500 + assert response.json()["detail"]["reason"] == "invalid_arg_bindings" + assert "expand_kwargs" in response.json()["detail"]["message"] + def test_arg_bindings_adapter_rejects_unknown_kind(self): """The discriminated union refuses serialized specs with an unrecognised kind.""" from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import get_arg_bindings_adapter diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index 374e7ffb48859..e0a1e26dd7298 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -219,16 +219,6 @@ def get_annotation_for(name: str, param: inspect.Parameter) -> Any: class _StubOperator(DecoratedOperator): custom_operator_name: str = "@task.stub" - # A mapped stub's arg *types* are uniform across map indexes (same function), but its - # arg *values* only resolve per map index at runtime, while the spec below is captured - # at parse time and the wire contract has no mapped-binding kind yet -- so .expand() - # is rejected rather than shipped with a wrong or empty spec. The task-sdk decorator - # machinery rejects direct .expand() at parse time for operator classes that opt out - # on Airflow >= 3.4 (older cores cannot enforce it, and never serialize a spec for the - # mapped stub); stubs called with arguments inside a mapped task group are rejected in - # __init__ below. - supports_expand: bool = False - def __init__( self, *, @@ -275,9 +265,10 @@ def __init__( # 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) - # supports_expand only blocks direct .expand() on the stub itself; a mapped task - # group still creates per-map-index instances of every task inside it, whose arg - # values resolve per map index at runtime -- after this parse-time capture. + # Direct .expand() on the stub needs no parse-time spec (ti_run derives per-map-index + # bindings from the serialized expand input), but a mapped task group creates + # per-map-index instances of the tasks inside it with no expand input of their own, + # so their arg values are unresolvable both here and server-side. in_mapped_group = getattr(self, "get_closest_mapped_task_group", lambda: None)() is not None if self._arg_bindings is not None and in_mapped_group: raise ValueError( diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index f45474a994bfe..fbec38ad104c5 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -28,7 +28,7 @@ from airflow.providers.common.compat.sdk import DAG, task_group from airflow.providers.standard.decorators.stub import _infer_value_schema, stub -from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS, AIRFLOW_V_3_4_PLUS +from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS def fn_ellipsis(): ... @@ -250,13 +250,10 @@ def test_arg_bindings_survive_dag_serialization_round_trip(self): }, ] - @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): + def test_expand_builds_mapped_stub_without_parse_time_bindings(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=[{}, {}]) + result = stub(fn_transform).expand(country=["uk", "fr"], extracted=[{}, {}]) + assert result.operator.is_mapped def test_stub_with_args_inside_mapped_task_group_rejected(self): @task_group diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 94a56701d5294..6dd7587870b41 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -735,6 +735,8 @@ class XComArgBinding(BaseModel): name: Annotated[str, Field(title="Name")] value_schema: ArgValueSchema | None = None task_id: Annotated[str, Field(title="Task Id")] + map_index: Annotated[int | None, Field(title="Map Index")] = -1 + element_index: Annotated[int | None, Field(title="Element Index")] = None class AssetEventDagRunReference(BaseModel): diff --git a/task-sdk/src/airflow/sdk/bases/decorator.py b/task-sdk/src/airflow/sdk/bases/decorator.py index fc3a01cf95a53..e45778725cc16 100644 --- a/task-sdk/src/airflow/sdk/bases/decorator.py +++ b/task-sdk/src/airflow/sdk/bases/decorator.py @@ -588,11 +588,6 @@ 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 2467b5f7ea990..b671959c50a00 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -4661,6 +4661,23 @@ "task_id": { "title": "Task Id", "type": "string" + }, + "map_index": { + "default": -1, + "title": "Map Index", + "type": "integer" + }, + "element_index": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Element Index" } }, "required": [ diff --git a/task-sdk/tests/task_sdk/bases/test_decorator.py b/task-sdk/tests/task_sdk/bases/test_decorator.py index 0578467aa26c1..860eb6e7b3312 100644 --- a/task-sdk/tests/task_sdk/bases/test_decorator.py +++ b/task-sdk/tests/task_sdk/bases/test_decorator.py @@ -383,29 +383,3 @@ def sync_task_fn(): return 42 assert not is_async_callable(sync_task_fn) - - -class DummyNoExpandDecoratedOperator(DecoratedOperator): - custom_operator_name = "@task.dummy_no_expand" - - supports_expand = False - - -class TestSupportsExpandOptOut: - """An operator class that sets ``supports_expand = False`` rejects dynamic task mapping at parse time.""" - - @pytest.fixture - def no_expand_task(self): - from airflow.sdk.bases.decorator import task_decorator_factory - - def fn(a): ... - - return task_decorator_factory(fn, decorated_operator_class=DummyNoExpandDecoratedOperator) - - def test_expand_rejected_with_operator_name(self, no_expand_task): - with pytest.raises(TypeError, match="@task.dummy_no_expand tasks do not support dynamic task"): - no_expand_task.expand(a=[1, 2]) - - def test_expand_kwargs_rejected(self, no_expand_task): - with pytest.raises(TypeError, match="do not support dynamic task mapping"): - no_expand_task.expand_kwargs([{"a": 1}]) From 980f1c78c4bdff9a9272e6c8bd512809ddaed7e4 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Sun, 26 Jul 2026 07:27:30 +0000 Subject: [PATCH 27/40] Keep stub temporal annotation normalization working on Python 3.10 On Python 3.10, isinstance(list[X], type) is True and issubclass on the alias silently consults the origin, so the plain-class branch swallowed parametrized generics before the origin/args reconstruction could rewrite their arguments; list[pendulum.DateTime] then degraded to no value schema at all. Python 3.11+ returns False there, which is why the regression only surfaced on the 3.10 CI jobs. Detecting parametrized generics first restores the normalization on every supported version. --- .../providers/standard/decorators/stub.py | 23 +++++++++++-------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index e0a1e26dd7298..9d63d30cc0dc1 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -76,18 +76,21 @@ def _normalize_temporal_annotation(annotation: Any) -> Any: Applied recursively through unions and containers, and only as a retry when direct schema generation fails, so temporal types carrying their own pydantic schema keep it. """ - if isinstance(annotation, type): - return next((base for base in _TEMPORAL_BASES if issubclass(annotation, base)), annotation) + # Parametrized generics must be detected before the plain-class branch: on Python + # 3.10, isinstance(list[X], type) is True and issubclass silently consults the + # origin, so the class branch would return list[X] unnormalized. origin = typing.get_origin(annotation) args = typing.get_args(annotation) - if origin is None or not args: - return annotation - normalized = tuple(_normalize_temporal_annotation(arg) for arg in args) - if normalized == args: - return annotation - if origin in (typing.Union, types.UnionType): - return typing.Union[normalized] # noqa: UP007 -- runtime construction from a tuple - return origin[normalized] + if origin is not None and args: + normalized = tuple(_normalize_temporal_annotation(arg) for arg in args) + if normalized == args: + return annotation + if origin in (typing.Union, types.UnionType): + return typing.Union[normalized] # noqa: UP007 -- runtime construction from a tuple + return origin[normalized] + if isinstance(annotation, type): + return next((base for base in _TEMPORAL_BASES if issubclass(annotation, base)), annotation) + return annotation def _infer_value_schema(annotation: Any) -> dict[str, Any] | None: From e73be7d97e3186dc18292e670c8e934da9f41f3c Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Sun, 26 Jul 2026 07:27:48 +0000 Subject: [PATCH 28/40] Regenerate ts-sdk supervisor types for the mapped stub arg bindings The map_index and element_index fields added to XComArgBinding in the supervisor schema must land together with the generated TypeScript output, which the check-ts-sdk-supervisor-schema static check enforces by regenerating and diffing the file. --- ts-sdk/src/generated/supervisor.ts | 32 +++++++++++++++++------------- 1 file changed, 18 insertions(+), 14 deletions(-) diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 757429e7f29fd..2bf5075dad997 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -254,6 +254,8 @@ export type TaskArgBinding = XComArgBinding | LiteralArgBinding; export type Kind = "xcom"; export type Name8 = string; export type TaskId1 = string; +export type MapIndex1 = number; +export type ElementIndex = number | null; export type Kind1 = "literal"; export type Name9 = string; export type FromDefault = boolean; @@ -337,7 +339,7 @@ export type Key5 = string; export type DagId6 = string; export type RunId5 = string; export type TaskId2 = string; -export type MapIndex1 = number | null; +export type MapIndex2 = number | null; export type Type25 = "DeleteXCom"; /** * Error types used in the API client. @@ -435,10 +437,10 @@ export type Type41 = "GetPreviousDagRun"; export type DagId12 = string; export type TaskId3 = string; export type LogicalDate4 = string | null; -export type MapIndex2 = number; +export type MapIndex3 = number; export type Type42 = "GetPreviousTI"; export type DagId13 = string; -export type MapIndex3 = number | null; +export type MapIndex4 = number | null; export type TaskIds = string[] | null; export type TaskGroupId = string | null; export type LogicalDates1 = string[] | null; @@ -455,7 +457,7 @@ export type TiId6 = string; export type Key8 = string; export type Type46 = "GetTaskStateStore"; export type DagId15 = string; -export type MapIndex4 = number | null; +export type MapIndex5 = number | null; export type TaskIds1 = string[] | null; export type TaskGroupId1 = string | null; export type LogicalDates2 = string[] | null; @@ -471,7 +473,7 @@ export type Key10 = string; export type DagId16 = string; export type RunId9 = string; export type TaskId4 = string; -export type MapIndex5 = number | null; +export type MapIndex6 = number | null; export type IncludePriorDates = boolean; export type Type50 = "GetXCom"; export type Key11 = string; @@ -528,7 +530,7 @@ export type StartDate5 = string | null; export type EndDate4 = string | null; export type State4 = string | null; export type TryNumber2 = number; -export type MapIndex6 = number | null; +export type MapIndex7 = number | null; export type Duration = number | null; export type Type60 = "PreviousTIResult"; export type Key14 = string; @@ -565,7 +567,7 @@ export type Key18 = string; export type DagId21 = string; export type RunId14 = string; export type TaskId9 = string; -export type MapIndex7 = number | null; +export type MapIndex8 = number | null; export type DagResult1 = boolean; export type MappedLength = number | null; export type Type71 = "SetXCom"; @@ -1054,6 +1056,8 @@ export interface XComArgBinding { name: Name8; value_schema?: ArgValueSchema | null; task_id: TaskId1; + map_index?: MapIndex1; + element_index?: ElementIndex; } /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema @@ -1233,7 +1237,7 @@ export interface DeleteXCom { dag_id: DagId6; run_id: RunId5; task_id: TaskId2; - map_index?: MapIndex1; + map_index?: MapIndex2; type?: Type25; } /** @@ -1401,7 +1405,7 @@ export interface GetPreviousTI { dag_id: DagId12; task_id: TaskId3; logical_date?: LogicalDate4; - map_index?: MapIndex2; + map_index?: MapIndex3; state?: TaskInstanceState | null; type?: Type42; } @@ -1411,7 +1415,7 @@ export interface GetPreviousTI { */ export interface GetTICount { dag_id: DagId13; - map_index?: MapIndex3; + map_index?: MapIndex4; task_ids?: TaskIds; task_group_id?: TaskGroupId; logical_dates?: LogicalDates1; @@ -1452,7 +1456,7 @@ export interface GetTaskStateStore { */ export interface GetTaskStates { dag_id: DagId15; - map_index?: MapIndex4; + map_index?: MapIndex5; task_ids?: TaskIds1; task_group_id?: TaskGroupId1; logical_dates?: LogicalDates2; @@ -1486,7 +1490,7 @@ export interface GetXCom { dag_id: DagId16; run_id: RunId9; task_id: TaskId4; - map_index?: MapIndex5; + map_index?: MapIndex6; include_prior_dates?: IncludePriorDates; type?: Type50; } @@ -1612,7 +1616,7 @@ export interface PreviousTIResponse { end_date?: EndDate4; state?: State4; try_number: TryNumber2; - map_index?: MapIndex6; + map_index?: MapIndex7; duration?: Duration; } /** @@ -1740,7 +1744,7 @@ export interface SetXCom { dag_id: DagId21; run_id: RunId14; task_id: TaskId9; - map_index?: MapIndex7; + map_index?: MapIndex8; dag_result?: DagResult1; mapped_length?: MappedLength; type?: Type71; From 31d9d6988ad5fa4216471f4735f722e32cc28798 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Sun, 26 Jul 2026 09:59:20 +0000 Subject: [PATCH 29/40] Assert mapped stub expansion via attributes Airflow 2.x also has The mapped-stub decorator test asserted MappedOperator.is_mapped, which does not exist on Airflow 2.x, so the provider compat suite failed on 2.11 with an AttributeError. Asserting on op_kwargs_expand_input and partial_kwargs keeps the test meaningful on every supported Airflow version and additionally pins down what "no parse-time bindings" means. --- .../standard/tests/unit/standard/decorators/test_stub.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index fbec38ad104c5..74d932148dc74 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -253,7 +253,13 @@ def test_arg_bindings_survive_dag_serialization_round_trip(self): def test_expand_builds_mapped_stub_without_parse_time_bindings(self): with DAG(dag_id="d"): result = stub(fn_transform).expand(country=["uk", "fr"], extracted=[{}, {}]) - assert result.operator.is_mapped + # op_kwargs_expand_input/partial_kwargs (not is_mapped) so the assertions also + # hold on the Airflow 2.x MappedOperator, which the provider still supports. + assert result.operator.op_kwargs_expand_input.value == { + "country": ["uk", "fr"], + "extracted": [{}, {}], + } + assert "_arg_bindings" not in result.operator.partial_kwargs def test_stub_with_args_inside_mapped_task_group_rejected(self): @task_group From e823dd8b5cb54fba9c98de53f4da1365702506a3 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 27 Jul 2026 14:09:00 +0000 Subject: [PATCH 30/40] Type and document the mapped stub arg-binding derivation The server-side derivation for mapped stub tasks typed the task as Any and left its non-obvious decisions undocumented: the expand_kwargs() gate, when NotFullyPopulated can actually fire, the partial()/expand() kwarg partition, and how a runtime consumes element_index. Narrowing to SerializedMappedOperator via the is_mapped() guard lets mypy check the mapped-only attribute access, and the comments capture the reasoning where it applies. Delivering value_schema on mapped bindings stays deferred, now tracked at https://github.com/apache/airflow/issues/70523. The serialization round-trip test asserted an either-or for the argless stub; deserialization never sets _arg_bindings for it, so assert exactly that. --- .../datamodels/task_arg_binding.py | 4 +++- .../execution_api/services/task_instances.py | 18 +++++++++++++++--- .../serialization/test_dag_serialization.py | 4 +--- 3 files changed, 19 insertions(+), 7 deletions(-) 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 index f34548e53a589..982c199868394 100644 --- 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 @@ -61,7 +61,9 @@ class XComArgBinding(BaseModel): element_index: int | None = None """When set, the pulled value is a sequence and this binding takes the element at - this index (the stub was expanded over an unmapped upstream's output).""" + this index (the stub was expanded over an unmapped upstream's output). The lang-SDK + side will GET the single row ``(task_id, map_index=-1)``, decode it, and take + ``value[element_index]``.""" class LiteralArgBinding(BaseModel): diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py index 26b4e192a4e6c..326d3399710f2 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py @@ -24,12 +24,14 @@ from fastapi import HTTPException, status from airflow.models.expandinput import NotFullyPopulated, SchedulerDictOfListsExpandInput +from airflow.serialization.definitions.mappedoperator import is_mapped from airflow.serialization.definitions.xcom_arg import SchedulerPlainXComArg, SchedulerXComArg if TYPE_CHECKING: from sqlalchemy.orm import Session from airflow.models.dagbag import DBDagBag + from airflow.serialization.definitions.mappedoperator import SerializedMappedOperator # Task type recorded on the TI row (``TaskInstance.operator``) for # ``airflow.providers.standard.decorators.stub._StubOperator``. Used to gate the @@ -46,7 +48,7 @@ def get_arg_bindings(dag_bag: DBDagBag, ti: Any, *, session: Session) -> list | return None if (task := dag.task_dict.get(ti.task_id)) is None: return None - if task.is_mapped: + if is_mapped(task): return _resolve_mapped_stub_arg_bindings(task, ti, session=session) return getattr(task, "_arg_bindings", None) @@ -61,7 +63,9 @@ def _unsupported_arg_bindings(detail: str) -> NoReturn: ) -def _resolve_mapped_stub_arg_bindings(task: Any, ti: Any, *, session: Session) -> list[dict[str, Any]]: +def _resolve_mapped_stub_arg_bindings( + task: SerializedMappedOperator, ti: Any, *, session: Session +) -> list[dict[str, Any]]: """ Build the per-map-index arg spec for a mapped (``.expand()``) stub task. @@ -70,9 +74,12 @@ def _resolve_mapped_stub_arg_bindings(task: Any, ti: Any, *, session: Session) - the map-index decomposition delegated to ``SchedulerDictOfListsExpandInput.resolve_expansion_sub_indexes``. Value schemas come from the stub function's annotations, which are not available - server-side, so mapped bindings omit them (runtimes fall back to decode-only checks). + server-side, so mapped bindings omit them (runtimes fall back to decode-only checks); + delivering them is tracked at https://github.com/apache/airflow/issues/70523. """ expand_input = task._get_specified_expand_input() + # TODO: Support the `expand_kwargs` path once https://github.com/apache/airflow/pull/69757 + # is merged and all the Lang-SDKs adapt it. if not isinstance(expand_input, SchedulerDictOfListsExpandInput): _unsupported_arg_bindings("expand_kwargs() is not supported on stub tasks") if ti.map_index < 0: @@ -82,12 +89,16 @@ def _resolve_mapped_stub_arg_bindings(task: Any, ti: Any, *, session: Session) - try: sub_indexes = expand_input.resolve_expansion_sub_indexes(ti.map_index, ti.run_id, session=session) except NotFullyPopulated as e: + # In the happy path this shouldn't happen at all, unless someone clears upstream + # TIs or XComs themselves during the DagRun. _unsupported_arg_bindings(f"upstream map lengths are not yet known for {sorted(e.missing)}") + # partial() kwargs spec = [ _bind_mapped_stub_arg(name, value, sub_index=None) for name, value in (task.partial_kwargs.get("op_kwargs") or {}).items() ] + # expand() kwargs spec += [ _bind_mapped_stub_arg(name, value, sub_index=sub_indexes[name]) for name, value in expand_value.items() @@ -113,6 +124,7 @@ def _bind_mapped_stub_arg(name: str, value: Any, *, sub_index: int | None) -> di " task outputs and literals are supported" ) if sub_index is not None: + # This kwarg was expanded over a literal collection written in the Dag file. items = list(value.items()) if isinstance(value, dict) else value try: value = items[sub_index] diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index e88445d75e5f8..889131a152032 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -3456,9 +3456,7 @@ def transform(country: str, extracted: dict): ... "task_id": "extract", }, ] - assert not hasattr(round_tripped.task_dict["extract"], "_arg_bindings") or ( - round_tripped.task_dict["extract"]._arg_bindings is None - ) + assert not hasattr(round_tripped.task_dict["extract"], "_arg_bindings") def test_handle_v1_serdag(): From 5f26a25161ba4a1340f48e1223b2f90acc6b91d1 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 27 Jul 2026 16:59:11 +0000 Subject: [PATCH 31/40] Reject stub args bound to a mapped upstream's aggregated output A stub argument fed a mapped task's combined output silently bound the unmapped XCom row (map_index=-1), which never exists for a mapped upstream, so the foreign runtime received nothing where Python TaskFlow delivers the aggregated list. The wire contract cannot express "pull all rows", so fail loudly like the other inexpressible constructs: at parse time for the unmapped call path, and server-side for already-serialized Dags. XComArgs inside partial() op_kwargs also deserialize to _XComRef and were never dereferenced, falling past the XComArg branches entirely; a partial() kwarg over an unmapped upstream failed spec validation instead of binding that XCom row. Dereferencing them fixes that and lets the mapped-upstream rejection see the real reference. --- .../execution_api/services/task_instances.py | 21 +++++++-- .../versions/head/test_task_instances.py | 46 +++++++++++++++++++ .../providers/standard/decorators/stub.py | 7 +++ .../unit/standard/decorators/test_stub.py | 8 ++++ 4 files changed, 78 insertions(+), 4 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py index 326d3399710f2..68c885dd14d46 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py @@ -26,11 +26,13 @@ from airflow.models.expandinput import NotFullyPopulated, SchedulerDictOfListsExpandInput from airflow.serialization.definitions.mappedoperator import is_mapped from airflow.serialization.definitions.xcom_arg import SchedulerPlainXComArg, SchedulerXComArg +from airflow.serialization.serialized_objects import _XComRef if TYPE_CHECKING: from sqlalchemy.orm import Session from airflow.models.dagbag import DBDagBag + from airflow.serialization.definitions.dag import SerializedDAG from airflow.serialization.definitions.mappedoperator import SerializedMappedOperator # Task type recorded on the TI row (``TaskInstance.operator``) for @@ -49,7 +51,7 @@ def get_arg_bindings(dag_bag: DBDagBag, ti: Any, *, session: Session) -> list | if (task := dag.task_dict.get(ti.task_id)) is None: return None if is_mapped(task): - return _resolve_mapped_stub_arg_bindings(task, ti, session=session) + return _resolve_mapped_stub_arg_bindings(task, ti, dag=dag, session=session) return getattr(task, "_arg_bindings", None) @@ -64,7 +66,7 @@ def _unsupported_arg_bindings(detail: str) -> NoReturn: def _resolve_mapped_stub_arg_bindings( - task: SerializedMappedOperator, ti: Any, *, session: Session + task: SerializedMappedOperator, ti: Any, *, dag: SerializedDAG, session: Session ) -> list[dict[str, Any]]: """ Build the per-map-index arg spec for a mapped (``.expand()``) stub task. @@ -93,9 +95,13 @@ def _resolve_mapped_stub_arg_bindings( # TIs or XComs themselves during the DagRun. _unsupported_arg_bindings(f"upstream map lengths are not yet known for {sorted(e.missing)}") - # partial() kwargs + # partial() kwargs. XComArgs inside partial() op_kwargs deserialize to _XComRef and + # are never dereferenced (set_task_dag_references only derefs the expand inputs), so + # resolve them here before binding. spec = [ - _bind_mapped_stub_arg(name, value, sub_index=None) + _bind_mapped_stub_arg( + name, value.deref(dag) if isinstance(value, _XComRef) else value, sub_index=None + ) for name, value in (task.partial_kwargs.get("op_kwargs") or {}).items() ] # expand() kwargs @@ -111,6 +117,13 @@ def _bind_mapped_stub_arg(name: str, value: Any, *, sub_index: int | None) -> di if isinstance(value, SchedulerPlainXComArg): if value.key != "return_value": _unsupported_arg_bindings(f"parameter {name!r} references the XCom key {value.key!r}") + if sub_index is None and value.operator.is_mapped: + # A partial() kwarg over a mapped upstream would bind the unmapped XCom row + # (map_index=-1), which never exists; the aggregated output is inexpressible. + _unsupported_arg_bindings( + f"parameter {name!r} references the aggregated output of the mapped task" + f" {value.operator.task_id!r}" + ) entry: dict[str, Any] = {"name": name, "kind": "xcom", "task_id": value.operator.task_id} if sub_index is not None: if value.operator.is_mapped: diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index 0401470b5ccdf..385b14906688b 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 @@ -588,6 +588,52 @@ def transform(country: str, extracted: dict): ... {"name": "extracted", "kind": "literal", "value": extracted}, ] + def test_ti_run_binds_partial_xcom_kwarg_over_unmapped_upstream(self, client, dag_maker): + """A partial() kwarg carrying an unmapped upstream's output binds that XCom for every index.""" + with dag_maker("test_mapped_stub_partial_xcom", serialized=True): + + @task.stub + def extract(): ... + + @task.stub + def transform(extracted: dict, country: str): ... + + transform.partial(extracted=extract()).expand(country=["uk", "fr"]) + + dr = dag_maker.create_dagrun() + ti = next(ti for ti in dr.get_task_instances() if ti.task_id == "transform" and ti.map_index == 1) + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "extracted", "kind": "xcom", "task_id": "extract"}, + {"name": "country", "kind": "literal", "value": "fr"}, + ] + + def test_ti_run_rejects_partial_kwarg_over_mapped_upstream(self, client, dag_maker): + """A partial() kwarg over a mapped upstream would bind the nonexistent unmapped XCom row.""" + with dag_maker("test_mapped_stub_partial_mapped_upstream", serialized=True): + + @task.stub + def seed(n: int): ... + + @task.stub + def transform(extracted: dict, country: str): ... + + transform.partial(extracted=seed.expand(n=[1, 2])).expand(country=["uk", "fr"]) + + dr = dag_maker.create_dagrun() + ti = next(ti for ti in dr.get_task_instances() if ti.task_id == "transform" and ti.map_index == 0) + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 500 + assert response.json()["detail"]["reason"] == "invalid_arg_bindings" + assert "aggregated output" in response.json()["detail"]["message"] + def test_ti_run_rejects_expand_kwargs_on_stub(self, client, dag_maker): """expand_kwargs() has no per-parameter spec to derive, so delivery fails structurally.""" with dag_maker("test_mapped_stub_expand_kwargs", serialized=True): diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index 9d63d30cc0dc1..1a7a451799cd3 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -189,6 +189,13 @@ def get_annotation_for(name: str, param: inspect.Parameter) -> Any: f"{value.key!r}; only an upstream task's return value can cross the language " "boundary -- indexing an output by a custom key is not supported" ) + if value.operator.is_mapped: + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} references the aggregated " + f"output of the mapped task {value.operator.task_id!r}; a foreign runtime " + "pulls single XCom rows, so a mapped upstream's combined output is not " + "supported -- use .expand() on the stub to consume it per element" + ) xcom_entry: dict[str, Any] = {"name": name, "kind": "xcom", "task_id": value.operator.task_id} if value_schema is not None: xcom_entry["value_schema"] = value_schema diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index 74d932148dc74..9321cf1861c81 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -221,6 +221,14 @@ def test_mapped_xcom_arg_rejected(self): with pytest.raises(ValueError, match="only direct upstream task outputs"): stub(fn_transform)("uk", extracted.map(lambda v: v)) + def test_mapped_upstream_aggregated_output_rejected(self): + def fn_produce(n: int): ... + + with DAG(dag_id="d"): + vals = stub(fn_produce).expand(n=[1, 2]) + with pytest.raises(ValueError, match="aggregated output of the mapped task"): + stub(fn_transform)("uk", vals) + def test_arg_bindings_survive_dag_serialization_round_trip(self): """The captured spec must survive whichever core serializer the provider runs against.""" try: From e894e43d5a8a5a047446412016f7764841a1f9cd Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 27 Jul 2026 17:08:19 +0000 Subject: [PATCH 32/40] Guard zero-length expansion when deriving mapped stub arg bindings Clearing only the upstream of a queued mapped stub and re-running it to an empty list records a TaskMap length of 0 while the expanded TI still exists; decomposing its map index then divided by zero and surfaced as an opaque catch-all 500. The task-sdk twin of this arithmetic guards mapped lengths below 1, so mirror it and route the failure through the structured invalid_arg_bindings error like every other undeliverable binding. --- .../execution_api/services/task_instances.py | 6 ++-- .../src/airflow/models/expandinput.py | 9 +++-- .../versions/head/test_task_instances.py | 35 +++++++++++++++++++ 3 files changed, 46 insertions(+), 4 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py index 68c885dd14d46..af1863a297f6f 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py @@ -91,9 +91,11 @@ def _resolve_mapped_stub_arg_bindings( try: sub_indexes = expand_input.resolve_expansion_sub_indexes(ti.map_index, ti.run_id, session=session) except NotFullyPopulated as e: - # In the happy path this shouldn't happen at all, unless someone clears upstream - # TIs or XComs themselves during the DagRun. + # Neither this nor the ValueError below can happen on the happy path: both take + # someone clearing upstream TIs or XComs themselves during the DagRun. _unsupported_arg_bindings(f"upstream map lengths are not yet known for {sorted(e.missing)}") + except ValueError as e: + _unsupported_arg_bindings(str(e)) # partial() kwargs. XComArgs inside partial() op_kwargs deserialize to _XComRef and # are never dereferenced (set_task_dag_references only derefs the expand inputs), so diff --git a/airflow-core/src/airflow/models/expandinput.py b/airflow-core/src/airflow/models/expandinput.py index 4b2db92757e59..ecc1b398eb735 100644 --- a/airflow-core/src/airflow/models/expandinput.py +++ b/airflow-core/src/airflow/models/expandinput.py @@ -163,14 +163,19 @@ def resolve_expansion_sub_indexes( expanded kwarg maps one-to-one, skipping the upstream length lookups. :raises NotFullyPopulated: if upstream map lengths are not all known yet. + :raises ValueError: if an expanded kwarg's recorded length is zero, e.g. an + upstream was cleared and re-ran to an empty list after this task instance + was expanded (the SDK twin guards the same case). """ if len(self.value) == 1: return dict.fromkeys(self.value, map_index) lengths = self._get_map_lengths(run_id, session=session) sub_indexes = {} for key in reversed(self.value): - sub_indexes[key] = map_index % lengths[key] - map_index //= lengths[key] + if (length := lengths[key]) < 1: + raise ValueError(f"cannot decompose map index over expanded kwarg {key!r} of length 0") + sub_indexes[key] = map_index % length + map_index //= length return sub_indexes def iter_references(self) -> Iterable[tuple[Operator, str]]: 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 385b14906688b..d1f6c8bb28725 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 @@ -634,6 +634,41 @@ def transform(extracted: dict, country: str): ... assert response.json()["detail"]["reason"] == "invalid_arg_bindings" assert "aggregated output" in response.json()["detail"]["message"] + def test_ti_run_rejects_zero_length_expansion_on_stub(self, client, dag_maker, session): + """An upstream re-run to an empty list after expansion fails structurally, not with a crash.""" + from airflow.models.taskmap import TaskMap + + with dag_maker("test_mapped_stub_zero_length", serialized=True, session=session): + + @task + def seed(): + return [0, 1] + + @task.stub + def combine(a: int, b: int): ... + + combine.expand(a=seed(), b=[1, 2]) + + dr = dag_maker.create_dagrun() + decision = dr.task_instance_scheduling_decisions(session=session) + (seed_ti,) = decision.schedulable_tis + seed_ti.state = TaskInstanceState.SUCCESS + session.add(TaskMap.from_task_instance_xcom(seed_ti, [0, 1])) + session.flush() + + decision = dr.task_instance_scheduling_decisions(session=session) + ti = next(t for t in decision.schedulable_tis if t.map_index == 0) + ti.set_state(State.QUEUED, session=session) + # Simulate the upstream being cleared and re-run to an empty list while this + # expanded TI is still queued. + session.execute(update(TaskMap).where(TaskMap.task_id == "seed").values(length=0, keys=None)) + session.commit() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 500 + assert response.json()["detail"]["reason"] == "invalid_arg_bindings" + assert "length 0" in response.json()["detail"]["message"] + def test_ti_run_rejects_expand_kwargs_on_stub(self, client, dag_maker): """expand_kwargs() has no per-parameter spec to derive, so delivery fails structurally.""" with dag_maker("test_mapped_stub_expand_kwargs", serialized=True): From c584b94f9bac726aa7169df4d1379dd1ce5ff690 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 27 Jul 2026 17:11:40 +0000 Subject: [PATCH 33/40] Skip stub arg-binding work for execution API clients that predate it ti_run derived arg bindings for every stub task regardless of the client's negotiated API version, so a stub Dag using a construct the derivation rejects (e.g. expand_kwargs) went from running with its args ignored to hard-failing with a 500 after a server upgrade -- even for clients whose responses have arg_bindings stripped anyway. The cadwyn migration only pops the field from successful responses; it cannot gate the computation, so consult the negotiated version before deriving. --- .../execution_api/routes/task_instances.py | 23 ++++++++++++++++++- .../v2026_10_30/test_task_instances.py | 20 ++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) 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 7951ce2eb107e..694a4f4729cf5 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -77,6 +77,7 @@ require_auth, ) from airflow.api_fastapi.execution_api.services.task_instances import STUB_TASK_TYPE, get_arg_bindings +from airflow.api_fastapi.execution_api.versions import bundle from airflow.configuration import conf from airflow.exceptions import InvalidPartitionKeyError, TaskNotFound from airflow.models.asset import AssetActive @@ -112,6 +113,22 @@ log = structlog.get_logger(__name__) tracer = trace.get_tracer(__name__) +# The first execution API version whose TIRunContext carries ``arg_bindings``. +ARG_BINDINGS_API_VERSION = "2026-10-30" + + +def _client_supports_arg_bindings() -> bool: + """ + Whether the request's negotiated API version can receive ``arg_bindings``. + + Clients on older versions never see the field (the version migration strips it from + the response), so the derivation -- and the structured failures it raises for + undeliverable mapped-stub specs -- must not run for them: a stub Dag that ran before + arg bindings existed keeps running against those clients. + """ + version = bundle.api_version_var.get(None) + return version is None or str(version) >= ARG_BINDINGS_API_VERSION + @ti_id_router.patch( "/{task_instance_id}/run", @@ -316,7 +333,11 @@ def ti_run( # 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(dag_bag, ti, session=session)): + if ( + ti.operator == STUB_TASK_TYPE + and _client_supports_arg_bindings() + and (arg_bindings := get_arg_bindings(dag_bag, ti, session=session)) + ): try: context.arg_bindings = get_arg_bindings_adapter().validate_python(arg_bindings) except ValidationError: diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py index a4b98bd10206e..dd29526c056c0 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py @@ -79,6 +79,26 @@ def test_old_version_strips_arg_bindings_even_when_set(self, old_ver_client, stu assert response.status_code == 200 assert "arg_bindings" not in response.json() + def test_old_version_skips_undeliverable_arg_bindings_derivation(self, old_ver_client, dag_maker): + """A stub whose bindings cannot be delivered must keep running for clients that never see them.""" + with dag_maker("test_arg_bindings_compat_expand_kwargs", serialized=True): + + @task.stub + def transform(country: str): ... + + transform.expand_kwargs([{"country": "uk"}]) + + dr = dag_maker.create_dagrun() + (ti,) = dr.get_task_instances() + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + # At head this Dag fails ti_run with a structured 500 (expand_kwargs is + # unsupported on stubs); a pre-arg-bindings client keeps the legacy behavior. + response = old_ver_client.patch(f"/execution/task-instances/{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 From 4a321755ed1bae072ba58f42dae73faa8f957cc0 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 27 Jul 2026 17:17:58 +0000 Subject: [PATCH 34/40] Capture mapped stub parameter metadata for the Dag serializer A mapped stub never instantiates at parse time, so ti_run derived its arg bindings blind to the stub signature: the spec came out in call-site dict order while the wire contract promises declaration order (a positional binder like the Go SDK's flat mode then receives swapped values), and parameters filled from signature defaults were silently dropped where the unmapped path ships from_default entries. Declaration order, defaults, and value schemas can only come from the real function, which exists nowhere but the Dag processor, so expose a classmethod the core serializer can call while serializing the mapped operator (wired up in a follow-up commit). Building the metadata also validates the mapping at parse time, so expand_kwargs() on a parameterful stub, partial() kwargs over a mapped upstream, and mappings that do not bind to the signature fail as Dag import errors instead of per-TI 500s at run time. --- .../providers/standard/decorators/stub.py | 210 +++++++++++++----- .../unit/standard/decorators/test_stub.py | 113 +++++++++- 2 files changed, 266 insertions(+), 57 deletions(-) diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index 1a7a451799cd3..edc46b649aa68 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -124,6 +124,77 @@ def _infer_value_schema(annotation: Any) -> dict[str, Any] | None: return schema or None +def _validate_stub_signature(signature: inspect.Signature, task_id: str) -> None: + for param in signature.parameters.values(): + if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD): + raise ValueError( + f"@task.stub task {task_id!r} must declare a fixed number of parameters for the " + f"foreign runtime to bind against; *{param.name} is not supported" + ) + if param.name in KNOWN_CONTEXT_KEYS: + raise ValueError( + f"@task.stub task {task_id!r} parameter {param.name!r} is an Airflow context key; " + "stub signatures declare only data parameters -- the lang-SDK runtime injects its " + "own task context natively (e.g. the Go SDK's sdk.TIRunContext parameter)" + ) + + +def _resolve_param_annotations(python_callable: Callable, signature: inspect.Signature) -> dict[str, Any]: + """Map each parameter to its parse-time-resolvable annotation (``Parameter.empty`` when not).""" + try: + hints = typing.get_type_hints(python_callable) + except (NameError, TypeError): + # Annotations that cannot be resolved at parse time (e.g. names behind + # TYPE_CHECKING with ``from __future__ import annotations``) degrade to "any". + hints = {} + + def resolve(name: str, param: inspect.Parameter) -> Any: + if name in hints: + return hints[name] + if isinstance(param.annotation, str): + return inspect.Parameter.empty + return param.annotation + + return {name: resolve(name, param) for name, param in signature.parameters.items()} + + +def _ensure_json_literal(value: Any, task_id: str, name: str) -> None: + try: + json.dumps(value, allow_nan=False) + except (TypeError, ValueError): + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} received a literal of type " + f"{type(value).__name__} that is not JSON-serializable, so it cannot be passed " + "to the foreign runtime" + ) + + +def _validate_xcom_value(value: Any, task_id: str, name: str, *, allow_mapped_upstream: bool = False) -> bool: + """Validate an XComArg argument, returning True when it is a bindable direct upstream output.""" + if isinstance(value, PlainXComArg): + if value.key != "return_value": + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} references the XCom key " + f"{value.key!r}; only an upstream task's return value can cross the language " + "boundary -- indexing an output by a custom key is not supported" + ) + if value.operator.is_mapped and not allow_mapped_upstream: + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} references the aggregated " + f"output of the mapped task {value.operator.task_id!r}; a foreign runtime " + "pulls single XCom rows, so a mapped upstream's combined output is not " + "supported -- use .expand() on the stub to consume it per element" + ) + return True + if isinstance(value, XComArg): + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} received a " + f"{type(value).__name__}; only direct upstream task outputs can cross the " + "language boundary -- .map()/.zip()/.concat() results are not supported" + ) + return False + + def _build_arg_bindings( python_callable: Callable, op_args: Collection[Any], @@ -146,75 +217,25 @@ def _build_arg_bindings( return None 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)" - ) + _validate_stub_signature(signature, task_id) bound = signature.bind(*op_args, **op_kwargs) explicitly_bound = set(bound.arguments) bound.apply_defaults() - try: - hints = typing.get_type_hints(python_callable) - except (NameError, TypeError): - # Annotations that cannot be resolved at parse time (e.g. names behind - # TYPE_CHECKING with ``from __future__ import annotations``) degrade to "any". - hints = {} - - def get_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 + annotations = _resolve_param_annotations(python_callable, signature) spec: list[dict[str, Any]] = [] - for name, param in signature.parameters.items(): + for name in signature.parameters: value = bound.arguments[name] - value_schema = _infer_value_schema(get_annotation_for(name, param)) - if isinstance(value, PlainXComArg): - if value.key != "return_value": - raise ValueError( - f"@task.stub task {task_id!r} parameter {name!r} references the XCom key " - f"{value.key!r}; only an upstream task's return value can cross the language " - "boundary -- indexing an output by a custom key is not supported" - ) - if value.operator.is_mapped: - raise ValueError( - f"@task.stub task {task_id!r} parameter {name!r} references the aggregated " - f"output of the mapped task {value.operator.task_id!r}; a foreign runtime " - "pulls single XCom rows, so a mapped upstream's combined output is not " - "supported -- use .expand() on the stub to consume it per element" - ) + value_schema = _infer_value_schema(annotations[name]) + if _validate_xcom_value(value, task_id, name): xcom_entry: dict[str, Any] = {"name": name, "kind": "xcom", "task_id": value.operator.task_id} if value_schema is not None: xcom_entry["value_schema"] = value_schema spec.append(xcom_entry) continue - if isinstance(value, XComArg): - raise ValueError( - f"@task.stub task {task_id!r} parameter {name!r} received a " - f"{type(value).__name__}; only direct upstream task outputs can cross the " - "language boundary -- .map()/.zip()/.concat() results are not supported" - ) - try: - json.dumps(value, allow_nan=False) - except (TypeError, ValueError): - raise ValueError( - f"@task.stub task {task_id!r} parameter {name!r} received a literal of type " - f"{type(value).__name__} that is not JSON-serializable, so it cannot be passed " - "to the foreign runtime" - ) + _ensure_json_literal(value, task_id, name) entry: dict[str, Any] = {"name": name, "kind": "literal", "value": value} if value_schema is not None: # Key omission (never ``None``) is the wire contract for "unconstrained": @@ -226,6 +247,64 @@ def get_annotation_for(name: str, param: inspect.Parameter) -> Any: return spec +def _build_mapped_arg_binding_params( + python_callable: Callable, + *, + partial_op_kwargs: Mapping[str, Any], + expand_input: Any, + task_id: str, +) -> list[dict[str, Any]] | None: + """ + Build the ordered per-parameter metadata for a mapped (``.expand()``) stub task. + + Per-map-index values only resolve at run time, so unlike ``_build_arg_bindings`` this + captures what the server-side derivation cannot recover from the serialized Dag alone: + the declaration order the wire contract promises, defaults for parameters no kwarg + covers, and each parameter's value schema. Returns ``None`` for parameterless stubs + (the legacy fan-out shape whose call args were always ignored keeps parsing). + """ + signature = inspect.signature(python_callable) + if not signature.parameters: + return None + _validate_stub_signature(signature, task_id) + if not isinstance(expand_input.value, Mapping): + # expand_kwargs() carries a list (or upstream XCom) of kwarg dicts whose + # parameter names are unknowable at parse time. + raise ValueError( + f"@task.stub task {task_id!r} does not support expand_kwargs(); the parameter " + "binding must be derivable at parse time, so use .expand() with explicit kwargs" + ) + expand_kwargs = expand_input.value + + try: + bound = signature.bind(**{**partial_op_kwargs, **expand_kwargs}) + except TypeError as e: + raise ValueError(f"@task.stub task {task_id!r} TaskFlow mapping does not bind to its signature: {e}") + bound.apply_defaults() + + annotations = _resolve_param_annotations(python_callable, signature) + + params: list[dict[str, Any]] = [] + for name in signature.parameters: + entry: dict[str, Any] = {"name": name} + if (value_schema := _infer_value_schema(annotations[name])) is not None: + entry["value_schema"] = value_schema + if name in expand_kwargs: + # The whole expanded collection ships through XCom/serialization; an upstream + # output is consumed per element, so a mapped upstream is fine here. + if not _validate_xcom_value(expand_kwargs[name], task_id, name, allow_mapped_upstream=True): + _ensure_json_literal(expand_kwargs[name], task_id, name) + elif name in partial_op_kwargs: + if not _validate_xcom_value(partial_op_kwargs[name], task_id, name): + _ensure_json_literal(partial_op_kwargs[name], task_id, name) + else: + default = bound.arguments[name] + _ensure_json_literal(default, task_id, name) + entry["default"] = default + params.append(entry) + return params + + class _StubOperator(DecoratedOperator): custom_operator_name: str = "@task.stub" @@ -292,6 +371,25 @@ def __init__( def get_serialized_fields(cls): return super().get_serialized_fields() | {"_arg_bindings"} + @classmethod + def get_mapped_serialized_fields(cls, mapped_op: Any) -> dict[str, Any]: + """ + Extra serialized fields for the mapped (``.expand()``) form of this operator. + + Called by the core Dag serializer (Airflow 3.4+) while ``python_callable`` is + still the real function; older cores never call it, so mapped stubs there keep + the legacy ignored-args behavior. + """ + params = _build_mapped_arg_binding_params( + mapped_op.python_callable, + partial_op_kwargs=mapped_op.partial_kwargs.get("op_kwargs") or {}, + expand_input=mapped_op._get_specified_expand_input(), + task_id=mapped_op.task_id, + ) + if params is None: + return {} + return {"_mapped_arg_binding_params": params} + def execute(self, context: Context) -> Any: raise RuntimeError( "@task.stub should not be executed directly -- we expected this to go to a remote worker. " diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index 9321cf1861c81..f1a5c0765c5d6 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -26,7 +26,7 @@ import pytest from airflow.providers.common.compat.sdk import DAG, task_group -from airflow.providers.standard.decorators.stub import _infer_value_schema, stub +from airflow.providers.standard.decorators.stub import _infer_value_schema, _StubOperator, stub from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS @@ -287,6 +287,117 @@ def group(n): group.expand(n=[1, 2]) +class TestMappedStubArgBindingParams: + """The serializer hook captures ordered per-parameter metadata for mapped stubs.""" + + def get_hook_fields(self, operator): + return _StubOperator.get_mapped_serialized_fields(operator) + + def test_params_follow_declaration_order_with_defaults_and_schemas(self): + with DAG(dag_id="d"): + # The partial() kwarg is declared *after* the expanded one: the captured + # order must come from the signature, not from the call sites. + result = stub(fn_transform).partial(extracted={"a": 1}).expand(country=["uk", "fr"]) + + assert self.get_hook_fields(result.operator) == { + "_mapped_arg_binding_params": [ + {"name": "country", "value_schema": {"type": "string"}}, + {"name": "extracted", "value_schema": {"type": "object", "additionalProperties": True}}, + { + "name": "retries_num", + "value_schema": {"type": "integer", "format": "int64"}, + "default": 3, + }, + ] + } + + def test_none_default_is_captured_by_key_presence(self): + def fn(x: str, y=None): ... + + with DAG(dag_id="d"): + result = stub(fn).expand(x=["a"]) + + params = self.get_hook_fields(result.operator)["_mapped_arg_binding_params"] + assert params[1] == {"name": "y", "default": None} + + def test_untyped_params_omit_value_schema(self): + with DAG(dag_id="d"): + result = stub(fn_untyped).expand(a=[1], b=[2]) + + assert self.get_hook_fields(result.operator) == { + "_mapped_arg_binding_params": [{"name": "a"}, {"name": "b"}] + } + + def test_parameterless_stub_captures_nothing(self): + with DAG(dag_id="d"): + result = stub(fn_extract).expand_kwargs([{}]) + + assert self.get_hook_fields(result.operator) == {} + + def test_expand_kwargs_rejected_for_parameterful_stub(self): + with DAG(dag_id="d"): + result = stub(fn_transform).expand_kwargs([{"country": "uk", "extracted": {}}]) + + with pytest.raises(ValueError, match="does not support expand_kwargs"): + self.get_hook_fields(result.operator) + + def test_missing_required_parameter_rejected(self): + with DAG(dag_id="d"): + result = stub(fn_transform).expand(country=["uk"]) + + with pytest.raises(ValueError, match="does not bind to its signature"): + self.get_hook_fields(result.operator) + + def test_partial_kwarg_over_mapped_upstream_rejected(self): + def fn_produce(n: int): ... + + with DAG(dag_id="d"): + vals = stub(fn_produce).expand(n=[1, 2]) + result = stub(fn_transform).partial(extracted=vals).expand(country=["uk"]) + + with pytest.raises(ValueError, match="aggregated output of the mapped task"): + self.get_hook_fields(result.operator) + + def test_expand_over_mapped_upstream_allowed(self): + def fn_produce(n: int): ... + + with DAG(dag_id="d"): + vals = stub(fn_produce).expand(n=[1, 2]) + result = stub(fn_transform).partial(country="uk").expand(extracted=vals) + + params = self.get_hook_fields(result.operator)["_mapped_arg_binding_params"] + assert [p["name"] for p in params] == ["country", "extracted", "retries_num"] + + def test_non_json_expand_literal_rejected(self): + with DAG(dag_id="d"): + result = stub(fn_transform).partial(country="uk").expand(extracted=[object()]) + + with pytest.raises(ValueError, match="not JSON-serializable"): + self.get_hook_fields(result.operator) + + def test_non_json_needed_default_rejected(self): + not_jsonable = object() + + def fn(x: str, y=not_jsonable): ... + + with DAG(dag_id="d"): + result = stub(fn).expand(x=["a"]) + + with pytest.raises(ValueError, match="not JSON-serializable"): + self.get_hook_fields(result.operator) + + def test_mapped_stub_inside_mapped_task_group_unconstructible(self): + """The SDK bans expansion inside an expanded group outright, so no hook guard is needed.""" + + @task_group + def group(n): + stub(fn_transform).partial(country="uk").expand(extracted=[{}]) + + with DAG(dag_id="d"): + with pytest.raises(NotImplementedError, match="expansion in an expanded task group"): + group.expand(n=[1, 2]) + + @pytest.mark.parametrize( ("annotation", "expected"), [ From 3e256b7bdaf8a4948f22e1c1157e1436b376eb5b Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 27 Jul 2026 17:27:57 +0000 Subject: [PATCH 35/40] Derive mapped stub arg bindings from parse-time parameter metadata The server-side derivation for mapped stubs was blind to the stub signature: it emitted the spec in call-site dict order while the wire contract promises declaration order (a positional binder like the Go SDK's flat mode then receives silently swapped values), dropped parameters filled from signature defaults where the unmapped path ships from_default entries, and could not attach value schemas. The Dag serializer now consults an optional operator-class hook while serializing a mapped operator -- the one point where operator_class and python_callable are still the real objects -- and stores the stub's per-parameter metadata under _mapped_arg_binding_params. ti_run walks that metadata in declaration order, fills expanded, partial, and defaulted parameters alike, and carries each parameter's value schema, closing the mapped/unmapped contract gap (https://github.com/apache/airflow/issues/70523). Dags serialized without the metadata (an older provider) keep the legacy ignored-args behavior instead of receiving order-uncertain bindings, and the old derivation's rejections stay as backstops for such Dags. --- .../execution_api/services/task_instances.py | 59 ++++--- .../src/airflow/serialization/schema.json | 29 +++ .../serialization/serialized_objects.py | 10 ++ .../versions/head/test_task_instances.py | 165 +++++++++++++++--- .../v2026_10_30/test_task_instances.py | 16 +- .../serialization/test_dag_serialization.py | 67 +++++++ 6 files changed, 295 insertions(+), 51 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py index af1863a297f6f..ce185dbeed129 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py @@ -67,27 +67,33 @@ def _unsupported_arg_bindings(detail: str) -> NoReturn: def _resolve_mapped_stub_arg_bindings( task: SerializedMappedOperator, ti: Any, *, dag: SerializedDAG, session: Session -) -> list[dict[str, Any]]: +) -> list[dict[str, Any]] | None: """ Build the per-map-index arg spec for a mapped (``.expand()``) stub task. - A mapped stub never instantiates at parse time, so no spec is captured in the - serialized Dag; it is derived here from the serialized expand input instead, with - the map-index decomposition delegated to - ``SchedulerDictOfListsExpandInput.resolve_expansion_sub_indexes``. - Value schemas come from the stub function's annotations, which are not available - server-side, so mapped bindings omit them (runtimes fall back to decode-only checks); - delivering them is tracked at https://github.com/apache/airflow/issues/70523. + A mapped stub never instantiates at parse time; the Dag serializer captures its + per-parameter metadata (declaration order, defaults, value schemas) from the stub + signature via ``get_mapped_serialized_fields``, and the map-index decomposition is + delegated to ``SchedulerDictOfListsExpandInput.resolve_expansion_sub_indexes``. + Dags serialized without the metadata (an older provider) resolve to ``None``: their + args were never deliverable, so they keep the legacy ignored-args behavior rather + than receive bindings whose order the server cannot know. """ + metadata = getattr(task, "_mapped_arg_binding_params", None) + if metadata is None: + return None + # The isinstance/map_index/unclaimed checks below re-reject what the provider now + # fails at parse time, for serialized Dags produced by other provider versions. expand_input = task._get_specified_expand_input() - # TODO: Support the `expand_kwargs` path once https://github.com/apache/airflow/pull/69757 - # is merged and all the Lang-SDKs adapt it. if not isinstance(expand_input, SchedulerDictOfListsExpandInput): _unsupported_arg_bindings("expand_kwargs() is not supported on stub tasks") if ti.map_index < 0: _unsupported_arg_bindings("the task instance has not been expanded to a map index") expand_value = expand_input.value + partial_op_kwargs = task.partial_kwargs.get("op_kwargs") or {} + if unclaimed := (set(expand_value) | set(partial_op_kwargs)) - {meta["name"] for meta in metadata}: + _unsupported_arg_bindings(f"kwargs {sorted(unclaimed)} are not in the captured parameter metadata") try: sub_indexes = expand_input.resolve_expansion_sub_indexes(ti.map_index, ti.run_id, session=session) except NotFullyPopulated as e: @@ -97,20 +103,25 @@ def _resolve_mapped_stub_arg_bindings( except ValueError as e: _unsupported_arg_bindings(str(e)) - # partial() kwargs. XComArgs inside partial() op_kwargs deserialize to _XComRef and - # are never dereferenced (set_task_dag_references only derefs the expand inputs), so - # resolve them here before binding. - spec = [ - _bind_mapped_stub_arg( - name, value.deref(dag) if isinstance(value, _XComRef) else value, sub_index=None - ) - for name, value in (task.partial_kwargs.get("op_kwargs") or {}).items() - ] - # expand() kwargs - spec += [ - _bind_mapped_stub_arg(name, value, sub_index=sub_indexes[name]) - for name, value in expand_value.items() - ] + spec = [] + for meta in metadata: # Declaration order, captured at parse time. + name = meta["name"] + if name in expand_value: + entry = _bind_mapped_stub_arg(name, expand_value[name], sub_index=sub_indexes[name]) + elif name in partial_op_kwargs: + value = partial_op_kwargs[name] + # XComArgs inside partial() op_kwargs deserialize to _XComRef and are never + # dereferenced (set_task_dag_references only derefs the expand inputs). + if isinstance(value, _XComRef): + value = value.deref(dag) + entry = _bind_mapped_stub_arg(name, value, sub_index=None) + elif "default" in meta: + entry = {"name": name, "kind": "literal", "value": meta["default"], "from_default": True} + else: + _unsupported_arg_bindings(f"parameter {name!r} has no expanded, partial, or default value") + if (value_schema := meta.get("value_schema")) is not None: + entry["value_schema"] = value_schema + spec.append(entry) return spec diff --git a/airflow-core/src/airflow/serialization/schema.json b/airflow-core/src/airflow/serialization/schema.json index fb954f0615b9d..39a33d1b7757b 100644 --- a/airflow-core/src/airflow/serialization/schema.json +++ b/airflow-core/src/airflow/serialization/schema.json @@ -184,6 +184,30 @@ ], "additionalProperties": false }, + "arg_binding_param": { + "$comment": "Per-parameter metadata of a mapped @task.stub task, in dict-encoded form and declaration order. The inner object stays open so future metadata fields keep validating on older cores", + "type": "object", + "properties": { + "__type": { + "type": "string", + "const": "dict" + }, + "__var": { + "type": "object", + "properties": { + "name": { "type": "string" }, + "value_schema": { "$ref": "#/definitions/typed_dict" }, + "default": {} + }, + "required": [ "name" ] + } + }, + "required": [ + "__type", + "__var" + ], + "additionalProperties": false + }, "color": { "type": "string", "pattern": "^#[a-fA-F0-9]{3,6}$" @@ -392,6 +416,11 @@ "$comment": "Only present on @task.stub tasks called with TaskFlow arguments", "type": "array", "items": { "$ref": "#/definitions/arg_binding" } + }, + "_mapped_arg_binding_params": { + "$comment": "Only present on mapped @task.stub tasks with parameters; ordered per-parameter binding metadata", + "type": "array", + "items": { "$ref": "#/definitions/arg_binding_param" } } }, "dependencies": { diff --git a/airflow-core/src/airflow/serialization/serialized_objects.py b/airflow-core/src/airflow/serialization/serialized_objects.py index 54bc3389c64ce..8f19dd9b8b20a 100644 --- a/airflow-core/src/airflow/serialization/serialized_objects.py +++ b/airflow-core/src/airflow/serialization/serialized_objects.py @@ -996,6 +996,16 @@ def serialize_mapped_operator(cls, op: MappedOperator) -> dict[str, Any]: ) del serialized_op["partial_kwargs"]["python_callable"] + # Optional per-class capability: an operator class may contribute extra serialized + # fields for its mapped form. This is the only point where operator_class is the + # real class and python_callable the real function (neither survives + # serialization), so signature-derived data must be captured here. Used by the + # standard provider's _StubOperator for its TaskFlow arg-binding metadata. + get_extra_fields = getattr(op.operator_class, "get_mapped_serialized_fields", None) + if get_extra_fields is not None: + for key, value in get_extra_fields(op).items(): + serialized_op[key] = cls.serialize(value) + serialized_op["_is_mapped"] = True return serialized_op 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 d1f6c8bb28725..0d61d4625f97a 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 @@ -486,7 +486,7 @@ def transform(country: str): ... ) assert response.status_code == 200 assert response.json()["arg_bindings"] == [ - {"name": "country", "kind": "literal", "value": country} + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": country} ] def test_ti_run_resolves_mapped_stub_over_unmapped_upstream(self, client, dag_maker): @@ -510,7 +510,13 @@ def transform(extracted: dict): ... response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) assert response.status_code == 200 assert response.json()["arg_bindings"] == [ - {"name": "extracted", "kind": "xcom", "task_id": "extract", "element_index": 1} + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "extract", + "element_index": 1, + } ] def test_ti_run_resolves_mapped_stub_over_mapped_upstream(self, client, dag_maker): @@ -534,7 +540,13 @@ def transform(extracted: dict): ... response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) assert response.status_code == 200 assert response.json()["arg_bindings"] == [ - {"name": "extracted", "kind": "xcom", "task_id": "seed", "map_index": 1} + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "seed", + "map_index": 1, + } ] def test_ti_run_decomposes_multi_kwarg_mapped_stub(self, client, dag_maker): @@ -559,8 +571,13 @@ def combine(a: str, b: int): ... ) assert response.status_code == 200 assert response.json()["arg_bindings"] == [ - {"name": "a", "kind": "literal", "value": a}, - {"name": "b", "kind": "literal", "value": b}, + {"name": "a", "kind": "literal", "value_schema": {"type": "string"}, "value": a}, + { + "name": "b", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": b, + }, ] def test_ti_run_binds_partial_kwargs_of_mapped_stub(self, client, dag_maker): @@ -584,8 +601,13 @@ def transform(country: str, extracted: dict): ... ) assert response.status_code == 200 assert response.json()["arg_bindings"] == [ - {"name": "country", "kind": "literal", "value": "uk"}, - {"name": "extracted", "kind": "literal", "value": extracted}, + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, + { + "name": "extracted", + "kind": "literal", + "value_schema": {"type": "object", "additionalProperties": True}, + "value": extracted, + }, ] def test_ti_run_binds_partial_xcom_kwarg_over_unmapped_upstream(self, client, dag_maker): @@ -608,21 +630,35 @@ def transform(extracted: dict, country: str): ... response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) assert response.status_code == 200 assert response.json()["arg_bindings"] == [ - {"name": "extracted", "kind": "xcom", "task_id": "extract"}, - {"name": "country", "kind": "literal", "value": "fr"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "extract", + }, + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "fr"}, ] def test_ti_run_rejects_partial_kwarg_over_mapped_upstream(self, client, dag_maker): - """A partial() kwarg over a mapped upstream would bind the nonexistent unmapped XCom row.""" - with dag_maker("test_mapped_stub_partial_mapped_upstream", serialized=True): + """ + A partial() kwarg over a mapped upstream would bind the nonexistent unmapped XCom row. - @task.stub - def seed(n: int): ... + The provider rejects this at parse time now; patching its capture hook simulates a + Dag serialized by another provider version, exercising the server-side backstop. + """ + from airflow.providers.standard.decorators.stub import _StubOperator - @task.stub - def transform(extracted: dict, country: str): ... + fabricated = {"_mapped_arg_binding_params": [{"name": "extracted"}, {"name": "country"}]} + with mock.patch.object(_StubOperator, "get_mapped_serialized_fields", return_value=fabricated): + with dag_maker("test_mapped_stub_partial_mapped_upstream", serialized=True): + + @task.stub + def seed(n: int): ... - transform.partial(extracted=seed.expand(n=[1, 2])).expand(country=["uk", "fr"]) + @task.stub + def transform(extracted: dict, country: str): ... + + transform.partial(extracted=seed.expand(n=[1, 2])).expand(country=["uk", "fr"]) dr = dag_maker.create_dagrun() ti = next(ti for ti in dr.get_task_instances() if ti.task_id == "transform" and ti.map_index == 0) @@ -634,6 +670,84 @@ def transform(extracted: dict, country: str): ... assert response.json()["detail"]["reason"] == "invalid_arg_bindings" assert "aggregated output" in response.json()["detail"]["message"] + def test_ti_run_orders_mapped_stub_spec_by_declaration_with_defaults(self, client, dag_maker): + """The spec follows the signature, not the call sites, and ships defaulted params.""" + with dag_maker("test_mapped_stub_declaration_order", serialized=True): + + @task.stub + def transform(country: str, extracted: dict, retries_num: int = 3): ... + + # The partial() kwarg is declared after the expanded one on purpose. + transform.partial(extracted={"a": 1}).expand(country=["uk", "fr"]) + + dr = dag_maker.create_dagrun() + ti = next(t for t in dr.get_task_instances() if t.map_index == 1) + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "fr"}, + { + "name": "extracted", + "kind": "literal", + "value_schema": {"type": "object", "additionalProperties": True}, + "value": {"a": 1}, + }, + { + "name": "retries_num", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 3, + "from_default": True, + }, + ] + + def test_ti_run_ignores_args_for_legacy_serialized_mapped_stub(self, client, dag_maker): + """A mapped stub serialized without parameter metadata keeps the ignored-args behavior.""" + from airflow.providers.standard.decorators.stub import _StubOperator + + with mock.patch.object(_StubOperator, "get_mapped_serialized_fields", return_value={}): + with dag_maker("test_mapped_stub_legacy", serialized=True): + + @task.stub + def transform(country: str): ... + + transform.expand(country=["uk", "fr"]) + + dr = dag_maker.create_dagrun() + ti = next(t for t in dr.get_task_instances() if t.map_index == 0) + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 200 + assert "arg_bindings" not in response.json() + + def test_ti_run_rejects_unexpanded_mapped_stub_ti(self, client, dag_maker): + """A mapped stub TI still at map_index=-1 cannot receive per-index bindings.""" + with dag_maker("test_mapped_stub_unexpanded", serialized=True): + + @task.stub + def extract(): ... + + @task.stub + def transform(extracted: dict): ... + + transform.expand(extracted=extract()) + + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance("transform") + assert ti.map_index == -1 + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 500 + assert response.json()["detail"]["reason"] == "invalid_arg_bindings" + assert "not been expanded" in response.json()["detail"]["message"] + def test_ti_run_rejects_zero_length_expansion_on_stub(self, client, dag_maker, session): """An upstream re-run to an empty list after expansion fails structurally, not with a crash.""" from airflow.models.taskmap import TaskMap @@ -670,13 +784,22 @@ def combine(a: int, b: int): ... assert "length 0" in response.json()["detail"]["message"] def test_ti_run_rejects_expand_kwargs_on_stub(self, client, dag_maker): - """expand_kwargs() has no per-parameter spec to derive, so delivery fails structurally.""" - with dag_maker("test_mapped_stub_expand_kwargs", serialized=True): + """ + expand_kwargs() has no per-parameter spec to derive, so delivery fails structurally. - @task.stub - def transform(country: str): ... + The provider rejects this at parse time now; patching its capture hook simulates a + Dag serialized by another provider version, exercising the server-side backstop. + """ + from airflow.providers.standard.decorators.stub import _StubOperator + + fabricated = {"_mapped_arg_binding_params": [{"name": "country"}]} + with mock.patch.object(_StubOperator, "get_mapped_serialized_fields", return_value=fabricated): + with dag_maker("test_mapped_stub_expand_kwargs", serialized=True): + + @task.stub + def transform(country: str): ... - transform.expand_kwargs([{"country": "uk"}]) + transform.expand_kwargs([{"country": "uk"}]) dr = dag_maker.create_dagrun() (ti,) = dr.get_task_instances() diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py index dd29526c056c0..b4dd521c760e2 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py @@ -81,20 +81,24 @@ def test_old_version_strips_arg_bindings_even_when_set(self, old_ver_client, stu def test_old_version_skips_undeliverable_arg_bindings_derivation(self, old_ver_client, dag_maker): """A stub whose bindings cannot be delivered must keep running for clients that never see them.""" - with dag_maker("test_arg_bindings_compat_expand_kwargs", serialized=True): + with dag_maker("test_arg_bindings_compat_unexpanded", serialized=True): @task.stub - def transform(country: str): ... + def extract(): ... + + @task.stub + def transform(extracted: dict): ... - transform.expand_kwargs([{"country": "uk"}]) + transform.expand(extracted=extract()) dr = dag_maker.create_dagrun() - (ti,) = dr.get_task_instances() + ti = dr.get_task_instance("transform") + assert ti.map_index == -1 ti.set_state(State.QUEUED) dag_maker.session.flush() - # At head this Dag fails ti_run with a structured 500 (expand_kwargs is - # unsupported on stubs); a pre-arg-bindings client keeps the legacy behavior. + # At head this TI fails ti_run with a structured 500 (it has not been expanded + # to a map index); a pre-arg-bindings client keeps the legacy behavior. response = old_ver_client.patch(f"/execution/task-instances/{ti.id}/run", json=RUN_PATCH_BODY) assert response.status_code == 200 assert "arg_bindings" not in response.json() diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index 889131a152032..a7459b204cc5c 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -3459,6 +3459,73 @@ def transform(country: str, extracted: dict): ... assert not hasattr(round_tripped.task_dict["extract"], "_arg_bindings") +def test_mapped_stub_param_metadata_round_trip(): + """The serializer collects the mapped stub's parameter metadata and it survives the round trip.""" + from airflow.sdk import task + + with DAG(dag_id="mapped_arg_binding_params_dag", schedule=None) as dag: + + @task.stub + def transform(country: str, extracted, retries_num: int = 3): ... + + @task + def plain(x): ... + + transform.partial(extracted={"a": 1}).expand(country=["uk", "fr"]) + plain.expand(x=[1, 2]) + + ser_dag = DagSerialization.to_dict(dag) + DagSerialization.validate_schema(ser_dag) + encoded_tasks = {t[Encoding.VAR]["task_id"]: t[Encoding.VAR] for t in ser_dag["dag"]["tasks"]} + assert encoded_tasks["transform"]["_mapped_arg_binding_params"] == [ + { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: { + "name": "country", + "value_schema": {Encoding.TYPE: DAT.DICT, Encoding.VAR: {"type": "string"}}, + }, + }, + {Encoding.TYPE: DAT.DICT, Encoding.VAR: {"name": "extracted"}}, + { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: { + "name": "retries_num", + "value_schema": { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: {"type": "integer", "format": "int64"}, + }, + "default": 3, + }, + }, + ], "metadata must serialize in declaration order" + assert "_mapped_arg_binding_params" not in encoded_tasks["plain"], ( + "only operator classes defining the capture hook contribute mapped fields" + ) + + round_tripped = DagSerialization.from_dict(ser_dag) + assert round_tripped.task_dict["transform"]._mapped_arg_binding_params == [ + {"name": "country", "value_schema": {"type": "string"}}, + {"name": "extracted"}, + {"name": "retries_num", "value_schema": {"type": "integer", "format": "int64"}, "default": 3}, + ] + assert not hasattr(round_tripped.task_dict["plain"], "_mapped_arg_binding_params") + + +def test_mapped_stub_capture_error_fails_serialization(): + """A capture-hook rejection surfaces as a Dag serialization (import) error.""" + from airflow.sdk import task + + with DAG(dag_id="mapped_arg_binding_params_invalid_dag", schedule=None) as dag: + + @task.stub + def transform(country: str): ... + + transform.expand_kwargs([{"country": "uk"}]) + + with pytest.raises(SerializationError, match="does not support expand_kwargs"): + DagSerialization.to_dict(dag) + + def test_handle_v1_serdag(): v1 = { "__version": 1, From f357db8507257dd3ce31cc1df0cac4a80f1c5157 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 27 Jul 2026 17:31:40 +0000 Subject: [PATCH 36/40] Cache stub value-schema generation across Dag re-parses TypeAdapter construction is one of pydantic's most expensive operations and ran fresh for every annotated stub parameter on every Dag file re-parse, which the Dag processor repeats continuously. Annotations are static, so cache the generated fragment per annotation for the process lifetime, deep-copying on the way out so embedded specs never alias the cache, and falling back to uncached generation for unhashable annotations. --- .../providers/standard/decorators/stub.py | 24 ++++++++++++++++--- .../unit/standard/decorators/test_stub.py | 12 ++++++++++ 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index edc46b649aa68..9eac3da2d11e9 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -18,12 +18,14 @@ from __future__ import annotations import ast +import copy import datetime import inspect import json import types import typing from collections.abc import Callable, Collection, Mapping +from functools import cache from typing import TYPE_CHECKING, Any try: @@ -112,16 +114,32 @@ def _infer_value_schema(annotation: Any) -> dict[str, Any] | None: # that can only ever be None constrains nothing worth shipping. return None try: - schema = TypeAdapter(annotation).json_schema(schema_generator=_ValueSchemaGenerator) + schema = _generate_value_schema(annotation) + except TypeError: + # Unhashable annotations cannot key the cache; generate directly. + schema = _generate_value_schema.__wrapped__(annotation) + # Deep-copy so callers embedding the fragment never alias the cached dict. + return copy.deepcopy(schema) if schema else None + + +@cache +def _generate_value_schema(annotation: Any) -> dict[str, Any] | None: + """ + Generate the schema for one annotation, cached for the process lifetime. + + TypeAdapter construction is one of pydantic's most expensive operations and + annotations are static, so re-parses of the same Dag file must not re-pay it. + """ + try: + return TypeAdapter(annotation).json_schema(schema_generator=_ValueSchemaGenerator) except (PydanticSchemaGenerationError, PydanticInvalidForJsonSchema): normalized = _normalize_temporal_annotation(annotation) if normalized is annotation: return None try: - schema = TypeAdapter(normalized).json_schema(schema_generator=_ValueSchemaGenerator) + return TypeAdapter(normalized).json_schema(schema_generator=_ValueSchemaGenerator) except (PydanticSchemaGenerationError, PydanticInvalidForJsonSchema): return None - return schema or None def _validate_stub_signature(signature: inspect.Signature, task_id: str) -> None: diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index f1a5c0765c5d6..b73c82f38d2f6 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -516,3 +516,15 @@ def test_infer_value_schema(annotation, expected): @mock.patch("airflow.providers.standard.decorators.stub.TypeAdapter", None) def test_infer_value_schema_without_pydantic(): assert _infer_value_schema(str) is None + + +def test_infer_value_schema_cache_returns_isolated_copies(): + first = _infer_value_schema(dict) + second = _infer_value_schema(dict) + assert first == second + assert first is not second, "callers embed and serialize the fragment, so it must not alias the cache" + + +def test_infer_value_schema_unhashable_annotation_generates_uncached(): + annotation = typing.Annotated[int, {"unhashable": True}] + assert _infer_value_schema(annotation) == {"type": "integer", "format": "int64"} From 1d9e5d4fd809f101b612e23ac6c89da81167a8c5 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 27 Jul 2026 17:33:03 +0000 Subject: [PATCH 37/40] Use XCOM_RETURN_KEY instead of hardcoded "return_value" in stub bindings Both binding validators duplicated the canonical XCom return-value key as a string literal, evading constant-based refactors and diverging from the neighboring code (serialization's xcom_arg already compares against XCOM_RETURN_KEY). The constant is importable in both contexts: common.compat.sdk re-exports it for the provider, airflow.models.xcom for core. --- .../api_fastapi/execution_api/services/task_instances.py | 3 ++- .../standard/src/airflow/providers/standard/decorators/stub.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py index ce185dbeed129..4a8c2589438d8 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py @@ -24,6 +24,7 @@ from fastapi import HTTPException, status from airflow.models.expandinput import NotFullyPopulated, SchedulerDictOfListsExpandInput +from airflow.models.xcom import XCOM_RETURN_KEY from airflow.serialization.definitions.mappedoperator import is_mapped from airflow.serialization.definitions.xcom_arg import SchedulerPlainXComArg, SchedulerXComArg from airflow.serialization.serialized_objects import _XComRef @@ -128,7 +129,7 @@ def _resolve_mapped_stub_arg_bindings( def _bind_mapped_stub_arg(name: str, value: Any, *, sub_index: int | None) -> dict[str, Any]: """Build one arg-binding dict; ``sub_index`` is set for expanded kwargs, None for partial ones.""" if isinstance(value, SchedulerPlainXComArg): - if value.key != "return_value": + if value.key != XCOM_RETURN_KEY: _unsupported_arg_bindings(f"parameter {name!r} references the XCom key {value.key!r}") if sub_index is None and value.operator.is_mapped: # A partial() kwarg over a mapped upstream would bind the unmapped XCom row diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index 9eac3da2d11e9..adf554095bac5 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -40,6 +40,7 @@ from airflow.providers.common.compat.sdk import ( KNOWN_CONTEXT_KEYS, + XCOM_RETURN_KEY, DecoratedOperator, PlainXComArg, TaskDecorator, @@ -190,7 +191,7 @@ def _ensure_json_literal(value: Any, task_id: str, name: str) -> None: def _validate_xcom_value(value: Any, task_id: str, name: str, *, allow_mapped_upstream: bool = False) -> bool: """Validate an XComArg argument, returning True when it is a bindable direct upstream output.""" if isinstance(value, PlainXComArg): - if value.key != "return_value": + if value.key != XCOM_RETURN_KEY: raise ValueError( f"@task.stub task {task_id!r} parameter {name!r} references the XCom key " f"{value.key!r}; only an upstream task's return value can cross the language " From 262d80321023622053117b713fa56b549bd1c398 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 28 Jul 2026 02:43:19 +0000 Subject: [PATCH 38/40] Move mapped stub arg-binding support to a follow-up branch Reviewing unmapped TaskFlow delivery and per-map-index derivation together made the PR hard to land, so this PR narrows to the unmapped contract: mapped (.expand()) stubs keep the released ignored-args behavior (they capture no parse-time spec, so ti_run naturally delivers no bindings), documented on the stub decorator. The mapped derivation -- the serializer capture hook, per-parameter metadata, map-index decomposition, and their tests -- moves wholesale to the stacked follow-up branch feature/lang-sdk/taskflow-stub-dag-mapped. --- .../execution_api/services/task_instances.py | 134 +------ .../src/airflow/models/expandinput.py | 27 -- .../src/airflow/serialization/schema.json | 29 -- .../serialization/serialized_objects.py | 10 - .../versions/head/test_task_instances.py | 338 +----------------- .../v2026_10_30/test_task_instances.py | 24 -- .../serialization/test_dag_serialization.py | 67 ---- .../tests_common/test_utils/version_compat.py | 1 - .../providers/standard/decorators/stub.py | 85 +---- .../unit/standard/decorators/test_stub.py | 114 +----- 10 files changed, 20 insertions(+), 809 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py index 4a8c2589438d8..e06db8ffa70be 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py @@ -18,23 +18,12 @@ from __future__ import annotations -import json -from typing import TYPE_CHECKING, Any, NoReturn - -from fastapi import HTTPException, status - -from airflow.models.expandinput import NotFullyPopulated, SchedulerDictOfListsExpandInput -from airflow.models.xcom import XCOM_RETURN_KEY -from airflow.serialization.definitions.mappedoperator import is_mapped -from airflow.serialization.definitions.xcom_arg import SchedulerPlainXComArg, SchedulerXComArg -from airflow.serialization.serialized_objects import _XComRef +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from sqlalchemy.orm import Session from airflow.models.dagbag import DBDagBag - from airflow.serialization.definitions.dag import SerializedDAG - from airflow.serialization.definitions.mappedoperator import SerializedMappedOperator # Task type recorded on the TI row (``TaskInstance.operator``) for # ``airflow.providers.standard.decorators.stub._StubOperator``. Used to gate the @@ -44,124 +33,17 @@ def get_arg_bindings(dag_bag: DBDagBag, ti: Any, *, session: Session) -> list | None: - """Extract or derive the stub task's TaskFlow arg spec from its Dag version.""" + """ + Extract the stub task's TaskFlow arg spec from its Dag version. + + Mapped (``.expand()``) stubs never capture a parse-time spec, so they resolve to + ``None`` here and keep the legacy ignored-args behavior; per-map-index delivery + lands in a follow-up. + """ if ti.dag_version_id is None: return None if (dag := dag_bag.get_dag(ti.dag_version_id, session=session)) is None: return None if (task := dag.task_dict.get(ti.task_id)) is None: return None - if is_mapped(task): - return _resolve_mapped_stub_arg_bindings(task, ti, dag=dag, session=session) return getattr(task, "_arg_bindings", None) - - -def _unsupported_arg_bindings(detail: str) -> NoReturn: - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={ - "reason": "invalid_arg_bindings", - "message": f"The stub task's TaskFlow arguments cannot be delivered: {detail}.", - }, - ) - - -def _resolve_mapped_stub_arg_bindings( - task: SerializedMappedOperator, ti: Any, *, dag: SerializedDAG, session: Session -) -> list[dict[str, Any]] | None: - """ - Build the per-map-index arg spec for a mapped (``.expand()``) stub task. - - A mapped stub never instantiates at parse time; the Dag serializer captures its - per-parameter metadata (declaration order, defaults, value schemas) from the stub - signature via ``get_mapped_serialized_fields``, and the map-index decomposition is - delegated to ``SchedulerDictOfListsExpandInput.resolve_expansion_sub_indexes``. - Dags serialized without the metadata (an older provider) resolve to ``None``: their - args were never deliverable, so they keep the legacy ignored-args behavior rather - than receive bindings whose order the server cannot know. - """ - metadata = getattr(task, "_mapped_arg_binding_params", None) - if metadata is None: - return None - # The isinstance/map_index/unclaimed checks below re-reject what the provider now - # fails at parse time, for serialized Dags produced by other provider versions. - expand_input = task._get_specified_expand_input() - if not isinstance(expand_input, SchedulerDictOfListsExpandInput): - _unsupported_arg_bindings("expand_kwargs() is not supported on stub tasks") - if ti.map_index < 0: - _unsupported_arg_bindings("the task instance has not been expanded to a map index") - - expand_value = expand_input.value - partial_op_kwargs = task.partial_kwargs.get("op_kwargs") or {} - if unclaimed := (set(expand_value) | set(partial_op_kwargs)) - {meta["name"] for meta in metadata}: - _unsupported_arg_bindings(f"kwargs {sorted(unclaimed)} are not in the captured parameter metadata") - try: - sub_indexes = expand_input.resolve_expansion_sub_indexes(ti.map_index, ti.run_id, session=session) - except NotFullyPopulated as e: - # Neither this nor the ValueError below can happen on the happy path: both take - # someone clearing upstream TIs or XComs themselves during the DagRun. - _unsupported_arg_bindings(f"upstream map lengths are not yet known for {sorted(e.missing)}") - except ValueError as e: - _unsupported_arg_bindings(str(e)) - - spec = [] - for meta in metadata: # Declaration order, captured at parse time. - name = meta["name"] - if name in expand_value: - entry = _bind_mapped_stub_arg(name, expand_value[name], sub_index=sub_indexes[name]) - elif name in partial_op_kwargs: - value = partial_op_kwargs[name] - # XComArgs inside partial() op_kwargs deserialize to _XComRef and are never - # dereferenced (set_task_dag_references only derefs the expand inputs). - if isinstance(value, _XComRef): - value = value.deref(dag) - entry = _bind_mapped_stub_arg(name, value, sub_index=None) - elif "default" in meta: - entry = {"name": name, "kind": "literal", "value": meta["default"], "from_default": True} - else: - _unsupported_arg_bindings(f"parameter {name!r} has no expanded, partial, or default value") - if (value_schema := meta.get("value_schema")) is not None: - entry["value_schema"] = value_schema - spec.append(entry) - return spec - - -def _bind_mapped_stub_arg(name: str, value: Any, *, sub_index: int | None) -> dict[str, Any]: - """Build one arg-binding dict; ``sub_index`` is set for expanded kwargs, None for partial ones.""" - if isinstance(value, SchedulerPlainXComArg): - if value.key != XCOM_RETURN_KEY: - _unsupported_arg_bindings(f"parameter {name!r} references the XCom key {value.key!r}") - if sub_index is None and value.operator.is_mapped: - # A partial() kwarg over a mapped upstream would bind the unmapped XCom row - # (map_index=-1), which never exists; the aggregated output is inexpressible. - _unsupported_arg_bindings( - f"parameter {name!r} references the aggregated output of the mapped task" - f" {value.operator.task_id!r}" - ) - entry: dict[str, Any] = {"name": name, "kind": "xcom", "task_id": value.operator.task_id} - if sub_index is not None: - if value.operator.is_mapped: - entry["map_index"] = sub_index - else: - entry["element_index"] = sub_index - return entry - if isinstance(value, SchedulerXComArg): - _unsupported_arg_bindings( - f"parameter {name!r} received a {type(value).__name__}; only direct upstream" - " task outputs and literals are supported" - ) - if sub_index is not None: - # This kwarg was expanded over a literal collection written in the Dag file. - items = list(value.items()) if isinstance(value, dict) else value - try: - value = items[sub_index] - except (IndexError, KeyError, TypeError): - _unsupported_arg_bindings(f"parameter {name!r} has no element at expansion index {sub_index}") - try: - json.dumps(value, allow_nan=False) - except (TypeError, ValueError): - _unsupported_arg_bindings( - f"parameter {name!r} carries a {type(value).__name__} value, which cannot cross" - " the language boundary" - ) - return {"name": name, "kind": "literal", "value": value} diff --git a/airflow-core/src/airflow/models/expandinput.py b/airflow-core/src/airflow/models/expandinput.py index ecc1b398eb735..0363bae92620a 100644 --- a/airflow-core/src/airflow/models/expandinput.py +++ b/airflow-core/src/airflow/models/expandinput.py @@ -151,33 +151,6 @@ def get_total_map_length(self, run_id: str, *, session: Session) -> int: lengths = self._get_map_lengths(run_id, session=session) return functools.reduce(operator.mul, (lengths[name] for name in self.value), 1) - def resolve_expansion_sub_indexes( - self, map_index: int, run_id: str, *, session: Session - ) -> dict[str, int]: - """ - Decompose a task instance's map index into one index per expanded kwarg. - - Server-side counterpart of the index decomposition in the SDK's - ``DictOfListsExpandInput._expand_mapped_field``: the cross-product of the - expanded kwargs is ordered with the last kwarg varying fastest. A single - expanded kwarg maps one-to-one, skipping the upstream length lookups. - - :raises NotFullyPopulated: if upstream map lengths are not all known yet. - :raises ValueError: if an expanded kwarg's recorded length is zero, e.g. an - upstream was cleared and re-ran to an empty list after this task instance - was expanded (the SDK twin guards the same case). - """ - if len(self.value) == 1: - return dict.fromkeys(self.value, map_index) - lengths = self._get_map_lengths(run_id, session=session) - sub_indexes = {} - for key in reversed(self.value): - if (length := lengths[key]) < 1: - raise ValueError(f"cannot decompose map index over expanded kwarg {key!r} of length 0") - sub_indexes[key] = map_index % length - map_index //= length - return sub_indexes - def iter_references(self) -> Iterable[tuple[Operator, str]]: from airflow.models.referencemixin import ReferenceMixin diff --git a/airflow-core/src/airflow/serialization/schema.json b/airflow-core/src/airflow/serialization/schema.json index 39a33d1b7757b..fb954f0615b9d 100644 --- a/airflow-core/src/airflow/serialization/schema.json +++ b/airflow-core/src/airflow/serialization/schema.json @@ -184,30 +184,6 @@ ], "additionalProperties": false }, - "arg_binding_param": { - "$comment": "Per-parameter metadata of a mapped @task.stub task, in dict-encoded form and declaration order. The inner object stays open so future metadata fields keep validating on older cores", - "type": "object", - "properties": { - "__type": { - "type": "string", - "const": "dict" - }, - "__var": { - "type": "object", - "properties": { - "name": { "type": "string" }, - "value_schema": { "$ref": "#/definitions/typed_dict" }, - "default": {} - }, - "required": [ "name" ] - } - }, - "required": [ - "__type", - "__var" - ], - "additionalProperties": false - }, "color": { "type": "string", "pattern": "^#[a-fA-F0-9]{3,6}$" @@ -416,11 +392,6 @@ "$comment": "Only present on @task.stub tasks called with TaskFlow arguments", "type": "array", "items": { "$ref": "#/definitions/arg_binding" } - }, - "_mapped_arg_binding_params": { - "$comment": "Only present on mapped @task.stub tasks with parameters; ordered per-parameter binding metadata", - "type": "array", - "items": { "$ref": "#/definitions/arg_binding_param" } } }, "dependencies": { diff --git a/airflow-core/src/airflow/serialization/serialized_objects.py b/airflow-core/src/airflow/serialization/serialized_objects.py index 8f19dd9b8b20a..54bc3389c64ce 100644 --- a/airflow-core/src/airflow/serialization/serialized_objects.py +++ b/airflow-core/src/airflow/serialization/serialized_objects.py @@ -996,16 +996,6 @@ def serialize_mapped_operator(cls, op: MappedOperator) -> dict[str, Any]: ) del serialized_op["partial_kwargs"]["python_callable"] - # Optional per-class capability: an operator class may contribute extra serialized - # fields for its mapped form. This is the only point where operator_class is the - # real class and python_callable the real function (neither survives - # serialization), so signature-derived data must be captured here. Used by the - # standard provider's _StubOperator for its TaskFlow arg-binding metadata. - get_extra_fields = getattr(op.operator_class, "get_mapped_serialized_fields", None) - if get_extra_fields is not None: - for key, value in get_extra_fields(op).items(): - serialized_op[key] = cls.serialize(value) - serialized_op["_is_mapped"] = True return serialized_op 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 0d61d4625f97a..e5dfea5360521 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -17,7 +17,6 @@ from __future__ import annotations -import itertools from datetime import datetime from types import SimpleNamespace from typing import TYPE_CHECKING @@ -464,257 +463,14 @@ def transform(country: str): ... "start_date": "2024-09-30T12:00:00Z", } - def test_ti_run_resolves_mapped_stub_literal_expand(self, client, dag_maker): - """Expanding a stub over a literal list resolves each map index to its element server-side.""" - with dag_maker("test_mapped_stub_literal", serialized=True): + def test_ti_run_returns_no_arg_bindings_for_mapped_stub(self, client, dag_maker): + """Mapped stubs keep the legacy ignored-args behavior until per-map-index delivery lands.""" + with dag_maker("test_mapped_stub_ignored_args", serialized=True): @task.stub def transform(country: str): ... - transform.expand(country=["uk", "fr", "de"]) - - dr = dag_maker.create_dagrun() - tis = {ti.map_index: ti for ti in dr.get_task_instances()} - assert set(tis) == {0, 1, 2} - for ti in tis.values(): - ti.set_state(State.QUEUED) - dag_maker.session.flush() - - for map_index, country in enumerate(["uk", "fr", "de"]): - response = client.patch( - f"/execution/task-instances/{tis[map_index].id}/run", json=self.RUN_PAYLOAD - ) - assert response.status_code == 200 - assert response.json()["arg_bindings"] == [ - {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": country} - ] - - def test_ti_run_resolves_mapped_stub_over_unmapped_upstream(self, client, dag_maker): - """Expanding over an unmapped upstream's output binds the whole XCom plus an element index.""" - with dag_maker("test_mapped_stub_unmapped_upstream", serialized=True): - - @task.stub - def extract(): ... - - @task.stub - def transform(extracted: dict): ... - - transform.expand(extracted=extract()) - - dr = dag_maker.create_dagrun() - ti = dr.get_task_instance("transform") - ti.map_index = 1 - ti.set_state(State.QUEUED) - dag_maker.session.flush() - - response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) - assert response.status_code == 200 - assert response.json()["arg_bindings"] == [ - { - "name": "extracted", - "kind": "xcom", - "value_schema": {"type": "object", "additionalProperties": True}, - "task_id": "extract", - "element_index": 1, - } - ] - - def test_ti_run_resolves_mapped_stub_over_mapped_upstream(self, client, dag_maker): - """Expanding over a mapped upstream binds the upstream XCom row at the same map index.""" - with dag_maker("test_mapped_stub_mapped_upstream", serialized=True): - - @task.stub - def seed(n: int): ... - - @task.stub - def transform(extracted: dict): ... - - transform.expand(extracted=seed.expand(n=[1, 2])) - - dr = dag_maker.create_dagrun() - ti = dr.get_task_instance("transform") - ti.map_index = 1 - ti.set_state(State.QUEUED) - dag_maker.session.flush() - - response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) - assert response.status_code == 200 - assert response.json()["arg_bindings"] == [ - { - "name": "extracted", - "kind": "xcom", - "value_schema": {"type": "object", "additionalProperties": True}, - "task_id": "seed", - "map_index": 1, - } - ] - - def test_ti_run_decomposes_multi_kwarg_mapped_stub(self, client, dag_maker): - """Cross-product expansion decomposes the map index per kwarg like the task-sdk does.""" - with dag_maker("test_mapped_stub_multi_kwarg", serialized=True): - - @task.stub - def combine(a: str, b: int): ... - - combine.expand(a=["x", "y"], b=[1, 2, 3]) - - dr = dag_maker.create_dagrun() - tis = {ti.map_index: ti for ti in dr.get_task_instances()} - assert set(tis) == set(range(6)) - for ti in tis.values(): - ti.set_state(State.QUEUED) - dag_maker.session.flush() - - for map_index, (a, b) in enumerate(itertools.product(["x", "y"], [1, 2, 3])): - response = client.patch( - f"/execution/task-instances/{tis[map_index].id}/run", json=self.RUN_PAYLOAD - ) - assert response.status_code == 200 - assert response.json()["arg_bindings"] == [ - {"name": "a", "kind": "literal", "value_schema": {"type": "string"}, "value": a}, - { - "name": "b", - "kind": "literal", - "value_schema": {"type": "integer", "format": "int64"}, - "value": b, - }, - ] - - def test_ti_run_binds_partial_kwargs_of_mapped_stub(self, client, dag_maker): - """partial() kwargs bind like an unmapped TaskFlow call alongside the expanded ones.""" - with dag_maker("test_mapped_stub_partial", serialized=True): - - @task.stub - def transform(country: str, extracted: dict): ... - - transform.partial(country="uk").expand(extracted=[{"a": 1}, {"b": 2}]) - - dr = dag_maker.create_dagrun() - tis = {ti.map_index: ti for ti in dr.get_task_instances()} - for ti in tis.values(): - ti.set_state(State.QUEUED) - dag_maker.session.flush() - - for map_index, extracted in enumerate([{"a": 1}, {"b": 2}]): - response = client.patch( - f"/execution/task-instances/{tis[map_index].id}/run", json=self.RUN_PAYLOAD - ) - assert response.status_code == 200 - assert response.json()["arg_bindings"] == [ - {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, - { - "name": "extracted", - "kind": "literal", - "value_schema": {"type": "object", "additionalProperties": True}, - "value": extracted, - }, - ] - - def test_ti_run_binds_partial_xcom_kwarg_over_unmapped_upstream(self, client, dag_maker): - """A partial() kwarg carrying an unmapped upstream's output binds that XCom for every index.""" - with dag_maker("test_mapped_stub_partial_xcom", serialized=True): - - @task.stub - def extract(): ... - - @task.stub - def transform(extracted: dict, country: str): ... - - transform.partial(extracted=extract()).expand(country=["uk", "fr"]) - - dr = dag_maker.create_dagrun() - ti = next(ti for ti in dr.get_task_instances() if ti.task_id == "transform" and ti.map_index == 1) - ti.set_state(State.QUEUED) - dag_maker.session.flush() - - response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) - assert response.status_code == 200 - assert response.json()["arg_bindings"] == [ - { - "name": "extracted", - "kind": "xcom", - "value_schema": {"type": "object", "additionalProperties": True}, - "task_id": "extract", - }, - {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "fr"}, - ] - - def test_ti_run_rejects_partial_kwarg_over_mapped_upstream(self, client, dag_maker): - """ - A partial() kwarg over a mapped upstream would bind the nonexistent unmapped XCom row. - - The provider rejects this at parse time now; patching its capture hook simulates a - Dag serialized by another provider version, exercising the server-side backstop. - """ - from airflow.providers.standard.decorators.stub import _StubOperator - - fabricated = {"_mapped_arg_binding_params": [{"name": "extracted"}, {"name": "country"}]} - with mock.patch.object(_StubOperator, "get_mapped_serialized_fields", return_value=fabricated): - with dag_maker("test_mapped_stub_partial_mapped_upstream", serialized=True): - - @task.stub - def seed(n: int): ... - - @task.stub - def transform(extracted: dict, country: str): ... - - transform.partial(extracted=seed.expand(n=[1, 2])).expand(country=["uk", "fr"]) - - dr = dag_maker.create_dagrun() - ti = next(ti for ti in dr.get_task_instances() if ti.task_id == "transform" and ti.map_index == 0) - ti.set_state(State.QUEUED) - dag_maker.session.flush() - - response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) - assert response.status_code == 500 - assert response.json()["detail"]["reason"] == "invalid_arg_bindings" - assert "aggregated output" in response.json()["detail"]["message"] - - def test_ti_run_orders_mapped_stub_spec_by_declaration_with_defaults(self, client, dag_maker): - """The spec follows the signature, not the call sites, and ships defaulted params.""" - with dag_maker("test_mapped_stub_declaration_order", serialized=True): - - @task.stub - def transform(country: str, extracted: dict, retries_num: int = 3): ... - - # The partial() kwarg is declared after the expanded one on purpose. - transform.partial(extracted={"a": 1}).expand(country=["uk", "fr"]) - - dr = dag_maker.create_dagrun() - ti = next(t for t in dr.get_task_instances() if t.map_index == 1) - ti.set_state(State.QUEUED) - dag_maker.session.flush() - - response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) - assert response.status_code == 200 - assert response.json()["arg_bindings"] == [ - {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "fr"}, - { - "name": "extracted", - "kind": "literal", - "value_schema": {"type": "object", "additionalProperties": True}, - "value": {"a": 1}, - }, - { - "name": "retries_num", - "kind": "literal", - "value_schema": {"type": "integer", "format": "int64"}, - "value": 3, - "from_default": True, - }, - ] - - def test_ti_run_ignores_args_for_legacy_serialized_mapped_stub(self, client, dag_maker): - """A mapped stub serialized without parameter metadata keeps the ignored-args behavior.""" - from airflow.providers.standard.decorators.stub import _StubOperator - - with mock.patch.object(_StubOperator, "get_mapped_serialized_fields", return_value={}): - with dag_maker("test_mapped_stub_legacy", serialized=True): - - @task.stub - def transform(country: str): ... - - transform.expand(country=["uk", "fr"]) + transform.expand(country=["uk", "fr"]) dr = dag_maker.create_dagrun() ti = next(t for t in dr.get_task_instances() if t.map_index == 0) @@ -725,92 +481,6 @@ def transform(country: str): ... assert response.status_code == 200 assert "arg_bindings" not in response.json() - def test_ti_run_rejects_unexpanded_mapped_stub_ti(self, client, dag_maker): - """A mapped stub TI still at map_index=-1 cannot receive per-index bindings.""" - with dag_maker("test_mapped_stub_unexpanded", serialized=True): - - @task.stub - def extract(): ... - - @task.stub - def transform(extracted: dict): ... - - transform.expand(extracted=extract()) - - dr = dag_maker.create_dagrun() - ti = dr.get_task_instance("transform") - assert ti.map_index == -1 - ti.set_state(State.QUEUED) - dag_maker.session.flush() - - response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) - assert response.status_code == 500 - assert response.json()["detail"]["reason"] == "invalid_arg_bindings" - assert "not been expanded" in response.json()["detail"]["message"] - - def test_ti_run_rejects_zero_length_expansion_on_stub(self, client, dag_maker, session): - """An upstream re-run to an empty list after expansion fails structurally, not with a crash.""" - from airflow.models.taskmap import TaskMap - - with dag_maker("test_mapped_stub_zero_length", serialized=True, session=session): - - @task - def seed(): - return [0, 1] - - @task.stub - def combine(a: int, b: int): ... - - combine.expand(a=seed(), b=[1, 2]) - - dr = dag_maker.create_dagrun() - decision = dr.task_instance_scheduling_decisions(session=session) - (seed_ti,) = decision.schedulable_tis - seed_ti.state = TaskInstanceState.SUCCESS - session.add(TaskMap.from_task_instance_xcom(seed_ti, [0, 1])) - session.flush() - - decision = dr.task_instance_scheduling_decisions(session=session) - ti = next(t for t in decision.schedulable_tis if t.map_index == 0) - ti.set_state(State.QUEUED, session=session) - # Simulate the upstream being cleared and re-run to an empty list while this - # expanded TI is still queued. - session.execute(update(TaskMap).where(TaskMap.task_id == "seed").values(length=0, keys=None)) - session.commit() - - response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) - assert response.status_code == 500 - assert response.json()["detail"]["reason"] == "invalid_arg_bindings" - assert "length 0" in response.json()["detail"]["message"] - - def test_ti_run_rejects_expand_kwargs_on_stub(self, client, dag_maker): - """ - expand_kwargs() has no per-parameter spec to derive, so delivery fails structurally. - - The provider rejects this at parse time now; patching its capture hook simulates a - Dag serialized by another provider version, exercising the server-side backstop. - """ - from airflow.providers.standard.decorators.stub import _StubOperator - - fabricated = {"_mapped_arg_binding_params": [{"name": "country"}]} - with mock.patch.object(_StubOperator, "get_mapped_serialized_fields", return_value=fabricated): - with dag_maker("test_mapped_stub_expand_kwargs", serialized=True): - - @task.stub - def transform(country: str): ... - - transform.expand_kwargs([{"country": "uk"}]) - - dr = dag_maker.create_dagrun() - (ti,) = dr.get_task_instances() - ti.set_state(State.QUEUED) - dag_maker.session.flush() - - response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) - assert response.status_code == 500 - assert response.json()["detail"]["reason"] == "invalid_arg_bindings" - assert "expand_kwargs" in response.json()["detail"]["message"] - def test_arg_bindings_adapter_rejects_unknown_kind(self): """The discriminated union refuses serialized specs with an unrecognised kind.""" from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import get_arg_bindings_adapter diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py index b4dd521c760e2..a4b98bd10206e 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py @@ -79,30 +79,6 @@ def test_old_version_strips_arg_bindings_even_when_set(self, old_ver_client, stu assert response.status_code == 200 assert "arg_bindings" not in response.json() - def test_old_version_skips_undeliverable_arg_bindings_derivation(self, old_ver_client, dag_maker): - """A stub whose bindings cannot be delivered must keep running for clients that never see them.""" - with dag_maker("test_arg_bindings_compat_unexpanded", serialized=True): - - @task.stub - def extract(): ... - - @task.stub - def transform(extracted: dict): ... - - transform.expand(extracted=extract()) - - dr = dag_maker.create_dagrun() - ti = dr.get_task_instance("transform") - assert ti.map_index == -1 - ti.set_state(State.QUEUED) - dag_maker.session.flush() - - # At head this TI fails ti_run with a structured 500 (it has not been expanded - # to a map index); a pre-arg-bindings client keeps the legacy behavior. - response = old_ver_client.patch(f"/execution/task-instances/{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 diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index a7459b204cc5c..889131a152032 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -3459,73 +3459,6 @@ def transform(country: str, extracted: dict): ... assert not hasattr(round_tripped.task_dict["extract"], "_arg_bindings") -def test_mapped_stub_param_metadata_round_trip(): - """The serializer collects the mapped stub's parameter metadata and it survives the round trip.""" - from airflow.sdk import task - - with DAG(dag_id="mapped_arg_binding_params_dag", schedule=None) as dag: - - @task.stub - def transform(country: str, extracted, retries_num: int = 3): ... - - @task - def plain(x): ... - - transform.partial(extracted={"a": 1}).expand(country=["uk", "fr"]) - plain.expand(x=[1, 2]) - - ser_dag = DagSerialization.to_dict(dag) - DagSerialization.validate_schema(ser_dag) - encoded_tasks = {t[Encoding.VAR]["task_id"]: t[Encoding.VAR] for t in ser_dag["dag"]["tasks"]} - assert encoded_tasks["transform"]["_mapped_arg_binding_params"] == [ - { - Encoding.TYPE: DAT.DICT, - Encoding.VAR: { - "name": "country", - "value_schema": {Encoding.TYPE: DAT.DICT, Encoding.VAR: {"type": "string"}}, - }, - }, - {Encoding.TYPE: DAT.DICT, Encoding.VAR: {"name": "extracted"}}, - { - Encoding.TYPE: DAT.DICT, - Encoding.VAR: { - "name": "retries_num", - "value_schema": { - Encoding.TYPE: DAT.DICT, - Encoding.VAR: {"type": "integer", "format": "int64"}, - }, - "default": 3, - }, - }, - ], "metadata must serialize in declaration order" - assert "_mapped_arg_binding_params" not in encoded_tasks["plain"], ( - "only operator classes defining the capture hook contribute mapped fields" - ) - - round_tripped = DagSerialization.from_dict(ser_dag) - assert round_tripped.task_dict["transform"]._mapped_arg_binding_params == [ - {"name": "country", "value_schema": {"type": "string"}}, - {"name": "extracted"}, - {"name": "retries_num", "value_schema": {"type": "integer", "format": "int64"}, "default": 3}, - ] - assert not hasattr(round_tripped.task_dict["plain"], "_mapped_arg_binding_params") - - -def test_mapped_stub_capture_error_fails_serialization(): - """A capture-hook rejection surfaces as a Dag serialization (import) error.""" - from airflow.sdk import task - - with DAG(dag_id="mapped_arg_binding_params_invalid_dag", schedule=None) as dag: - - @task.stub - def transform(country: str): ... - - transform.expand_kwargs([{"country": "uk"}]) - - with pytest.raises(SerializationError, match="does not support expand_kwargs"): - DagSerialization.to_dict(dag) - - def test_handle_v1_serdag(): v1 = { "__version": 1, 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 d96b9dce07b4d..7eb25dec2b3cb 100644 --- a/devel-common/src/tests_common/test_utils/version_compat.py +++ b/devel-common/src/tests_common/test_utils/version_compat.py @@ -42,7 +42,6 @@ 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/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index adf554095bac5..72f54421bff32 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -188,7 +188,7 @@ def _ensure_json_literal(value: Any, task_id: str, name: str) -> None: ) -def _validate_xcom_value(value: Any, task_id: str, name: str, *, allow_mapped_upstream: bool = False) -> bool: +def _validate_xcom_value(value: Any, task_id: str, name: str) -> bool: """Validate an XComArg argument, returning True when it is a bindable direct upstream output.""" if isinstance(value, PlainXComArg): if value.key != XCOM_RETURN_KEY: @@ -197,7 +197,7 @@ def _validate_xcom_value(value: Any, task_id: str, name: str, *, allow_mapped_up f"{value.key!r}; only an upstream task's return value can cross the language " "boundary -- indexing an output by a custom key is not supported" ) - if value.operator.is_mapped and not allow_mapped_upstream: + if value.operator.is_mapped: raise ValueError( f"@task.stub task {task_id!r} parameter {name!r} references the aggregated " f"output of the mapped task {value.operator.task_id!r}; a foreign runtime " @@ -266,64 +266,6 @@ def _build_arg_bindings( return spec -def _build_mapped_arg_binding_params( - python_callable: Callable, - *, - partial_op_kwargs: Mapping[str, Any], - expand_input: Any, - task_id: str, -) -> list[dict[str, Any]] | None: - """ - Build the ordered per-parameter metadata for a mapped (``.expand()``) stub task. - - Per-map-index values only resolve at run time, so unlike ``_build_arg_bindings`` this - captures what the server-side derivation cannot recover from the serialized Dag alone: - the declaration order the wire contract promises, defaults for parameters no kwarg - covers, and each parameter's value schema. Returns ``None`` for parameterless stubs - (the legacy fan-out shape whose call args were always ignored keeps parsing). - """ - signature = inspect.signature(python_callable) - if not signature.parameters: - return None - _validate_stub_signature(signature, task_id) - if not isinstance(expand_input.value, Mapping): - # expand_kwargs() carries a list (or upstream XCom) of kwarg dicts whose - # parameter names are unknowable at parse time. - raise ValueError( - f"@task.stub task {task_id!r} does not support expand_kwargs(); the parameter " - "binding must be derivable at parse time, so use .expand() with explicit kwargs" - ) - expand_kwargs = expand_input.value - - try: - bound = signature.bind(**{**partial_op_kwargs, **expand_kwargs}) - except TypeError as e: - raise ValueError(f"@task.stub task {task_id!r} TaskFlow mapping does not bind to its signature: {e}") - bound.apply_defaults() - - annotations = _resolve_param_annotations(python_callable, signature) - - params: list[dict[str, Any]] = [] - for name in signature.parameters: - entry: dict[str, Any] = {"name": name} - if (value_schema := _infer_value_schema(annotations[name])) is not None: - entry["value_schema"] = value_schema - if name in expand_kwargs: - # The whole expanded collection ships through XCom/serialization; an upstream - # output is consumed per element, so a mapped upstream is fine here. - if not _validate_xcom_value(expand_kwargs[name], task_id, name, allow_mapped_upstream=True): - _ensure_json_literal(expand_kwargs[name], task_id, name) - elif name in partial_op_kwargs: - if not _validate_xcom_value(partial_op_kwargs[name], task_id, name): - _ensure_json_literal(partial_op_kwargs[name], task_id, name) - else: - default = bound.arguments[name] - _ensure_json_literal(default, task_id, name) - entry["default"] = default - params.append(entry) - return params - - class _StubOperator(DecoratedOperator): custom_operator_name: str = "@task.stub" @@ -390,25 +332,6 @@ def __init__( def get_serialized_fields(cls): return super().get_serialized_fields() | {"_arg_bindings"} - @classmethod - def get_mapped_serialized_fields(cls, mapped_op: Any) -> dict[str, Any]: - """ - Extra serialized fields for the mapped (``.expand()``) form of this operator. - - Called by the core Dag serializer (Airflow 3.4+) while ``python_callable`` is - still the real function; older cores never call it, so mapped stubs there keep - the legacy ignored-args behavior. - """ - params = _build_mapped_arg_binding_params( - mapped_op.python_callable, - partial_op_kwargs=mapped_op.partial_kwargs.get("op_kwargs") or {}, - expand_input=mapped_op._get_specified_expand_input(), - task_id=mapped_op.task_id, - ) - if params is None: - return {} - return {"_mapped_arg_binding_params": params} - def execute(self, context: Context) -> Any: raise RuntimeError( "@task.stub should not be executed directly -- we expected this to go to a remote worker. " @@ -432,6 +355,10 @@ def stub( outputs or JSON-serializable literals; the resulting argument-binding spec (parameter names, value schemas, and values, in declaration order) is delivered to the foreign runtime, which binds the values onto the native task function. + + Mapped (``.expand()``) stubs do not receive TaskFlow arguments yet -- their call args + keep the legacy ignored behavior; per-map-index delivery is part of + https://github.com/apache/airflow/issues/66937 and lands in a follow-up. """ return task_decorator_factory( decorated_operator_class=_StubOperator, diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index b73c82f38d2f6..68d804de4fc17 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -26,7 +26,7 @@ import pytest from airflow.providers.common.compat.sdk import DAG, task_group -from airflow.providers.standard.decorators.stub import _infer_value_schema, _StubOperator, stub +from airflow.providers.standard.decorators.stub import _infer_value_schema, stub from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS @@ -259,6 +259,7 @@ def test_arg_bindings_survive_dag_serialization_round_trip(self): ] def test_expand_builds_mapped_stub_without_parse_time_bindings(self): + """Mapped stubs capture no spec: their call args keep the legacy ignored behavior for now.""" with DAG(dag_id="d"): result = stub(fn_transform).expand(country=["uk", "fr"], extracted=[{}, {}]) # op_kwargs_expand_input/partial_kwargs (not is_mapped) so the assertions also @@ -287,117 +288,6 @@ def group(n): group.expand(n=[1, 2]) -class TestMappedStubArgBindingParams: - """The serializer hook captures ordered per-parameter metadata for mapped stubs.""" - - def get_hook_fields(self, operator): - return _StubOperator.get_mapped_serialized_fields(operator) - - def test_params_follow_declaration_order_with_defaults_and_schemas(self): - with DAG(dag_id="d"): - # The partial() kwarg is declared *after* the expanded one: the captured - # order must come from the signature, not from the call sites. - result = stub(fn_transform).partial(extracted={"a": 1}).expand(country=["uk", "fr"]) - - assert self.get_hook_fields(result.operator) == { - "_mapped_arg_binding_params": [ - {"name": "country", "value_schema": {"type": "string"}}, - {"name": "extracted", "value_schema": {"type": "object", "additionalProperties": True}}, - { - "name": "retries_num", - "value_schema": {"type": "integer", "format": "int64"}, - "default": 3, - }, - ] - } - - def test_none_default_is_captured_by_key_presence(self): - def fn(x: str, y=None): ... - - with DAG(dag_id="d"): - result = stub(fn).expand(x=["a"]) - - params = self.get_hook_fields(result.operator)["_mapped_arg_binding_params"] - assert params[1] == {"name": "y", "default": None} - - def test_untyped_params_omit_value_schema(self): - with DAG(dag_id="d"): - result = stub(fn_untyped).expand(a=[1], b=[2]) - - assert self.get_hook_fields(result.operator) == { - "_mapped_arg_binding_params": [{"name": "a"}, {"name": "b"}] - } - - def test_parameterless_stub_captures_nothing(self): - with DAG(dag_id="d"): - result = stub(fn_extract).expand_kwargs([{}]) - - assert self.get_hook_fields(result.operator) == {} - - def test_expand_kwargs_rejected_for_parameterful_stub(self): - with DAG(dag_id="d"): - result = stub(fn_transform).expand_kwargs([{"country": "uk", "extracted": {}}]) - - with pytest.raises(ValueError, match="does not support expand_kwargs"): - self.get_hook_fields(result.operator) - - def test_missing_required_parameter_rejected(self): - with DAG(dag_id="d"): - result = stub(fn_transform).expand(country=["uk"]) - - with pytest.raises(ValueError, match="does not bind to its signature"): - self.get_hook_fields(result.operator) - - def test_partial_kwarg_over_mapped_upstream_rejected(self): - def fn_produce(n: int): ... - - with DAG(dag_id="d"): - vals = stub(fn_produce).expand(n=[1, 2]) - result = stub(fn_transform).partial(extracted=vals).expand(country=["uk"]) - - with pytest.raises(ValueError, match="aggregated output of the mapped task"): - self.get_hook_fields(result.operator) - - def test_expand_over_mapped_upstream_allowed(self): - def fn_produce(n: int): ... - - with DAG(dag_id="d"): - vals = stub(fn_produce).expand(n=[1, 2]) - result = stub(fn_transform).partial(country="uk").expand(extracted=vals) - - params = self.get_hook_fields(result.operator)["_mapped_arg_binding_params"] - assert [p["name"] for p in params] == ["country", "extracted", "retries_num"] - - def test_non_json_expand_literal_rejected(self): - with DAG(dag_id="d"): - result = stub(fn_transform).partial(country="uk").expand(extracted=[object()]) - - with pytest.raises(ValueError, match="not JSON-serializable"): - self.get_hook_fields(result.operator) - - def test_non_json_needed_default_rejected(self): - not_jsonable = object() - - def fn(x: str, y=not_jsonable): ... - - with DAG(dag_id="d"): - result = stub(fn).expand(x=["a"]) - - with pytest.raises(ValueError, match="not JSON-serializable"): - self.get_hook_fields(result.operator) - - def test_mapped_stub_inside_mapped_task_group_unconstructible(self): - """The SDK bans expansion inside an expanded group outright, so no hook guard is needed.""" - - @task_group - def group(n): - stub(fn_transform).partial(country="uk").expand(extracted=[{}]) - - with DAG(dag_id="d"): - with pytest.raises(NotImplementedError, match="expansion in an expanded task group"): - group.expand(n=[1, 2]) - - @pytest.mark.parametrize( ("annotation", "expected"), [ From dca8397fbf6d82a5c122bc28d4d9d70f42d628be Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 28 Jul 2026 02:47:50 +0000 Subject: [PATCH 39/40] Drop mapped-only fields from the arg-binding wire contract XComArgBinding carried map_index and element_index for the per-map-index delivery that now lands in the stacked follow-up branch; the unmapped path never sets either, so this PR ships the contract without them. The task-sdk client models, supervisor schema snapshot, and ts-sdk types are regenerated accordingly; the follow-up re-adds the fields with its derivation. --- .../datamodels/task_arg_binding.py | 9 -- .../airflow/sdk/api/datamodels/_generated.py | 2 - .../sdk/execution_time/schema/schema.json | 89 ++++++++----------- ts-sdk/src/generated/supervisor.ts | 32 +++---- 4 files changed, 50 insertions(+), 82 deletions(-) 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 index 982c199868394..69b6453bb1c8e 100644 --- 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 @@ -56,15 +56,6 @@ class XComArgBinding(BaseModel): task_id: str """Upstream task id whose ``return_value`` XCom is pulled.""" - map_index: int = -1 - """Map index of the upstream XCom row to pull; -1 is the unmapped row.""" - - element_index: int | None = None - """When set, the pulled value is a sequence and this binding takes the element at - this index (the stub was expanded over an unmapped upstream's output). The lang-SDK - side will GET the single row ``(task_id, map_index=-1)``, decode it, and take - ``value[element_index]``.""" - class LiteralArgBinding(BaseModel): """One positional stub-task argument carrying an inline literal from the Dag file.""" diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 739ac5ef66595..d97d38613834c 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -735,8 +735,6 @@ class XComArgBinding(BaseModel): name: Annotated[str, Field(title="Name")] value_schema: ArgValueSchema | None = None task_id: Annotated[str, Field(title="Task Id")] - map_index: Annotated[int | None, Field(title="Map Index")] = -1 - element_index: Annotated[int | None, Field(title="Element Index")] = None class AssetEventDagRunReference(BaseModel): 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 4c2f503dc828e..8e47af32cf5d4 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -4396,6 +4396,42 @@ "title": "VariableResult", "type": "object" }, + "XComArgBinding": { + "description": "One positional stub-task argument pulled from an upstream task's XCom.", + "properties": { + "kind": { + "const": "xcom", + "title": "Kind", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "value_schema": { + "anyOf": [ + { + "$ref": "#/$defs/ArgValueSchema" + }, + { + "type": "null" + } + ], + "default": null + }, + "task_id": { + "title": "Task Id", + "type": "string" + } + }, + "required": [ + "kind", + "name", + "task_id" + ], + "title": "XComArgBinding", + "type": "object" + }, "XComCountResponse": { "properties": { "len": { @@ -4635,59 +4671,6 @@ ], "title": "TaskArgBinding" }, - "XComArgBinding": { - "description": "One positional stub-task argument pulled from an upstream task's XCom.", - "properties": { - "kind": { - "const": "xcom", - "title": "Kind", - "type": "string" - }, - "name": { - "title": "Name", - "type": "string" - }, - "value_schema": { - "anyOf": [ - { - "$ref": "#/$defs/ArgValueSchema" - }, - { - "type": "null" - } - ], - "default": null - }, - "task_id": { - "title": "Task Id", - "type": "string" - }, - "map_index": { - "default": -1, - "title": "Map Index", - "type": "integer" - }, - "element_index": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "default": null, - "title": "Element Index" - } - }, - "required": [ - "kind", - "name", - "task_id" - ], - "title": "XComArgBinding", - "type": "object" - }, "AssetEventDagRunReference": { "additionalProperties": false, "description": "Schema for AssetEvent model used in DagRun.", diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 31d44a52b18c5..af944e4e565d7 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -254,8 +254,6 @@ export type TaskArgBinding = XComArgBinding | LiteralArgBinding; export type Kind = "xcom"; export type Name8 = string; export type TaskId1 = string; -export type MapIndex1 = number; -export type ElementIndex = number | null; export type Kind1 = "literal"; export type Name9 = string; export type FromDefault = boolean; @@ -339,7 +337,7 @@ export type Key5 = string; export type DagId6 = string; export type RunId5 = string; export type TaskId2 = string; -export type MapIndex2 = number | null; +export type MapIndex1 = number | null; export type Type25 = "DeleteXCom"; /** * Error types used in the API client. @@ -437,10 +435,10 @@ export type Type41 = "GetPreviousDagRun"; export type DagId12 = string; export type TaskId3 = string; export type LogicalDate4 = string | null; -export type MapIndex3 = number; +export type MapIndex2 = number; export type Type42 = "GetPreviousTI"; export type DagId13 = string; -export type MapIndex4 = number | null; +export type MapIndex3 = number | null; export type TaskIds = string[] | null; export type TaskGroupId = string | null; export type LogicalDates1 = string[] | null; @@ -457,7 +455,7 @@ export type TiId6 = string; export type Key8 = string; export type Type46 = "GetTaskStateStore"; export type DagId15 = string; -export type MapIndex5 = number | null; +export type MapIndex4 = number | null; export type TaskIds1 = string[] | null; export type TaskGroupId1 = string | null; export type LogicalDates2 = string[] | null; @@ -473,7 +471,7 @@ export type Key10 = string; export type DagId16 = string; export type RunId9 = string; export type TaskId4 = string; -export type MapIndex6 = number | null; +export type MapIndex5 = number | null; export type IncludePriorDates = boolean; export type Type50 = "GetXCom"; export type Key11 = string; @@ -530,7 +528,7 @@ export type StartDate5 = string | null; export type EndDate4 = string | null; export type State4 = string | null; export type TryNumber2 = number; -export type MapIndex7 = number | null; +export type MapIndex6 = number | null; export type Duration = number | null; export type Type60 = "PreviousTIResult"; export type Key14 = string; @@ -567,7 +565,7 @@ export type Key18 = string; export type DagId21 = string; export type RunId14 = string; export type TaskId9 = string; -export type MapIndex8 = number | null; +export type MapIndex7 = number | null; export type DagResult1 = boolean; export type MappedLength = number | null; export type Type71 = "SetXCom"; @@ -1063,8 +1061,6 @@ export interface XComArgBinding { name: Name8; value_schema?: ArgValueSchema | null; task_id: TaskId1; - map_index?: MapIndex1; - element_index?: ElementIndex; } /** * One positional stub-task argument carrying an inline literal from the Dag file. @@ -1237,7 +1233,7 @@ export interface DeleteXCom { dag_id: DagId6; run_id: RunId5; task_id: TaskId2; - map_index?: MapIndex2; + map_index?: MapIndex1; type?: Type25; } /** @@ -1405,7 +1401,7 @@ export interface GetPreviousTI { dag_id: DagId12; task_id: TaskId3; logical_date?: LogicalDate4; - map_index?: MapIndex3; + map_index?: MapIndex2; state?: TaskInstanceState | null; type?: Type42; } @@ -1415,7 +1411,7 @@ export interface GetPreviousTI { */ export interface GetTICount { dag_id: DagId13; - map_index?: MapIndex4; + map_index?: MapIndex3; task_ids?: TaskIds; task_group_id?: TaskGroupId; logical_dates?: LogicalDates1; @@ -1456,7 +1452,7 @@ export interface GetTaskStateStore { */ export interface GetTaskStates { dag_id: DagId15; - map_index?: MapIndex5; + map_index?: MapIndex4; task_ids?: TaskIds1; task_group_id?: TaskGroupId1; logical_dates?: LogicalDates2; @@ -1490,7 +1486,7 @@ export interface GetXCom { dag_id: DagId16; run_id: RunId9; task_id: TaskId4; - map_index?: MapIndex6; + map_index?: MapIndex5; include_prior_dates?: IncludePriorDates; type?: Type50; } @@ -1616,7 +1612,7 @@ export interface PreviousTIResponse { end_date?: EndDate4; state?: State4; try_number: TryNumber2; - map_index?: MapIndex7; + map_index?: MapIndex6; duration?: Duration; } /** @@ -1744,7 +1740,7 @@ export interface SetXCom { dag_id: DagId21; run_id: RunId14; task_id: TaskId9; - map_index?: MapIndex8; + map_index?: MapIndex7; dag_result?: DagResult1; mapped_length?: MappedLength; type?: Type71; From 2c949d8c4e6a118b371636b7539012ed74e49bd3 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 28 Jul 2026 02:54:18 +0000 Subject: [PATCH 40/40] Support dynamic task mapping on stub tasks Restore the mapped (.expand()) stub arg-binding support split out of the unmapped PR: the Dag serializer captures per-parameter metadata (declaration order, defaults, value schemas) from the stub signature via the get_mapped_serialized_fields hook, ti_run derives per-map-index bindings from it with the map-index decomposition on SchedulerDictOfListsExpandInput, and XComArgBinding regains the map_index/element_index delivery fields across the task-sdk and ts-sdk generated models. --- .../datamodels/task_arg_binding.py | 9 + .../execution_api/services/task_instances.py | 134 ++++++- .../src/airflow/models/expandinput.py | 27 ++ .../src/airflow/serialization/schema.json | 29 ++ .../serialization/serialized_objects.py | 10 + .../versions/head/test_task_instances.py | 338 +++++++++++++++++- .../v2026_10_30/test_task_instances.py | 24 ++ .../serialization/test_dag_serialization.py | 67 ++++ .../tests_common/test_utils/version_compat.py | 1 + .../providers/standard/decorators/stub.py | 85 ++++- .../unit/standard/decorators/test_stub.py | 114 +++++- .../airflow/sdk/api/datamodels/_generated.py | 2 + .../sdk/execution_time/schema/schema.json | 89 +++-- ts-sdk/src/generated/supervisor.ts | 32 +- 14 files changed, 891 insertions(+), 70 deletions(-) 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 index 69b6453bb1c8e..982c199868394 100644 --- 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 @@ -56,6 +56,15 @@ class XComArgBinding(BaseModel): task_id: str """Upstream task id whose ``return_value`` XCom is pulled.""" + map_index: int = -1 + """Map index of the upstream XCom row to pull; -1 is the unmapped row.""" + + element_index: int | None = None + """When set, the pulled value is a sequence and this binding takes the element at + this index (the stub was expanded over an unmapped upstream's output). The lang-SDK + side will GET the single row ``(task_id, map_index=-1)``, decode it, and take + ``value[element_index]``.""" + class LiteralArgBinding(BaseModel): """One positional stub-task argument carrying an inline literal from the Dag file.""" diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py index e06db8ffa70be..4a8c2589438d8 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py @@ -18,12 +18,23 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +import json +from typing import TYPE_CHECKING, Any, NoReturn + +from fastapi import HTTPException, status + +from airflow.models.expandinput import NotFullyPopulated, SchedulerDictOfListsExpandInput +from airflow.models.xcom import XCOM_RETURN_KEY +from airflow.serialization.definitions.mappedoperator import is_mapped +from airflow.serialization.definitions.xcom_arg import SchedulerPlainXComArg, SchedulerXComArg +from airflow.serialization.serialized_objects import _XComRef if TYPE_CHECKING: from sqlalchemy.orm import Session from airflow.models.dagbag import DBDagBag + from airflow.serialization.definitions.dag import SerializedDAG + from airflow.serialization.definitions.mappedoperator import SerializedMappedOperator # Task type recorded on the TI row (``TaskInstance.operator``) for # ``airflow.providers.standard.decorators.stub._StubOperator``. Used to gate the @@ -33,17 +44,124 @@ def get_arg_bindings(dag_bag: DBDagBag, ti: Any, *, session: Session) -> list | None: - """ - Extract the stub task's TaskFlow arg spec from its Dag version. - - Mapped (``.expand()``) stubs never capture a parse-time spec, so they resolve to - ``None`` here and keep the legacy ignored-args behavior; per-map-index delivery - lands in a follow-up. - """ + """Extract or derive the stub task's TaskFlow arg spec from its Dag version.""" if ti.dag_version_id is None: return None if (dag := dag_bag.get_dag(ti.dag_version_id, session=session)) is None: return None if (task := dag.task_dict.get(ti.task_id)) is None: return None + if is_mapped(task): + return _resolve_mapped_stub_arg_bindings(task, ti, dag=dag, session=session) return getattr(task, "_arg_bindings", None) + + +def _unsupported_arg_bindings(detail: str) -> NoReturn: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={ + "reason": "invalid_arg_bindings", + "message": f"The stub task's TaskFlow arguments cannot be delivered: {detail}.", + }, + ) + + +def _resolve_mapped_stub_arg_bindings( + task: SerializedMappedOperator, ti: Any, *, dag: SerializedDAG, session: Session +) -> list[dict[str, Any]] | None: + """ + Build the per-map-index arg spec for a mapped (``.expand()``) stub task. + + A mapped stub never instantiates at parse time; the Dag serializer captures its + per-parameter metadata (declaration order, defaults, value schemas) from the stub + signature via ``get_mapped_serialized_fields``, and the map-index decomposition is + delegated to ``SchedulerDictOfListsExpandInput.resolve_expansion_sub_indexes``. + Dags serialized without the metadata (an older provider) resolve to ``None``: their + args were never deliverable, so they keep the legacy ignored-args behavior rather + than receive bindings whose order the server cannot know. + """ + metadata = getattr(task, "_mapped_arg_binding_params", None) + if metadata is None: + return None + # The isinstance/map_index/unclaimed checks below re-reject what the provider now + # fails at parse time, for serialized Dags produced by other provider versions. + expand_input = task._get_specified_expand_input() + if not isinstance(expand_input, SchedulerDictOfListsExpandInput): + _unsupported_arg_bindings("expand_kwargs() is not supported on stub tasks") + if ti.map_index < 0: + _unsupported_arg_bindings("the task instance has not been expanded to a map index") + + expand_value = expand_input.value + partial_op_kwargs = task.partial_kwargs.get("op_kwargs") or {} + if unclaimed := (set(expand_value) | set(partial_op_kwargs)) - {meta["name"] for meta in metadata}: + _unsupported_arg_bindings(f"kwargs {sorted(unclaimed)} are not in the captured parameter metadata") + try: + sub_indexes = expand_input.resolve_expansion_sub_indexes(ti.map_index, ti.run_id, session=session) + except NotFullyPopulated as e: + # Neither this nor the ValueError below can happen on the happy path: both take + # someone clearing upstream TIs or XComs themselves during the DagRun. + _unsupported_arg_bindings(f"upstream map lengths are not yet known for {sorted(e.missing)}") + except ValueError as e: + _unsupported_arg_bindings(str(e)) + + spec = [] + for meta in metadata: # Declaration order, captured at parse time. + name = meta["name"] + if name in expand_value: + entry = _bind_mapped_stub_arg(name, expand_value[name], sub_index=sub_indexes[name]) + elif name in partial_op_kwargs: + value = partial_op_kwargs[name] + # XComArgs inside partial() op_kwargs deserialize to _XComRef and are never + # dereferenced (set_task_dag_references only derefs the expand inputs). + if isinstance(value, _XComRef): + value = value.deref(dag) + entry = _bind_mapped_stub_arg(name, value, sub_index=None) + elif "default" in meta: + entry = {"name": name, "kind": "literal", "value": meta["default"], "from_default": True} + else: + _unsupported_arg_bindings(f"parameter {name!r} has no expanded, partial, or default value") + if (value_schema := meta.get("value_schema")) is not None: + entry["value_schema"] = value_schema + spec.append(entry) + return spec + + +def _bind_mapped_stub_arg(name: str, value: Any, *, sub_index: int | None) -> dict[str, Any]: + """Build one arg-binding dict; ``sub_index`` is set for expanded kwargs, None for partial ones.""" + if isinstance(value, SchedulerPlainXComArg): + if value.key != XCOM_RETURN_KEY: + _unsupported_arg_bindings(f"parameter {name!r} references the XCom key {value.key!r}") + if sub_index is None and value.operator.is_mapped: + # A partial() kwarg over a mapped upstream would bind the unmapped XCom row + # (map_index=-1), which never exists; the aggregated output is inexpressible. + _unsupported_arg_bindings( + f"parameter {name!r} references the aggregated output of the mapped task" + f" {value.operator.task_id!r}" + ) + entry: dict[str, Any] = {"name": name, "kind": "xcom", "task_id": value.operator.task_id} + if sub_index is not None: + if value.operator.is_mapped: + entry["map_index"] = sub_index + else: + entry["element_index"] = sub_index + return entry + if isinstance(value, SchedulerXComArg): + _unsupported_arg_bindings( + f"parameter {name!r} received a {type(value).__name__}; only direct upstream" + " task outputs and literals are supported" + ) + if sub_index is not None: + # This kwarg was expanded over a literal collection written in the Dag file. + items = list(value.items()) if isinstance(value, dict) else value + try: + value = items[sub_index] + except (IndexError, KeyError, TypeError): + _unsupported_arg_bindings(f"parameter {name!r} has no element at expansion index {sub_index}") + try: + json.dumps(value, allow_nan=False) + except (TypeError, ValueError): + _unsupported_arg_bindings( + f"parameter {name!r} carries a {type(value).__name__} value, which cannot cross" + " the language boundary" + ) + return {"name": name, "kind": "literal", "value": value} diff --git a/airflow-core/src/airflow/models/expandinput.py b/airflow-core/src/airflow/models/expandinput.py index 0363bae92620a..ecc1b398eb735 100644 --- a/airflow-core/src/airflow/models/expandinput.py +++ b/airflow-core/src/airflow/models/expandinput.py @@ -151,6 +151,33 @@ def get_total_map_length(self, run_id: str, *, session: Session) -> int: lengths = self._get_map_lengths(run_id, session=session) return functools.reduce(operator.mul, (lengths[name] for name in self.value), 1) + def resolve_expansion_sub_indexes( + self, map_index: int, run_id: str, *, session: Session + ) -> dict[str, int]: + """ + Decompose a task instance's map index into one index per expanded kwarg. + + Server-side counterpart of the index decomposition in the SDK's + ``DictOfListsExpandInput._expand_mapped_field``: the cross-product of the + expanded kwargs is ordered with the last kwarg varying fastest. A single + expanded kwarg maps one-to-one, skipping the upstream length lookups. + + :raises NotFullyPopulated: if upstream map lengths are not all known yet. + :raises ValueError: if an expanded kwarg's recorded length is zero, e.g. an + upstream was cleared and re-ran to an empty list after this task instance + was expanded (the SDK twin guards the same case). + """ + if len(self.value) == 1: + return dict.fromkeys(self.value, map_index) + lengths = self._get_map_lengths(run_id, session=session) + sub_indexes = {} + for key in reversed(self.value): + if (length := lengths[key]) < 1: + raise ValueError(f"cannot decompose map index over expanded kwarg {key!r} of length 0") + sub_indexes[key] = map_index % length + map_index //= length + return sub_indexes + def iter_references(self) -> Iterable[tuple[Operator, str]]: from airflow.models.referencemixin import ReferenceMixin diff --git a/airflow-core/src/airflow/serialization/schema.json b/airflow-core/src/airflow/serialization/schema.json index fb954f0615b9d..39a33d1b7757b 100644 --- a/airflow-core/src/airflow/serialization/schema.json +++ b/airflow-core/src/airflow/serialization/schema.json @@ -184,6 +184,30 @@ ], "additionalProperties": false }, + "arg_binding_param": { + "$comment": "Per-parameter metadata of a mapped @task.stub task, in dict-encoded form and declaration order. The inner object stays open so future metadata fields keep validating on older cores", + "type": "object", + "properties": { + "__type": { + "type": "string", + "const": "dict" + }, + "__var": { + "type": "object", + "properties": { + "name": { "type": "string" }, + "value_schema": { "$ref": "#/definitions/typed_dict" }, + "default": {} + }, + "required": [ "name" ] + } + }, + "required": [ + "__type", + "__var" + ], + "additionalProperties": false + }, "color": { "type": "string", "pattern": "^#[a-fA-F0-9]{3,6}$" @@ -392,6 +416,11 @@ "$comment": "Only present on @task.stub tasks called with TaskFlow arguments", "type": "array", "items": { "$ref": "#/definitions/arg_binding" } + }, + "_mapped_arg_binding_params": { + "$comment": "Only present on mapped @task.stub tasks with parameters; ordered per-parameter binding metadata", + "type": "array", + "items": { "$ref": "#/definitions/arg_binding_param" } } }, "dependencies": { diff --git a/airflow-core/src/airflow/serialization/serialized_objects.py b/airflow-core/src/airflow/serialization/serialized_objects.py index 54bc3389c64ce..8f19dd9b8b20a 100644 --- a/airflow-core/src/airflow/serialization/serialized_objects.py +++ b/airflow-core/src/airflow/serialization/serialized_objects.py @@ -996,6 +996,16 @@ def serialize_mapped_operator(cls, op: MappedOperator) -> dict[str, Any]: ) del serialized_op["partial_kwargs"]["python_callable"] + # Optional per-class capability: an operator class may contribute extra serialized + # fields for its mapped form. This is the only point where operator_class is the + # real class and python_callable the real function (neither survives + # serialization), so signature-derived data must be captured here. Used by the + # standard provider's _StubOperator for its TaskFlow arg-binding metadata. + get_extra_fields = getattr(op.operator_class, "get_mapped_serialized_fields", None) + if get_extra_fields is not None: + for key, value in get_extra_fields(op).items(): + serialized_op[key] = cls.serialize(value) + serialized_op["_is_mapped"] = True return serialized_op 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 e5dfea5360521..0d61d4625f97a 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 @@ -17,6 +17,7 @@ from __future__ import annotations +import itertools from datetime import datetime from types import SimpleNamespace from typing import TYPE_CHECKING @@ -463,14 +464,257 @@ def transform(country: str): ... "start_date": "2024-09-30T12:00:00Z", } - def test_ti_run_returns_no_arg_bindings_for_mapped_stub(self, client, dag_maker): - """Mapped stubs keep the legacy ignored-args behavior until per-map-index delivery lands.""" - with dag_maker("test_mapped_stub_ignored_args", serialized=True): + def test_ti_run_resolves_mapped_stub_literal_expand(self, client, dag_maker): + """Expanding a stub over a literal list resolves each map index to its element server-side.""" + with dag_maker("test_mapped_stub_literal", serialized=True): @task.stub def transform(country: str): ... - transform.expand(country=["uk", "fr"]) + transform.expand(country=["uk", "fr", "de"]) + + dr = dag_maker.create_dagrun() + tis = {ti.map_index: ti for ti in dr.get_task_instances()} + assert set(tis) == {0, 1, 2} + for ti in tis.values(): + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + for map_index, country in enumerate(["uk", "fr", "de"]): + response = client.patch( + f"/execution/task-instances/{tis[map_index].id}/run", json=self.RUN_PAYLOAD + ) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": country} + ] + + def test_ti_run_resolves_mapped_stub_over_unmapped_upstream(self, client, dag_maker): + """Expanding over an unmapped upstream's output binds the whole XCom plus an element index.""" + with dag_maker("test_mapped_stub_unmapped_upstream", serialized=True): + + @task.stub + def extract(): ... + + @task.stub + def transform(extracted: dict): ... + + transform.expand(extracted=extract()) + + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance("transform") + ti.map_index = 1 + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "extract", + "element_index": 1, + } + ] + + def test_ti_run_resolves_mapped_stub_over_mapped_upstream(self, client, dag_maker): + """Expanding over a mapped upstream binds the upstream XCom row at the same map index.""" + with dag_maker("test_mapped_stub_mapped_upstream", serialized=True): + + @task.stub + def seed(n: int): ... + + @task.stub + def transform(extracted: dict): ... + + transform.expand(extracted=seed.expand(n=[1, 2])) + + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance("transform") + ti.map_index = 1 + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "seed", + "map_index": 1, + } + ] + + def test_ti_run_decomposes_multi_kwarg_mapped_stub(self, client, dag_maker): + """Cross-product expansion decomposes the map index per kwarg like the task-sdk does.""" + with dag_maker("test_mapped_stub_multi_kwarg", serialized=True): + + @task.stub + def combine(a: str, b: int): ... + + combine.expand(a=["x", "y"], b=[1, 2, 3]) + + dr = dag_maker.create_dagrun() + tis = {ti.map_index: ti for ti in dr.get_task_instances()} + assert set(tis) == set(range(6)) + for ti in tis.values(): + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + for map_index, (a, b) in enumerate(itertools.product(["x", "y"], [1, 2, 3])): + response = client.patch( + f"/execution/task-instances/{tis[map_index].id}/run", json=self.RUN_PAYLOAD + ) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "a", "kind": "literal", "value_schema": {"type": "string"}, "value": a}, + { + "name": "b", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": b, + }, + ] + + def test_ti_run_binds_partial_kwargs_of_mapped_stub(self, client, dag_maker): + """partial() kwargs bind like an unmapped TaskFlow call alongside the expanded ones.""" + with dag_maker("test_mapped_stub_partial", serialized=True): + + @task.stub + def transform(country: str, extracted: dict): ... + + transform.partial(country="uk").expand(extracted=[{"a": 1}, {"b": 2}]) + + dr = dag_maker.create_dagrun() + tis = {ti.map_index: ti for ti in dr.get_task_instances()} + for ti in tis.values(): + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + for map_index, extracted in enumerate([{"a": 1}, {"b": 2}]): + response = client.patch( + f"/execution/task-instances/{tis[map_index].id}/run", json=self.RUN_PAYLOAD + ) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, + { + "name": "extracted", + "kind": "literal", + "value_schema": {"type": "object", "additionalProperties": True}, + "value": extracted, + }, + ] + + def test_ti_run_binds_partial_xcom_kwarg_over_unmapped_upstream(self, client, dag_maker): + """A partial() kwarg carrying an unmapped upstream's output binds that XCom for every index.""" + with dag_maker("test_mapped_stub_partial_xcom", serialized=True): + + @task.stub + def extract(): ... + + @task.stub + def transform(extracted: dict, country: str): ... + + transform.partial(extracted=extract()).expand(country=["uk", "fr"]) + + dr = dag_maker.create_dagrun() + ti = next(ti for ti in dr.get_task_instances() if ti.task_id == "transform" and ti.map_index == 1) + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "extract", + }, + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "fr"}, + ] + + def test_ti_run_rejects_partial_kwarg_over_mapped_upstream(self, client, dag_maker): + """ + A partial() kwarg over a mapped upstream would bind the nonexistent unmapped XCom row. + + The provider rejects this at parse time now; patching its capture hook simulates a + Dag serialized by another provider version, exercising the server-side backstop. + """ + from airflow.providers.standard.decorators.stub import _StubOperator + + fabricated = {"_mapped_arg_binding_params": [{"name": "extracted"}, {"name": "country"}]} + with mock.patch.object(_StubOperator, "get_mapped_serialized_fields", return_value=fabricated): + with dag_maker("test_mapped_stub_partial_mapped_upstream", serialized=True): + + @task.stub + def seed(n: int): ... + + @task.stub + def transform(extracted: dict, country: str): ... + + transform.partial(extracted=seed.expand(n=[1, 2])).expand(country=["uk", "fr"]) + + dr = dag_maker.create_dagrun() + ti = next(ti for ti in dr.get_task_instances() if ti.task_id == "transform" and ti.map_index == 0) + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 500 + assert response.json()["detail"]["reason"] == "invalid_arg_bindings" + assert "aggregated output" in response.json()["detail"]["message"] + + def test_ti_run_orders_mapped_stub_spec_by_declaration_with_defaults(self, client, dag_maker): + """The spec follows the signature, not the call sites, and ships defaulted params.""" + with dag_maker("test_mapped_stub_declaration_order", serialized=True): + + @task.stub + def transform(country: str, extracted: dict, retries_num: int = 3): ... + + # The partial() kwarg is declared after the expanded one on purpose. + transform.partial(extracted={"a": 1}).expand(country=["uk", "fr"]) + + dr = dag_maker.create_dagrun() + ti = next(t for t in dr.get_task_instances() if t.map_index == 1) + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "fr"}, + { + "name": "extracted", + "kind": "literal", + "value_schema": {"type": "object", "additionalProperties": True}, + "value": {"a": 1}, + }, + { + "name": "retries_num", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 3, + "from_default": True, + }, + ] + + def test_ti_run_ignores_args_for_legacy_serialized_mapped_stub(self, client, dag_maker): + """A mapped stub serialized without parameter metadata keeps the ignored-args behavior.""" + from airflow.providers.standard.decorators.stub import _StubOperator + + with mock.patch.object(_StubOperator, "get_mapped_serialized_fields", return_value={}): + with dag_maker("test_mapped_stub_legacy", serialized=True): + + @task.stub + def transform(country: str): ... + + transform.expand(country=["uk", "fr"]) dr = dag_maker.create_dagrun() ti = next(t for t in dr.get_task_instances() if t.map_index == 0) @@ -481,6 +725,92 @@ def transform(country: str): ... assert response.status_code == 200 assert "arg_bindings" not in response.json() + def test_ti_run_rejects_unexpanded_mapped_stub_ti(self, client, dag_maker): + """A mapped stub TI still at map_index=-1 cannot receive per-index bindings.""" + with dag_maker("test_mapped_stub_unexpanded", serialized=True): + + @task.stub + def extract(): ... + + @task.stub + def transform(extracted: dict): ... + + transform.expand(extracted=extract()) + + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance("transform") + assert ti.map_index == -1 + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 500 + assert response.json()["detail"]["reason"] == "invalid_arg_bindings" + assert "not been expanded" in response.json()["detail"]["message"] + + def test_ti_run_rejects_zero_length_expansion_on_stub(self, client, dag_maker, session): + """An upstream re-run to an empty list after expansion fails structurally, not with a crash.""" + from airflow.models.taskmap import TaskMap + + with dag_maker("test_mapped_stub_zero_length", serialized=True, session=session): + + @task + def seed(): + return [0, 1] + + @task.stub + def combine(a: int, b: int): ... + + combine.expand(a=seed(), b=[1, 2]) + + dr = dag_maker.create_dagrun() + decision = dr.task_instance_scheduling_decisions(session=session) + (seed_ti,) = decision.schedulable_tis + seed_ti.state = TaskInstanceState.SUCCESS + session.add(TaskMap.from_task_instance_xcom(seed_ti, [0, 1])) + session.flush() + + decision = dr.task_instance_scheduling_decisions(session=session) + ti = next(t for t in decision.schedulable_tis if t.map_index == 0) + ti.set_state(State.QUEUED, session=session) + # Simulate the upstream being cleared and re-run to an empty list while this + # expanded TI is still queued. + session.execute(update(TaskMap).where(TaskMap.task_id == "seed").values(length=0, keys=None)) + session.commit() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 500 + assert response.json()["detail"]["reason"] == "invalid_arg_bindings" + assert "length 0" in response.json()["detail"]["message"] + + def test_ti_run_rejects_expand_kwargs_on_stub(self, client, dag_maker): + """ + expand_kwargs() has no per-parameter spec to derive, so delivery fails structurally. + + The provider rejects this at parse time now; patching its capture hook simulates a + Dag serialized by another provider version, exercising the server-side backstop. + """ + from airflow.providers.standard.decorators.stub import _StubOperator + + fabricated = {"_mapped_arg_binding_params": [{"name": "country"}]} + with mock.patch.object(_StubOperator, "get_mapped_serialized_fields", return_value=fabricated): + with dag_maker("test_mapped_stub_expand_kwargs", serialized=True): + + @task.stub + def transform(country: str): ... + + transform.expand_kwargs([{"country": "uk"}]) + + dr = dag_maker.create_dagrun() + (ti,) = dr.get_task_instances() + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 500 + assert response.json()["detail"]["reason"] == "invalid_arg_bindings" + assert "expand_kwargs" in response.json()["detail"]["message"] + def test_arg_bindings_adapter_rejects_unknown_kind(self): """The discriminated union refuses serialized specs with an unrecognised kind.""" from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import get_arg_bindings_adapter diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py index a4b98bd10206e..b4dd521c760e2 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py @@ -79,6 +79,30 @@ def test_old_version_strips_arg_bindings_even_when_set(self, old_ver_client, stu assert response.status_code == 200 assert "arg_bindings" not in response.json() + def test_old_version_skips_undeliverable_arg_bindings_derivation(self, old_ver_client, dag_maker): + """A stub whose bindings cannot be delivered must keep running for clients that never see them.""" + with dag_maker("test_arg_bindings_compat_unexpanded", serialized=True): + + @task.stub + def extract(): ... + + @task.stub + def transform(extracted: dict): ... + + transform.expand(extracted=extract()) + + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance("transform") + assert ti.map_index == -1 + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + # At head this TI fails ti_run with a structured 500 (it has not been expanded + # to a map index); a pre-arg-bindings client keeps the legacy behavior. + response = old_ver_client.patch(f"/execution/task-instances/{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 diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index 889131a152032..a7459b204cc5c 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -3459,6 +3459,73 @@ def transform(country: str, extracted: dict): ... assert not hasattr(round_tripped.task_dict["extract"], "_arg_bindings") +def test_mapped_stub_param_metadata_round_trip(): + """The serializer collects the mapped stub's parameter metadata and it survives the round trip.""" + from airflow.sdk import task + + with DAG(dag_id="mapped_arg_binding_params_dag", schedule=None) as dag: + + @task.stub + def transform(country: str, extracted, retries_num: int = 3): ... + + @task + def plain(x): ... + + transform.partial(extracted={"a": 1}).expand(country=["uk", "fr"]) + plain.expand(x=[1, 2]) + + ser_dag = DagSerialization.to_dict(dag) + DagSerialization.validate_schema(ser_dag) + encoded_tasks = {t[Encoding.VAR]["task_id"]: t[Encoding.VAR] for t in ser_dag["dag"]["tasks"]} + assert encoded_tasks["transform"]["_mapped_arg_binding_params"] == [ + { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: { + "name": "country", + "value_schema": {Encoding.TYPE: DAT.DICT, Encoding.VAR: {"type": "string"}}, + }, + }, + {Encoding.TYPE: DAT.DICT, Encoding.VAR: {"name": "extracted"}}, + { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: { + "name": "retries_num", + "value_schema": { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: {"type": "integer", "format": "int64"}, + }, + "default": 3, + }, + }, + ], "metadata must serialize in declaration order" + assert "_mapped_arg_binding_params" not in encoded_tasks["plain"], ( + "only operator classes defining the capture hook contribute mapped fields" + ) + + round_tripped = DagSerialization.from_dict(ser_dag) + assert round_tripped.task_dict["transform"]._mapped_arg_binding_params == [ + {"name": "country", "value_schema": {"type": "string"}}, + {"name": "extracted"}, + {"name": "retries_num", "value_schema": {"type": "integer", "format": "int64"}, "default": 3}, + ] + assert not hasattr(round_tripped.task_dict["plain"], "_mapped_arg_binding_params") + + +def test_mapped_stub_capture_error_fails_serialization(): + """A capture-hook rejection surfaces as a Dag serialization (import) error.""" + from airflow.sdk import task + + with DAG(dag_id="mapped_arg_binding_params_invalid_dag", schedule=None) as dag: + + @task.stub + def transform(country: str): ... + + transform.expand_kwargs([{"country": "uk"}]) + + with pytest.raises(SerializationError, match="does not support expand_kwargs"): + DagSerialization.to_dict(dag) + + def test_handle_v1_serdag(): v1 = { "__version": 1, 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/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index 72f54421bff32..adf554095bac5 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -188,7 +188,7 @@ def _ensure_json_literal(value: Any, task_id: str, name: str) -> None: ) -def _validate_xcom_value(value: Any, task_id: str, name: str) -> bool: +def _validate_xcom_value(value: Any, task_id: str, name: str, *, allow_mapped_upstream: bool = False) -> bool: """Validate an XComArg argument, returning True when it is a bindable direct upstream output.""" if isinstance(value, PlainXComArg): if value.key != XCOM_RETURN_KEY: @@ -197,7 +197,7 @@ def _validate_xcom_value(value: Any, task_id: str, name: str) -> bool: f"{value.key!r}; only an upstream task's return value can cross the language " "boundary -- indexing an output by a custom key is not supported" ) - if value.operator.is_mapped: + if value.operator.is_mapped and not allow_mapped_upstream: raise ValueError( f"@task.stub task {task_id!r} parameter {name!r} references the aggregated " f"output of the mapped task {value.operator.task_id!r}; a foreign runtime " @@ -266,6 +266,64 @@ def _build_arg_bindings( return spec +def _build_mapped_arg_binding_params( + python_callable: Callable, + *, + partial_op_kwargs: Mapping[str, Any], + expand_input: Any, + task_id: str, +) -> list[dict[str, Any]] | None: + """ + Build the ordered per-parameter metadata for a mapped (``.expand()``) stub task. + + Per-map-index values only resolve at run time, so unlike ``_build_arg_bindings`` this + captures what the server-side derivation cannot recover from the serialized Dag alone: + the declaration order the wire contract promises, defaults for parameters no kwarg + covers, and each parameter's value schema. Returns ``None`` for parameterless stubs + (the legacy fan-out shape whose call args were always ignored keeps parsing). + """ + signature = inspect.signature(python_callable) + if not signature.parameters: + return None + _validate_stub_signature(signature, task_id) + if not isinstance(expand_input.value, Mapping): + # expand_kwargs() carries a list (or upstream XCom) of kwarg dicts whose + # parameter names are unknowable at parse time. + raise ValueError( + f"@task.stub task {task_id!r} does not support expand_kwargs(); the parameter " + "binding must be derivable at parse time, so use .expand() with explicit kwargs" + ) + expand_kwargs = expand_input.value + + try: + bound = signature.bind(**{**partial_op_kwargs, **expand_kwargs}) + except TypeError as e: + raise ValueError(f"@task.stub task {task_id!r} TaskFlow mapping does not bind to its signature: {e}") + bound.apply_defaults() + + annotations = _resolve_param_annotations(python_callable, signature) + + params: list[dict[str, Any]] = [] + for name in signature.parameters: + entry: dict[str, Any] = {"name": name} + if (value_schema := _infer_value_schema(annotations[name])) is not None: + entry["value_schema"] = value_schema + if name in expand_kwargs: + # The whole expanded collection ships through XCom/serialization; an upstream + # output is consumed per element, so a mapped upstream is fine here. + if not _validate_xcom_value(expand_kwargs[name], task_id, name, allow_mapped_upstream=True): + _ensure_json_literal(expand_kwargs[name], task_id, name) + elif name in partial_op_kwargs: + if not _validate_xcom_value(partial_op_kwargs[name], task_id, name): + _ensure_json_literal(partial_op_kwargs[name], task_id, name) + else: + default = bound.arguments[name] + _ensure_json_literal(default, task_id, name) + entry["default"] = default + params.append(entry) + return params + + class _StubOperator(DecoratedOperator): custom_operator_name: str = "@task.stub" @@ -332,6 +390,25 @@ def __init__( def get_serialized_fields(cls): return super().get_serialized_fields() | {"_arg_bindings"} + @classmethod + def get_mapped_serialized_fields(cls, mapped_op: Any) -> dict[str, Any]: + """ + Extra serialized fields for the mapped (``.expand()``) form of this operator. + + Called by the core Dag serializer (Airflow 3.4+) while ``python_callable`` is + still the real function; older cores never call it, so mapped stubs there keep + the legacy ignored-args behavior. + """ + params = _build_mapped_arg_binding_params( + mapped_op.python_callable, + partial_op_kwargs=mapped_op.partial_kwargs.get("op_kwargs") or {}, + expand_input=mapped_op._get_specified_expand_input(), + task_id=mapped_op.task_id, + ) + if params is None: + return {} + return {"_mapped_arg_binding_params": params} + def execute(self, context: Context) -> Any: raise RuntimeError( "@task.stub should not be executed directly -- we expected this to go to a remote worker. " @@ -355,10 +432,6 @@ def stub( outputs or JSON-serializable literals; the resulting argument-binding spec (parameter names, value schemas, and values, in declaration order) is delivered to the foreign runtime, which binds the values onto the native task function. - - Mapped (``.expand()``) stubs do not receive TaskFlow arguments yet -- their call args - keep the legacy ignored behavior; per-map-index delivery is part of - https://github.com/apache/airflow/issues/66937 and lands in a follow-up. """ return task_decorator_factory( decorated_operator_class=_StubOperator, diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index 68d804de4fc17..b73c82f38d2f6 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -26,7 +26,7 @@ import pytest from airflow.providers.common.compat.sdk import DAG, task_group -from airflow.providers.standard.decorators.stub import _infer_value_schema, stub +from airflow.providers.standard.decorators.stub import _infer_value_schema, _StubOperator, stub from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS @@ -259,7 +259,6 @@ def test_arg_bindings_survive_dag_serialization_round_trip(self): ] def test_expand_builds_mapped_stub_without_parse_time_bindings(self): - """Mapped stubs capture no spec: their call args keep the legacy ignored behavior for now.""" with DAG(dag_id="d"): result = stub(fn_transform).expand(country=["uk", "fr"], extracted=[{}, {}]) # op_kwargs_expand_input/partial_kwargs (not is_mapped) so the assertions also @@ -288,6 +287,117 @@ def group(n): group.expand(n=[1, 2]) +class TestMappedStubArgBindingParams: + """The serializer hook captures ordered per-parameter metadata for mapped stubs.""" + + def get_hook_fields(self, operator): + return _StubOperator.get_mapped_serialized_fields(operator) + + def test_params_follow_declaration_order_with_defaults_and_schemas(self): + with DAG(dag_id="d"): + # The partial() kwarg is declared *after* the expanded one: the captured + # order must come from the signature, not from the call sites. + result = stub(fn_transform).partial(extracted={"a": 1}).expand(country=["uk", "fr"]) + + assert self.get_hook_fields(result.operator) == { + "_mapped_arg_binding_params": [ + {"name": "country", "value_schema": {"type": "string"}}, + {"name": "extracted", "value_schema": {"type": "object", "additionalProperties": True}}, + { + "name": "retries_num", + "value_schema": {"type": "integer", "format": "int64"}, + "default": 3, + }, + ] + } + + def test_none_default_is_captured_by_key_presence(self): + def fn(x: str, y=None): ... + + with DAG(dag_id="d"): + result = stub(fn).expand(x=["a"]) + + params = self.get_hook_fields(result.operator)["_mapped_arg_binding_params"] + assert params[1] == {"name": "y", "default": None} + + def test_untyped_params_omit_value_schema(self): + with DAG(dag_id="d"): + result = stub(fn_untyped).expand(a=[1], b=[2]) + + assert self.get_hook_fields(result.operator) == { + "_mapped_arg_binding_params": [{"name": "a"}, {"name": "b"}] + } + + def test_parameterless_stub_captures_nothing(self): + with DAG(dag_id="d"): + result = stub(fn_extract).expand_kwargs([{}]) + + assert self.get_hook_fields(result.operator) == {} + + def test_expand_kwargs_rejected_for_parameterful_stub(self): + with DAG(dag_id="d"): + result = stub(fn_transform).expand_kwargs([{"country": "uk", "extracted": {}}]) + + with pytest.raises(ValueError, match="does not support expand_kwargs"): + self.get_hook_fields(result.operator) + + def test_missing_required_parameter_rejected(self): + with DAG(dag_id="d"): + result = stub(fn_transform).expand(country=["uk"]) + + with pytest.raises(ValueError, match="does not bind to its signature"): + self.get_hook_fields(result.operator) + + def test_partial_kwarg_over_mapped_upstream_rejected(self): + def fn_produce(n: int): ... + + with DAG(dag_id="d"): + vals = stub(fn_produce).expand(n=[1, 2]) + result = stub(fn_transform).partial(extracted=vals).expand(country=["uk"]) + + with pytest.raises(ValueError, match="aggregated output of the mapped task"): + self.get_hook_fields(result.operator) + + def test_expand_over_mapped_upstream_allowed(self): + def fn_produce(n: int): ... + + with DAG(dag_id="d"): + vals = stub(fn_produce).expand(n=[1, 2]) + result = stub(fn_transform).partial(country="uk").expand(extracted=vals) + + params = self.get_hook_fields(result.operator)["_mapped_arg_binding_params"] + assert [p["name"] for p in params] == ["country", "extracted", "retries_num"] + + def test_non_json_expand_literal_rejected(self): + with DAG(dag_id="d"): + result = stub(fn_transform).partial(country="uk").expand(extracted=[object()]) + + with pytest.raises(ValueError, match="not JSON-serializable"): + self.get_hook_fields(result.operator) + + def test_non_json_needed_default_rejected(self): + not_jsonable = object() + + def fn(x: str, y=not_jsonable): ... + + with DAG(dag_id="d"): + result = stub(fn).expand(x=["a"]) + + with pytest.raises(ValueError, match="not JSON-serializable"): + self.get_hook_fields(result.operator) + + def test_mapped_stub_inside_mapped_task_group_unconstructible(self): + """The SDK bans expansion inside an expanded group outright, so no hook guard is needed.""" + + @task_group + def group(n): + stub(fn_transform).partial(country="uk").expand(extracted=[{}]) + + with DAG(dag_id="d"): + with pytest.raises(NotImplementedError, match="expansion in an expanded task group"): + group.expand(n=[1, 2]) + + @pytest.mark.parametrize( ("annotation", "expected"), [ diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index d97d38613834c..739ac5ef66595 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -735,6 +735,8 @@ class XComArgBinding(BaseModel): name: Annotated[str, Field(title="Name")] value_schema: ArgValueSchema | None = None task_id: Annotated[str, Field(title="Task Id")] + map_index: Annotated[int | None, Field(title="Map Index")] = -1 + element_index: Annotated[int | None, Field(title="Element Index")] = None class AssetEventDagRunReference(BaseModel): 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 8e47af32cf5d4..4c2f503dc828e 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -4396,42 +4396,6 @@ "title": "VariableResult", "type": "object" }, - "XComArgBinding": { - "description": "One positional stub-task argument pulled from an upstream task's XCom.", - "properties": { - "kind": { - "const": "xcom", - "title": "Kind", - "type": "string" - }, - "name": { - "title": "Name", - "type": "string" - }, - "value_schema": { - "anyOf": [ - { - "$ref": "#/$defs/ArgValueSchema" - }, - { - "type": "null" - } - ], - "default": null - }, - "task_id": { - "title": "Task Id", - "type": "string" - } - }, - "required": [ - "kind", - "name", - "task_id" - ], - "title": "XComArgBinding", - "type": "object" - }, "XComCountResponse": { "properties": { "len": { @@ -4671,6 +4635,59 @@ ], "title": "TaskArgBinding" }, + "XComArgBinding": { + "description": "One positional stub-task argument pulled from an upstream task's XCom.", + "properties": { + "kind": { + "const": "xcom", + "title": "Kind", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "value_schema": { + "anyOf": [ + { + "$ref": "#/$defs/ArgValueSchema" + }, + { + "type": "null" + } + ], + "default": null + }, + "task_id": { + "title": "Task Id", + "type": "string" + }, + "map_index": { + "default": -1, + "title": "Map Index", + "type": "integer" + }, + "element_index": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Element Index" + } + }, + "required": [ + "kind", + "name", + "task_id" + ], + "title": "XComArgBinding", + "type": "object" + }, "AssetEventDagRunReference": { "additionalProperties": false, "description": "Schema for AssetEvent model used in DagRun.", diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index af944e4e565d7..31d44a52b18c5 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -254,6 +254,8 @@ export type TaskArgBinding = XComArgBinding | LiteralArgBinding; export type Kind = "xcom"; export type Name8 = string; export type TaskId1 = string; +export type MapIndex1 = number; +export type ElementIndex = number | null; export type Kind1 = "literal"; export type Name9 = string; export type FromDefault = boolean; @@ -337,7 +339,7 @@ export type Key5 = string; export type DagId6 = string; export type RunId5 = string; export type TaskId2 = string; -export type MapIndex1 = number | null; +export type MapIndex2 = number | null; export type Type25 = "DeleteXCom"; /** * Error types used in the API client. @@ -435,10 +437,10 @@ export type Type41 = "GetPreviousDagRun"; export type DagId12 = string; export type TaskId3 = string; export type LogicalDate4 = string | null; -export type MapIndex2 = number; +export type MapIndex3 = number; export type Type42 = "GetPreviousTI"; export type DagId13 = string; -export type MapIndex3 = number | null; +export type MapIndex4 = number | null; export type TaskIds = string[] | null; export type TaskGroupId = string | null; export type LogicalDates1 = string[] | null; @@ -455,7 +457,7 @@ export type TiId6 = string; export type Key8 = string; export type Type46 = "GetTaskStateStore"; export type DagId15 = string; -export type MapIndex4 = number | null; +export type MapIndex5 = number | null; export type TaskIds1 = string[] | null; export type TaskGroupId1 = string | null; export type LogicalDates2 = string[] | null; @@ -471,7 +473,7 @@ export type Key10 = string; export type DagId16 = string; export type RunId9 = string; export type TaskId4 = string; -export type MapIndex5 = number | null; +export type MapIndex6 = number | null; export type IncludePriorDates = boolean; export type Type50 = "GetXCom"; export type Key11 = string; @@ -528,7 +530,7 @@ export type StartDate5 = string | null; export type EndDate4 = string | null; export type State4 = string | null; export type TryNumber2 = number; -export type MapIndex6 = number | null; +export type MapIndex7 = number | null; export type Duration = number | null; export type Type60 = "PreviousTIResult"; export type Key14 = string; @@ -565,7 +567,7 @@ export type Key18 = string; export type DagId21 = string; export type RunId14 = string; export type TaskId9 = string; -export type MapIndex7 = number | null; +export type MapIndex8 = number | null; export type DagResult1 = boolean; export type MappedLength = number | null; export type Type71 = "SetXCom"; @@ -1061,6 +1063,8 @@ export interface XComArgBinding { name: Name8; value_schema?: ArgValueSchema | null; task_id: TaskId1; + map_index?: MapIndex1; + element_index?: ElementIndex; } /** * One positional stub-task argument carrying an inline literal from the Dag file. @@ -1233,7 +1237,7 @@ export interface DeleteXCom { dag_id: DagId6; run_id: RunId5; task_id: TaskId2; - map_index?: MapIndex1; + map_index?: MapIndex2; type?: Type25; } /** @@ -1401,7 +1405,7 @@ export interface GetPreviousTI { dag_id: DagId12; task_id: TaskId3; logical_date?: LogicalDate4; - map_index?: MapIndex2; + map_index?: MapIndex3; state?: TaskInstanceState | null; type?: Type42; } @@ -1411,7 +1415,7 @@ export interface GetPreviousTI { */ export interface GetTICount { dag_id: DagId13; - map_index?: MapIndex3; + map_index?: MapIndex4; task_ids?: TaskIds; task_group_id?: TaskGroupId; logical_dates?: LogicalDates1; @@ -1452,7 +1456,7 @@ export interface GetTaskStateStore { */ export interface GetTaskStates { dag_id: DagId15; - map_index?: MapIndex4; + map_index?: MapIndex5; task_ids?: TaskIds1; task_group_id?: TaskGroupId1; logical_dates?: LogicalDates2; @@ -1486,7 +1490,7 @@ export interface GetXCom { dag_id: DagId16; run_id: RunId9; task_id: TaskId4; - map_index?: MapIndex5; + map_index?: MapIndex6; include_prior_dates?: IncludePriorDates; type?: Type50; } @@ -1612,7 +1616,7 @@ export interface PreviousTIResponse { end_date?: EndDate4; state?: State4; try_number: TryNumber2; - map_index?: MapIndex6; + map_index?: MapIndex7; duration?: Duration; } /** @@ -1740,7 +1744,7 @@ export interface SetXCom { dag_id: DagId21; run_id: RunId14; task_id: TaskId9; - map_index?: MapIndex7; + map_index?: MapIndex8; dag_result?: DagResult1; mapped_length?: MappedLength; type?: Type71;