phanikumv commented on code in PR #69757:
URL: https://github.com/apache/airflow/pull/69757#discussion_r3673479795
##########
providers/standard/src/airflow/providers/standard/decorators/stub.py:
##########
@@ -18,19 +18,277 @@
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,
+ MappedOperator,
+ 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"
+ )
Review Comment:
Before the json.dumps, walk the value for embedded XCom references using the
same helper that the regular TaskFlow uses..
--
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]
