Pass native datetime and UUID literals to stub tasks

airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py

@@ -37,7 +37,43 @@
     "ArgValueSchema", Annotated[dict[str, JsonValue], Field(title="ArgValueSchema")]
 )
 """JSON-schema fragment constraining the value a stub-task argument binds to; generated
-by pydantic from the stub annotation, carried verbatim, unknown keywords ignored."""
+by pydantic from the stub annotation, carried verbatim, unknown keywords ignored.
+
+``format`` carries the part of the contract ``type`` alone cannot: which native type a
+lang SDK should decode the value into. Every SDK is expected to follow the same table,
+so a Dag author sees one behaviour regardless of the task's language:
+
+=========================  ========================  ====================================  ==================  =================================
+Python annotation          JSON-schema signal        Wire spelling                         Native target       Inline literal handling
+=========================  ========================  ====================================  ==================  =================================
+``datetime``               string + ``date-time``    ``2024-01-02T03:04:05Z``              timestamp           converted
+``date``                   string + ``date``         ``2024-01-02``                         date                converted
+``time``                   string + ``time``         ``03:04:05``                           time of day         converted
+``timedelta``              string + ``duration``     ``P1DT2H3M4S`` or ``-PT1M30S``         duration            converted
+``UUID``                   string + ``uuid``         ``6ba7b810-9dad-...-...``              UUID                converted
+``int``                    integer + ``int64``       ``42``                                 64-bit integer      JSON-native
+``float``                  number + ``double``       ``1.5``                                64-bit float        JSON-native
+``Enum`` (string value)    string + ``enum``         ``"value"``                            enum                member rejected; pass ``.value``
+``str``-backed ``Enum``    string + ``enum``         ``"value"``                            enum                JSON-native
+``int``-backed ``Enum``    integer + ``enum``        ``1``                                  enum                JSON-native
+``Decimal``                number or string          ``1.2`` or ``"1.20"``                   decimal             rejected; pass number/string
+``Path``                   string + ``path``         ``"/tmp/example"``                     path or string      rejected; pass string
+``set[datetime]``          array + ``uniqueItems``   ``["2024-01-02T03:04:05Z"]``           set of timestamps   converted in stable order
+=========================  ========================  ====================================  ==================  =================================
+
+``Inline literal handling`` describes a value captured directly from the Python Dag. The
+schema also accompanies XCom bindings, where the value is produced later. Schema generation
+does not itself serialize a literal: unsupported Python objects are rejected before a binding
+is emitted even when pydantic can describe their eventual JSON representation.
+
+Timestamps always carry an explicit offset -- a naive ``datetime`` is pinned to Airflow's
+default timezone at serialization time -- because an offset-less timestamp means different
+instants to different runtimes (UTC in Go, worker-local in JavaScript, unparsable in Java).
+
+A ``string`` target is always acceptable for any of the string formats: an SDK that does
+not model a format hands the raw text to the task and lets it parse. Unions serialize as
+``anyOf``, and a ``null`` branch means the argument may arrive absent, so the native
+parameter has to be nullable."""
 
 
 class _ArgBindingBase(BaseModel):

airflow-core/src/airflow/serialization/stub_arg_bindings.py

@@ -31,13 +31,16 @@
 import json
 import types
 import typing
+import uuid
+from collections.abc import Mapping, Sequence
 from functools import cache
 from inspect import Parameter, Signature, signature
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, NamedTuple
 
 from pydantic import PydanticUserError, TypeAdapter
 from pydantic.json_schema import GenerateJsonSchema
 
+from airflow._shared.timezones.timezone import coerce_datetime
 from airflow.models.xcom import XCOM_RETURN_KEY
 from airflow.sdk import XComArg
 from airflow.sdk.definitions.context import KNOWN_CONTEXT_KEYS
@@ -66,6 +69,7 @@ def float_schema(self, schema):
 
 # Most-derived first: datetime subclasses date, so it must be matched before date.
 _TEMPORAL_BASES = (datetime.datetime, datetime.date, datetime.time, datetime.timedelta)
+_NATIVE_VALUE_BASES = (*_TEMPORAL_BASES, uuid.UUID)
 
 
 def _normalize_temporal_annotation(annotation: Any) -> Any:
@@ -108,36 +112,189 @@ def _infer_value_schema(annotation: Any) -> dict[str, Any] | 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
+    wire_form = _get_wire_form(annotation)
+    # Deep-copy so callers embedding the fragment never alias the cached dict.
+    return copy.deepcopy(wire_form.schema) if wire_form else None
+
+
+class _ValueWireForm(NamedTuple):
+    """The schema describing an annotation's JSON form, and the adapter that renders values into it."""
+
+    adapter: TypeAdapter
+    schema: dict[str, Any]
+
+
+def _get_wire_form(annotation: Any) -> _ValueWireForm | None:
     try:
