feat(sidekick/rust): support grpc_client option (#7492)

For googleapis/google-cloud-rust#5991 .

Author: joshuatants

doc/config-schema.md

@@ -570,6 +570,7 @@ This document describes the schema for the librarian.yaml.
 | `include_bidi_streaming_methods` | bool (optional) | Indicates whether to include gRPC bi-directional streaming methods. |
 | `include_server_streaming_methods` | bool (optional) | Indicates whether to include gRPC server-side streaming methods. |
 | `allow_streaming_any_types` | list of string | Is a list of protobuf field/message IDs with google.protobuf.Any permitted in streaming RPCs (their fields will be dropped in prost conversion). |
+| `grpc_client` | string | Is the Rust type used for the inner gRPC client in generated transports. Defaults to "gaxi::grpc::Client". |
 
 ## RustDocumentationOverride Configuration
 
@@ -583,6 +584,7 @@ This document describes the schema for the librarian.yaml.
 
 | Field | Type | Description |
 | :--- | :--- | :--- |
+| `grpc_client` | string | Is the Rust type used for the inner gRPC client in generated transports. This overrides the crate-level setting. Defaults to "gaxi::grpc::Client". |
 | `disabled_rustdoc_warnings` | yaml.StringSlice | Specifies rustdoc lints to disable. An empty slice explicitly enables all warnings. |
 | `detailed_tracing_attributes` | bool (optional) | Indicates whether to include detailed tracing attributes. This overrides the crate-level setting. |
 | `lro_stub_options` | bool (optional) | Indicates whether to include LRO poller options in generated stub traits. This overrides the crate-level setting. |

internal/config/language.go

@@ -133,12 +133,20 @@ type RustDefault struct {
 	// AllowStreamingAnyTypes is a list of protobuf field/message IDs with google.protobuf.Any
 	// permitted in streaming RPCs (their fields will be dropped in prost conversion).
 	AllowStreamingAnyTypes []string `yaml:"allow_streaming_any_types,omitempty"`
+
+	// GrpcClient is the Rust type used for the inner gRPC client in generated transports.
+	// Defaults to "gaxi::grpc::Client".
+	GrpcClient string `yaml:"grpc_client,omitempty"`
 }
 
 // RustModule defines a generation target within a veneer crate.
 // Each module specifies what proto source to use, which template to apply,
 // and where to output the generated code.
 type RustModule struct {
+	// GrpcClient is the Rust type used for the inner gRPC client in generated transports.
+	// This overrides the crate-level setting. Defaults to "gaxi::grpc::Client".
+	GrpcClient string `yaml:"grpc_client,omitempty"`
+
 	// DisabledRustdocWarnings specifies rustdoc lints to disable. An empty slice explicitly enables all warnings.
 	DisabledRustdocWarnings yaml.StringSlice `yaml:"disabled_rustdoc_warnings,omitempty"`
 

internal/librarian/rust/codec.go

@@ -171,6 +171,9 @@ func buildCodec(library *config.Library, releaseLevel string) map[string]string
 	if rust.QuickstartServiceOverride != "" {
 		codec["quickstart-service-override"] = rust.QuickstartServiceOverride
 	}
+	if rust.GrpcClient != "" {
+		codec["grpc-client"] = rust.GrpcClient
+	}
 	return codec
 }
 
@@ -368,5 +371,15 @@ func buildModuleCodec(library *config.Library, module *config.RustModule) map[st
 	if module.InternalBuilders {
 		codec["internal-builders"] = "true"
 	}
+	grpcClient := ""
+	if library.Rust != nil {
+		grpcClient = library.Rust.GrpcClient
+	}
+	if module.GrpcClient != "" {
+		grpcClient = module.GrpcClient
+	}
+	if grpcClient != "" {
+		codec["grpc-client"] = grpcClient
+	}
 	return codec
 }

internal/librarian/rust/codec_test.go

@@ -1439,6 +1439,7 @@ func TestBuildCodec(t *testing.T) {
 								Ignore:  true,
 							},
 						},
+						GrpcClient: "crate::storage::bidi::GrpcClient",
 					},
 					ModulePath:                "gcs",
 					TemplateOverride:          "custom-template",
@@ -1467,6 +1468,7 @@ func TestBuildCodec(t *testing.T) {
 				"generate-rpc-samples":        "true",
 				"name-overrides":              "foo=bar",
 				"quickstart-service-override": "OverriddenService",
+				"grpc-client":                 "crate::storage::bidi::GrpcClient",
 				"default-features":            "feature1,feature2",
 				"disabled-rustdoc-warnings":   "warning1,warning2",
 				"disabled-clippy-warnings":    "clippy1,clippy2",
@@ -1535,3 +1537,52 @@ func TestModuleToModelConfig_SkippedIds(t *testing.T) {
 		}
 	})
 }
