jason810496 commented on code in PR #69757: URL: https://github.com/apache/airflow/pull/69757#discussion_r3663929498
########## airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py: ########## @@ -0,0 +1,65 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Business logic backing the task-instance execution routes.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from sqlalchemy.orm import Session + + from airflow.models.dagbag import DBDagBag + +# Task types (``TaskInstance.operator``, the operator class name) whose tasks carry a +# lang-SDK ``arg_bindings`` spec. Used to gate the serialized-Dag lookup so regular tasks +# never pay for it. The gate matches exact class names; a new lang-SDK operator adds its +# name here. +LANG_SDK_OPERATORS = frozenset({"_StubOperator"}) + + +def client_supports_arg_bindings() -> bool: + """ + Whether the request's negotiated API version can receive ``arg_bindings``. + + Clients on older versions never see the field (the version migration strips it from + the response), so the derivation must not run for them: a stub Dag that ran before + arg bindings existed keeps running against those clients. The version that introduced + the field is the single source of truth via Cadwyn's ``is_applied`` side-effect check. + """ Review Comment: ```suggestion Clients on older versions never see the field (the version migration strips it from the response), so the derivation must not run for them. Additionally, we shouldn't distinguish by date, we should check the `VersionChangeWithSideEffects` subclass `.is_applied`. Reference: https://docs.cadwyn.dev/concepts/version_changes/?h=versionchangewithsideeffects#version-changes-with-side-effects """ ``` ########## providers/standard/src/airflow/providers/standard/decorators/stub.py: ########## @@ -18,19 +18,262 @@ from __future__ import annotations import ast -from collections.abc import Callable +import copy +import datetime +import inspect +import json +import types +import typing +from collections.abc import Callable, Collection, Mapping +from functools import cache from typing import TYPE_CHECKING, Any +try: + from pydantic import PydanticUserError, TypeAdapter + from pydantic.json_schema import GenerateJsonSchema +except ImportError: + # Airflow 3 always ships pydantic but Airflow 2.x base installs do not; without it, + # stub args carry no value schemas and runtimes keep their decode-only fallback. + GenerateJsonSchema = object # type: ignore[assignment,misc] + TypeAdapter = None # type: ignore[assignment,misc] + PydanticUserError = None # type: ignore[assignment,misc] + from airflow.providers.common.compat.sdk import ( + KNOWN_CONTEXT_KEYS, + XCOM_RETURN_KEY, DecoratedOperator, + PlainXComArg, TaskDecorator, + XComArg, task_decorator_factory, ) if TYPE_CHECKING: from airflow.providers.common.compat.sdk import Context +class _ValueSchemaGenerator(GenerateJsonSchema): + """ + Pydantic's stock JSON-schema generation plus OpenAPI's fixed-width numeric formats. + + A foreign runtime decodes numbers into machine types, which the bare + ``integer``/``number`` type names cannot convey; ``format`` is an annotation per + JSON schema, so runtimes that don't know these names simply skip them. + """ + + def int_schema(self, schema): + return {**super().int_schema(schema), "format": "int64"} + + def float_schema(self, schema): + return {**super().float_schema(schema), "format": "double"} + + +# Most-derived first: datetime subclasses date, so it must be matched before date. +_TEMPORAL_BASES = (datetime.datetime, datetime.date, datetime.time, datetime.timedelta) + + +def _normalize_temporal_annotation(annotation: Any) -> Any: + """ + Map temporal subclasses (e.g. ``pendulum.DateTime``) to their stdlib base. + + Applied recursively through unions and containers, and only as a retry when direct + schema generation fails, so temporal types carrying their own pydantic schema keep it. + """ + # Parametrized generics must be detected before the plain-class branch: on Python + # 3.10, isinstance(list[X], type) is True and issubclass silently consults the + # origin, so the class branch would return list[X] unnormalized. + origin = typing.get_origin(annotation) + args = typing.get_args(annotation) + if origin is not None and args: + normalized = tuple(_normalize_temporal_annotation(arg) for arg in args) + if normalized == args: + return annotation + if origin in (typing.Union, types.UnionType): + return typing.Union[normalized] # noqa: UP007 -- runtime construction from a tuple + return origin[normalized] + if isinstance(annotation, type): + return next((base for base in _TEMPORAL_BASES if issubclass(annotation, base)), annotation) + return annotation + + +def _infer_value_schema(annotation: Any) -> dict[str, Any] | None: + """ + Build the JSON-schema fragment for one stub parameter annotation, via pydantic. + + The pydantic-generated schema ships verbatim, so runtimes must treat it as + open-vocabulary JSON schema. Returns ``None`` when the annotation constrains nothing + (missing, ``Any``, bare ``None``) or pydantic cannot generate a schema for it; the + binding then omits ``value_schema`` and the foreign runtime falls back to a + decode-only check. + """ + if TypeAdapter is None: + return None + if annotation is inspect.Parameter.empty or annotation is None or annotation is Any: + return None + if annotation is type(None): + # get_type_hints normalizes a bare ``None`` annotation to NoneType; a parameter + # that can only ever be None constrains nothing worth shipping. + return None + try: + schema = _generate_value_schema(annotation) + except TypeError: + # Unhashable annotations cannot key the cache; generate directly. Any pydantic + # failure inside the body degrades to None there, so this retry never re-raises. + schema = _generate_value_schema.__wrapped__(annotation) + # Deep-copy so callers embedding the fragment never alias the cached dict. + return copy.deepcopy(schema) if schema else None + + +@cache +def _generate_value_schema(annotation: Any) -> dict[str, Any] | None: + """ + Generate the schema for one annotation, cached for the process lifetime. + + TypeAdapter construction is one of pydantic's most expensive operations and + annotations are static, so re-parses of the same Dag file must not re-pay it. + """ + # Reached only when pydantic is installed (``_infer_value_schema`` guards on + # ``TypeAdapter is None``), so ``PydanticUserError`` is a real exception class here. + # It is the base of PydanticSchemaGenerationError and PydanticInvalidForJsonSchema and + # covers annotations pydantic rejects outright (e.g. bare ClassVar); TypeError catches + # the exotic generics pydantic chokes on with a plain TypeError. Either way, "pydantic + # cannot schema this" degrades to no schema rather than failing Dag parsing. + try: + return TypeAdapter(annotation).json_schema(schema_generator=_ValueSchemaGenerator) + except (PydanticUserError, TypeError): + normalized = _normalize_temporal_annotation(annotation) + if normalized is annotation: + return None + try: + return TypeAdapter(normalized).json_schema(schema_generator=_ValueSchemaGenerator) + except (PydanticUserError, TypeError): + return None + + +def _validate_stub_signature(signature: inspect.Signature, task_id: str) -> None: + for param in signature.parameters.values(): + if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD): + raise ValueError( + f"@task.stub task {task_id!r} must declare a fixed number of parameters for the " + f"foreign runtime to bind against; *{param.name} is not supported" + ) + if param.name in KNOWN_CONTEXT_KEYS: + raise ValueError( + f"@task.stub task {task_id!r} parameter {param.name!r} is an Airflow context key; " + "stub signatures declare only data parameters -- the lang-SDK runtime injects its " + "own task context natively (e.g. the Go SDK's sdk.TIRunContext parameter)" + ) + + +def _resolve_param_annotations(python_callable: Callable, signature: inspect.Signature) -> dict[str, Any]: + """Map each parameter to its parse-time-resolvable annotation (``Parameter.empty`` when not).""" + try: + hints = typing.get_type_hints(python_callable) + except (NameError, TypeError): + # Annotations that cannot be resolved at parse time (e.g. names behind + # TYPE_CHECKING with ``from __future__ import annotations``) degrade to "any". + hints = {} + + def resolve(name: str, param: inspect.Parameter) -> Any: + if name in hints: + return hints[name] + if isinstance(param.annotation, str): + return inspect.Parameter.empty + return param.annotation + + return {name: resolve(name, param) for name, param in signature.parameters.items()} + + +def _ensure_json_literal(value: Any, task_id: str, name: str) -> None: + try: + json.dumps(value, allow_nan=False) + except (TypeError, ValueError): + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} received a literal of type " + f"{type(value).__name__} that is not JSON-serializable, so it cannot be passed " + "to the foreign runtime; pass it in its JSON form instead (e.g. datetimes, dates, " + "and times as ISO 8601 strings)" + ) + + +def _validate_xcom_value(value: Any, task_id: str, name: str) -> bool: + """Validate an XComArg argument, returning True when it is a bindable direct upstream output.""" + if isinstance(value, PlainXComArg): + if value.key != XCOM_RETURN_KEY: + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} references the XCom key " + f"{value.key!r}; only an upstream task's return value can cross the language " + "boundary -- indexing an output by a custom key is not supported" + ) + if value.operator.is_mapped: + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} references the aggregated " + f"output of the mapped task {value.operator.task_id!r}; a foreign runtime " + "pulls single XCom rows, so a mapped upstream's combined output is not " + "supported" + ) + return True + if isinstance(value, XComArg): + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} received a " + f"{type(value).__name__}; only direct upstream task outputs can cross the " + "language boundary -- .map()/.zip()/.concat() results are not supported" + ) + return False + + +def _build_arg_bindings( + python_callable: Callable, + op_args: Collection[Any], + 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. Review Comment: ```suggestion always the stub function's parameter name, so a foreign runtime can bind by name in addition to the existing positional order. ``` ########## providers/standard/pyproject.toml: ########## @@ -60,7 +60,8 @@ requires-python = ">=3.10" # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build`` dependencies = [ "apache-airflow>=2.11.0", - "apache-airflow-providers-common-compat>=1.14.1", + # 1.18.0 first ships KNOWN_CONTEXT_KEYS and PlainXComArg, which the stub decorator imports. + "apache-airflow-providers-common-compat>=1.18.0", # use next version Review Comment: ```suggestion "apache-airflow-providers-common-compat>=1.18.0", # use next version ``` ########## providers/standard/src/airflow/providers/standard/decorators/stub.py: ########## @@ -18,19 +18,262 @@ from __future__ import annotations import ast -from collections.abc import Callable +import copy +import datetime +import inspect +import json +import types +import typing +from collections.abc import Callable, Collection, Mapping +from functools import cache from typing import TYPE_CHECKING, Any +try: + from pydantic import PydanticUserError, TypeAdapter + from pydantic.json_schema import GenerateJsonSchema +except ImportError: + # Airflow 3 always ships pydantic but Airflow 2.x base installs do not; without it, + # stub args carry no value schemas and runtimes keep their decode-only fallback. + GenerateJsonSchema = object # type: ignore[assignment,misc] + TypeAdapter = None # type: ignore[assignment,misc] + PydanticUserError = None # type: ignore[assignment,misc] + from airflow.providers.common.compat.sdk import ( + KNOWN_CONTEXT_KEYS, + XCOM_RETURN_KEY, DecoratedOperator, + PlainXComArg, TaskDecorator, + XComArg, task_decorator_factory, ) if TYPE_CHECKING: from airflow.providers.common.compat.sdk import Context +class _ValueSchemaGenerator(GenerateJsonSchema): + """ + Pydantic's stock JSON-schema generation plus OpenAPI's fixed-width numeric formats. + + A foreign runtime decodes numbers into machine types, which the bare + ``integer``/``number`` type names cannot convey; ``format`` is an annotation per + JSON schema, so runtimes that don't know these names simply skip them. + """ + + def int_schema(self, schema): + return {**super().int_schema(schema), "format": "int64"} + + def float_schema(self, schema): + return {**super().float_schema(schema), "format": "double"} + + +# Most-derived first: datetime subclasses date, so it must be matched before date. +_TEMPORAL_BASES = (datetime.datetime, datetime.date, datetime.time, datetime.timedelta) + + +def _normalize_temporal_annotation(annotation: Any) -> Any: + """ + Map temporal subclasses (e.g. ``pendulum.DateTime``) to their stdlib base. + + Applied recursively through unions and containers, and only as a retry when direct + schema generation fails, so temporal types carrying their own pydantic schema keep it. + """ + # Parametrized generics must be detected before the plain-class branch: on Python + # 3.10, isinstance(list[X], type) is True and issubclass silently consults the + # origin, so the class branch would return list[X] unnormalized. + origin = typing.get_origin(annotation) + args = typing.get_args(annotation) + if origin is not None and args: + normalized = tuple(_normalize_temporal_annotation(arg) for arg in args) + if normalized == args: + return annotation + if origin in (typing.Union, types.UnionType): + return typing.Union[normalized] # noqa: UP007 -- runtime construction from a tuple + return origin[normalized] + if isinstance(annotation, type): + return next((base for base in _TEMPORAL_BASES if issubclass(annotation, base)), annotation) + return annotation + + +def _infer_value_schema(annotation: Any) -> dict[str, Any] | None: + """ + Build the JSON-schema fragment for one stub parameter annotation, via pydantic. + + The pydantic-generated schema ships verbatim, so runtimes must treat it as + open-vocabulary JSON schema. Returns ``None`` when the annotation constrains nothing + (missing, ``Any``, bare ``None``) or pydantic cannot generate a schema for it; the + binding then omits ``value_schema`` and the foreign runtime falls back to a + decode-only check. + """ + if TypeAdapter is None: + return None + if annotation is inspect.Parameter.empty or annotation is None or annotation is Any: + return None + if annotation is type(None): + # get_type_hints normalizes a bare ``None`` annotation to NoneType; a parameter + # that can only ever be None constrains nothing worth shipping. + return None + try: + schema = _generate_value_schema(annotation) + except TypeError: + # Unhashable annotations cannot key the cache; generate directly. Any pydantic + # failure inside the body degrades to None there, so this retry never re-raises. + schema = _generate_value_schema.__wrapped__(annotation) + # Deep-copy so callers embedding the fragment never alias the cached dict. + return copy.deepcopy(schema) if schema else None + + +@cache +def _generate_value_schema(annotation: Any) -> dict[str, Any] | None: + """ + Generate the schema for one annotation, cached for the process lifetime. + + TypeAdapter construction is one of pydantic's most expensive operations and + annotations are static, so re-parses of the same Dag file must not re-pay it. + """ + # Reached only when pydantic is installed (``_infer_value_schema`` guards on + # ``TypeAdapter is None``), so ``PydanticUserError`` is a real exception class here. + # It is the base of PydanticSchemaGenerationError and PydanticInvalidForJsonSchema and + # covers annotations pydantic rejects outright (e.g. bare ClassVar); TypeError catches + # the exotic generics pydantic chokes on with a plain TypeError. Either way, "pydantic + # cannot schema this" degrades to no schema rather than failing Dag parsing. + try: + return TypeAdapter(annotation).json_schema(schema_generator=_ValueSchemaGenerator) + except (PydanticUserError, TypeError): + normalized = _normalize_temporal_annotation(annotation) + if normalized is annotation: + return None + try: + return TypeAdapter(normalized).json_schema(schema_generator=_ValueSchemaGenerator) + except (PydanticUserError, TypeError): + return None + + +def _validate_stub_signature(signature: inspect.Signature, task_id: str) -> None: + for param in signature.parameters.values(): + if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD): + raise ValueError( + f"@task.stub task {task_id!r} must declare a fixed number of parameters for the " + f"foreign runtime to bind against; *{param.name} is not supported" + ) + if param.name in KNOWN_CONTEXT_KEYS: + raise ValueError( + f"@task.stub task {task_id!r} parameter {param.name!r} is an Airflow context key; " + "stub signatures declare only data parameters -- the lang-SDK runtime injects its " + "own task context natively (e.g. the Go SDK's sdk.TIRunContext parameter)" + ) + + +def _resolve_param_annotations(python_callable: Callable, signature: inspect.Signature) -> dict[str, Any]: + """Map each parameter to its parse-time-resolvable annotation (``Parameter.empty`` when not).""" + try: + hints = typing.get_type_hints(python_callable) + except (NameError, TypeError): + # Annotations that cannot be resolved at parse time (e.g. names behind + # TYPE_CHECKING with ``from __future__ import annotations``) degrade to "any". + hints = {} + + def resolve(name: str, param: inspect.Parameter) -> Any: + if name in hints: + return hints[name] + if isinstance(param.annotation, str): + return inspect.Parameter.empty + return param.annotation + + return {name: resolve(name, param) for name, param in signature.parameters.items()} + + +def _ensure_json_literal(value: Any, task_id: str, name: str) -> None: + try: + json.dumps(value, allow_nan=False) + except (TypeError, ValueError): + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} received a literal of type " + f"{type(value).__name__} that is not JSON-serializable, so it cannot be passed " + "to the foreign runtime; pass it in its JSON form instead (e.g. datetimes, dates, " + "and times as ISO 8601 strings)" + ) Review Comment: We handled the datetime and dates. ```suggestion 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; pass it in its JSON form instead. ) ``` ########## providers/standard/src/airflow/providers/standard/decorators/stub.py: ########## @@ -75,7 +318,30 @@ def __init__( f"Functions passed to @task.stub must be an empty function (`pass`, or `...` only) (got {stmt})" ) - ... + # Bind the TaskFlow call to the *original* signature (DecoratedOperator mangles context + # key defaults, which stubs reject anyway) and persist the ordered arg spec so the + # execution API can hand it to the foreign runtime via StartupDetails. + self._arg_bindings = _build_arg_bindings(python_callable, self.op_args, self.op_kwargs, self.task_id) + + # Direct .expand() on the stub needs no parse-time spec (ti_run derives per-map-index + # bindings from the serialized expand input), but a mapped task group creates + # per-map-index instances of the tasks inside it with no expand input of their own, + # so their arg values are unresolvable both here and server-side. + in_mapped_group = self.get_closest_mapped_task_group() is not None + if self._arg_bindings is not None and in_mapped_group: + raise ValueError( + f"@task.stub task {self.task_id!r} passes TaskFlow call arguments inside a mapped " + "task group; the captured spec cannot carry values that resolve per map index at " + "runtime, so stub tasks with arguments are not supported under a task group's " + ".expand()" + ) Review Comment: Should we move the `in_mapped_group = self.get_closest_mapped_task_group() is not None` validation into the `_build_arg_bindings`? Or at least before the `_build_arg_bindings` call. -- 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]
