jason810496 commented on code in PR #69757:
URL: https://github.com/apache/airflow/pull/69757#discussion_r3638459163
##########
airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py:
##########
@@ -51,9 +51,11 @@
AddTeamNameField,
AddVariableKeysEndpoint,
)
+from airflow.api_fastapi.execution_api.versions.v2026_07_30 import
AddArgBindingsToTIRunContext
Review Comment:
> What Airflow version are we targeting this for release? If its 3.4 then
this should be further in the future, say 2026_10_30.
It's 3.4, I will make it 2026_10_30 then.
##########
providers/standard/src/airflow/providers/standard/decorators/stub.py:
##########
@@ -18,22 +18,170 @@
from __future__ import annotations
import ast
-from collections.abc import Callable
-from typing import TYPE_CHECKING, Any
+import inspect
+import json
+import types
+import typing
+from collections.abc import Callable, Collection, Mapping, Sequence
+from typing import TYPE_CHECKING, Any, Union
from airflow.providers.common.compat.sdk import (
+ KNOWN_CONTEXT_KEYS,
DecoratedOperator,
+ PlainXComArg,
TaskDecorator,
+ XComArg,
task_decorator_factory,
)
if TYPE_CHECKING:
from airflow.providers.common.compat.sdk import Context
+def _infer_data_type(annotation: Any) -> str:
+ """
+ Map a stub function parameter annotation to the language-neutral arg-type
vocabulary.
+
+ The returned name is one of the execution API's ``ArgBindingDataType``
values
+
(``string``/``integer``/``number``/``boolean``/``object``/``array``/``any``);
the foreign
+ runtime type-checks the bound value against it. Anything we cannot
classify confidently
+ maps to ``any`` so binding falls back to a decode-only check.
+ """
+ if annotation is inspect.Parameter.empty or annotation is None or
annotation is Any:
+ return "any"
+ origin = typing.get_origin(annotation)
+ if origin is not None:
+ if origin is Union or origin is types.UnionType:
+ members = [a for a in typing.get_args(annotation) if a is not
type(None)]
+ if len(members) == 1:
+ return _infer_data_type(members[0])
+ return "any"
+ annotation = origin
+ if not isinstance(annotation, type):
+ return "any"
+ # bool subclasses int, and str/bytes are Sequences -- order matters.
+ if issubclass(annotation, bool):
+ return "boolean"
+ if issubclass(annotation, int):
+ return "integer"
+ if issubclass(annotation, float):
+ return "number"
+ if issubclass(annotation, str):
+ return "string"
+ if issubclass(annotation, bytes):
+ return "any"
+ if issubclass(annotation, (dict, Mapping)):
+ return "object"
+ if issubclass(annotation, (list, tuple, set, frozenset, Sequence)):
+ return "array"
+ return "any"
+
+
+def _build_arg_bindings(
+ python_callable: Callable,
+ op_args: Collection[Any],
+ op_kwargs: Mapping[str, Any],
+ task_id: str,
+) -> list[dict[str, Any]] | None:
+ """
+ Bind the TaskFlow call arguments to the stub signature and build the
ordered arg spec.
+
+ Each spec entry is a plain dict matching one variant of the execution API's
+ ``TaskArgBinding`` union: an ``XComArgBinding`` (``kind="xcom"``) for
upstream TaskFlow
+ outputs, or a ``LiteralArgBinding`` (``kind="literal"``) for everything
else. ``name`` is
+ always the stub function's parameter name, so a foreign runtime can bind
by name (e.g. the
+ Go SDK's ``sdk.TaskInput`` struct fields) in addition to the existing
positional order.
+ Returns ``None`` for argless calls: the binding contract (including the
signature checks
+ below) applies only once a TaskFlow call actually passes arguments, so
pre-TaskFlow stub
+ Dags whose call arguments were always ignored keep parsing.
+ """
+ if not op_args and not op_kwargs:
+ return None
+
+ signature = inspect.signature(python_callable)
+
+ for param in signature.parameters.values():
+ if param.kind in (inspect.Parameter.VAR_POSITIONAL,
inspect.Parameter.VAR_KEYWORD):
+ raise ValueError(
+ f"@task.stub task {task_id!r} must declare a fixed number of
parameters for the "
+ f"foreign runtime to bind against; *{param.name} is not
supported"
+ )
+ if param.name in KNOWN_CONTEXT_KEYS:
+ raise ValueError(
+ f"@task.stub task {task_id!r} parameter {param.name!r} is an
Airflow context key; "
+ "stub signatures declare only data parameters -- the lang-SDK
runtime injects its "
+ "own task context natively (e.g. the Go SDK's sdk.TIRunContext
parameter)"
+ )
+
+ bound = signature.bind(*op_args, **op_kwargs)
+ explicitly_bound = set(bound.arguments)
+ bound.apply_defaults()
+
+ try:
+ hints = typing.get_type_hints(python_callable)
+ except (NameError, TypeError):
+ # Annotations that cannot be resolved at parse time (e.g. names behind
+ # TYPE_CHECKING with ``from __future__ import annotations``) degrade
to "any".
+ hints = {}
+
+ def get_annotation_for(name: str, param: inspect.Parameter) -> Any:
+ if name in hints:
+ return hints[name]
+ if isinstance(param.annotation, str):
+ return inspect.Parameter.empty
+ return param.annotation
+
+ spec: list[dict[str, Any]] = []
+ for name, param in signature.parameters.items():
+ value = bound.arguments[name]
+ data_type = _infer_data_type(get_annotation_for(name, param))
+ if isinstance(value, PlainXComArg):
+ if value.key != "return_value":
+ raise ValueError(
+ f"@task.stub task {task_id!r} parameter {name!r}
references the XCom key "
+ f"{value.key!r}; only an upstream task's return value can
cross the language "
+ "boundary -- indexing an output by a custom key is not
supported"
+ )
+ spec.append(
+ {
+ "name": name,
+ "kind": "xcom",
+ "data_type": data_type,
+ "task_id": value.operator.task_id,
+ }
+ )
+ continue
+ if isinstance(value, XComArg):
+ raise ValueError(
+ f"@task.stub task {task_id!r} parameter {name!r} received a "
+ f"{type(value).__name__}; only direct upstream task outputs
can cross the "
+ "language boundary -- .map()/.zip()/.concat() results are not
supported"
+ )
+ try:
+ json.dumps(value, allow_nan=False)
+ except (TypeError, ValueError):
+ raise ValueError(
+ f"@task.stub task {task_id!r} parameter {name!r} received a
literal of type "
+ f"{type(value).__name__} that is not JSON-serializable, so it
cannot be passed "
+ "to the foreign runtime"
+ )
+ entry: dict[str, Any] = {"name": name, "kind": "literal", "data_type":
data_type, "value": value}
+ if name not in explicitly_bound:
+ entry["from_default"] = True
+ spec.append(entry)
+ return spec
+
+
class _StubOperator(DecoratedOperator):
custom_operator_name: str = "@task.stub"
+ # Mapped stubs would need per-map-index arg specs, which the foreign
runtime cannot
Review Comment:
I will double check this part.
##########
airflow-core/tests/unit/serialization/test_dag_serialization.py:
##########
Review Comment:
I will update the `airflow-core/src/airflow/serialization/schema.json` to
show the new `args_binding` field (the native Dag TaskFlow will leverage the
`args_binding` field as well). However, I don't think it's not necessary to
bump the "serialization version".
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py:
##########
@@ -370,6 +371,87 @@ async def workload_token(request: Request) -> TIToken:
assert extras["scope"] == "execution"
assert extras["sub"] == str(ti.id)
+ def test_ti_run_returns_arg_bindings_for_stub_task(self, client,
dag_maker):
+ """A stub task's TaskFlow arg spec is extracted from the serialized
Dag and returned."""
+ with dag_maker("test_arg_bindings_dag", serialized=True):
+
+ @task.stub
+ def extract(): ...
+
+ @task.stub
+ def transform(country: str, extracted: dict, limit: int = 10): ...
+
+ transform("uk", extract())
+
+ dr = dag_maker.create_dagrun()
+ tis = {ti.task_id: ti for ti in dr.get_task_instances()}
+ for ti in tis.values():
+ ti.set_state(State.QUEUED)
+ dag_maker.session.flush()
+
+ payload = {
+ "state": "running",
+ "hostname": "random-hostname",
+ "unixname": "random-unixname",
+ "pid": 100,
+ "start_date": "2024-09-30T12:00:00Z",
+ }
+
+ response =
client.patch(f"/execution/task-instances/{tis['transform'].id}/run",
json=payload)
+ assert response.status_code == 200
+ assert response.json()["arg_bindings"] == [
+ {"name": "country", "kind": "literal", "data_type": "string",
"value": "uk"},
+ {"name": "extracted", "kind": "xcom", "data_type": "object",
"task_id": "extract"},
Review Comment:
Actually, all the "data_type" here only comes from the **annotation** at the
Stub Operator level. So what exactly the upstream task return doesn't really
matter IMO.
There comes up another case that I needs to resolve. The case you give, user
might annotate the argument as `dict | bool`, but I haven't deal with the union
type annotation yet.
##########
providers/standard/src/airflow/providers/standard/decorators/stub.py:
##########
Review Comment:
I will look into whether it would be better to change the dag parsing side
or to keep it as is.
##########
airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_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.
+"""
+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 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"
Review Comment:
Good call. I will make these JSON schema to make the "type" distinguishable
across language.
https://json-schema.org/understanding-json-schema/reference/type#format
--
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]
