jason810496 commented on code in PR #69757: URL: https://github.com/apache/airflow/pull/69757#discussion_r3619530055
########## airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py: ########## @@ -0,0 +1,112 @@ +# 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. + +Each binding is one variant of a union discriminated on ``kind``: an ``XComArgBinding`` +pulls the value from an upstream task's XCom, a ``LiteralArgBinding`` carries an inline +value from the Dag file. Both variants still emit plain named structs +(``$defs/XComArgBinding``, ``$defs/LiteralArgBinding``) for the foreign-language SDKs +consuming the supervisor schema. +""" + +from __future__ import annotations + +from enum import Enum +from functools import cache +from typing import Annotated, Literal + +from pydantic import Field, JsonValue, TypeAdapter +from typing_extensions import TypeAliasType + +from airflow.api_fastapi.core_api.base import BaseModel + + +class ArgBindingDataType(str, Enum): + """Language-neutral value type a stub-task argument binds to in the foreign runtime.""" + + STRING = "string" + INTEGER = "integer" + NUMBER = "number" + BOOLEAN = "boolean" + OBJECT = "object" + ARRAY = "array" + ANY = "any" + + +class XComArgBinding(BaseModel): + """One positional stub-task argument pulled from an upstream task's XCom.""" + + # No default on purpose: a required ``kind`` stays non-nullable through the OpenAPI + # round trip, which discriminated-union codegen needs (a defaulted field turns + # ``Literal`` into ``Literal | None`` in the generated task-sdk models). Review Comment: ```suggestion ``` ########## airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py: ########## @@ -0,0 +1,112 @@ +# 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. + +Each binding is one variant of a union discriminated on ``kind``: an ``XComArgBinding`` +pulls the value from an upstream task's XCom, a ``LiteralArgBinding`` carries an inline +value from the Dag file. Both variants still emit plain named structs +(``$defs/XComArgBinding``, ``$defs/LiteralArgBinding``) for the foreign-language SDKs +consuming the supervisor schema. Review Comment: ```suggestion ``` ########## ts-sdk/src/generated/supervisor.ts: ########## @@ -22,6 +22,20 @@ // // Re-run with: pnpm run generate:supervisor +/** Review Comment: We should not bump `ts-sdk/src/generated/supervisor.ts` in this PR. ########## go-sdk/pkg/binding/binding.go: ########## @@ -0,0 +1,701 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// Package binding turns a task function's parameter list into the concrete +// argument values it is called with at execution time. +// +// Three kinds of parameter are supported: +// +// - Injectable runtime values: context.Context, sdk.TIRunContext, +// *slog.Logger, and any interface whose method set is a subset of +// sdk.Client. These are filled by type, in any position. +// - Data parameters: everything else (except TaskInput structs, below), in +// declaration order. They receive the positional arguments the Python +// stub Dag captured at parse time from the TaskFlow call +// (“transform("uk", extract())“) and delivered in StartupDetails. A +// literal argument decodes directly; an XCom argument is pulled from the +// named upstream task in the current dag run first. +// - TaskInput structs: a struct that anonymously embeds sdk.TaskInput opts +// into per-field, name-based binding instead of consuming one positional +// slot as a whole-value decode target. Each exported field binds by name +// against the Dag's TaskFlow call arguments: an `arg:"<name>"` tag names +// the argument to claim, and a field with no tag claims its own Go field +// name, verbatim. At most one such parameter is allowed per function. +// +// A TaskInput struct's fields are resolved first, by name, claiming entries +// out of the argument spec; the remaining unclaimed entries are then +// distributed, in their original relative order, onto the plain flat data +// parameters in declaration order -- so flat parameters and a TaskInput +// struct can coexist in the same function signature regardless of where each +// sits, and (with no TaskInput struct present) this reduces to exactly +// today's positional-only behaviour. Review Comment: We should make it either using flat-base or TaskInput struct-base. Not supporting flat-base plus TaskInput struct-base, which is too ambiguous. ########## airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py: ########## @@ -141,3 +145,28 @@ def remove_partition_date_from_dag_run(response: ResponseInfo) -> None: # type: """Strip ``partition_date`` from the nested ``dag_run`` payload for older clients.""" if "dag_run" in response.body and isinstance(response.body["dag_run"], dict): response.body["dag_run"].pop("partition_date", None) + + +class AddArgBindingsToTIRunContext(VersionChange): + """ + Add the ``arg_bindings`` positional-argument binding spec for stub (foreign-runtime) tasks. + + Each entry is a discriminated union of ``XComArgBinding`` and ``LiteralArgBinding`` keyed + on ``kind``. ``data_type`` is declared as the ``ArgBindingDataType`` enum rather than an + inline ``Literal``; the wire representation (a JSON string) is unchanged, so no migration + instruction is needed -- this version has not been released with the ``arg_bindings`` field + in any other shape. Review Comment: ```suggestion ``` ########## go-sdk/pkg/binding/binding.go: ########## @@ -0,0 +1,701 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// Package binding turns a task function's parameter list into the concrete +// argument values it is called with at execution time. +// +// Three kinds of parameter are supported: +// +// - Injectable runtime values: context.Context, sdk.TIRunContext, +// *slog.Logger, and any interface whose method set is a subset of +// sdk.Client. These are filled by type, in any position. +// - Data parameters: everything else (except TaskInput structs, below), in +// declaration order. They receive the positional arguments the Python +// stub Dag captured at parse time from the TaskFlow call +// (“transform("uk", extract())“) and delivered in StartupDetails. A +// literal argument decodes directly; an XCom argument is pulled from the +// named upstream task in the current dag run first. +// - TaskInput structs: a struct that anonymously embeds sdk.TaskInput opts +// into per-field, name-based binding instead of consuming one positional +// slot as a whole-value decode target. Each exported field binds by name +// against the Dag's TaskFlow call arguments: an `arg:"<name>"` tag names +// the argument to claim, and a field with no tag claims its own Go field +// name, verbatim. At most one such parameter is allowed per function. +// +// A TaskInput struct's fields are resolved first, by name, claiming entries +// out of the argument spec; the remaining unclaimed entries are then +// distributed, in their original relative order, onto the plain flat data +// parameters in declaration order -- so flat parameters and a TaskInput +// struct can coexist in the same function signature regardless of where each +// sits, and (with no TaskInput struct present) this reduces to exactly +// today's positional-only behaviour. +// +// Conceptually, flat data parameters are positional-argument binding: order +// matters, and every parameter must be filled or Resolve fails the task +// before its body runs. A TaskInput struct is closer to keyword-argument +// binding: fields match by name, and a field whose name has no corresponding +// TaskFlow call argument is simply left at its Go zero value instead of +// failing the task -- the same way an unpassed keyword argument falls back +// to its default in a kwargs-style call. +// +// Analyze inspects a function once at registration and returns a Plan; Resolve +// builds the call arguments for each execution from that Plan and the +// per-execution argument spec, failing loudly on arity or type mismatches of +// flat data parameters. A declared type of "any" (or a Go parameter typed +// any) opts that argument out of the type check; the decode step still fails +// loudly on unusable values. +// +// Mapped upstream fan-in is out of scope: XCom arguments always pull the +// unmapped upstream instance (map_index is never sent). +package binding + +import ( + "bytes" + "context" + "encoding/json" + "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" +) + +// 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" +) Review Comment: We should reference these from `go-sdk/pkg/execution/genmodels/models.gen.go` instead of hand-written to prevent drift ########## airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py: ########## @@ -110,6 +111,36 @@ log = structlog.get_logger(__name__) tracer = trace.get_tracer(__name__) +# Task type recorded on the TI row (``TaskInstance.operator``) for +# ``airflow.providers.standard.decorators.stub._StubOperator``. Used to gate the +# serialized-dag lookup for ``arg_bindings`` so regular tasks never pay for it. +_STUB_TASK_TYPE = "_StubOperator" + + +def _get_arg_bindings(dag_version_id: UUID | None, task_id: str, *, session) -> list[dict] | None: + """Extract the stub task's serialized positional-arg spec from the serialized Dag blob.""" + # Imported here on purpose: only the Multi-Lang stub-task path touches the + # serialized-dag machinery, so keep it off the module's top-level imports. + from airflow.models.dag_version import DagVersion + from airflow.serialization.enums import Encoding + from airflow.serialization.serialized_objects import BaseSerialization + + if dag_version_id is None: + return None + dag_version = session.get(DagVersion, dag_version_id) + if dag_version is None or dag_version.serialized_dag is None: + return None + data = dag_version.serialized_dag.data + if not data: Review Comment: ```suggestion if not (data := dag_version.serialized_dag.data): ``` ########## airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py: ########## @@ -0,0 +1,112 @@ +# 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. + +Each binding is one variant of a union discriminated on ``kind``: an ``XComArgBinding`` +pulls the value from an upstream task's XCom, a ``LiteralArgBinding`` carries an inline +value from the Dag file. Both variants still emit plain named structs +(``$defs/XComArgBinding``, ``$defs/LiteralArgBinding``) for the foreign-language SDKs +consuming the supervisor schema. +""" + +from __future__ import annotations + +from enum import Enum +from functools import cache +from typing import Annotated, Literal + +from pydantic import Field, JsonValue, TypeAdapter +from typing_extensions import TypeAliasType + +from airflow.api_fastapi.core_api.base import BaseModel + + +class ArgBindingDataType(str, Enum): + """Language-neutral value type a stub-task argument binds to in the foreign runtime.""" + + STRING = "string" + INTEGER = "integer" + NUMBER = "number" + BOOLEAN = "boolean" + OBJECT = "object" + ARRAY = "array" + ANY = "any" + + +class XComArgBinding(BaseModel): + """One positional stub-task argument pulled from an upstream task's XCom.""" + + # No default on purpose: a required ``kind`` stays non-nullable through the OpenAPI + # round trip, which discriminated-union codegen needs (a defaulted field turns + # ``Literal`` into ``Literal | None`` in the generated task-sdk models). + kind: Literal["xcom"] + + name: str + """The stub function's parameter name this binding fills, in declaration order.""" + + data_type: ArgBindingDataType = ArgBindingDataType.ANY + """Declared type from the stub function's annotation; runtimes type-check against it.""" + + task_id: str + """Upstream task id to pull the XCom from.""" + + key: str = "return_value" + """XCom key to pull.""" Review Comment: The `key` will always be `return_value` across TaskFlow, we can remove the `key` field from end to end. ########## go-sdk/pkg/binding/binding.go: ########## @@ -0,0 +1,701 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// Package binding turns a task function's parameter list into the concrete +// argument values it is called with at execution time. +// +// Three kinds of parameter are supported: +// +// - Injectable runtime values: context.Context, sdk.TIRunContext, +// *slog.Logger, and any interface whose method set is a subset of +// sdk.Client. These are filled by type, in any position. +// - Data parameters: everything else (except TaskInput structs, below), in +// declaration order. They receive the positional arguments the Python +// stub Dag captured at parse time from the TaskFlow call +// (“transform("uk", extract())“) and delivered in StartupDetails. A +// literal argument decodes directly; an XCom argument is pulled from the +// named upstream task in the current dag run first. +// - TaskInput structs: a struct that anonymously embeds sdk.TaskInput opts +// into per-field, name-based binding instead of consuming one positional +// slot as a whole-value decode target. Each exported field binds by name +// against the Dag's TaskFlow call arguments: an `arg:"<name>"` tag names +// the argument to claim, and a field with no tag claims its own Go field +// name, verbatim. At most one such parameter is allowed per function. +// +// A TaskInput struct's fields are resolved first, by name, claiming entries +// out of the argument spec; the remaining unclaimed entries are then +// distributed, in their original relative order, onto the plain flat data +// parameters in declaration order -- so flat parameters and a TaskInput +// struct can coexist in the same function signature regardless of where each +// sits, and (with no TaskInput struct present) this reduces to exactly +// today's positional-only behaviour. +// +// Conceptually, flat data parameters are positional-argument binding: order +// matters, and every parameter must be filled or Resolve fails the task +// before its body runs. A TaskInput struct is closer to keyword-argument +// binding: fields match by name, and a field whose name has no corresponding +// TaskFlow call argument is simply left at its Go zero value instead of +// failing the task -- the same way an unpassed keyword argument falls back +// to its default in a kwargs-style call. +// +// Analyze inspects a function once at registration and returns a Plan; Resolve +// builds the call arguments for each execution from that Plan and the +// per-execution argument spec, failing loudly on arity or type mismatches of +// flat data parameters. A declared type of "any" (or a Go parameter typed +// any) opts that argument out of the type check; the decode step still fails +// loudly on unusable values. +// +// Mapped upstream fan-in is out of scope: XCom arguments always pull the +// unmapped upstream instance (map_index is never sent). +package binding + +import ( + "bytes" + "context" + "encoding/json" + "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" +) + +// 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: an XComArg or a LiteralArg. A sealed sum type mirroring +// the wire model's XComArgBinding/LiteralArgBinding split, kept runtime-neutral +// so this package stays decoupled from the generated coordinator schema types. +type Arg interface { + // ArgName is the stub function's parameter name this binding fills; used + // to match a TaskInput struct field's `arg:` tag (or its verbatim + // field-name fallback). + ArgName() string + // DeclaredType is the Dag-declared language-neutral type for the argument; + // empty is treated as DataTypeAny. + DeclaredType() DataType + // sealedArg restricts implementations to this package, keeping + // resolveOne's type switch the single exhaustive consumer. + sealedArg() +} + +// XComArg sources the argument from an upstream task's XCom. +type XComArg struct { + // Name is the stub function's parameter name this binding fills. + Name string + // TaskID is the upstream task to pull from. + TaskID string + // Key is the XCom key to pull; empty means the return-value key. + Key string + // DataType is the declared type to check the Go parameter against. + DataType DataType +} + +// LiteralArg carries an inline value from the Dag file. +type LiteralArg struct { + // Name is the stub function's parameter name this binding fills. + Name string + // Value is the literal value from the Dag file. + Value any + // DataType is the declared type to check the Go parameter against. + DataType DataType +} Review Comment: We should reference these from go-sdk/pkg/execution/genmodels/models.gen.go instead of hand-written to prevent drift if possible. ########## go-sdk/dags/go_examples.py: ########## @@ -107,3 +119,90 @@ def concurrent_xcom_dag(): concurrent_xcom_dag() + + [email protected](queue="golang") +def make_config(): ... + + [email protected](queue="golang") +def make_numbers(): ... + + [email protected](queue="golang") +def via_flat_args( + name: str, + count: int, + ratio: float, + enabled: bool, + tags: list, + config: dict, + numbers: list, + note: str | None = None, +): ... + + +# Capitalized parameters on purpose: with no ``arg:`` tags on the Go side, each +# struct field binds the argument spelled exactly like its Go field name. [email protected](queue="golang") +def via_struct_no_tags(RegionCode: str, Threshold: float): ... + + [email protected](queue="golang") +def via_struct_arg_tag(region_code: str, threshold: float): ... + + [email protected](queue="golang") +def via_struct_unmatched_arg(region_code: str): ... + + +@dag(dag_id="taskflow_binding_dag") +def taskflow_binding_dag(): + """ + Stress the TaskFlow argument-binding surface beyond ``simple_dag``'s transform. + + Conceptually, the flat parameter list is *positional-argument* binding: order + matters, and every parameter must be filled or the task fails before it runs. + ``sdk.TaskInput`` structs are closer to *keyword-argument* binding: fields match + by name, and (see ``via_struct_unmatched_arg`` below) a field whose name has no + corresponding TaskFlow call argument simply stays at its zero value instead of + failing the task -- the same way an unpassed keyword argument falls back to a + default in kwargs-style calls. + + ``via_flat_args``'s one mixed positional/keyword call carries literals of every + scalar type plus an array literal, and fans in XComs from *two* upstream Go + tasks: ``make_config`` returns an object that binds onto a strictly-decoded Go + struct, ``make_numbers`` an array that binds onto ``[]int``. ``note`` is not + passed, so its ``None`` default is captured and arrives in Go as a nil + ``*string``. The Go ``via_flat_args`` (``go-sdk/example/bundle/taskflowbinding``) + verifies every bound value and fails the task on any mismatch. + + Three further tasks demonstrate the Go SDK's ``sdk.TaskInput`` struct injection + mode, one field-binding mode at a time: + + * ``via_struct_no_tags``: no ``arg:`` tags at all -- each struct field binds + the argument spelled exactly like its Go field name, hence this stub's + capitalized ``RegionCode``/``Threshold`` parameters. + * ``via_struct_arg_tag``: every field names its argument via an explicit + ``arg:`` tag -- ``Region`` is genuinely renamed to ``region_code``, and + ``Threshold`` is tagged ``threshold`` to pull the snake_case argument its + verbatim field name would miss. + * ``via_struct_unmatched_arg``: the Go struct declares a field with no + corresponding argument in this TaskFlow call at all -- it stays at its Go + zero value rather than failing the task. + """ + via_flat_args( + "summary", + 3, + 2.5, + True, + ["metrics", "hourly"], + config=make_config(), + numbers=make_numbers(), + ) + via_struct_no_tags(RegionCode="eu-west-1", Threshold=0.75) + via_struct_arg_tag(region_code="eu-west-1", threshold=0.75) + via_struct_unmatched_arg(region_code="eu-west-1") Review Comment: For via_struct_no_tags, via_struct_arg_tag and via_struct_unmatched_arg tasks, we should make the parameters reference the value from XCom as well, not just exercising the literal case only. We need to make the changes e2e. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