+
+func TestBuildModuleCodec_GrpcClient(t *testing.T) {
+	t.Run("module level grpc_client", func(t *testing.T) {
+		library := &config.Library{
+			Name: "google-cloud-storage",
+		}
+		module := &config.RustModule{
+			GrpcClient: "crate::storage::bidi::GrpcClient",
+		}
+		got := buildModuleCodec(library, module)
+		if got["grpc-client"] != "crate::storage::bidi::GrpcClient" {
+			t.Errorf("expected grpc-client to be %q, got %q", "crate::storage::bidi::GrpcClient", got["grpc-client"])
+		}
+	})
+
+	t.Run("crate level fallback grpc_client", func(t *testing.T) {
+		library := &config.Library{
+			Name: "google-cloud-storage",
+			Rust: &config.RustCrate{
+				RustDefault: config.RustDefault{
+					GrpcClient: "crate::storage::bidi::GrpcClient",
+				},
+			},
+		}
+		module := &config.RustModule{}
+		got := buildModuleCodec(library, module)
+		if got["grpc-client"] != "crate::storage::bidi::GrpcClient" {
+			t.Errorf("expected grpc-client to be %q, got %q", "crate::storage::bidi::GrpcClient", got["grpc-client"])
+		}
+	})
+
+	t.Run("module overrides crate level grpc_client", func(t *testing.T) {
+		library := &config.Library{
+			Name: "google-cloud-storage",
+			Rust: &config.RustCrate{
+				RustDefault: config.RustDefault{
+					GrpcClient: "crate::storage::default::GrpcClient",
+				},
+			},
+		}
+		module := &config.RustModule{
+			GrpcClient: "crate::storage::bidi::GrpcClient",
+		}
+		got := buildModuleCodec(library, module)
+		if got["grpc-client"] != "crate::storage::bidi::GrpcClient" {
+			t.Errorf("expected grpc-client to be %q, got %q", "crate::storage::bidi::GrpcClient", got["grpc-client"])
+		}
+	})
+}

internal/sidekick/rust/annotate_model.go

@@ -97,6 +97,8 @@ type modelAnnotations struct {
 	LroStubOptions bool
 	// If true, the generated builders's visibility should be restricted to the crate.
 	InternalBuilders bool
+	// The Rust type used for the inner gRPC client in generated transports.
+	GrpcClient string
 	// The service to use for the package-level quickstart sample.
 	// Rust generation may decide not to generate some services,
 	// e.g. if the methods have no bindings. On occasion the service
@@ -366,6 +368,7 @@ func annotateModel(model *api.API, codec *codec) (*modelAnnotations, error) {
 		DetailedTracingAttributes: codec.detailedTracingAttributes,
 		LroStubOptions:            codec.lroStubOptions,
 		InternalBuilders:          codec.internalBuilders,
+		GrpcClient:                codec.grpcClient,
 		QuickstartService:         quickstartService,
 	}
 

internal/sidekick/rust/annotate_model_test.go

@@ -188,6 +188,38 @@ func TestInternalBuildersAnnotation(t *testing.T) {
 	}
 }
 