-        schema = _generate_value_schema(annotation)
+        return _build_wire_form(annotation)
     except TypeError:
-        # Unhashable annotations cannot key the cache; generate directly. Any pydantic
+        # Unhashable annotations cannot key the cache; build 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
+        return _build_wire_form.__wrapped__(annotation)
 
 
 @cache
-def _generate_value_schema(annotation: Any) -> dict[str, Any] | None:
+def _build_wire_form(annotation: Any) -> _ValueWireForm | None:
     """
-    Generate the schema for one annotation, cached for the process lifetime.
+    Build the adapter and schema for one annotation together, cached for the process lifetime.
 
     TypeAdapter construction is one of pydantic's most expensive operations and
     annotations are static, so re-serializations of the same Dag must not re-pay it.
+
+    Pairing them keeps an explicitly supported native value from being rendered in a
+    spelling its own ``value_schema`` does not describe, including when the
+    temporal-normalization retry below settles on a different annotation.
     """
     # PydanticUserError/TypeError cover annotations pydantic can't schema; either way,
     # that degrades to no schema rather than failing Dag serialization.
-    try:
-        return TypeAdapter(annotation).json_schema(schema_generator=_ValueSchemaGenerator)
-    except (PydanticUserError, TypeError):
-        normalized = _normalize_temporal_annotation(annotation)
-        if normalized is annotation:
-            return None
+    for candidate in (annotation, _normalize_temporal_annotation(annotation)):
         try:
-            return TypeAdapter(normalized).json_schema(schema_generator=_ValueSchemaGenerator)
+            adapter = TypeAdapter(candidate)
+            return _ValueWireForm(adapter, adapter.json_schema(schema_generator=_ValueSchemaGenerator))
         except (PydanticUserError, TypeError):
