cool9850311 commented on code in PR #73420:
URL: https://github.com/apache/airflow/pull/73420#discussion_r4071905242
##########
go-sdk/pkg/execution/client.go:
##########
@@ -253,3 +325,193 @@ func (c *CoordinatorClient) PushXCom(
_, err := c.comm.Communicate(ctx, msg)
return err
}
+
+// GetTaskState requests a task state value from the supervisor.
+func (c *CoordinatorClient) GetTaskState(ctx context.Context, key string)
(any, error) {
+ resp, err := c.comm.Communicate(
+ ctx,
+ genmodels.GetTaskStateStore{TIID: c.tiID, Key: key},
+ )
+ if err != nil {
+ return nil, translateApiError(err, errCodeTaskStoreNotFound,
sdk.TaskStateNotFound, key)
+ }
+
+ var result genmodels.TaskStateStoreResult
+ if err := decodeBody(resp, &result); err != nil {
+ return nil, fmt.Errorf("decoding task state result: %w", err)
+ }
+
+ return result.Value, nil
+}
+
+// UnmarshalJSONTaskState gets a task state value and unmarshals it into
pointer.
+func (c *CoordinatorClient) UnmarshalJSONTaskState(
+ ctx context.Context,
+ key string,
+ pointer any,
+) error {
+ val, err := c.GetTaskState(ctx, key)
+ if err != nil {
+ return err
+ }
+ // The wire form is msgpack, so the value arrives as a decoded Go value
+ // (map, slice, number) rather than the JSON text UnmarshalJSONVariable
+ // gets. Round-tripping it through JSON is what lets encoding/json fill
+ // the caller's typed pointer.
+ b, err := json.Marshal(val)
+ if err != nil {
+ return fmt.Errorf("marshaling task state value: %w", err)
+ }
+ return json.Unmarshal(b, pointer)
+}
+
+// SetTaskState asks the supervisor to store a task state value, expiring it
+// according to the deployment's default retention.
+func (c *CoordinatorClient) SetTaskState(ctx context.Context, key string,
value any) error {
+ expiry, err := resolveDefaultExpiry(time.Now())
+ if err != nil {
+ return err
+ }
+ return c.sendSetTaskState(ctx, key, value, expiry)
+}
+
+// SetTaskStateWithRetention stores a task state value with a caller-chosen
lifetime.
+func (c *CoordinatorClient) SetTaskStateWithRetention(
+ ctx context.Context,
+ key string,
+ value any,
+ retention time.Duration,
+) error {
+ var expiry any
+ switch {
+ // Must precede any arithmetic on now: NeverExpire is the maximum
+ // time.Duration, so adding it overflows.
+ case retention == sdk.NeverExpire:
+ expiry = nil
+ case retention <= 0:
+ return fmt.Errorf(
+ "task state retention must be positive or
sdk.NeverExpire, got %s: "+
+ "use SetTaskState to follow the deployment
default, or DeleteTaskState to drop key %q",
+ retention, key,
+ )
+ default:
+ expiry = time.Now().UTC().Add(retention)
+ }
+ return c.sendSetTaskState(ctx, key, value, expiry)
+}
+
+// sendSetTaskState writes one SetTaskStateStore frame on behalf of both
setters.
+func (c *CoordinatorClient) sendSetTaskState(
+ ctx context.Context,
+ key string,
+ value any,
+ expiry any,
+) error {
+ if value == nil {
+ return fmt.Errorf("cannot set task state key %q to nil", key)
+ }
+ if err := validateJSONRepresentable(reflect.ValueOf(value)); err != nil
{
+ return fmt.Errorf("cannot set task state key %q: %w", key, err)
+ }
+
+ // TODO: warn when the serialized value exceeds the deployment's
+ // [state_store] max_value_storage_bytes, matching Python's
+ // airflow.sdk.execution_time.context task store setter.
+
+ _, err := c.comm.Communicate(ctx, genmodels.SetTaskStateStore{
+ TIID: c.tiID,
+ Key: key,
+ Value: value,
+ ExpiresAt: expiry,
+ })
+ return err
+}
+
+// DeleteTaskState asks the supervisor to delete a task state value.
+func (c *CoordinatorClient) DeleteTaskState(ctx context.Context, key string)
error {
+ _, err := c.comm.Communicate(ctx, genmodels.DeleteTaskStateStore{TIID:
c.tiID, Key: key})
+ return err
+}
+
+// ClearTaskState asks the supervisor to delete every task state value for this
+// task instance.
+func (c *CoordinatorClient) ClearTaskState(ctx context.Context) error {
+ _, err := c.comm.Communicate(ctx, genmodels.ClearTaskStateStore{TIID:
c.tiID})
+ return err
+}
+
+// timeType is rejected by validateJSONRepresentable: msgpack encodes a
+// time.Time as its timestamp extension, which the supervisor decodes to a
+// datetime and Pydantic then refuses as a task state value.
+var timeType = reflect.TypeFor[time.Time]()
+
+// validateJSONRepresentable reports whether v survives the round trip into the
+// supervisor's JsonValue: string, number, bool, list, and object, nested
+// freely. It mirrors the Pydantic validation Python applies to the same value,
+// so a bad value is rejected here with a useful message instead of costing a
+// round trip and coming back as an opaque API error.
+//
+// A Go struct is allowed because msgpack encodes it as a map, which arrives as
+// a JSON object; the types rejected below are the ones that arrive as
+// something JSON has no spelling for.
+func validateJSONRepresentable(v reflect.Value) error {
+ if v.Type() == timeType {
+ return fmt.Errorf(
+ "time.Time is not JSON representable; store
value.Format(time.RFC3339) " +
+ "and parse it back with time.Parse",
+ )
+ }
+ switch v.Kind() {
+ case reflect.String, reflect.Bool,
+ reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32,
reflect.Int64,
+ reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32,
reflect.Uint64:
+ return nil
+ case reflect.Float32, reflect.Float64:
+ if f := v.Float(); math.IsNaN(f) || math.IsInf(f, 0) {
+ return fmt.Errorf(
+ "value must be a finite number; NaN and Inf are
not JSON representable",
+ )
+ }
+ return nil
+ case reflect.Interface, reflect.Pointer:
+ if v.IsNil() {
+ return nil
+ }
+ return validateJSONRepresentable(v.Elem())
+ case reflect.Slice, reflect.Array:
+ // A byte slice encodes to msgpack binary, which arrives as
Python bytes.
+ if v.Type().Elem().Kind() == reflect.Uint8 && v.Kind() ==
reflect.Slice {
Review Comment:
Fixed.
--
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]