+func TestGrpcClientAnnotation(t *testing.T) {
+	for _, test := range []struct {
+		Options map[string]string
+		Want    string
+	}{
+		{
+			Options: map[string]string{},
+			Want:    "gaxi::grpc::Client",
+		},
+		{
+			Options: map[string]string{
+				"grpc-client": "crate::storage::bidi::GrpcClient",
+			},
+			Want: "crate::storage::bidi::GrpcClient",
+		},
+	} {
+		model := newTestAnnotateModelAPI()
+		codec := newTestCodec(t, libconfig.SpecProtobuf, "", test.Options)
+		got, err := annotateModel(model, codec)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got.GrpcClient != test.Want {
+			t.Errorf("mismatch in GrpcClient, want=%v, got=%v", test.Want, got.GrpcClient)
+		}
+		svcAnn := model.Services[0].Codec.(*serviceAnnotations)
+		if svcAnn.GrpcClient != test.Want {
+			t.Errorf("mismatch in service GrpcClient, want=%v, got=%v", test.Want, svcAnn.GrpcClient)
+		}
+	}
+}
+
 func TestQuickstartServiceAnnotation(t *testing.T) {
 	t.Run("survives filtering", func(t *testing.T) {
 		model := newTestAnnotateModelAPI()
@@ -375,6 +407,7 @@ func TestPackageNames(t *testing.T) {
 		GenerateSetterSamples:     true,
 		GenerateRpcSamples:        true,
 		DetailedTracingAttributes: true,
+		GrpcClient:                "gaxi::grpc::Client",
 	}
 	if diff := cmp.Diff(want, got, cmpopts.IgnoreFields(modelAnnotations{}, "BoilerPlate")); diff != "" {
 		t.Errorf("mismatch (-want +got):\n%s", diff)

internal/sidekick/rust/annotate_service.go

@@ -56,6 +56,8 @@ type serviceAnnotations struct {
 	DetailedTracingAttributes bool
 	// If true, the generated builders's visibility should be restricted to the crate.
 	InternalBuilders bool
+	// The Rust type used for the inner gRPC client in generated transports.
+	GrpcClient string
 }
 
 // HasGrpc returns true if the service has at least one generated gRPC method.
@@ -218,6 +220,7 @@ func (c *codec) annotateService(s *api.Service) (*serviceAnnotations, error) {
 		Incomplete:                slices.ContainsFunc(s.Methods, func(m *api.Method) bool { return !c.generateMethod(m) }),
 		DetailedTracingAttributes: c.detailedTracingAttributes,
 		InternalBuilders:          c.internalBuilders,
+		GrpcClient:                c.grpcClient,
 	}
 	s.Codec = ann
 	return ann, nil

internal/sidekick/rust/annotate_service_test.go

@@ -89,6 +89,7 @@ func TestServiceAnnotationsPerServiceFeatures(t *testing.T) {
 		ModuleName:         "resource_service",
 		PerServiceFeatures: true,
 		Incomplete:         true,
+		GrpcClient:         "gaxi::grpc::Client",
 	}
 	if diff := cmp.Diff(wantService, service.Codec, cmpopts.IgnoreFields(serviceAnnotations{}, "Methods")); diff != "" {
 		t.Errorf("mismatch (-want +got):\n%s", diff)
@@ -245,6 +246,7 @@ func TestServiceAnnotationsLROTypes(t *testing.T) {
 			resource,
 			empty,
 		},
+		GrpcClient: "gaxi::grpc::Client",
 	}
 	if !wantService.HasLROs() {
 		t.Errorf("HasLRO should be true. The service has several LROs.")
@@ -275,6 +277,7 @@ func TestServiceAnnotationsNameOverrides(t *testing.T) {
 		Name:       "Renamed",
 		ModuleName: "renamed",
 		Incomplete: true,
+		GrpcClient: "gaxi::grpc::Client",
 	}
 	if diff := cmp.Diff(wantService, service.Codec, serviceFilter); diff != "" {
 		t.Errorf("mismatch (-want +got):\n%s", diff)
@@ -312,6 +315,7 @@ func TestServiceAnnotations(t *testing.T) {
 		PackageModuleName: "test::v1",
 		ModuleName:        "resource_service",
 		Incomplete:        true,
+		GrpcClient:        "gaxi::grpc::Client",
 	}
 	if diff := cmp.Diff(wantService, service.Codec, cmpopts.IgnoreFields(serviceAnnotations{}, "Methods")); diff != "" {
 		t.Errorf("mismatch (-want +got):\n%s", diff)

internal/sidekick/rust/codec.go

@@ -76,6 +76,7 @@ func newCodec(specificationFormat string, options map[string]string) (*codec, er
 		systemParameters:        sysParams,
 		serializeEnumsAsStrings: specificationFormat != libconfig.SpecProtobuf,
 		bytesUseUrlSafeAlphabet: specificationFormat == libconfig.SpecDiscovery,
+		grpcClient:              "gaxi::grpc::Client",
 	}
 
 	for key, definition := range options {
@@ -212,6 +213,8 @@ func newCodec(specificationFormat string, options map[string]string) (*codec, er
 				return nil, fmt.Errorf("cannot convert `include-rpc-status-conversion` value %q to boolean: %w", definition, err)
 			}
 			codec.includeRpcStatusConversion = value
+		case key == "grpc-client":
+			codec.grpcClient = definition
 		default:
 			return nil, fmt.Errorf("unknown Rust codec option %q", key)
 		}
@@ -372,6 +375,8 @@ type codec struct {
 	internalBuilders bool
 	// Overrides the default heuristically selected service for the package-level quickstart.
 	quickstartServiceOverride string
+	// The Rust type used for the inner gRPC client in generated transports.
+	grpcClient string
 }
 
 type systemParameter struct {

internal/sidekick/rust/codec_test.go

@@ -369,6 +369,15 @@ func TestParseOptions(t *testing.T) {
 				c.quickstartServiceOverride = "OverriddenService"
 			},
 		},
+		{
+			Format: libconfig.SpecProtobuf,
+			Options: map[string]string{
+				"grpc-client": "custom::GrpcClient",
+			},
+			Update: func(c *codec) {
+				c.grpcClient = "custom::GrpcClient"
+			},
+		},
 	} {
 		t.Run(fmt.Sprintf("case_%d", i), func(t *testing.T) {
 			want, err := newCodec(test.Format, map[string]string{})