[pkg/ottl] Add PMapGetSetter and StandardPMapGetSetter

Author: rockdaboot

pkg/ottl/expression.go

@@ -406,6 +406,24 @@ func (g StandardFunctionGetter[K]) Get(args Arguments) (Expr[K], error) {
 	return Expr[K]{exprFunc: fn}, nil
 }
 
+type PMapGetSetter[K any] interface {
+	Get(ctx context.Context, tCtx K) (pcommon.Map, error)
+	Set(ctx context.Context, tCtx K, val pcommon.Map) error
+}
+
+type StandardPMapGetSetter[K any] struct {
+	Getter func(ctx context.Context, tCtx K) (pcommon.Map, error)
+	Setter func(ctx context.Context, tCtx K, val any) error
+}
+
+func (path StandardPMapGetSetter[K]) Get(ctx context.Context, tCtx K) (pcommon.Map, error) {
+	return path.Getter(ctx, tCtx)
+}
+
+func (path StandardPMapGetSetter[K]) Set(ctx context.Context, tCtx K, val pcommon.Map) error {
+	return path.Setter(ctx, tCtx, val)
+}
+
 // PMapGetter is a Getter that must return a pcommon.Map.
 type PMapGetter[K any] interface {
 	// Get retrieves a pcommon.Map value.

pkg/ottl/functions.go

@@ -590,6 +590,16 @@ func (p *Parser[K]) buildArg(argVal value, argType reflect.Type) (any, error) {
 			return nil, err
 		}
 		return StandardIntLikeGetter[K]{Getter: arg.Get}, nil
+	case strings.HasPrefix(name, "PMapGetSetter"):
+		if argVal.Literal == nil || argVal.Literal.Path == nil {
+			return nil, errors.New("must be a path")
+		}
+		pathGetSetter, err := p.buildGetSetterFromPath(argVal.Literal.Path)
+		if err != nil {
+			return nil, err
+		}
+		stdMapGetter := StandardPMapGetter[K]{Getter: pathGetSetter.Get}
+		return StandardPMapGetSetter[K]{Getter: stdMapGetter.Get, Setter: pathGetSetter.Set}, nil
 	case strings.HasPrefix(name, "PMapGetter"):
 		arg, err := p.newGetter(argVal)
 		if err != nil {

pkg/ottl/functions_test.go

@@ -24,6 +24,11 @@ func Test_NewFunctionCall_invalid(t *testing.T) {
 			&errorFunctionArguments{},
 			functionThatHasAnError,
 		),
+		createFactory[any](
+			"testing_pmapgetsetter",
+			&pMapGetSetterArguments{},
+			functionWithPMapGetSetter,
+		),
 		createFactory[any](
 			"testing_getsetter",
 			&getSetterArguments{},
@@ -112,6 +117,19 @@ func Test_NewFunctionCall_invalid(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "not accessor (pmap)",
+			inv: editor{
+				Function: "testing_pmapgetsetter",
+				Arguments: []argument{
+					{
+						Value: value{
+							String: ottltest.Strp("not path"),
+						},
+					},
+				},
+			},
+		},
 		{
 			name: "not accessor",
 			inv: editor{
@@ -391,6 +409,30 @@ func Test_NewFunctionCall_invalid(t *testing.T) {
 				Function: "non_pointer",
 			},
 		},
+		{
+			name: "path parts not all used (pmap)",
+			inv: editor{
+				Function: "testing_pmapgetsetter",
+				Arguments: []argument{
+					{
+						Value: value{
+							Literal: &mathExprLiteral{
+								Path: &path{
+									Fields: []field{
+										{
+											Name: "name",
+										},
+										{
+											Name: "not-used",
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
 		{
 			name: "path parts not all used",
 			inv: editor{
@@ -415,6 +457,35 @@ func Test_NewFunctionCall_invalid(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "Keys not allowed (pmap)",
+			inv: editor{
+				Function: "testing_pmapgetsetter",
+				Arguments: []argument{
+					{
+						Value: value{
+							Literal: &mathExprLiteral{
+								Path: &path{
+									Fields: []field{
+										{
+											Name: "name",
+											Keys: []key{
+												{
+													String: ottltest.Strp("foo"),
+												},
+												{
+													String: ottltest.Strp("bar"),
+												},
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
 		{
 			name: "Keys not allowed",
 			inv: editor{
@@ -806,6 +877,28 @@ func Test_NewFunctionCall(t *testing.T) {
 			},
 			want: 2,
 		},
+		{
+			name: "pmapgetsetter arg",
+			inv: editor{
+				Function: "testing_pmapgetsetter",
+				Arguments: []argument{
+					{
+						Value: value{
+							Literal: &mathExprLiteral{
+								Path: &path{
+									Fields: []field{
+										{
+											Name: "name",
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			want: nil,
+		},
 		{
 			name: "pmapgetter slice arg",
 			inv: editor{
@@ -1452,6 +1545,36 @@ func Test_NewFunctionCall(t *testing.T) {
 			},
 			want: nil,
 		},
+		{
+			name: "Complex Indexing (pmap)",
+			inv: editor{
+				Function: "testing_pmapgetsetter",
+				Arguments: []argument{
+					{
+						Value: value{
+							Literal: &mathExprLiteral{
+								Path: &path{
+									Fields: []field{
+										{
+											Name: "attributes",
+											Keys: []key{
+												{
+													String: ottltest.Strp("foo"),
+												},
+												{
+													String: ottltest.Strp("bar"),
+												},
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			want: nil,
+		},
 		{
 			name: "Complex Indexing",
 			inv: editor{
@@ -1482,6 +1605,28 @@ func Test_NewFunctionCall(t *testing.T) {
 			},
 			want: nil,
 		},
+		{
+			name: "path that allows keys but none have been specified (pmap)",
+			inv: editor{
+				Function: "testing_pmapgetsetter",
+				Arguments: []argument{
+					{
+						Value: value{
+							Literal: &mathExprLiteral{
+								Path: &path{
+									Fields: []field{
+										{
+											Name: "attributes",
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			want: nil,
+		},
 		{
 			name: "path that allows keys but none have been specified",
 			inv: editor{
@@ -1888,6 +2033,16 @@ func functionWithByteSliceLikeGetter(ByteSliceLikeGetter[any]) (ExprFunc[any], e
 	}, nil
 }
 
+type pMapGetSetterArguments struct {
+	PMapGetSetterArg PMapGetSetter[any]
+}
+
+func functionWithPMapGetSetter(PMapGetSetter[any]) (ExprFunc[any], error) {
+	return func(context.Context, any) (any, error) {
+		return "anything", nil
+	}, nil
+}
+
 type pMapGetterArguments struct {
 	PMapArg PMapGetter[any]
 }
@@ -2109,6 +2264,11 @@ func defaultFunctionsForTests() map[string]Factory[any] {
 			&setterArguments{},
 			functionWithSetter,
 		),
+		createFactory[any](
+			"testing_pmapgetsetter",
+			&pMapGetSetterArguments{},
+			functionWithPMapGetSetter,
+		),
 		createFactory[any](
 			"testing_getsetter",
 			&getSetterArguments{},