-            return None
+            continue
+    return None
+
+
+def _unwrap_annotated(annotation: Any) -> Any:
+    while typing.get_origin(annotation) is typing.Annotated:
+        annotation = typing.get_args(annotation)[0]
+    return annotation
+
+
+def _get_annotation_native_base(annotation: Any) -> type[Any] | None:
+    annotation = _unwrap_annotated(annotation)
+    if not isinstance(annotation, type):
+        return None
+    return next((base for base in _NATIVE_VALUE_BASES if issubclass(annotation, base)), None)
+
+
+def _get_value_native_base(value: Any) -> type[Any] | None:
+    return next((base for base in _NATIVE_VALUE_BASES if isinstance(value, base)), None)
+
+
+def _has_native_value_annotation(annotation: Any) -> bool:
+    annotation = _unwrap_annotated(annotation)
+    if _get_annotation_native_base(annotation) is not None:
+        return True
+    return any(
+        arg is not Ellipsis and _has_native_value_annotation(arg) for arg in typing.get_args(annotation)
+    )
+
+
+def _annotation_accepts_value(annotation: Any, value: Any) -> bool:
+    annotation = _unwrap_annotated(annotation)
+    if annotation is Any:
+        return False
+
+    annotation_native_base = _get_annotation_native_base(annotation)
+    if annotation_native_base is not None:
+        return annotation_native_base is _get_value_native_base(value)
+
+    origin = typing.get_origin(annotation)
+    if origin in (typing.Union, types.UnionType):
+        return any(_annotation_accepts_value(member, value) for member in typing.get_args(annotation))
+    if origin is typing.Literal:
+        return value in typing.get_args(annotation)
+
+    runtime_type = origin or annotation
+    try:
+        return isinstance(value, runtime_type)
+    except TypeError:
+        return False
+
+
+def _select_union_member(annotation: Any, value: Any) -> Any:
+    unwrapped = _unwrap_annotated(annotation)
+    if typing.get_origin(unwrapped) not in (typing.Union, types.UnionType):
+        return annotation
+    return next(
+        (member for member in typing.get_args(unwrapped) if _annotation_accepts_value(member, value)),
+        annotation,
+    )
+
+
+def _origin_accepts_value(origin: Any, value: Any) -> bool:
+    try:
+        return isinstance(value, origin)
+    except TypeError:
+        return False
+
+
+def _get_json_sort_key(value: Any) -> str:
+    return json.dumps(value, allow_nan=False, ensure_ascii=False, separators=(",", ":"), sort_keys=True)
+
+
+def _to_json_value(value: Any, annotation: Any) -> Any:
+    """
+    Render supported native values in the JSON form their ``value_schema`` advertises.
+
+    Temporal and UUID values are converted directly or recursively inside lists, tuples,
+    mapping values, and sets. Each native leaf uses the same adapter that produced its schema,
+    giving every lang SDK one spelling per format without opting unrelated pydantic-supported
+    types into serialization. Sets containing supported leaves become deterministically ordered
+    JSON arrays so repeated Dag serialization remains stable.
+
+    Unsupported values and containers pass through untouched, leaving the
+    JSON-serializability check to reject them.
+    """
+    if annotation is Parameter.empty or annotation is None or annotation is Any:
+        return value
+
+    annotation = _select_union_member(annotation, value)
+    annotation_native_base = _get_annotation_native_base(annotation)
+    value_native_base = _get_value_native_base(value)
+    if annotation_native_base is not None:
+        if annotation_native_base is not value_native_base:
+            return value
+        if value_native_base is datetime.datetime:
+            # A naive timestamp is ambiguous once it leaves Python: Go would read it as UTC,
+            # JavaScript as the worker's local time, and Java would refuse to parse it. Pin
+            # the offset here, using the same default timezone the rest of Airflow applies.
+            value = coerce_datetime(value)
+        wire_form = _get_wire_form(annotation)
+        if wire_form is None:
+            return value
+        # warnings=False: a value that does not match its annotation is the JSON-literal
+        # check's business, not a serializer warning's.
+        return wire_form.adapter.dump_python(value, mode="json", warnings=False)
+
+    unwrapped = _unwrap_annotated(annotation)
+    origin = typing.get_origin(unwrapped)
+    args = typing.get_args(unwrapped)
+    if not args or not _origin_accepts_value(origin, value):
+        return value
+
+    if isinstance(value, dict) and isinstance(origin, type) and issubclass(origin, Mapping):
+        value_annotation = args[1]
+        return {key: _to_json_value(item, value_annotation) for key, item in value.items()}
+
+    if origin is set and isinstance(value, set):
+        item_annotation = args[0]
+        if not _has_native_value_annotation(item_annotation):
+            return value
+        rendered = [_to_json_value(item, item_annotation) for item in value]
+        try:
+            return sorted(rendered, key=_get_json_sort_key)
+        except (TypeError, ValueError):
+            return value
+
+    if not isinstance(origin, type) or not issubclass(origin, Sequence):
+        return value
+    if isinstance(value, list):
+        item_annotation = args[0]
+        return [_to_json_value(item, item_annotation) for item in value]
+    if not isinstance(value, tuple):
+        return value
+    if origin is tuple:
+        if len(args) == 2 and args[1] is Ellipsis:
+            return tuple(_to_json_value(item, args[0]) for item in value)
+        if len(args) != len(value):
+            return value
+        return tuple(
+            _to_json_value(item, item_annotation) for item, item_annotation in zip(value, args, strict=True)
+        )
+    return tuple(_to_json_value(item, args[0]) for item in value)
 
 
 def _validate_stub_signature(sig: Signature, task_id: str) -> None:
@@ -174,20 +331,24 @@ def _resolve(name: str, param: Parameter) -> Any:
     return {name: _resolve(name, param) for name, param in sig.parameters.items()}
 
 
-def _ensure_json_literal(value: Any, task_id: str, name: str) -> None:
+def _reject_nested_xcom(value: Any, task_id: str, name: str) -> None:
     if next(XComArg.iter_xcom_references(value), None) is not None:
         raise ValueError(
             f"@task.stub task {task_id!r} parameter {name!r} received a collection with an "
             "upstream task output nested inside it; only a direct XComArg argument can cross "
             "the language boundary -- pass the upstream output as its own argument"
         )
+
+
+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"
+            "to the foreign runtime; annotate the stub parameter with the value's type so it "
+            "can be serialized, or pass it in its JSON form instead"
         )
 
 
@@ -274,6 +435,8 @@ def build_arg_bindings(op: DecoratedOperator) -> list[dict[str, Any]] | None:
                 xcom_entry["value_schema"] = value_schema
             spec.append(xcom_entry)
             continue
+        _reject_nested_xcom(value, task_id, name)
+        value = _to_json_value(value, annotations[name])
         _ensure_json_literal(value, task_id, name)
         entry: dict[str, Any] = {"name": name, "kind": "literal", "value": value}
         if value_schema is not None: