From 3cf3b28c13a6872398f974f1552ad4571563749c Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 10 Jul 2026 07:24:18 +0000 Subject: [PATCH 1/5] 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 67c7ca053562b..28e2812ec6a0b 100644 --- a/task-sdk/pyproject.toml +++ b/task-sdk/pyproject.toml @@ -253,6 +253,7 @@ disable-timestamp=true enable-version-header=true enum-field-as-literal='one' # When a single enum member, make it output a `Literal["..."]` input-file-type='openapi' +set-default-enum-member=true # `= DataType.ANY` not `= "any"`, keeping mypy happy with enum-typed fields output-model-type='pydantic_v2.BaseModel' output-datetime-class='AwareDatetime' target-python-version='3.10' diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 72aebeab3e49b..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 01e3c8970b5cb..ae8624fad897c 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -1,6 +1,6 @@ { "$schema": "https://json-schema.org/draft/2020-12/schema", - "api_version": "2026-06-16", + "api_version": "2026-07-30", "description": "Apache Airflow SDK Supervisor Schema", "$defs": { "AssetAliasReferenceAssetEventDagRun": { @@ -1395,6 +1395,19 @@ "title": "DagRunType", "type": "string" }, + "DataType": { + "enum": [ + "string", + "integer", + "number", + "boolean", + "object", + "array", + "any" + ], + "title": "DataType", + "type": "string" + }, "DeferTask": { "additionalProperties": false, "description": "Update a task instance state to deferred.", @@ -2998,6 +3011,14 @@ "type": "object" }, "JsonValue": {}, + "Kind": { + "enum": [ + "xcom", + "literal" + ], + "title": "Kind", + "type": "string" + }, "LazyDeserializedDAG": { "description": "Lazily build information from the serialized DAG structure.\n\nAn object that will present \"enough\" of the DAG like interface to update DAG db models etc, without having\nto deserialize the full DAG and Task hierarchy.", "properties": { @@ -4797,6 +4818,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": { @@ -4878,6 +4959,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 4766c0e8a7949..c09f9de82cc3c 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"; /** @@ -386,10 +396,10 @@ export type Extra8 = { } | null; export type Type30 = "GetAssetEventByAssetAlias"; export type Name11 = string; -export type Key6 = string; +export type Key7 = string; export type Type31 = "GetAssetStateStoreByName"; export type Uri8 = string; -export type Key7 = string; +export type Key8 = string; export type Type32 = "GetAssetStateStoreByUri"; export type AliasName1 = string; export type Type33 = "GetAssetsByAlias"; @@ -417,7 +427,7 @@ export type LogicalDate3 = string; export type State3 = string | null; export type Type41 = "GetPreviousDagRun"; export type DagId12 = string; -export type TaskId2 = string; +export type TaskId3 = string; export type LogicalDate4 = string | null; export type MapIndex2 = number; export type Type42 = "GetPreviousTI"; @@ -436,7 +446,7 @@ export type TiId5 = string; export type TryNumber1 = number; export type Type45 = "GetTaskRescheduleStartDate"; export type TiId6 = string; -export type Key8 = string; +export type Key9 = string; export type Type46 = "GetTaskStateStore"; export type DagId15 = string; export type MapIndex4 = number | null; @@ -445,34 +455,34 @@ export type TaskGroupId1 = string | null; export type LogicalDates2 = string[] | null; export type RunIds2 = string[] | null; export type Type47 = "GetTaskStates"; -export type Key9 = string; +export type Key10 = string; export type Type48 = "GetVariable"; export type Prefix = string | null; export type Limit2 = number; export type Offset = number; export type Type49 = "GetVariableKeys"; -export type Key10 = string; +export type Key11 = string; export type DagId16 = string; export type RunId9 = string; -export type TaskId3 = string; +export type TaskId4 = string; export type MapIndex5 = number | null; export type IncludePriorDates = boolean; export type Type50 = "GetXCom"; -export type Key11 = string; +export type Key12 = string; export type DagId17 = string; export type RunId10 = string; -export type TaskId4 = string; +export type TaskId5 = string; export type Type51 = "GetXComCount"; -export type Key12 = string; +export type Key13 = string; export type DagId18 = string; export type RunId11 = string; -export type TaskId5 = string; +export type TaskId6 = string; export type Offset1 = number; export type Type52 = "GetXComSequenceItem"; -export type Key13 = string; +export type Key14 = string; export type DagId19 = string; export type RunId12 = string; -export type TaskId6 = string; +export type TaskId7 = string; export type Start = number | null; export type Stop = number | null; export type Step = number | null; @@ -494,6 +504,11 @@ export type AssignedUsers1 = HITLUser[] | null; export type Type54 = "HITLDetailRequestResult"; export type InactiveAssets = AssetProfile[] | null; export type Type55 = "InactiveAssetsResult"; +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "Kind". + */ +export type Kind1 = "xcom" | "literal"; export type Name12 = string | null; export type Type56 = "MaskSecret"; export type Ok = boolean; @@ -504,7 +519,7 @@ export type StartDate4 = string | null; export type EndDate3 = string | null; export type Type58 = "PrevSuccessfulDagRunResult"; export type Type59 = "PreviousDagRunResult"; -export type TaskId7 = string; +export type TaskId8 = string; export type DagId20 = string; export type RunId13 = string; export type LogicalDate5 = string | null; @@ -515,7 +530,7 @@ export type TryNumber2 = number; export type MapIndex6 = number | null; export type Duration = number | null; export type Type60 = "PreviousTIResult"; -export type Key14 = string; +export type Key15 = string; export type Value1 = string | null; export type Description = string | null; export type Type61 = "PutVariable"; @@ -533,22 +548,22 @@ export type Type64 = "RetryTask"; export type Type65 = "SentFDs"; export type Fds = number[]; export type Name13 = string; -export type Key15 = string; +export type Key16 = string; export type Type66 = "SetAssetStateStoreByName"; export type Uri9 = string; -export type Key16 = string; +export type Key17 = string; export type Type67 = "SetAssetStateStoreByUri"; export type Type68 = "SetRenderedFields"; export type RenderedMapIndex3 = string; export type Type69 = "SetRenderedMapIndex"; export type TiId8 = string; -export type Key17 = string; +export type Key18 = string; export type ExpiresAt = string | null; export type Type70 = "SetTaskStateStore"; -export type Key18 = string; +export type Key19 = string; export type DagId21 = string; export type RunId14 = string; -export type TaskId8 = string; +export type TaskId9 = string; export type MapIndex7 = number | null; export type DagResult1 = boolean; export type MappedLength = number | null; @@ -608,12 +623,12 @@ export type Type83 = "ValidateInletsAndOutlets"; export type Keys = string[]; export type TotalEntries = number; export type Type84 = "VariableKeysResult"; -export type Key19 = string; +export type Key20 = string; export type Value2 = string | null; export type Type85 = "VariableResult"; export type Len = number; export type Type86 = "XComCountResponse"; -export type Key20 = string; +export type Key21 = string; export type Type87 = "XComResult"; export type Type88 = "XComSequenceIndexResult"; export type Root = JsonValue[]; @@ -999,6 +1014,7 @@ export interface TIRunContext { xcom_keys_to_clear?: XcomKeysToClear; should_retry?: ShouldRetry; start_date?: StartDate2; + stub_args?: StubArgs; } /** * Variable schema for responses with fields that are needed for Runtime. @@ -1026,6 +1042,22 @@ export interface ConnectionResponse { port: Port1; extra: Extra6; } +/** + * One positional argument of a stub (foreign-runtime) task, in declaration order. + * + * A deliberately flat shape (``kind`` discriminates instead of a union) so the JSON schema + * generates a plain struct in the foreign-language SDKs consuming the supervisor schema. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "StubTaskArg". + */ +export interface StubTaskArg { + kind: Kind; + data_type?: DataType; + task_id?: TaskId1; + key?: Key1; + value?: unknown; +} /** * Email notification request for task failures/retries. * @@ -1146,7 +1178,7 @@ export interface DeferTask { */ export interface DeleteAssetStateStoreByName { name: Name8; - key: Key1; + key: Key2; type?: Type21; } /** @@ -1155,7 +1187,7 @@ export interface DeleteAssetStateStoreByName { */ export interface DeleteAssetStateStoreByUri { uri: Uri5; - key: Key2; + key: Key3; type?: Type22; } /** @@ -1164,7 +1196,7 @@ export interface DeleteAssetStateStoreByUri { */ export interface DeleteTaskStateStore { ti_id: TiId2; - key: Key3; + key: Key4; type?: Type23; } /** @@ -1172,7 +1204,7 @@ export interface DeleteTaskStateStore { * via the `definition` "DeleteVariable". */ export interface DeleteVariable { - key: Key4; + key: Key5; type?: Type24; } /** @@ -1180,10 +1212,10 @@ export interface DeleteVariable { * via the `definition` "DeleteXCom". */ export interface DeleteXCom { - key: Key5; + key: Key6; dag_id: DagId6; run_id: RunId5; - task_id: TaskId1; + task_id: TaskId2; map_index?: MapIndex1; type?: Type25; } @@ -1245,7 +1277,7 @@ export interface GetAssetEventByAssetAlias { */ export interface GetAssetStateStoreByName { name: Name11; - key: Key6; + key: Key7; type?: Type31; } /** @@ -1254,7 +1286,7 @@ export interface GetAssetStateStoreByName { */ export interface GetAssetStateStoreByUri { uri: Uri8; - key: Key7; + key: Key8; type?: Type32; } /** @@ -1346,7 +1378,7 @@ export interface GetPreviousDagRun { */ export interface GetPreviousTI { dag_id: DagId12; - task_id: TaskId2; + task_id: TaskId3; logical_date?: LogicalDate4; map_index?: MapIndex2; state?: TaskInstanceState | null; @@ -1390,7 +1422,7 @@ export interface GetTaskRescheduleStartDate { */ export interface GetTaskStateStore { ti_id: TiId6; - key: Key8; + key: Key9; type?: Type46; } /** @@ -1411,7 +1443,7 @@ export interface GetTaskStates { * via the `definition` "GetVariable". */ export interface GetVariable { - key: Key9; + key: Key10; type?: Type48; } /** @@ -1429,10 +1461,10 @@ export interface GetVariableKeys { * via the `definition` "GetXCom". */ export interface GetXCom { - key: Key10; + key: Key11; dag_id: DagId16; run_id: RunId9; - task_id: TaskId3; + task_id: TaskId4; map_index?: MapIndex5; include_prior_dates?: IncludePriorDates; type?: Type50; @@ -1444,10 +1476,10 @@ export interface GetXCom { * via the `definition` "GetXComCount". */ export interface GetXComCount { - key: Key11; + key: Key12; dag_id: DagId17; run_id: RunId10; - task_id: TaskId4; + task_id: TaskId5; type?: Type51; } /** @@ -1455,10 +1487,10 @@ export interface GetXComCount { * via the `definition` "GetXComSequenceItem". */ export interface GetXComSequenceItem { - key: Key12; + key: Key13; dag_id: DagId18; run_id: RunId11; - task_id: TaskId5; + task_id: TaskId6; offset: Offset1; type?: Type52; } @@ -1467,10 +1499,10 @@ export interface GetXComSequenceItem { * via the `definition` "GetXComSequenceSlice". */ export interface GetXComSequenceSlice { - key: Key13; + key: Key14; dag_id: DagId19; run_id: RunId12; - task_id: TaskId6; + task_id: TaskId7; start: Start; stop: Stop; step: Step; @@ -1551,7 +1583,7 @@ export interface PreviousDagRunResult { * via the `definition` "PreviousTIResponse". */ export interface PreviousTIResponse { - task_id: TaskId7; + task_id: TaskId8; dag_id: DagId20; run_id: RunId13; logical_date?: LogicalDate5; @@ -1577,7 +1609,7 @@ export interface PreviousTIResult { * via the `definition` "PutVariable". */ export interface PutVariable { - key: Key14; + key: Key15; value: Value1; description: Description; type?: Type61; @@ -1629,7 +1661,7 @@ export interface SentFDs { */ export interface SetAssetStateStoreByName { name: Name13; - key: Key15; + key: Key16; value: JsonValue; type?: Type66; } @@ -1639,7 +1671,7 @@ export interface SetAssetStateStoreByName { */ export interface SetAssetStateStoreByUri { uri: Uri9; - key: Key16; + key: Key17; value: JsonValue; type?: Type67; } @@ -1672,7 +1704,7 @@ export interface SetRenderedMapIndex { */ export interface SetTaskStateStore { ti_id: TiId8; - key: Key17; + key: Key18; value: JsonValue; expires_at: ExpiresAt; type?: Type70; @@ -1682,11 +1714,11 @@ export interface SetTaskStateStore { * via the `definition` "SetXCom". */ export interface SetXCom { - key: Key18; + key: Key19; value: JsonValue; dag_id: DagId21; run_id: RunId14; - task_id: TaskId8; + task_id: TaskId9; map_index?: MapIndex7; dag_result?: DagResult1; mapped_length?: MappedLength; @@ -1843,7 +1875,7 @@ export interface VariableKeysResult { * via the `definition` "VariableResult". */ export interface VariableResult { - key: Key19; + key: Key20; value?: Value2; type?: Type85; } @@ -1862,7 +1894,7 @@ export interface XComCountResponse { * via the `definition` "XComResult". */ export interface XComResult { - key: Key20; + key: Key21; value: JsonValue; type?: Type87; } @@ -1888,4 +1920,4 @@ export interface XComSequenceSliceResult { * (e.g. bundle metadata) and runs the migrator accordingly. * Exposed so the SDK author / operator can confirm which schema * version their build is pinned to. */ -export const SUPERVISOR_API_VERSION = "2026-06-16" as const; +export const SUPERVISOR_API_VERSION = "2026-07-30" as const; From 8bb8e4f8f5a18d46126dda77fa9eb2d88a4f5095 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Sat, 11 Jul 2026 12:15:51 +0000 Subject: [PATCH 2/5] 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 ae8624fad897c..66797d25d81bb 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -4818,66 +4818,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": { @@ -4960,11 +4900,11 @@ "default": null, "title": "Start Date" }, - "stub_args": { + "arg_bindings": { "anyOf": [ { "items": { - "$ref": "#/$defs/StubTaskArg" + "$ref": "#/$defs/TaskArgBinding" }, "type": "array" }, @@ -4973,7 +4913,7 @@ } ], "default": null, - "title": "Stub Args" + "title": "Arg Bindings" } }, "required": [ @@ -4983,6 +4923,66 @@ "title": "TIRunContext", "type": "object" }, + "TaskArgBinding": { + "description": "One positional argument of a stub (foreign-runtime) task, in declaration order.\n\nA deliberately flat shape (``kind`` discriminates instead of a union) so the JSON schema\ngenerates a plain struct in the foreign-language SDKs consuming the supervisor schema.", + "properties": { + "kind": { + "enum": [ + "xcom", + "literal" + ], + "title": "Kind", + "type": "string" + }, + "data_type": { + "default": "any", + "enum": [ + "string", + "integer", + "number", + "boolean", + "object", + "array", + "any" + ], + "title": "Data Type", + "type": "string" + }, + "task_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Task Id" + }, + "key": { + "default": "return_value", + "title": "Key", + "type": "string" + }, + "value": { + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ], + "default": null + } + }, + "required": [ + "kind" + ], + "title": "TaskArgBinding", + "type": "object" + }, "TaskInstance": { "description": "Schema for TaskInstance model with minimal required fields needed for Runtime.", "properties": { diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py index 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 c09f9de82cc3c..0d3aac7e95a77 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; @@ -1014,7 +1014,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. @@ -1049,9 +1049,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 e74cddaebb472ab79ee479e04942f961922a23ff Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Sun, 12 Jul 2026 05:42:46 +0000 Subject: [PATCH 3/5] 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 1503642e7aa9db87f4ec0394f1805c49fd5e3a46 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 13 Jul 2026 04:18:03 +0000 Subject: [PATCH 4/5] 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 bec0e8daa5cf8986a02a1baeb2074581cbd02f5e Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Mon, 13 Jul 2026 04:18:56 +0000 Subject: [PATCH 5/5] 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") +}