Add support for caller to determine origin for union_replace

ygnmi/gnmi.go

@@ -37,6 +37,15 @@ import (
 	closer "github.com/openconfig/gocloser"
 )
 
+const (
+	// OriginOverride is the key to custom opt that sets the path origin.
+	OriginOverride = "origin-override"
+	// cliOrigin is the  path origin string for CLI-originated payloads.
+	cliOrigin = "cli"
+	// openconfigOrigin isgNMI the gNMI path origin string for OpenConfig payloads.
+	openconfigOrigin = "openconfig"
+)
+
 // subscribe create a gNMI SubscribeClient for the given query.
 func subscribe[T any](ctx context.Context, c *Client, q AnyQuery[T], mode gpb.SubscriptionList_Mode, o *opt) (_ gpb.GNMI_SubscribeClient, rerr error) {
 	var queryPaths []*gpb.Path
@@ -498,85 +507,98 @@ const (
 	unionreplacePath
 )
 
+type cliASCIIConfig string
+
 // populateSetRequest fills a SetResponse for a val and operation type.
-func populateSetRequest(req *gpb.SetRequest, path *gpb.Path, val interface{}, op setOperation, preferShadowPath, isLeaf bool, compressInfo *CompressionInfo, opts ...Option) error {
+func populateSetRequest(req *gpb.SetRequest, path *gpb.Path, val any, op setOperation, preferShadowPath, isLeaf bool, compressInfo *CompressionInfo, opts ...Option) error {
 	if req == nil {
 		return fmt.Errorf("cannot populate a nil SetRequest")
 	}
 	opt := resolveOpts(opts)
 
-	switch op {
-	case deletePath:
+	if op == deletePath {
 		req.Delete = append(req.Delete, path)
-	case replacePath, updatePath, unionreplacePath:
-		var typedVal *gpb.TypedValue
-		var err error
-		if s, ok := val.(*string); ok && path.Origin == "cli" {
-			typedVal = &gpb.TypedValue{Value: &gpb.TypedValue_AsciiVal{AsciiVal: *s}}
-		} else if s, ok := val.(string); ok && strings.HasSuffix(path.Origin, "_cli") {
-			typedVal = &gpb.TypedValue{Value: &gpb.TypedValue_AsciiVal{AsciiVal: s}}
-		} else if opt.preferProto {
-			typedVal, err = ygot.EncodeTypedValue(val, gpb.Encoding_JSON_IETF, &ygot.RFC7951JSONConfig{AppendModuleName: opt.appendModuleName, PreferShadowPath: preferShadowPath})
-		} else {
-			// Since the GoStructs are generated using preferOperationalState, we
-			// need to turn on preferShadowPath to prefer marshalling config paths.
-			var b []byte
-			b, err = ygot.Marshal7951(val, ygot.JSONIndent("  "), &ygot.RFC7951JSONConfig{AppendModuleName: opt.appendModuleName, PreferShadowPath: preferShadowPath})
-
-			// Respect the encoding option.
-			switch opt.encoding {
-			case gpb.Encoding_JSON:
-				typedVal = &gpb.TypedValue{Value: &gpb.TypedValue_JsonVal{JsonVal: b}}
-			case gpb.Encoding_JSON_IETF:
-				fallthrough
-			default:
-				typedVal = &gpb.TypedValue{Value: &gpb.TypedValue_JsonIetfVal{JsonIetfVal: b}}
-			}
-		}
+		return nil
+	}
 
-		if err != nil && opt.setFallback && path.Origin != "openconfig" {
-			if m, ok := val.(proto.Message); ok {
-				any, err := anypb.New(m)
-				if err != nil {
-					return fmt.Errorf("failed to marshal proto: %v", err)
-				}
-				typedVal = &gpb.TypedValue{Value: &gpb.TypedValue_AnyVal{AnyVal: any}}
-			} else {
-				b, err := json.Marshal(val)
-				if err != nil {
-					return fmt.Errorf("failed to marshal json: %v", err)
-				}
-				typedVal = &gpb.TypedValue{Value: &gpb.TypedValue_JsonVal{JsonVal: b}}
-			}
-		} else if err != nil {
-			return fmt.Errorf("failed to encode set request: %v", err)
-		}
+	var (
+		typedVal           *gpb.TypedValue
+		err                error
+		isCLIUnionReplace  bool
+		cliUnionReplaceVal string
+	)
+
+	if s, ok := val.(cliASCIIConfig); ok && op == unionreplacePath {
+		isCLIUnionReplace = true
+		cliUnionReplaceVal = string(s)
+	}
 
-		var modifyTypedValueFn func(*gpb.TypedValue) error
-		if !isLeaf && compressInfo != nil && len(compressInfo.PostRelPath) > 0 && len(typedVal.GetJsonIetfVal()) > 0 {
-			// When the path struct points to a node that's compressed out,
-			// then we know that the type is a node lower than it should be
-			// as far as the JSON is concerned.
-			modifyTypedValueFn = func(tv *gpb.TypedValue) error { return wrapJSONIETF(tv, compressInfo.PostRelPath) }
+	switch {
+	case isCLIUnionReplace:
+		typedVal = &gpb.TypedValue{Value: &gpb.TypedValue_AsciiVal{AsciiVal: cliUnionReplaceVal}}
+	case path.Origin == cliOrigin:
+		s, ok := val.(*string)
+		if !ok {
+			return fmt.Errorf("replace/update for CLI origin must have string pointer value, got %T", val)
 		}
-		if modifyTypedValueFn != nil {
-			if err := modifyTypedValueFn(typedVal); err != nil {
-				return fmt.Errorf("failed to modify TypedValue: %v", err)
-			}
+		typedVal = &gpb.TypedValue{Value: &gpb.TypedValue_AsciiVal{AsciiVal: *s}}
+	case opt.preferProto:
+		typedVal, err = ygot.EncodeTypedValue(val, gpb.Encoding_JSON_IETF, &ygot.RFC7951JSONConfig{AppendModuleName: opt.appendModuleName, PreferShadowPath: preferShadowPath})
+	default:
+		// Since the GoStructs are generated using preferOperationalState, we
+		// need to turn on preferShadowPath to prefer marshalling config paths.
+		var b []byte
+		b, err = ygot.Marshal7951(val, ygot.JSONIndent("  "), &ygot.RFC7951JSONConfig{AppendModuleName: opt.appendModuleName, PreferShadowPath: preferShadowPath})
+
+		// Respect the encoding option.
+		switch opt.encoding {
+		case gpb.Encoding_JSON:
+			typedVal = &gpb.TypedValue{Value: &gpb.TypedValue_JsonVal{JsonVal: b}}
+		case gpb.Encoding_JSON_IETF:
+			fallthrough
+		default:
+			typedVal = &gpb.TypedValue{Value: &gpb.TypedValue_JsonIetfVal{JsonIetfVal: b}}
 		}
-		update := &gpb.Update{
-			Path: path,
-			Val:  typedVal,
+	}
+
+	if err != nil && opt.setFallback && path.Origin != openconfigOrigin {
+		if m, ok := val.(proto.Message); ok {
+			any, err := anypb.New(m)
+			if err != nil {
+				return fmt.Errorf("failed to marshal proto: %v", err)
+			}
+			typedVal = &gpb.TypedValue{Value: &gpb.TypedValue_AnyVal{AnyVal: any}}
+		} else {
+			b, err := json.Marshal(val)
+			if err != nil {
+				return fmt.Errorf("failed to marshal json: %v", err)
+			}
+			typedVal = &gpb.TypedValue{Value: &gpb.TypedValue_JsonVal{JsonVal: b}}
 		}
+	} else if err != nil {
+		return fmt.Errorf("failed to encode set request: %v", err)
+	}
 
-		switch op {
-		case replacePath:
-			req.Replace = append(req.Replace, update)
-		case updatePath:
-			req.Update = append(req.Update, update)
-		case unionreplacePath:
-			req.UnionReplace = append(req.UnionReplace, update)
+	if !isLeaf && compressInfo != nil && len(compressInfo.PostRelPath) > 0 && len(typedVal.GetJsonIetfVal()) > 0 {
+		// When the path struct points to a node that's compressed out,
+		// then we know that the type is a node lower than it should be
+		// as far as the JSON is concerned.
+		if err := wrapJSONIETF(typedVal, compressInfo.PostRelPath); err != nil {
+			return fmt.Errorf("failed to modify TypedValue: %v", err)
 		}
+	}
+	update := &gpb.Update{
+		Path: path,
+		Val:  typedVal,
+	}
+
+	switch op {
+	case replacePath:
+		req.Replace = append(req.Replace, update)
+	case updatePath:
+		req.Update = append(req.Update, update)
+	case unionreplacePath:
+		req.UnionReplace = append(req.UnionReplace, update)
 	default:
 		return fmt.Errorf("unknown set operation: %v", op)
 	}
@@ -622,13 +644,14 @@ func prettySetRequest(setRequest *gpb.SetRequest) string {
 		writePath(update.Path)
 		writeVal(update.Val)
 	}
+	for i, update := range setRequest.UnionReplace {
+		fmt.Fprintf(&buf, "-------union_replace path/value pair #%d------\n", i)
+		writePath(update.Path)
+		writeVal(update.Val)
+	}
 	return buf.String()
 }
 
-const (
-	// OriginOverride is the key to custom opt that sets the path origin.
-	OriginOverride = "origin-override"
-)
 
 func resolvePath(q PathStruct) (*gpb.Path, error) {
 	path, opts, err := ResolvePath(q)
@@ -646,7 +669,7 @@ func resolvePath(q PathStruct) (*gpb.Path, error) {
 
 	// TODO: remove when fixed https://github.com/openconfig/ygot/issues/615
 	if path.Origin == "" && (len(path.Elem) == 0 || path.Elem[0].Name != "meta") {
-		path.Origin = "openconfig"
+		path.Origin = openconfigOrigin
 	}
 
 	return path, nil

ygnmi/gnmi_test.go

@@ -23,6 +23,112 @@ import (
 	"google.golang.org/protobuf/testing/protocmp"
 )
 
+func TestPopulateSetRequest(t *testing.T) {
+	path := &gpb.Path{Elem: []*gpb.PathElem{{Name: "a"}, {Name: "b"}}}
+	stringVal := "foo"
+	tests := []struct {
+		desc             string
+		path             *gpb.Path
+		val              any
+		op               setOperation
+		preferShadowPath bool
+		isLeaf           bool
+		compressInfo     *CompressionInfo
+		opts             []Option
+		want             *gpb.SetRequest
+		wantErr          bool
+	}{{
+		desc: "delete",
+		path: path,
+		op:   deletePath,
+		want: &gpb.SetRequest{
+			Delete: []*gpb.Path{path},
+		},
+	}, {
+		desc:   "replace-leaf",
+		path:   path,
+		val:    &stringVal,
+		op:     replacePath,
+		isLeaf: true,
+		want: &gpb.SetRequest{
+			Replace: []*gpb.Update{{
+				Path: path,
+				Val:  &gpb.TypedValue{Value: &gpb.TypedValue_JsonIetfVal{JsonIetfVal: []byte(`"foo"`)}},
+			}},
+		},
+	}, {
+		desc:   "update-leaf",
+		path:   path,
+		val:    &stringVal,
+		op:     updatePath,
+		isLeaf: true,
+		want: &gpb.SetRequest{
+			Update: []*gpb.Update{{
+				Path: path,
+				Val:  &gpb.TypedValue{Value: &gpb.TypedValue_JsonIetfVal{JsonIetfVal: []byte(`"foo"`)}},
+			}},
+		},
+	}, {
+		desc:   "union-replace-leaf",
+		path:   path,
+		val:    &stringVal,
+		op:     unionreplacePath,
+		isLeaf: true,
+		want: &gpb.SetRequest{
+			UnionReplace: []*gpb.Update{{
+				Path: path,
+				Val:  &gpb.TypedValue{Value: &gpb.TypedValue_JsonIetfVal{JsonIetfVal: []byte(`"foo"`)}}}},
+		},
+	}, {
+		desc:   "union-replace-cli",
+		path:   &gpb.Path{Origin: "cli"},
+		val:    cliASCIIConfig("foo"),
+		op:     unionreplacePath,
+		isLeaf: true,
+		want: &gpb.SetRequest{
+			UnionReplace: []*gpb.Update{{
+				Path: &gpb.Path{Origin: "cli"},
+				Val:  &gpb.TypedValue{Value: &gpb.TypedValue_AsciiVal{AsciiVal: "foo"}},
+			}},
+		},
+	}, {
+		desc:   "cli-origin-replace",
+		path:   &gpb.Path{Origin: "cli", Elem: []*gpb.PathElem{{Name: "a"}}},
+		val:    &stringVal,
+		op:     replacePath,
+		isLeaf: true,
+		want: &gpb.SetRequest{
+			Replace: []*gpb.Update{{
+				Path: &gpb.Path{Origin: "cli", Elem: []*gpb.PathElem{{Name: "a"}}},
+				Val:  &gpb.TypedValue{Value: &gpb.TypedValue_AsciiVal{AsciiVal: "foo"}},
+			}},
+		},
+	}, {
+		desc:    "cli-origin-replace-wrong-type",
+		path:    &gpb.Path{Origin: "cli", Elem: []*gpb.PathElem{{Name: "a"}}},
+		val:     stringVal,
+		op:      replacePath,
+		isLeaf:  true,
+		wantErr: true,
+	}}
+
+	for _, tt := range tests {
+		t.Run(tt.desc, func(t *testing.T) {
+			req := &gpb.SetRequest{}
+			err := populateSetRequest(req, tt.path, tt.val, tt.op, tt.preferShadowPath, tt.isLeaf, tt.compressInfo, tt.opts...)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("populateSetRequest() error = %v, wantErr %v", err, tt.wantErr)
+			}
+			if err != nil {
+				return
+			}
+			if diff := cmp.Diff(tt.want, req, protocmp.Transform()); diff != "" {
+				t.Errorf("populateSetRequest() returned diff (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
+
 func TestWrapJSONIETF(t *testing.T) {
 	tests := []struct {
 		desc               string
@@ -188,11 +294,11 @@ func (m *MockPathStructWithOrigin) schemaPath() []string {
 	return nil
 }
 
-func (m *MockPathStructWithOrigin) getKeys() map[string]interface{} {
+func (m *MockPathStructWithOrigin) getKeys() map[string]any {
 	return nil
 }
 
-func (m *MockPathStructWithOrigin) CustomData() map[string]interface{} {
+func (m *MockPathStructWithOrigin) CustomData() map[string]any {
 	return nil
 }
 
@@ -244,11 +350,11 @@ func (m *MockPathStructWithoutOrigin) schemaPath() []string {
 	return nil
 }
 
-func (m *MockPathStructWithoutOrigin) getKeys() map[string]interface{} {
+func (m *MockPathStructWithoutOrigin) getKeys() map[string]any {
 	return nil
 }
 
-func (m *MockPathStructWithoutOrigin) CustomData() map[string]interface{} {
+func (m *MockPathStructWithoutOrigin) CustomData() map[string]any {
 	return nil
 }
 

ygnmi/ygnmi.go

@@ -775,6 +775,21 @@ func (sb *SetBatch) Set(ctx context.Context, c *Client, opts ...Option) (*Result
 	return responseToResult(resp), err
 }
 
+// String returns a string representation of the queued operations in the SetBatch.
+func (sb *SetBatch) String() string {
+	req := &gpb.SetRequest{}
+	for _, op := range sb.ops {
+		path, err := resolvePath(op.path)
+		if err != nil {
+			return fmt.Sprintf("error resolving path %v: %v", op.path, err)
+		}
+		if err := populateSetRequest(req, path, op.val, op.mode, op.shadowpath, op.isLeaf, op.compressInfo); err != nil {
+			return fmt.Sprintf("error populating set request for path %v: %v", op.path, err)
+		}
+	}
+	return prettySetRequest(req)
+}
+
 // BatchUpdate stores an update operation in the SetBatch.
 func BatchUpdate[T any](sb *SetBatch, q ConfigQuery[T], val T) {
 	var setVal interface{} = val
@@ -835,10 +850,11 @@ func BatchUnionReplace[T any](sb *SetBatch, q ConfigQuery[T], val T) {
 // https://github.com/openconfig/reference/blob/master/rpc/gnmi/gnmi-union_replace.md
 func BatchUnionReplaceCLI(sb *SetBatch, nos, ascii string) {
 	ps := NewDeviceRootBase()
-	ps.PutCustomData(OriginOverride, nos+"_cli")
+	ps.PutCustomData(OriginOverride, nos)
+
 	sb.ops = append(sb.ops, &batchOp{
 		path:         ps,
-		val:          ascii,
+		val:          cliASCIIConfig(ascii),
 		mode:         unionreplacePath,
 		shadowpath:   false,
 		compressInfo: nil,