Bind Python temporal and UUID arguments to their native Go types

airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_taskflow_binding.py

@@ -83,6 +83,7 @@ def test_all_tasks_succeeded(completed_run: _CompletedRun):
         "via_flat_map",
         "via_struct_map",
         "via_plain_map",
+        "via_temporal_args",
     ):
         assert completed_run.ti_states.get(task_id) == "success", completed_run.ti_states
 
@@ -168,3 +169,13 @@ def test_via_plain_map_decodes_dict_into_a_typed_map(completed_run: _CompletedRu
     """``via_plain_map`` passes one dict literal onto a Go ``map[string]string``
     parameter, so it decodes with no struct involved at all."""
     assert completed_run.xcom("via_plain_map") == {"team": "data", "tier": "gold"}
+
+
+def test_via_temporal_args_binds_formatted_strings_onto_native_go_types(
+    completed_run: _CompletedRun,
+):
+    assert completed_run.xcom("via_temporal_args") == {
+        "when": "2024-01-02T03:04:05Z",
+        "window_seconds": 300,
+        "trace_id": "6ba7b810-9dad-11d1-80b4-00c04fd430c8",
+    }

go-sdk/README.md

@@ -169,6 +169,33 @@ data parameters is rejected at registration. A sole struct parameter also falls
 decoding when it gets exactly one passed argument no field claims, so a task can still take an
 upstream object as a single argument.
 
+### Temporal and UUID arguments
+
+The stub's Python annotation decides which Go type an argument binds onto. The annotation describes
+what the Go task receives, not what the Dag passes: a TaskFlow call argument has to be
+JSON-serializable to cross into another runtime, so pass the formatted string (`window="PT5M"`)
+rather than a native `timedelta`, which fails when the Dag is serialized.
+
+| Python annotation | Go parameter type |
+| --- | --- |
+| `datetime`, `pendulum.DateTime` | `time.Time` |
+| `date` | `time.Time` (midnight) |
+| `time` | `time.Time` (zero date) |
+| `timedelta`, `pendulum.Duration` | `time.Duration` |
+| `UUID` | `uuid.UUID` |
+
+Any of these also binds onto a plain Go `string` if the task would rather parse the value itself —
+the escape hatch for annotations the table doesn't cover, alongside any type that decodes itself
+from JSON. They nest too: `list[timedelta]` binds onto `[]time.Duration`, `dict[str, datetime]` onto
+`map[string]time.Time`, and a struct field is reached just as readily, with a malformed value
+reported by its position, key or field name.
+
+The annotation is authoritative: a Go type the declared Python type cannot fill fails the task
+before its body runs rather than binding a value the author never meant — including a plain `int`
+declared as a Go `time.Duration`, which would otherwise be read as a count of nanoseconds. A
+nullable annotation (`datetime | None`) needs a pointer on the Go side so a `None` has somewhere to
+land.
+
 ### Reading the task runtime context
 
 Declare an `sdk.TIRunContext` parameter on a task to read the identifiers and scheduling timestamps of the

go-sdk/cmd/airflow-go-pack/pack_integration_test.go

@@ -166,6 +166,7 @@ dags:
       - "via_flat_map"
       - "via_struct_map"
       - "via_plain_map"
+      - "via_temporal_args"
 `
 	assert.Equal(t, expectedManifest, string(metadata))
 

go-sdk/dags/go_examples.py

@@ -54,7 +54,8 @@
 
 from __future__ import annotations
 
-from datetime import timedelta
+from datetime import datetime, timedelta
+from uuid import UUID
 
 from airflow.sdk import dag, task
 
@@ -165,6 +166,11 @@ def via_struct_map(payload: dict): ...
 def via_plain_map(labels: dict): ...
 
 
+# Annotations select native Go types.
+@task.stub(queue="golang")
+def via_temporal_args(when: datetime, window: timedelta, trace_id: UUID): ...
+
+
 @dag(dag_id="taskflow_binding_dag")
 def taskflow_binding_dag():
     """
@@ -184,6 +190,7 @@ def taskflow_binding_dag():
       default no Go field claims.
     * ``via_flat_map`` / ``via_struct_map`` / ``via_plain_map``: one dict bound
       whole into a struct, onto a struct's map field, and into a plain Go map.
+    * ``via_temporal_args``: formatted strings bound to native temporal and UUID types.
 
     Each ``via_struct_*`` call mixes a ``threshold`` literal with ``make_region``'s
     XCom, so struct fields are proven against both argument sources. The Go
@@ -205,6 +212,12 @@ def taskflow_binding_dag():
     via_flat_map(config={"region": "eu-west-1", "count": 3})
     via_struct_map(payload={"region": "eu-west-1", "count": 3})
     via_plain_map(labels={"team": "data", "tier": "gold"})
+    # Cross-runtime arguments use JSON strings.
+    via_temporal_args(
+        when="2024-01-02T03:04:05Z",
+        window="PT5M",
+        trace_id="6ba7b810-9dad-11d1-80b4-00c04fd430c8",
+    )
 
 
 taskflow_binding_dag()

go-sdk/example/bundle/main.go

@@ -57,6 +57,7 @@ func (m *myBundle) RegisterDags(dagbag v1.Registry) error {
 	bindingDag.AddTaskWithName("via_flat_map", taskflowbinding.ViaFlatMap)
 	bindingDag.AddTaskWithName("via_struct_map", taskflowbinding.ViaStructMap)
 	bindingDag.AddTaskWithName("via_plain_map", taskflowbinding.ViaPlainMap)
+	bindingDag.AddTaskWithName("via_temporal_args", taskflowbinding.ViaTemporalArgs)
 
 	return nil
 }

go-sdk/example/bundle/taskflowbinding/taskflowbinding.go

@@ -22,6 +22,9 @@ import (
 	"fmt"
 	"log/slog"
 	"reflect"
+	"time"
+
+	"github.com/google/uuid"
 
 	"github.com/apache/airflow/go-sdk/sdk"
 )
@@ -264,3 +267,35 @@ func ViaPlainMap(
 	log.InfoContext(ctx, "Bound dict into a plain map", "labels", fmt.Sprint(labels))
 	return map[string]any{"team": labels["team"], "tier": labels["tier"]}, nil
 }
+
+// ViaTemporalArgs validates native argument binding.
+func ViaTemporalArgs(
+	ctx sdk.TIRunContext,
+	log *slog.Logger,
+	when time.Time,
+	window time.Duration,
+	traceID uuid.UUID,
+) (any, error) {
+	wantWhen := time.Date(2024, 1, 2, 3, 4, 5, 0, time.UTC)
+	if !when.Equal(wantWhen) {
+		return nil, fmt.Errorf("datetime bound incorrectly: when=%s, want %s", when, wantWhen)
+	}
+	if window != 5*time.Minute {
+		return nil, fmt.Errorf("timedelta bound incorrectly: window=%s, want 5m", window)
+	}
+	wantTrace := uuid.MustParse("6ba7b810-9dad-11d1-80b4-00c04fd430c8")
+	if traceID != wantTrace {
+		return nil, fmt.Errorf("UUID bound incorrectly: trace_id=%s, want %s", traceID, wantTrace)
+	}
+
+	log.InfoContext(ctx, "Bound temporal arguments",
+		"when", when.Format(time.RFC3339),
+		"window", window.String(),
+		"trace_id", traceID.String(),
+	)
+	return map[string]any{
+		"when":           when.UTC().Format(time.RFC3339),
+		"window_seconds": window.Seconds(),
+		"trace_id":       traceID.String(),
+	}, nil
+}

go-sdk/example/bundle/taskflowbinding/taskflowbinding_test.go

@@ -21,7 +21,9 @@ import (
 	"context"
 	"log/slog"
 	"testing"
+	"time"
 
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -178,3 +180,35 @@ func TestViaPlainMapRejectsWrongBinding(t *testing.T) {
 	_, err := ViaPlainMap(ctx, slog.Default(), map[string]string{"team": "wrong"})
 	assert.ErrorContains(t, err, "plain map bound incorrectly")
 }
+
+func TestViaTemporalArgs(t *testing.T) {
+	ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{})
+	got, err := ViaTemporalArgs(
+		ctx, slog.Default(),
+		time.Date(2024, 1, 2, 3, 4, 5, 0, time.UTC),
+		5*time.Minute,
+		uuid.MustParse("6ba7b810-9dad-11d1-80b4-00c04fd430c8"),
+	)
+	require.NoError(t, err)
+
+	summary, ok := got.(map[string]any)
+	require.True(t, ok, "ViaTemporalArgs should return a map summary, got %T", got)
+	assert.Equal(t, "2024-01-02T03:04:05Z", summary["when"])
+	assert.Equal(t, 300.0, summary["window_seconds"])
+	assert.Equal(t, "6ba7b810-9dad-11d1-80b4-00c04fd430c8", summary["trace_id"])
+}
+
+func TestViaTemporalArgsRejectsWrongBinding(t *testing.T) {
+	ctx := sdk.NewTIRunContext(context.Background(), sdk.TaskInstance{}, sdk.DagRun{})
+	trace := uuid.MustParse("6ba7b810-9dad-11d1-80b4-00c04fd430c8")
+	when := time.Date(2024, 1, 2, 3, 4, 5, 0, time.UTC)
+
+	_, err := ViaTemporalArgs(ctx, slog.Default(), when.Add(time.Hour), 5*time.Minute, trace)
+	assert.ErrorContains(t, err, "datetime bound incorrectly")
+
+	_, err = ViaTemporalArgs(ctx, slog.Default(), when, 300, trace)
+	assert.ErrorContains(t, err, "timedelta bound incorrectly")
+
+	_, err = ViaTemporalArgs(ctx, slog.Default(), when, 5*time.Minute, uuid.Nil)
+	assert.ErrorContains(t, err, "UUID bound incorrectly")
+}