This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 175d0278691 Decode Go SDK date-times from the supervisor in UTC
(#75737)
175d0278691 is described below
commit 175d0278691b728b29fc61cdd8e4384514b90815
Author: ChickenBenny <[email protected]>
AuthorDate: Tue Nov 10 11:55:09 2026 +0800
Decode Go SDK date-times from the supervisor in UTC (#75737)
* Give Go SDK tasks Dag run dates in UTC
* Add a UTC-decoding Time type for Go SDK genmodels
* Retype Go SDK genmodels time.Time fields to the UTC-decoding Time
* Retype nullable Go SDK genmodels date-times to *Time
* Test that generated Go SDK date-time fields decode in UTC
* Test Go SDK date-time string decoding and the gen date-time pass
* Address review: date-time lists, unmatched defs, Time JSON
---------
Co-authored-by: Tzu-ping Chung <[email protected]>
---
go-sdk/pkg/execution/client.go | 8 +-
go-sdk/pkg/execution/client_test.go | 20 +-
go-sdk/pkg/execution/genmodels/gen.go | 5 +-
go-sdk/pkg/execution/genmodels/gen/main.go | 225 +++++++++++++++++++++-
go-sdk/pkg/execution/genmodels/gen/main_test.go | 169 +++++++++++++++++
go-sdk/pkg/execution/genmodels/models.gen.go | 102 +++++-----
go-sdk/pkg/execution/genmodels/time.go | 84 +++++++++
go-sdk/pkg/execution/genmodels/time_test.go | 236 ++++++++++++++++++++++++
go-sdk/pkg/execution/integration_test.go | 11 +-
go-sdk/pkg/execution/messages.go | 12 --
go-sdk/pkg/execution/messages_test.go | 33 ++--
go-sdk/pkg/execution/task_runner.go | 24 +--
go-sdk/pkg/execution/trigger_runner.go | 11 +-
go-sdk/sdk/context.go | 5 +-
14 files changed, 820 insertions(+), 125 deletions(-)
diff --git a/go-sdk/pkg/execution/client.go b/go-sdk/pkg/execution/client.go
index f05f894f960..d49f843c19e 100644
--- a/go-sdk/pkg/execution/client.go
+++ b/go-sdk/pkg/execution/client.go
@@ -88,7 +88,7 @@ func NewCoordinatorClient(comm *CoordinatorComm, tiID string)
*CoordinatorClient
// supervisor always passes the setting, so an absent or malformed value is a
// misconfiguration and fails the write rather than silently retaining the key
// for a period nobody configured.
-func resolveDefaultExpiry(now time.Time) (any, error) {
+func resolveDefaultExpiry(now time.Time) (*genmodels.Time, error) {
raw, ok := os.LookupEnv(defaultRetentionDaysEnv)
if !ok {
return nil, fmt.Errorf(
@@ -117,7 +117,7 @@ func resolveDefaultExpiry(now time.Time) (any, error) {
"state_store",
)
}
- return expiry, nil
+ return genmodels.TimePtr(expiry), nil
}
// parseRetentionDays accepts "7.0" because Python's getint does.
@@ -502,7 +502,7 @@ func (s taskStateStore) Clear(ctx context.Context) error {
// resolveExpiry turns a retention into the wire expires_at. A nil retention -
// no sdk.WithRetention - follows the deployment default.
-func resolveExpiry(retention *time.Duration, now time.Time) (any, error) {
+func resolveExpiry(retention *time.Duration, now time.Time) (*genmodels.Time,
error) {
if retention == nil {
return resolveDefaultExpiry(now)
}
@@ -517,7 +517,7 @@ func resolveExpiry(retention *time.Duration, now time.Time)
(any, error) {
r,
)
default:
- return now.UTC().Add(r), nil
+ return genmodels.TimePtr(now.UTC().Add(r)), nil
}
}
diff --git a/go-sdk/pkg/execution/client_test.go
b/go-sdk/pkg/execution/client_test.go
index 3c293bfae72..f8a5b76c1f4 100644
--- a/go-sdk/pkg/execution/client_test.go
+++ b/go-sdk/pkg/execution/client_test.go
@@ -544,16 +544,24 @@ func TestResolveDefaultExpiry(t *testing.T) {
name string
env string
unset bool
- want any
+ want *genmodels.Time
wantErr string
// The timestamp only wraps for day counts that need a 64-bit
int.
needs64BitInt bool
}{
{name: "unset is an error", unset: true, wantErr: "is not set"},
- {name: "honours supervisor value", env: "7", want:
now.UTC().AddDate(0, 0, 7)},
+ {
+ name: "honours supervisor value",
+ env: "7",
+ want: genmodels.TimePtr(now.UTC().AddDate(0, 0, 7)),
+ },
{name: "zero days never expires", env: "0", want: nil},
// Python's config parser accepts a whole-number float spelling.
- {name: "whole float accepted", env: "7.0", want:
now.UTC().AddDate(0, 0, 7)},
+ {
+ name: "whole float accepted",
+ env: "7.0",
+ want: genmodels.TimePtr(now.UTC().AddDate(0, 0, 7)),
+ },
{name: "unparsable is an error", env: "abc", wantErr: "failed
to convert value to int"},
// conf.get returns "" for an empty setting, which must not
read as unset.
{name: "empty is an error", env: "", wantErr: "failed to
convert value to int"},
@@ -595,10 +603,6 @@ func TestResolveDefaultExpiry(t *testing.T) {
return
}
require.NoError(t, err)
- if tc.want == nil {
- assert.Nil(t, got, "a nil expiry must be
untyped so msgpack encodes null")
- return
- }
assert.Equal(t, tc.want, got)
})
}
@@ -1277,7 +1281,7 @@ func TestCoordinatorClientTriggerDagRun(t *testing.T) {
genmodels.TriggerDagRun{
DagID: "downstream",
RunID: "run_1",
- LogicalDate: logicalDate,
+ LogicalDate:
genmodels.TimePtr(logicalDate),
Conf: &genmodels.Conf{"k": "v"},
},
)
diff --git a/go-sdk/pkg/execution/genmodels/gen.go
b/go-sdk/pkg/execution/genmodels/gen.go
index 900207090c7..9715b94563b 100644
--- a/go-sdk/pkg/execution/genmodels/gen.go
+++ b/go-sdk/pkg/execution/genmodels/gen.go
@@ -32,8 +32,9 @@
// go-jsonschema's dead anyOf-branch typedefs, widens concrete int/float/bool
// fields whose schema default their Go zero value does not satisfy to pointers
// (so an unset value is omitted on the wire and the supervisor reapplies the
-// default, while an explicit 0/false still encodes), and emits
-// discriminators.gen.go and defaults.gen.go.
+// default, while an explicit 0/false still encodes), retypes every date-time
+// field to the UTC-decoding Time in time.go, and emits discriminators.gen.go
and
+// defaults.gen.go.
package genmodels
//go:generate go run github.com/atombender/[email protected] --only-models
--struct-name-from-title --tags msgpack --capitalization ID --capitalization
URI --capitalization TI -p genmodels -o models.gen.go
../../../schema/supervisor-schema.json
diff --git a/go-sdk/pkg/execution/genmodels/gen/main.go
b/go-sdk/pkg/execution/genmodels/gen/main.go
index 7a6f3bbc6e0..d734e1c49f8 100644
--- a/go-sdk/pkg/execution/genmodels/gen/main.go
+++ b/go-sdk/pkg/execution/genmodels/gen/main.go
@@ -16,7 +16,7 @@
// under the License.
// Command gen post-processes the genmodels package after go-jsonschema, doing
-// five things it cannot:
+// six things it cannot:
//
// 1. Strips the dead `<Type>_<N>` typedefs emitted for each branch of a
nullable
// `anyOf` (e.g. AssetEventResponsePartitionKey_0); the structs use `any`,
so
@@ -39,6 +39,10 @@
// on a nullable interface{}/enum field), seeding it before decode since
// msgpack applies no defaults. Pointer-widened fields (see 2) are skipped:
// their default is simply nil.
+// 6. Retypes every schema date-time to the package's UTC-decoding Time, since
+// msgpack decodes a bare time.Time in the worker's local zone: time.Time
+// becomes Time, and a nullable date-time, which go-jsonschema emits as
+// interface{}, becomes *Time.
//
// Constants and cases key on the generated struct names (read from
models.gen.go)
// matched to the schema's wire values case-insensitively, since go-jsonschema
may
@@ -108,6 +112,13 @@ func main() {
if err := pointerizeFields(*modelsPath, pointerized); err != nil {
log.Fatalf("gen: pointerizing %s: %v", *modelsPath, err)
}
+ nullableDateTimes, err := findNullableDateTimes(doc, structByKey)
+ if err != nil {
+ log.Fatalf("gen: %v", err)
+ }
+ if err := retypeDateTimes(*modelsPath, nullableDateTimes); err != nil {
+ log.Fatalf("gen: retyping date-times in %s: %v", *modelsPath,
err)
+ }
if err := genlicense.EnsureHeader(*modelsPath); err != nil {
log.Fatalf("gen: adding license header to %s: %v", *modelsPath,
err)
}
@@ -354,9 +365,221 @@ func identUseCounts(file *ast.File) map[string]int {
return counts
}
+// timeTypeName is the hand-written UTC-decoding type in the genmodels package.
+const timeTypeName = "Time"
+
+func (p schemaProp) isDateTime() bool { return p.Format == "date-time" }
+
+func (p schemaProp) isNull() bool { return p.Type == "null" }
+
+// hasDateTime reports whether p carries a date-time anywhere, so a shape
+// findNullableDateTimes does not recognise fails generation instead of being
skipped.
+func (p schemaProp) hasDateTime() bool {
+ if p.isDateTime() || (p.Items != nil && p.Items.hasDateTime()) {
+ return true
+ }
+ for _, branch := range p.AnyOf {
+ if branch.hasDateTime() {
+ return true
+ }
+ }
+ return false
+}
+
+// nonNullBranch returns X for an anyOf [X, null] prop.
+func (p schemaProp) nonNullBranch() (schemaProp, bool) {
+ if len(p.AnyOf) != 2 {
+ return schemaProp{}, false
+ }
+ value, null := p.AnyOf[0], p.AnyOf[1]
+ if value.isNull() {
+ value, null = null, value
+ }
+ return value, null.isNull()
+}
+
+// findNullableDateTimes returns structName -> wire tag for every
+// anyOf [date-time, null] prop. go-jsonschema renders those as interface{},
+// which the AST cannot tell apart from other nullable scalars; every other
+// date-time it renders as time.Time, which retypeDateTimes swaps by type
alone.
+func findNullableDateTimes(
+ doc *schemaDoc,
+ structByKey map[string]string,
+) (map[string]map[string]bool, error) {
+ set := map[string]map[string]bool{}
+ for defName, def := range doc.Defs {
+ structName, ok := structByKey[strings.ToUpper(defName)]
+ if !ok {
+ for propName, prop := range def.Properties {
+ if prop.hasDateTime() {
+ return nil, fmt.Errorf(
+ "%s.%s: date-time in a $def
with no generated struct",
+ defName,
+ propName,
+ )
+ }
+ }
+ continue
+ }
+ for propName, prop := range def.Properties {
+ if !prop.hasDateTime() {
+ continue
+ }
+ value, nullable := prop.nonNullBranch()
+ switch {
+ case prop.isDateTime(),
+ prop.Items != nil && prop.Items.isDateTime(),
+ nullable && value.Items != nil &&
value.Items.isDateTime():
+ continue
+ case nullable && value.isDateTime():
+ if set[structName] == nil {
+ set[structName] = map[string]bool{}
+ }
+ set[structName][propName] = true
+ default:
+ return nil, fmt.Errorf(
+ "%s.%s: unsupported date-time schema
shape",
+ defName,
+ propName,
+ )
+ }
+ }
+ }
+ return set, nil
+}
+
+// retypeDateTimes swaps every time.Time for Time and each nullable date-time
+// field from interface{} to *Time. Dropping the time import makes a time.Time
+// it misses fail to compile; an already retyped file is left untouched.
+func retypeDateTimes(path string, nullable map[string]map[string]bool) error {
+ src, err := os.ReadFile(path)
+ if err != nil {
+ return err
+ }
+ fset := token.NewFileSet()
+ file, err := parser.ParseFile(fset, path, src, parser.ParseComments)
+ if err != nil {
+ return err
+ }
+
+ changed := false
+ var retypeErr error
+ ast.Inspect(file, func(n ast.Node) bool {
+ switch n := n.(type) {
+ case *ast.TypeSpec:
+ if st, ok := n.Type.(*ast.StructType); ok &&
nullable[n.Name.Name] != nil {
+ if err := retypeNullable(st,
nullable[n.Name.Name], &changed); err != nil {
+ retypeErr = fmt.Errorf("%s.%w",
n.Name.Name, err)
+ }
+ }
+ case *ast.Field:
+ if isTimeTime(n.Type) {
+ n.Type = ast.NewIdent(timeTypeName)
+ changed = true
+ }
+ case *ast.ArrayType:
+ if isTimeTime(n.Elt) {
+ n.Elt = ast.NewIdent(timeTypeName)
+ changed = true
+ }
+ case *ast.StarExpr:
+ if isTimeTime(n.X) {
+ n.X = ast.NewIdent(timeTypeName)
+ changed = true
+ }
+ }
+ return retypeErr == nil
+ })
+ if retypeErr != nil {
+ return retypeErr
+ }
+ if !changed {
+ return nil
+ }
+ dropImport(file, "time")
+
+ var buf bytes.Buffer
+ if err := format.Node(&buf, fset, file); err != nil {
+ return err
+ }
+ return os.WriteFile(path, buf.Bytes(), 0o644)
+}
+
+// retypeNullable rewrites the struct's nullable date-time fields from
+// interface{} to *Time; one already *Time is left alone.
+func retypeNullable(st *ast.StructType, tags map[string]bool, changed *bool)
error {
+ for _, field := range st.Fields.List {
+ if field.Tag == nil || len(field.Names) != 1 ||
!tags[msgpackName(field.Tag.Value)] {
+ continue
+ }
+ switch {
+ case isEmptyInterface(field.Type):
+ field.Type = &ast.StarExpr{X:
ast.NewIdent(timeTypeName)}
+ *changed = true
+ case isStarIdent(field.Type, timeTypeName):
+ default:
+ return fmt.Errorf("%s: nullable date-time is not
interface{}", field.Names[0].Name)
+ }
+ }
+ return nil
+}
+
+func isTimeTime(e ast.Expr) bool {
+ sel, ok := e.(*ast.SelectorExpr)
+ return ok && isIdent(sel.X, "time") && sel.Sel.Name == "Time"
+}
+
+func isIdent(e ast.Expr, name string) bool {
+ ident, ok := e.(*ast.Ident)
+ return ok && ident.Name == name
+}
+
+func isStarIdent(e ast.Expr, name string) bool {
+ star, ok := e.(*ast.StarExpr)
+ return ok && isIdent(star.X, name)
+}
+
+func isEmptyInterface(e ast.Expr) bool {
+ iface, ok := e.(*ast.InterfaceType)
+ return ok && (iface.Methods == nil || len(iface.Methods.List) == 0)
+}
+
+// dropImport removes the import of path from file, and the import declaration
+// with it when that was its only spec.
+func dropImport(file *ast.File, path string) {
+ quoted := strconv.Quote(path)
+ for i, decl := range file.Decls {
+ gen, ok := decl.(*ast.GenDecl)
+ if !ok || gen.Tok != token.IMPORT {
+ continue
+ }
+ for j, spec := range gen.Specs {
+ if spec.(*ast.ImportSpec).Path.Value != quoted {
+ continue
+ }
+ gen.Specs = append(gen.Specs[:j], gen.Specs[j+1:]...)
+ if len(gen.Specs) == 0 {
+ file.Decls = append(file.Decls[:i],
file.Decls[i+1:]...)
+ }
+ imports := file.Imports[:0]
+ for _, imp := range file.Imports {
+ if imp.Path.Value != quoted {
+ imports = append(imports, imp)
+ }
+ }
+ file.Imports = imports
+ return
+ }
+ }
+}
+
type schemaProp struct {
Const string `json:"const"`
Default json.RawMessage `json:"default"`
+ Type any `json:"type"` // a string, or a list for a
multi-type prop
+ Format string `json:"format"`
+ Items *schemaProp `json:"items"`
+ AnyOf []schemaProp `json:"anyOf"`
}
type schemaDoc struct {
diff --git a/go-sdk/pkg/execution/genmodels/gen/main_test.go
b/go-sdk/pkg/execution/genmodels/gen/main_test.go
new file mode 100644
index 00000000000..8adb242fc6a
--- /dev/null
+++ b/go-sdk/pkg/execution/genmodels/gen/main_test.go
@@ -0,0 +1,169 @@
+// 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.
+
+package main
+
+import (
+ "encoding/json"
+ "go/format"
+ "os"
+ "path/filepath"
+ "strings"
+ "testing"
+
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+func loadTestSchema(t *testing.T, defs string) *schemaDoc {
+ t.Helper()
+ var doc schemaDoc
+ require.NoError(t, json.Unmarshal([]byte(`{"$defs": `+defs+`}`), &doc))
+ return &doc
+}
+
+const (
+ dateTimeJSON = `{"type": "string", "format": "date-time"}`
+ nullJSON = `{"type": "null"}`
+ dateTimeListJSON = `{"type": "array", "items": ` + dateTimeJSON + `}`
+)
+
+func TestFindNullableDateTimes(t *testing.T) {
+ doc := loadTestSchema(t, `{"GetAssetByUri": {"properties": {
+ "value": `+dateTimeJSON+`,
+ "nullable": {"anyOf": [`+dateTimeJSON+`, `+nullJSON+`]},
+ "null_first": {"anyOf": [`+nullJSON+`, `+dateTimeJSON+`]},
+ "list": `+dateTimeListJSON+`,
+ "nullable_list": {"anyOf": [`+dateTimeListJSON+`,
`+nullJSON+`]},
+ "nullable_text": {"anyOf": [{"type": "string"}, `+nullJSON+`]}
+ }}}`)
+ structByKey := map[string]string{"GETASSETBYURI": "GetAssetByURI"}
+
+ got, err := findNullableDateTimes(doc, structByKey)
+
+ require.NoError(t, err)
+ assert.Equal(t, map[string]map[string]bool{
+ "GetAssetByURI": {"nullable": true, "null_first": true},
+ }, got)
+}
+
+func TestFindNullableDateTimes_RejectsUnsupportedShape(t *testing.T) {
+ tests := map[string]string{
+ "date-time or string": `{"anyOf": [` + dateTimeJSON + `,
{"type": "string"}]}`,
+ "three branches": `{"anyOf": [` + dateTimeJSON + `,
{"type": "string"}, ` + nullJSON + `]}`,
+ "date-time nested in list": `{"anyOf": [{"type": "array",
"items": {"anyOf": [` + dateTimeJSON + `, ` + nullJSON + `]}}, ` + nullJSON +
`]}`,
+ }
+ for name, prop := range tests {
+ t.Run(name, func(t *testing.T) {
+ doc := loadTestSchema(t, `{"DagRun": {"properties":
{"end_date": `+prop+`}}}`)
+
+ _, err := findNullableDateTimes(doc,
map[string]string{"DAGRUN": "DagRun"})
+
+ require.ErrorContains(t, err, "DagRun.end_date:
unsupported date-time schema shape")
+ })
+ }
+}
+
+func TestFindNullableDateTimes_RejectsDefWithoutStruct(t *testing.T) {
+ doc := loadTestSchema(t, `{
+ "DagRunState": {"type": "string", "enum": ["queued",
"running"]},
+ "DagRun": {"properties": {"end_date": {"anyOf":
[`+dateTimeJSON+`, `+nullJSON+`]}}}
+ }`)
+
+ _, err := findNullableDateTimes(doc, map[string]string{"DAGRUNINFO":
"DagRunInfo"})
+
+ require.ErrorContains(t, err, "DagRun.end_date: date-time in a $def
with no generated struct")
+}
+
+const unretypedModels = `package genmodels
+
+import "time"
+
+type DagRun struct {
+ RunAfter time.Time ` + "`msgpack:\"run_after\"`" + `
+ QueuedAt *time.Time ` + "`msgpack:\"queued_at,omitempty\"`" + `
+ LogicalDate interface{} ` + "`msgpack:\"logical_date\"`" + `
+ Note interface{} ` + "`msgpack:\"note\"`" + `
+}
+
+type LogicalDates []time.Time
+`
+
+const retypedModels = `package genmodels
+
+type DagRun struct {
+ RunAfter Time ` + "`msgpack:\"run_after\"`" + `
+ QueuedAt *Time ` + "`msgpack:\"queued_at,omitempty\"`" + `
+ LogicalDate *Time ` + "`msgpack:\"logical_date\"`" + `
+ Note interface{} ` + "`msgpack:\"note\"`" + `
+}
+
+type LogicalDates []Time
+`
+
+var nullableLogicalDate = map[string]map[string]bool{"DagRun":
{"logical_date": true}}
+
+func writeTestModels(t *testing.T, src string) string {
+ t.Helper()
+ path := filepath.Join(t.TempDir(), "models.gen.go")
+ require.NoError(t, os.WriteFile(path, []byte(src), 0o644))
+ return path
+}
+
+func readFormatted(t *testing.T, path string) string {
+ t.Helper()
+ src, err := os.ReadFile(path)
+ require.NoError(t, err)
+ formatted, err := format.Source(src)
+ require.NoError(t, err)
+ return string(formatted)
+}
+
+func TestRetypeDateTimes(t *testing.T) {
+ path := writeTestModels(t, unretypedModels)
+
+ require.NoError(t, retypeDateTimes(path, nullableLogicalDate))
+
+ want, err := format.Source([]byte(retypedModels))
+ require.NoError(t, err)
+ assert.Equal(t, string(want), readFormatted(t, path))
+}
+
+func TestRetypeDateTimes_RerunIsNoOp(t *testing.T) {
+ // Not gofmt-clean, so a rewrite would show up as reformatted bytes.
+ src := strings.Replace(
+ retypedModels,
+ "type LogicalDates []Time",
+ "type LogicalDates []Time",
+ 1,
+ )
+ path := writeTestModels(t, src)
+
+ require.NoError(t, retypeDateTimes(path, nullableLogicalDate))
+
+ got, err := os.ReadFile(path)
+ require.NoError(t, err)
+ assert.Equal(t, src, string(got), "an already retyped file must not be
rewritten")
+}
+
+func TestRetypeDateTimes_RejectsNullableThatIsNotInterface(t *testing.T) {
+ path := writeTestModels(t, unretypedModels)
+
+ err := retypeDateTimes(path, map[string]map[string]bool{"DagRun":
{"run_after": true}})
+
+ require.ErrorContains(t, err, "DagRun.RunAfter: nullable date-time is
not interface{}")
+}
diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go
b/go-sdk/pkg/execution/genmodels/models.gen.go
index a4c41c88041..32009919b8b 100644
--- a/go-sdk/pkg/execution/genmodels/models.gen.go
+++ b/go-sdk/pkg/execution/genmodels/models.gen.go
@@ -18,8 +18,6 @@
package genmodels
-import "time"
-
type ArgBindings []TaskArgBinding
type ArgValueSchema map[string]JsonValue
@@ -57,7 +55,7 @@ type AssetEventDagRunReference struct {
SourceTaskID interface{} `msgpack:"source_task_id"`
// Timestamp corresponds to the JSON schema field "timestamp".
- Timestamp time.Time `msgpack:"timestamp"`
+ Timestamp Time `msgpack:"timestamp"`
}
// Asset event schema with fields that are needed for Runtime.
@@ -90,7 +88,7 @@ type AssetEventResponse struct {
SourceTaskID interface{} `msgpack:"source_task_id,omitempty"`
// Timestamp corresponds to the JSON schema field "timestamp".
- Timestamp time.Time `msgpack:"timestamp"`
+ Timestamp Time `msgpack:"timestamp"`
}
// Response to GetAssetEvent request.
@@ -445,7 +443,7 @@ type DagResult struct {
IsPaused bool `msgpack:"is_paused"`
// NextDagrun corresponds to the JSON schema field "next_dagrun".
- NextDagrun interface{} `msgpack:"next_dagrun"`
+ NextDagrun *Time `msgpack:"next_dagrun"`
// Owners corresponds to the JSON schema field "owners".
Owners interface{} `msgpack:"owners"`
@@ -476,28 +474,28 @@ type DagRun struct {
DagID string `msgpack:"dag_id"`
// DataIntervalEnd corresponds to the JSON schema field
"data_interval_end".
- DataIntervalEnd interface{} `msgpack:"data_interval_end"`
+ DataIntervalEnd *Time `msgpack:"data_interval_end"`
// DataIntervalStart corresponds to the JSON schema field
"data_interval_start".
- DataIntervalStart interface{} `msgpack:"data_interval_start"`
+ DataIntervalStart *Time `msgpack:"data_interval_start"`
// EndDate corresponds to the JSON schema field "end_date".
- EndDate interface{} `msgpack:"end_date"`
+ EndDate *Time `msgpack:"end_date"`
// LogicalDate corresponds to the JSON schema field "logical_date".
- LogicalDate interface{} `msgpack:"logical_date"`
+ LogicalDate *Time `msgpack:"logical_date"`
// Note corresponds to the JSON schema field "note".
Note interface{} `msgpack:"note,omitempty"`
// PartitionDate corresponds to the JSON schema field "partition_date".
- PartitionDate interface{} `msgpack:"partition_date,omitempty"`
+ PartitionDate *Time `msgpack:"partition_date,omitempty"`
// PartitionKey corresponds to the JSON schema field "partition_key".
PartitionKey interface{} `msgpack:"partition_key"`
// RunAfter corresponds to the JSON schema field "run_after".
- RunAfter time.Time `msgpack:"run_after"`
+ RunAfter Time `msgpack:"run_after"`
// RunID corresponds to the JSON schema field "run_id".
RunID string `msgpack:"run_id"`
@@ -506,7 +504,7 @@ type DagRun struct {
RunType DagRunType `msgpack:"run_type"`
// StartDate corresponds to the JSON schema field "start_date".
- StartDate interface{} `msgpack:"start_date"`
+ StartDate *Time `msgpack:"start_date"`
// State corresponds to the JSON schema field "state".
State DagRunState `msgpack:"state"`
@@ -524,16 +522,16 @@ type DagRunAssetReference struct {
DagID string `msgpack:"dag_id"`
// DataIntervalEnd corresponds to the JSON schema field
"data_interval_end".
- DataIntervalEnd interface{} `msgpack:"data_interval_end"`
+ DataIntervalEnd *Time `msgpack:"data_interval_end"`
// DataIntervalStart corresponds to the JSON schema field
"data_interval_start".
- DataIntervalStart interface{} `msgpack:"data_interval_start"`
+ DataIntervalStart *Time `msgpack:"data_interval_start"`
// EndDate corresponds to the JSON schema field "end_date".
- EndDate interface{} `msgpack:"end_date"`
+ EndDate *Time `msgpack:"end_date"`
// LogicalDate corresponds to the JSON schema field "logical_date".
- LogicalDate interface{} `msgpack:"logical_date"`
+ LogicalDate *Time `msgpack:"logical_date"`
// PartitionKey corresponds to the JSON schema field "partition_key".
PartitionKey interface{} `msgpack:"partition_key"`
@@ -542,7 +540,7 @@ type DagRunAssetReference struct {
RunID string `msgpack:"run_id"`
// StartDate corresponds to the JSON schema field "start_date".
- StartDate interface{} `msgpack:"start_date"`
+ StartDate *Time `msgpack:"start_date"`
// State corresponds to the JSON schema field "state".
State string `msgpack:"state"`
@@ -572,28 +570,28 @@ type DagRunResult struct {
DagID string `msgpack:"dag_id"`
// DataIntervalEnd corresponds to the JSON schema field
"data_interval_end".
- DataIntervalEnd interface{} `msgpack:"data_interval_end"`
+ DataIntervalEnd *Time `msgpack:"data_interval_end"`
// DataIntervalStart corresponds to the JSON schema field
"data_interval_start".
- DataIntervalStart interface{} `msgpack:"data_interval_start"`
+ DataIntervalStart *Time `msgpack:"data_interval_start"`
// EndDate corresponds to the JSON schema field "end_date".
- EndDate interface{} `msgpack:"end_date"`
+ EndDate *Time `msgpack:"end_date"`
// LogicalDate corresponds to the JSON schema field "logical_date".
- LogicalDate interface{} `msgpack:"logical_date"`
+ LogicalDate *Time `msgpack:"logical_date"`
// Note corresponds to the JSON schema field "note".
Note interface{} `msgpack:"note,omitempty"`
// PartitionDate corresponds to the JSON schema field "partition_date".
- PartitionDate interface{} `msgpack:"partition_date,omitempty"`
+ PartitionDate *Time `msgpack:"partition_date,omitempty"`
// PartitionKey corresponds to the JSON schema field "partition_key".
PartitionKey interface{} `msgpack:"partition_key"`
// RunAfter corresponds to the JSON schema field "run_after".
- RunAfter time.Time `msgpack:"run_after"`
+ RunAfter Time `msgpack:"run_after"`
// RunID corresponds to the JSON schema field "run_id".
RunID string `msgpack:"run_id"`
@@ -602,7 +600,7 @@ type DagRunResult struct {
RunType DagRunType `msgpack:"run_type"`
// StartDate corresponds to the JSON schema field "start_date".
- StartDate interface{} `msgpack:"start_date"`
+ StartDate *Time `msgpack:"start_date"`
// State corresponds to the JSON schema field "state".
State DagRunState `msgpack:"state"`
@@ -829,13 +827,13 @@ type GetAssetByURI struct {
type GetAssetEventByAsset struct {
// After corresponds to the JSON schema field "after".
- After interface{} `msgpack:"after,omitempty"`
+ After *Time `msgpack:"after,omitempty"`
// Ascending corresponds to the JSON schema field "ascending".
Ascending *bool `msgpack:"ascending,omitempty"`
// Before corresponds to the JSON schema field "before".
- Before interface{} `msgpack:"before,omitempty"`
+ Before *Time `msgpack:"before,omitempty"`
// Extra corresponds to the JSON schema field "extra".
Extra *Extra `msgpack:"extra,omitempty"`
@@ -862,7 +860,7 @@ type GetAssetEventByAsset struct {
type GetAssetEventByAssetAlias struct {
// After corresponds to the JSON schema field "after".
- After interface{} `msgpack:"after,omitempty"`
+ After *Time `msgpack:"after,omitempty"`
// AliasName corresponds to the JSON schema field "alias_name".
AliasName string `msgpack:"alias_name"`
@@ -871,7 +869,7 @@ type GetAssetEventByAssetAlias struct {
Ascending *bool `msgpack:"ascending,omitempty"`
// Before corresponds to the JSON schema field "before".
- Before interface{} `msgpack:"before,omitempty"`
+ Before *Time `msgpack:"before,omitempty"`
// Extra corresponds to the JSON schema field "extra".
Extra *Extra `msgpack:"extra,omitempty"`
@@ -997,7 +995,7 @@ type GetPreviousDagRun struct {
DagID string `msgpack:"dag_id"`
// LogicalDate corresponds to the JSON schema field "logical_date".
- LogicalDate time.Time `msgpack:"logical_date"`
+ LogicalDate Time `msgpack:"logical_date"`
// State corresponds to the JSON schema field "state".
State interface{} `msgpack:"state,omitempty"`
@@ -1012,7 +1010,7 @@ type GetPreviousTI struct {
DagID string `msgpack:"dag_id"`
// LogicalDate corresponds to the JSON schema field "logical_date".
- LogicalDate interface{} `msgpack:"logical_date,omitempty"`
+ LogicalDate *Time `msgpack:"logical_date,omitempty"`
// MapIndex corresponds to the JSON schema field "map_index".
MapIndex *int `msgpack:"map_index,omitempty"`
@@ -1272,7 +1270,7 @@ type HITLDetailResponseResult struct {
ParamsInput *ParamsInput `msgpack:"params_input,omitempty"`
// RespondedAt corresponds to the JSON schema field "responded_at".
- RespondedAt interface{} `msgpack:"responded_at"`
+ RespondedAt *Time `msgpack:"responded_at"`
// RespondedByUser corresponds to the JSON schema field
"responded_by_user".
RespondedByUser *HITLUser `msgpack:"responded_by_user,omitempty"`
@@ -1318,7 +1316,7 @@ type LazyDeserializedDAG struct {
Data Data `msgpack:"data"`
// LastLoaded corresponds to the JSON schema field "last_loaded".
- LastLoaded interface{} `msgpack:"last_loaded,omitempty"`
+ LastLoaded *Time `msgpack:"last_loaded,omitempty"`
}
// One positional stub-task argument carrying an inline literal from the Dag
file.
@@ -1339,7 +1337,7 @@ type LiteralArgBinding struct {
ValueSchema *ArgValueSchema `msgpack:"value_schema,omitempty"`
}
-type LogicalDates []time.Time
+type LogicalDates []Time
// Pinned loop definition and enclosing iteration for a task execution.
type LoopContext struct {
@@ -1411,16 +1409,16 @@ type ParamsInput map[string]interface{}
type PrevSuccessfulDagRunResult struct {
// DataIntervalEnd corresponds to the JSON schema field
"data_interval_end".
- DataIntervalEnd interface{} `msgpack:"data_interval_end,omitempty"`
+ DataIntervalEnd *Time `msgpack:"data_interval_end,omitempty"`
// DataIntervalStart corresponds to the JSON schema field
"data_interval_start".
- DataIntervalStart interface{} `msgpack:"data_interval_start,omitempty"`
+ DataIntervalStart *Time `msgpack:"data_interval_start,omitempty"`
// EndDate corresponds to the JSON schema field "end_date".
- EndDate interface{} `msgpack:"end_date,omitempty"`
+ EndDate *Time `msgpack:"end_date,omitempty"`
// StartDate corresponds to the JSON schema field "start_date".
- StartDate interface{} `msgpack:"start_date,omitempty"`
+ StartDate *Time `msgpack:"start_date,omitempty"`
// Type corresponds to the JSON schema field "type".
Type string `msgpack:"type,omitempty"`
@@ -1444,10 +1442,10 @@ type PreviousTIResponse struct {
Duration interface{} `msgpack:"duration,omitempty"`
// EndDate corresponds to the JSON schema field "end_date".
- EndDate interface{} `msgpack:"end_date,omitempty"`
+ EndDate *Time `msgpack:"end_date,omitempty"`
// LogicalDate corresponds to the JSON schema field "logical_date".
- LogicalDate interface{} `msgpack:"logical_date,omitempty"`
+ LogicalDate *Time `msgpack:"logical_date,omitempty"`
// MapIndex corresponds to the JSON schema field "map_index".
MapIndex interface{} `msgpack:"map_index,omitempty"`
@@ -1462,7 +1460,7 @@ type PreviousTIResponse struct {
RunID string `msgpack:"run_id"`
// StartDate corresponds to the JSON schema field "start_date".
- StartDate interface{} `msgpack:"start_date,omitempty"`
+ StartDate *Time `msgpack:"start_date,omitempty"`
// State corresponds to the JSON schema field "state".
State interface{} `msgpack:"state,omitempty"`
@@ -1502,10 +1500,10 @@ type RenderedFields map[string]JsonValue
// Update a task instance state to reschedule/up_for_reschedule.
type RescheduleTask struct {
// EndDate corresponds to the JSON schema field "end_date".
- EndDate time.Time `msgpack:"end_date"`
+ EndDate Time `msgpack:"end_date"`
// RescheduleDate corresponds to the JSON schema field
"reschedule_date".
- RescheduleDate time.Time `msgpack:"reschedule_date"`
+ RescheduleDate Time `msgpack:"reschedule_date"`
// State corresponds to the JSON schema field "state".
State string `msgpack:"state,omitempty"`
@@ -1522,7 +1520,7 @@ type ResendLoggingFD struct {
// Update a task instance state to up_for_retry.
type RetryTask struct {
// EndDate corresponds to the JSON schema field "end_date".
- EndDate time.Time `msgpack:"end_date"`
+ EndDate Time `msgpack:"end_date"`
// RenderedMapIndex corresponds to the JSON schema field
"rendered_map_index".
RenderedMapIndex interface{} `msgpack:"rendered_map_index,omitempty"`
@@ -1608,7 +1606,7 @@ type SetRenderedMapIndex struct {
type SetTaskStateStore struct {
// ExpiresAt corresponds to the JSON schema field "expires_at".
- ExpiresAt interface{} `msgpack:"expires_at"`
+ ExpiresAt *Time `msgpack:"expires_at"`
// Key corresponds to the JSON schema field "key".
Key string `msgpack:"key"`
@@ -1673,7 +1671,7 @@ type StartupDetails struct {
SentryIntegration string `msgpack:"sentry_integration"`
// StartDate corresponds to the JSON schema field "start_date".
- StartDate time.Time `msgpack:"start_date"`
+ StartDate Time `msgpack:"start_date"`
// TI corresponds to the JSON schema field "ti".
TI TaskInstance `msgpack:"ti"`
@@ -1691,7 +1689,7 @@ type States []string
// registering asset events.
type SucceedTask struct {
// EndDate corresponds to the JSON schema field "end_date".
- EndDate time.Time `msgpack:"end_date"`
+ EndDate Time `msgpack:"end_date"`
// LoopDecision corresponds to the JSON schema field "loop_decision".
LoopDecision *LoopDecisionPayload `msgpack:"loop_decision,omitempty"`
@@ -1743,7 +1741,7 @@ type TIRunContext struct {
// FirstTaskRescheduleStartDate corresponds to the JSON schema field
// "first_task_reschedule_start_date".
- FirstTaskRescheduleStartDate interface{}
`msgpack:"first_task_reschedule_start_date,omitempty"`
+ FirstTaskRescheduleStartDate *Time
`msgpack:"first_task_reschedule_start_date,omitempty"`
// Loop corresponds to the JSON schema field "loop".
Loop *LoopContext `msgpack:"loop,omitempty"`
@@ -1764,7 +1762,7 @@ type TIRunContext struct {
ShouldRetry bool `msgpack:"should_retry,omitempty"`
// StartDate corresponds to the JSON schema field "start_date".
- StartDate interface{} `msgpack:"start_date,omitempty"`
+ StartDate *Time `msgpack:"start_date,omitempty"`
// TaskRescheduleCount corresponds to the JSON schema field
// "task_reschedule_count".
@@ -1886,7 +1884,7 @@ type TaskOutlets []AssetProfile
// Response containing the first reschedule date for a task instance.
type TaskRescheduleStartDate struct {
// StartDate corresponds to the JSON schema field "start_date".
- StartDate interface{} `msgpack:"start_date"`
+ StartDate *Time `msgpack:"start_date"`
// Type corresponds to the JSON schema field "type".
Type string `msgpack:"type,omitempty"`
@@ -1900,7 +1898,7 @@ type TaskRescheduleStartDate struct {
// - anything else = FAILED
type TaskState struct {
// EndDate corresponds to the JSON schema field "end_date".
- EndDate interface{} `msgpack:"end_date,omitempty"`
+ EndDate *Time `msgpack:"end_date,omitempty"`
// LoopDecision corresponds to the JSON schema field "loop_decision".
LoopDecision *LoopDecisionPayload `msgpack:"loop_decision,omitempty"`
@@ -1967,7 +1965,7 @@ type TriggerDagRun struct {
DagID string `msgpack:"dag_id"`
// LogicalDate corresponds to the JSON schema field "logical_date".
- LogicalDate interface{} `msgpack:"logical_date,omitempty"`
+ LogicalDate *Time `msgpack:"logical_date,omitempty"`
// Note corresponds to the JSON schema field "note".
Note interface{} `msgpack:"note,omitempty"`
@@ -1979,7 +1977,7 @@ type TriggerDagRun struct {
ResetDagRun interface{} `msgpack:"reset_dag_run,omitempty"`
// RunAfter corresponds to the JSON schema field "run_after".
- RunAfter interface{} `msgpack:"run_after,omitempty"`
+ RunAfter *Time `msgpack:"run_after,omitempty"`
// RunID corresponds to the JSON schema field "run_id".
RunID string `msgpack:"run_id"`
diff --git a/go-sdk/pkg/execution/genmodels/time.go
b/go-sdk/pkg/execution/genmodels/time.go
new file mode 100644
index 00000000000..7dbbd154a4e
--- /dev/null
+++ b/go-sdk/pkg/execution/genmodels/time.go
@@ -0,0 +1,84 @@
+// 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.
+
+package genmodels
+
+import (
+ "time"
+
+ "github.com/vmihailenco/msgpack/v5"
+)
+
+// Time is a schema date-time. The supervisor's msgpack timestamp carries no
+// zone and msgpack decodes it with time.Unix, i.e. in the worker's local zone,
+// while Airflow means UTC.
+//
+// A string must carry an offset (RFC 3339). A naive one is rejected rather
than
+// read as UTC: it names no zone, and Airflow refuses to store naive datetimes.
+//
+// A defined type inherits none of time.Time's methods, so Time restates the
+// ones it needs: EncodeMsgpack (else it encodes as an empty map), IsZero (so
+// omitempty treats Time and *Time as it does time.Time: without it a Time is
+// always dropped and a *Time to the zero time is kept), String (for fmt) and
+// MarshalText/UnmarshalText (else encoding/json renders it as {}).
+type Time time.Time
+
+// EncodeMsgpack encodes t as a msgpack timestamp.
+func (t Time) EncodeMsgpack(enc *msgpack.Encoder) error {
+ return enc.EncodeTime(time.Time(t))
+}
+
+// IsZero reports whether t is the zero time, as time.Time.IsZero does.
+func (t Time) IsZero() bool {
+ return time.Time(t).IsZero()
+}
+
+// String formats t as time.Time.String does.
+func (t Time) String() string {
+ return time.Time(t).String()
+}
+
+// MarshalText formats t as RFC 3339, as time.Time.MarshalText does.
+func (t Time) MarshalText() ([]byte, error) {
+ return time.Time(t).MarshalText()
+}
+
+// UnmarshalText parses an RFC 3339 string into t in UTC.
+func (t *Time) UnmarshalText(data []byte) error {
+ var v time.Time
+ if err := v.UnmarshalText(data); err != nil {
+ return err
+ }
+ *t = Time(v.UTC())
+ return nil
+}
+
+// DecodeMsgpack decodes a msgpack timestamp into t in UTC.
+func (t *Time) DecodeMsgpack(dec *msgpack.Decoder) error {
+ v, err := dec.DecodeTime()
+ if err != nil {
+ return err
+ }
+ *t = Time(v.UTC())
+ return nil
+}
+
+// TimePtr returns t as a *Time, for setting a nullable date-time field.
+func TimePtr(t time.Time) *Time {
+ v := Time(t)
+ return &v
+}
diff --git a/go-sdk/pkg/execution/genmodels/time_test.go
b/go-sdk/pkg/execution/genmodels/time_test.go
new file mode 100644
index 00000000000..4212d1f6388
--- /dev/null
+++ b/go-sdk/pkg/execution/genmodels/time_test.go
@@ -0,0 +1,236 @@
+// 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.
+
+package genmodels
+
+import (
+ "encoding/json"
+ "fmt"
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+ "github.com/vmihailenco/msgpack/v5"
+)
+
+// withLocalZone runs the test with a non-UTC time.Local, so a timestamp that
+// msgpack decodes in the local zone is distinguishable from one in UTC.
+func withLocalZone(t *testing.T) {
+ t.Helper()
+ orig := time.Local
+ time.Local = time.FixedZone("UTC+8", 8*60*60)
+ t.Cleanup(func() { time.Local = orig })
+}
+
+func TestTimeDecodeMsgpack_UTC(t *testing.T) {
+ withLocalZone(t)
+ want := time.Date(2026, 1, 2, 3, 4, 5, 6, time.UTC)
+ raw, err := msgpack.Marshal(want)
+ require.NoError(t, err)
+
+ var got Time
+ require.NoError(t, msgpack.Unmarshal(raw, &got))
+
+ assert.Equal(t, want, time.Time(got))
+ assert.Equal(t, time.UTC, time.Time(got).Location())
+}
+
+func TestTimeDecodeMsgpack_String(t *testing.T) {
+ withLocalZone(t)
+ want := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
+ tests := map[string]struct {
+ wire string
+ wantErr bool
+ }{
+ "RFC 3339 Z": {wire: "2026-01-02T03:04:05Z"},
+ "offset +00:00": {wire: "2026-01-02T03:04:05+00:00"},
+ "offset +08:00": {wire: "2026-01-02T11:04:05+08:00"},
+ "naive is refused": {wire: "2026-01-02T03:04:05", wantErr:
true},
+ }
+ for name, tc := range tests {
+ t.Run(name, func(t *testing.T) {
+ raw, err := msgpack.Marshal(tc.wire)
+ require.NoError(t, err)
+
+ var got Time
+ err = msgpack.Unmarshal(raw, &got)
+
+ if tc.wantErr {
+ require.Error(t, err)
+ return
+ }
+ require.NoError(t, err)
+ assert.Equal(t, want, time.Time(got))
+ assert.Equal(t, time.UTC, time.Time(got).Location())
+ })
+ }
+}
+
+func TestTimeDecodeMsgpack_Nil(t *testing.T) {
+ raw, err := msgpack.Marshal(map[string]any{"value": nil, "nullable":
nil})
+ require.NoError(t, err)
+
+ var got struct {
+ Value Time `msgpack:"value"`
+ Nullable *Time `msgpack:"nullable"`
+ }
+ got.Nullable = new(Time)
+ require.NoError(t, msgpack.Unmarshal(raw, &got))
+
+ assert.True(t, time.Time(got.Value).IsZero())
+ assert.Nil(t, got.Nullable)
+}
+
+func TestTimeEncodeMsgpack_OmitEmpty(t *testing.T) {
+ type omitEmpty struct {
+ Value Time `msgpack:"value,omitempty"`
+ }
+ want := time.Date(2026, 1, 2, 3, 4, 5, 6, time.UTC)
+ tests := map[string]struct {
+ value Time
+ wantKeys []string
+ }{
+ "set is encoded": {value: Time(want), wantKeys:
[]string{"value"}},
+ "zero is omitted": {value: Time{}, wantKeys: []string{}},
+ }
+ for name, tc := range tests {
+ t.Run(name, func(t *testing.T) {
+ raw, err := msgpack.Marshal(omitEmpty{Value: tc.value})
+ require.NoError(t, err)
+
+ var got map[string]any
+ require.NoError(t, msgpack.Unmarshal(raw, &got))
+
+ keys := []string{}
+ for k := range got {
+ keys = append(keys, k)
+ }
+ assert.Equal(t, tc.wantKeys, keys)
+ })
+ }
+}
+
+func TestTimeString(t *testing.T) {
+ want := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
+
+ assert.Equal(t, want.String(), Time(want).String())
+ assert.Equal(t, want.String(), fmt.Sprint(Time(want)))
+}
+
+func TestTimeMarshalJSON(t *testing.T) {
+ v := Time(time.Date(2026, 1, 2, 3, 4, 5, 6, time.UTC))
+
+ got, err := json.Marshal(struct {
+ Value Time `json:"value"`
+ Nullable *Time `json:"nullable"`
+ }{Value: v, Nullable: &v})
+
+ require.NoError(t, err)
+ assert.JSONEq(
+ t,
+ `{"value": "2026-01-02T03:04:05.000000006Z", "nullable":
"2026-01-02T03:04:05.000000006Z"}`,
+ string(got),
+ )
+}
+
+func TestTimeUnmarshalJSON(t *testing.T) {
+ withLocalZone(t)
+ want := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
+ tests := map[string]struct {
+ wire string
+ wantErr bool
+ }{
+ "RFC 3339 Z": {wire: `"2026-01-02T03:04:05Z"`},
+ "offset +08:00": {wire: `"2026-01-02T11:04:05+08:00"`},
+ "naive is refused": {wire: `"2026-01-02T03:04:05"`, wantErr:
true},
+ }
+ for name, tc := range tests {
+ t.Run(name, func(t *testing.T) {
+ var got Time
+ err := json.Unmarshal([]byte(tc.wire), &got)
+
+ if tc.wantErr {
+ require.Error(t, err)
+ return
+ }
+ require.NoError(t, err)
+ assert.Equal(t, want, time.Time(got))
+ assert.Equal(t, time.UTC, time.Time(got).Location())
+ })
+ }
+}
+
+func TestTimeEncodeMsgpack_Timestamp(t *testing.T) {
+ want := time.Date(2026, 1, 2, 3, 4, 5, 6, time.UTC)
+ nullable := Time(want)
+ raw, err := msgpack.Marshal(struct {
+ Value Time `msgpack:"value"`
+ Nullable *Time `msgpack:"nullable"`
+ Null *Time `msgpack:"null"`
+ }{Value: Time(want), Nullable: &nullable})
+ require.NoError(t, err)
+
+ var got map[string]any
+ require.NoError(t, msgpack.Unmarshal(raw, &got))
+
+ for _, key := range []string{"value", "nullable"} {
+ ts, ok := got[key].(time.Time)
+ require.True(t, ok, "%s must encode as a timestamp, got %T",
key, got[key])
+ assert.True(t, want.Equal(ts), key)
+ }
+ assert.Nil(t, got["null"])
+}
+
+func TestGeneratedDateTimeFieldsDecodeUTC(t *testing.T) {
+ withLocalZone(t)
+ runAfter := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
+ logical := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
+ raw, err := msgpack.Marshal(map[string]any{
+ "run_after": runAfter,
+ "logical_date": logical,
+ "end_date": nil,
+ })
+ require.NoError(t, err)
+
+ var got DagRun
+ require.NoError(t, msgpack.Unmarshal(raw, &got))
+
+ assert.Equal(t, runAfter, time.Time(got.RunAfter))
+ require.NotNil(t, got.LogicalDate)
+ assert.Equal(t, logical, time.Time(*got.LogicalDate))
+ assert.Nil(t, got.EndDate)
+}
+
+func TestGeneratedDateTimeListDecodesUTC(t *testing.T) {
+ withLocalZone(t)
+ dates := []time.Time{
+ time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC),
+ time.Date(2026, 1, 2, 0, 0, 0, 0, time.UTC),
+ }
+ raw, err := msgpack.Marshal(map[string]any{"dag_id": "d",
"logical_dates": dates})
+ require.NoError(t, err)
+
+ var got GetDRCount
+ require.NoError(t, msgpack.Unmarshal(raw, &got))
+
+ require.NotNil(t, got.LogicalDates)
+ require.Len(t, *got.LogicalDates, len(dates))
+ for i, want := range dates {
+ assert.Equal(t, want, time.Time((*got.LogicalDates)[i]))
+ }
+}
diff --git a/go-sdk/pkg/execution/integration_test.go
b/go-sdk/pkg/execution/integration_test.go
index a93ff793a31..121997a1457 100644
--- a/go-sdk/pkg/execution/integration_test.go
+++ b/go-sdk/pkg/execution/integration_test.go
@@ -521,14 +521,13 @@ func TestRunTaskInjectsAirflowContext(t *testing.T) {
MapIndex: ptr(-1),
},
BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"},
- // The supervisor nests scheduling timestamps under dag_run; the
- // generated nullable date-time fields hold time.Time values
directly.
+ // The supervisor nests scheduling timestamps under dag_run.
TIContext: genmodels.TIRunContext{
DagRun: genmodels.DagRun{
- LogicalDate: logical,
- DataIntervalStart: start,
- DataIntervalEnd: end,
- RunAfter: runAfter,
+ LogicalDate: genmodels.TimePtr(logical),
+ DataIntervalStart: genmodels.TimePtr(start),
+ DataIntervalEnd: genmodels.TimePtr(end),
+ RunAfter: genmodels.Time(runAfter),
RunType: genmodels.DagRunTypeManual,
Conf: &conf,
},
diff --git a/go-sdk/pkg/execution/messages.go b/go-sdk/pkg/execution/messages.go
index 3f45c1c3fc1..e494cf36edb 100644
--- a/go-sdk/pkg/execution/messages.go
+++ b/go-sdk/pkg/execution/messages.go
@@ -19,7 +19,6 @@ package execution
import (
"fmt"
- "time"
"github.com/vmihailenco/msgpack/v5"
@@ -177,17 +176,6 @@ func ifaceStringPtr(v any) *string {
return &s
}
-// ifaceTimePtr returns a pointer to the time carried by a nullable date-time
-// field decoded as any, or nil when missing/null. msgpack decodes the
-// supervisor's timestamp extension straight into time.Time.
-func ifaceTimePtr(v any) *time.Time {
- t, ok := v.(time.Time)
- if !ok {
- return nil
- }
- return &t
-}
-
// ifaceInt converts a numeric value decoded as any (msgpack yields various
// int/uint/float widths) into int, returning def when nil or non-numeric.
func ifaceInt(v any, def int) int {
diff --git a/go-sdk/pkg/execution/messages_test.go
b/go-sdk/pkg/execution/messages_test.go
index 32fe80d3508..b94e26b14fc 100644
--- a/go-sdk/pkg/execution/messages_test.go
+++ b/go-sdk/pkg/execution/messages_test.go
@@ -124,12 +124,10 @@ func TestDecodeStartupDetails(t *testing.T) {
assert.Equal(t, "dags/tutorial.go", details.DagRelPath)
assert.Equal(t, "example_dags", details.BundleInfo.Name)
assert.Equal(t, "1.0.0", ifaceString(details.BundleInfo.Version))
- // msgpack decodes the timestamp extension into the local zone; compare
the
- // instant in UTC rather than the zoned representation.
- assert.Equal(t, startDate, details.StartDate.UTC())
- assert.Equal(t, logical,
ifaceTimePtr(details.TIContext.DagRun.LogicalDate).UTC())
- assert.Equal(t, intervalStart,
ifaceTimePtr(details.TIContext.DagRun.DataIntervalStart).UTC())
- assert.Equal(t, intervalEnd,
ifaceTimePtr(details.TIContext.DagRun.DataIntervalEnd).UTC())
+ assert.Equal(t, startDate, time.Time(details.StartDate))
+ assert.Equal(t, logical,
time.Time(*details.TIContext.DagRun.LogicalDate))
+ assert.Equal(t, intervalStart,
time.Time(*details.TIContext.DagRun.DataIntervalStart))
+ assert.Equal(t, intervalEnd,
time.Time(*details.TIContext.DagRun.DataIntervalEnd))
// max_tries / should_retry drive the retry decision and live at the top
// level of ti_context, not under dag_run.
assert.Equal(t, 3, details.TIContext.MaxTries)
@@ -150,10 +148,10 @@ func TestDecodeStartupDetails_MissingOptionalTimestamps(t
*testing.T) {
body, err := decodeIncomingBody(raw)
require.NoError(t, err)
details := body.(*genmodels.StartupDetails)
- assert.True(t, details.StartDate.IsZero())
- assert.Nil(t, ifaceTimePtr(details.TIContext.DagRun.LogicalDate))
- assert.Nil(t, ifaceTimePtr(details.TIContext.DagRun.DataIntervalStart))
- assert.Nil(t, ifaceTimePtr(details.TIContext.DagRun.DataIntervalEnd))
+ assert.True(t, time.Time(details.StartDate).IsZero())
+ assert.Nil(t, details.TIContext.DagRun.LogicalDate)
+ assert.Nil(t, details.TIContext.DagRun.DataIntervalStart)
+ assert.Nil(t, details.TIContext.DagRun.DataIntervalEnd)
// map_index is a pointer now, so an omitted key decodes to nil (the
natural
// "unmapped") rather than being seeded to the -1 sentinel.
assert.Nil(t, details.TI.MapIndex)
@@ -388,7 +386,7 @@ func TestSucceedTaskWire(t *testing.T) {
endDate := time.Date(2024, 1, 15, 10, 30, 0, 123456789, time.UTC)
m := decodeToMap(t, genmodels.SucceedTask{
Type: genmodels.TypeSucceedTask,
- EndDate: endDate,
+ EndDate: genmodels.Time(endDate),
TaskOutlets: &genmodels.TaskOutlets{},
OutletEvents: &genmodels.OutletEvents{},
})
@@ -407,7 +405,7 @@ func TestSucceedTaskWire(t *testing.T) {
func TestTaskStateWire(t *testing.T) {
endDate := time.Date(2024, 1, 15, 10, 30, 0, 0, time.UTC)
m := decodeToMap(t, genmodels.TaskState{
- Type: genmodels.TypeTaskState, State:
genmodels.TaskStateStateFailed, EndDate: endDate,
+ Type: genmodels.TypeTaskState, State:
genmodels.TaskStateStateFailed, EndDate: genmodels.TimePtr(endDate),
})
assert.Equal(t, "TaskState", m["type"])
assert.Equal(t, "failed", m["state"])
@@ -417,7 +415,7 @@ func TestRetryTaskWire(t *testing.T) {
endDate := time.Date(2024, 1, 15, 10, 30, 0, 0, time.UTC)
m := decodeToMap(t, genmodels.RetryTask{
Type: genmodels.TypeRetryTask,
- EndDate: endDate,
+ EndDate: genmodels.Time(endDate),
RetryReason: "boom",
})
assert.Equal(t, "RetryTask", m["type"])
@@ -566,15 +564,6 @@ func TestIfaceStringHelpers(t *testing.T) {
}
}
-func TestIfaceTimePtr(t *testing.T) {
- assert.Nil(t, ifaceTimePtr(nil))
- assert.Nil(t, ifaceTimePtr("not a time"))
- now := time.Date(2024, 1, 15, 0, 0, 0, 0, time.UTC)
- if p := ifaceTimePtr(now); assert.NotNil(t, p) {
- assert.Equal(t, now, *p)
- }
-}
-
func TestIsNilRaw(t *testing.T) {
assert.True(t, isNilRaw(nil))
assert.True(t, isNilRaw(msgpack.RawMessage{}))
diff --git a/go-sdk/pkg/execution/task_runner.go
b/go-sdk/pkg/execution/task_runner.go
index d01b770e836..3981e1ec799 100644
--- a/go-sdk/pkg/execution/task_runner.go
+++ b/go-sdk/pkg/execution/task_runner.go
@@ -61,7 +61,7 @@ func RunTask(
)
return genmodels.TaskState{
State: genmodels.TaskStateStateRemoved,
- EndDate: time.Now().UTC(),
+ EndDate: genmodels.TimePtr(time.Now().UTC()),
}
}
@@ -85,10 +85,10 @@ func RunTask(
sdk.DagRun{
DagID: details.TI.DagID,
RunID: details.TI.RunID,
- LogicalDate: ifaceTimePtr(dagRun.LogicalDate),
- DataIntervalStart:
ifaceTimePtr(dagRun.DataIntervalStart),
- DataIntervalEnd: ifaceTimePtr(dagRun.DataIntervalEnd),
- RunAfter: dagRun.RunAfter.UTC(), // msgpack
decodes it in the local zone.
+ LogicalDate: (*time.Time)(dagRun.LogicalDate),
+ DataIntervalStart:
(*time.Time)(dagRun.DataIntervalStart),
+ DataIntervalEnd: (*time.Time)(dagRun.DataIntervalEnd),
+ RunAfter: time.Time(dagRun.RunAfter),
RunType: dagRun.RunType,
Conf: conf,
},
@@ -200,13 +200,13 @@ func taskInstanceOf(details *genmodels.StartupDetails)
sdk.TaskInstance {
func failTask(shouldRetry bool, err error) any {
if shouldRetry {
return genmodels.RetryTask{
- EndDate: time.Now().UTC(),
+ EndDate: genmodels.Time(time.Now().UTC()),
RetryReason: err.Error(),
}
}
return genmodels.TaskState{
State: genmodels.TaskStateStateFailed,
- EndDate: time.Now().UTC(),
+ EndDate: genmodels.TimePtr(time.Now().UTC()),
}
}
@@ -323,13 +323,13 @@ func executeTask(
)
if shouldRetry {
result = genmodels.RetryTask{
- EndDate: time.Now().UTC(),
+ EndDate:
genmodels.Time(time.Now().UTC()),
RetryReason: fmt.Sprintf("panic: %v",
r),
}
} else {
result = genmodels.TaskState{
State: genmodels.TaskStateStateFailed,
- EndDate: time.Now().UTC(),
+ EndDate:
genmodels.TimePtr(time.Now().UTC()),
}
}
}
@@ -342,13 +342,13 @@ func executeTask(
// UP_FOR_RETRY via RetryTask; otherwise it terminates as
FAILED.
if shouldRetry {
return genmodels.RetryTask{
- EndDate: time.Now().UTC(),
+ EndDate: genmodels.Time(time.Now().UTC()),
RetryReason: err.Error(),
}
}
return genmodels.TaskState{
State: genmodels.TaskStateStateFailed,
- EndDate: time.Now().UTC(),
+ EndDate: genmodels.TimePtr(time.Now().UTC()),
}
}
@@ -361,7 +361,7 @@ func executeTask(
// empty collections rather than None.
func succeedTask() genmodels.SucceedTask {
return genmodels.SucceedTask{
- EndDate: time.Now().UTC(),
+ EndDate: genmodels.Time(time.Now().UTC()),
TaskOutlets: &genmodels.TaskOutlets{},
OutletEvents: &genmodels.OutletEvents{},
}
diff --git a/go-sdk/pkg/execution/trigger_runner.go
b/go-sdk/pkg/execution/trigger_runner.go
index 8c85ace9a27..fc8b67ded47 100644
--- a/go-sdk/pkg/execution/trigger_runner.go
+++ b/go-sdk/pkg/execution/trigger_runner.go
@@ -169,10 +169,10 @@ func runTriggerDagRun(
ResetDagRun: spec.ResetDagRun,
}
if !logicalDate.IsZero() {
- msg.LogicalDate = logicalDate
+ msg.LogicalDate = genmodels.TimePtr(logicalDate)
}
if !spec.RunAfter.IsZero() {
- msg.RunAfter = spec.RunAfter.UTC()
+ msg.RunAfter = genmodels.TimePtr(spec.RunAfter.UTC())
}
if spec.Conf != nil {
conf := genmodels.Conf(spec.Conf)
@@ -193,13 +193,16 @@ func runTriggerDagRun(
)
return genmodels.TaskState{
State: genmodels.TaskStateStateSkipped,
- EndDate: time.Now().UTC(),
+ EndDate: genmodels.TimePtr(time.Now().UTC()),
}
}
logger.ErrorContext(ctx, "Dag Run already exists, marking task
as failed.",
"dag_id", spec.DagID,
)
- return genmodels.TaskState{State:
genmodels.TaskStateStateFailed, EndDate: time.Now().UTC()}
+ return genmodels.TaskState{
+ State: genmodels.TaskStateStateFailed,
+ EndDate: genmodels.TimePtr(time.Now().UTC()),
+ }
}
logger.InfoContext(ctx, "Dag Run triggered successfully.",
"trigger_dag_id", spec.DagID)
if err := client.PushXCom(ctx, ti, triggerRunIDXComKey, runID); err !=
nil {
diff --git a/go-sdk/sdk/context.go b/go-sdk/sdk/context.go
index eb0263a10d6..b9e7b98c71a 100644
--- a/go-sdk/sdk/context.go
+++ b/go-sdk/sdk/context.go
@@ -91,8 +91,9 @@ type DagRunType = genmodels.DagRunType
// DagRun identifies the Dag run the current task instance belongs to and
// carries its scheduling timestamps, how it was created, and its conf. The
-// *time.Time fields are nil when the supervisor did not provide a value (for
-// example, a manually triggered run without a logical date).
+// time fields are in UTC. The *time.Time fields are nil when the supervisor
did
+// not provide a value (for example, a manually triggered run without a logical
+// date).
type DagRun struct {
DagID string
RunID string
