Context RequestValues for ygnmi queries

internal/gnmitestutil/gnmi.go

@@ -20,6 +20,7 @@ import (
 	"io"
 
 	"github.com/openconfig/gnmi/testing/fake/gnmi"
+	"github.com/openconfig/ygnmi/ygnmi"
 	"github.com/pkg/errors"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/local"
@@ -32,7 +33,7 @@ import (
 type FakeGNMI struct {
 	agent         *gnmi.Agent
 	stub          *Stubber
-	clientWrapper *clientWithGetter
+	clientWrapper *clientWrap
 }
 
 // StartGNMI launches a new fake GNMI server on the given port
@@ -53,7 +54,7 @@ func StartGNMI(port int) (*FakeGNMI, error) {
 	return &FakeGNMI{
 		agent:         agent,
 		stub:          stub,
-		clientWrapper: &clientWithGetter{stub: stub},
+		clientWrapper: &clientWrap{stub: stub},
 	}, nil
 }
 
@@ -79,22 +80,30 @@ func (g *FakeGNMI) Requests() []*gpb.SubscribeRequest {
 	return g.agent.Requests()
 }
 
+// LastRequestContextValues returns the request-scoped values from the last
+// ygnmi request sent to the gNMI server.
+func (g *FakeGNMI) LastRequestContextValues() *ygnmi.RequestValues {
+	return g.clientWrapper.requestValues
+}
+
 // GetRequests returns the set of GetRequests sent to the gNMI server.
 //
 // They're ordered in reverse time of request.
 func (g *FakeGNMI) GetRequests() []*gpb.GetRequest {
 	return g.clientWrapper.getRequests
 }
 
-// clientWithGetter adds gNMI Get functionality to a GNMI client.
-type clientWithGetter struct {
+// clientWrap adds gNMI Get functionality to a GNMI client.
+type clientWrap struct {
 	gpb.GNMIClient
-	stub        *Stubber
-	getRequests []*gpb.GetRequest
+	stub          *Stubber
+	getRequests   []*gpb.GetRequest
+	requestValues *ygnmi.RequestValues
 }
 
 // Get is a fake implementation of gnmi.Get, it returns the GetResponse contained in the stub.
-func (g *clientWithGetter) Get(ctx context.Context, req *gpb.GetRequest, _ ...grpc.CallOption) (*gpb.GetResponse, error) {
+func (g *clientWrap) Get(ctx context.Context, req *gpb.GetRequest, _ ...grpc.CallOption) (*gpb.GetResponse, error) {
+	g.requestValues = ygnmi.FromContext(ctx)
 	g.getRequests = append([]*gpb.GetRequest{req}, g.getRequests...)
 	if len(g.stub.getResponses) == 0 {
 		return nil, io.EOF
@@ -106,6 +115,12 @@ func (g *clientWithGetter) Get(ctx context.Context, req *gpb.GetRequest, _ ...gr
 	return resp, err
 }
 
+// Subscribe is a wrapper of the fake implementation of gnmi.Subscribe.
+func (g *clientWrap) Subscribe(ctx context.Context, opts ...grpc.CallOption) (gpb.GNMI_SubscribeClient, error) {
+	g.requestValues = ygnmi.FromContext(ctx)
+	return g.GNMIClient.Subscribe(ctx, opts...)
+}
+
 // Stubber is a handle to add stubbed responses.
 type Stubber struct {
 	gen          *fpb.FixedGenerator

internal/uexampleoc/uexampleocpath/uexampleocpath.go

@@ -265,7 +265,7 @@ func (b *Batch) Query() ygnmi.SingletonQuery[*oc.Device] {
 		false,
 		false,
 		false,
-		true,
+		false,
 		false,
 		ygnmi.NewDeviceRootBase(),
 		nil,

pathgen/gnmigen.go

@@ -132,7 +132,7 @@ func defaultStateTmplStruct(pathStructName, fakeRootName string, compressBehavio
 		SchemaStructPkgAccessor: "oc.",
 		IsState:                 true,
 		IsShadowPath:            compressBehaviour == genutil.PreferIntendedConfig,
-		IsCompressedSchema:      true,
+		IsCompressedSchema:      compressBehaviour.CompressEnabled(),
 		MethodName:              "State",
 		SingletonTypeName:       singletonQueryTypeName,
 		WildcardTypeName:        wildcardQueryTypeName,
@@ -230,7 +230,6 @@ func GNMIFieldGenerator(pathStructName, _ string, compressBehaviour genutil.Comp
 // to distinguish between config and state queries.
 func GNMIInitGenerator(pathStructName, fakeRootName string, compressBehaviour genutil.CompressBehaviour, _ *ygen.ParsedDirectory, node *NodeData, wildcard bool) (string, error) {
 	tmplStruct := defaultStateTmplStruct(pathStructName, fakeRootName, compressBehaviour, node)
-	tmplStruct.IsCompressedSchema = false
 	if err := modifyQueryType(node, &tmplStruct); err != nil {
 		return "", err
 	}

ygnmi/context.go

@@ -0,0 +1,54 @@
+// Copyright 2023 Google Inc.
+//
+// Licensed 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 ygnmi
+
+import (
+	"context"
+)
+
+// RequestValues contains request-scoped values for ygnmi queries.
+type RequestValues struct {
+	// CompressedConfigQuery is a key type that means that the query is
+	// uninterested in /state paths.
+	CompressedConfigQuery bool
+	// CompressedStateQuery is a key type that means that the query is
+	// uninterested in /config paths.
+	CompressedStateQuery bool
+}
+
+// FromContext extracts certain ygnmi request-scoped values, if present.
+func FromContext(ctx context.Context) *RequestValues {
+	compConfig, _ := ctx.Value(compressedConfigQuery{}).(bool)
+	compState, _ := ctx.Value(compressedStateQuery{}).(bool)
+	return &RequestValues{
+		CompressedConfigQuery: compConfig,
+		CompressedStateQuery:  compState,
+	}
+}
+
+// NewContext returns a new Context carrying ygnmi request-scoped values.
+func NewContext(ctx context.Context, q UntypedQuery) context.Context {
+	if q.isCompressedSchema() {
+		if q.IsState() {
+			return context.WithValue(ctx, compressedStateQuery{}, true)
+		} else {
+			return context.WithValue(ctx, compressedConfigQuery{}, true)
+		}
+	}
+	return ctx
+}
+
+type compressedConfigQuery struct{}
+type compressedStateQuery struct{}

ygnmi/context_test.go

@@ -0,0 +1,91 @@
+// Copyright 2023 Google Inc.
+//
+// Licensed 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 ygnmi_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/openconfig/ygnmi/exampleoc/exampleocpath"
+	"github.com/openconfig/ygnmi/internal/uexampleoc/uexampleocpath"
+	"github.com/openconfig/ygnmi/ygnmi"
+)
+
+func TestFromContext(t *testing.T) {
+	tests := []struct {
+		desc              string
+		inContext         context.Context
+		wantRequestValues *ygnmi.RequestValues
+	}{{
+		desc:      "compress-config",
+		inContext: ygnmi.NewContext(context.Background(), exampleocpath.Root().Parent().Config()),
+		wantRequestValues: &ygnmi.RequestValues{
+			CompressedConfigQuery: true,
+			CompressedStateQuery:  false,
+		},
+	}, {
+		desc:      "compress-config-leaf",
+		inContext: ygnmi.NewContext(context.Background(), exampleocpath.Root().Parent().Child().Five().Config()),
+		wantRequestValues: &ygnmi.RequestValues{
+			CompressedConfigQuery: true,
+			CompressedStateQuery:  false,
+		},
+	}, {
+		desc:      "compress-state",
+		inContext: ygnmi.NewContext(context.Background(), exampleocpath.Root().Parent().State()),
+		wantRequestValues: &ygnmi.RequestValues{
+			CompressedConfigQuery: false,
+			CompressedStateQuery:  true,
+		},
+	}, {
+		desc:      "compress-state-leaf",
+		inContext: ygnmi.NewContext(context.Background(), exampleocpath.Root().Parent().Child().Five().State()),
+		wantRequestValues: &ygnmi.RequestValues{
+			CompressedConfigQuery: false,
+			CompressedStateQuery:  true,
+		},
+	}, {
+		desc:      "uncompressed-container",
+		inContext: ygnmi.NewContext(context.Background(), uexampleocpath.Root().Parent()),
+		wantRequestValues: &ygnmi.RequestValues{
+			CompressedConfigQuery: false,
+			CompressedStateQuery:  false,
+		},
+	}, {
+		desc:      "uncompressed-config-container",
+		inContext: ygnmi.NewContext(context.Background(), uexampleocpath.Root().Parent().Child().Config()),
+		wantRequestValues: &ygnmi.RequestValues{
+			CompressedConfigQuery: false,
+			CompressedStateQuery:  false,
+		},
+	}, {
+		desc:      "uncompressed-leaf",
+		inContext: ygnmi.NewContext(context.Background(), uexampleocpath.Root().Parent().Child().Config().Five()),
+		wantRequestValues: &ygnmi.RequestValues{
+			CompressedConfigQuery: false,
+			CompressedStateQuery:  false,
+		},
+	}}
+
+	for _, tt := range tests {
+		t.Run(tt.desc, func(t *testing.T) {
+			got := ygnmi.FromContext(tt.inContext)
+			if diff := cmp.Diff(tt.wantRequestValues, got); diff != "" {
+				t.Errorf("(-want, +got):\n%s", diff)
+			}
+		})
+	}
+}

ygnmi/gnmi.go

@@ -71,6 +71,8 @@ func subscribe[T any](ctx context.Context, c *Client, q AnyQuery[T], mode gpb.Su
 		return nil, fmt.Errorf("using gnmi.Get is only valid for ONCE subscriptions")
 	}
 
+	ctx = NewContext(ctx, q)
+
 	var sub gpb.GNMI_SubscribeClient
 	var err error
 	if o.useGet {