From a39e33bc78ff3aad0d1f335870281ccb7d40420d Mon Sep 17 00:00:00 2001
From: Carlos O'Ryan <coryan@google.com>
Date: Fri, 24 Jan 2025 10:23:43 -0500
Subject: [PATCH] Adddress review comments

---
 generator/internal/parser/mixin.go    | 34 +++++++--------------
 generator/internal/parser/protobuf.go | 43 +++++++++++++--------------
 2 files changed, 31 insertions(+), 46 deletions(-)

diff --git a/generator/internal/parser/mixin.go b/generator/internal/parser/mixin.go
index 3fe75e59b..410c8a351 100644
--- a/generator/internal/parser/mixin.go
+++ b/generator/internal/parser/mixin.go
@@ -40,15 +40,21 @@ const (
 type mixinMethods map[string]bool
 
 // loadMixins loads file descriptors for configured mixins.
-func loadMixins(serviceConfig *serviceconfig.Service) (mixinMethods, []*descriptorpb.FileDescriptorProto) {
+func loadMixins(serviceConfig *serviceconfig.Service, withLongrunning bool) (mixinMethods, []*descriptorpb.FileDescriptorProto) {
 	var files []*descriptorpb.FileDescriptorProto
 	var enabledMixinMethods mixinMethods = make(map[string]bool)
-	apis := serviceConfig.GetApis()
-	if len(apis) < 2 {
+	var apiNames []string
+	for _, api := range serviceConfig.GetApis() {
+		apiNames = append(apiNames, api.GetName())
+	}
+	if withLongrunning {
+		apiNames = append(apiNames, longrunningService)
+	}
+	if len(apiNames) < 2 {
 		return enabledMixinMethods, files
 	}
-	for _, api := range apis {
-		switch api.GetName() {
+	for _, apiName := range apiNames {
+		switch apiName {
 		case locationService:
 			files = append(files, protodesc.ToFileDescriptorProto(location.File_google_cloud_location_locations_proto))
 		case iamService:
@@ -76,24 +82,6 @@ func loadMixinMethods(serviceConfig *serviceconfig.Service) mixinMethods {
 	return enabledMixinMethods
 }
 
-// appendLongrunningMixin appends the longrunning mixin. Many services have
-// LROs, which need the mixin to be usable, but do not include the mixin in the
-// service config YAML file.
-func appendLongrunningMixin(methods mixinMethods, files []*descriptorpb.FileDescriptorProto) (mixinMethods, []*descriptorpb.FileDescriptorProto) {
-	desc := protodesc.ToFileDescriptorProto(longrunningpb.File_google_longrunning_operations_proto)
-	for _, f := range files {
-		if f.GetName() == desc.GetName() {
-			return methods, files
-		}
-	}
-	files = append(files, desc)
-	methods[".google.protobuf.Operations.ListOperations"] = true
-	methods[".google.protobuf.Operations.GetOperation"] = true
-	methods[".google.protobuf.Operations.DeleteOperation"] = true
-	methods[".google.protobuf.Operations.CancelOperation"] = true
-	return methods, files
-}
-
 // Apply `serviceConfig` overrides to `targetMethod`.
 //
 // The service config file may include overrides to mixin method definitions.
diff --git a/generator/internal/parser/protobuf.go b/generator/internal/parser/protobuf.go
index c18711199..1ed41c393 100644
--- a/generator/internal/parser/protobuf.go
+++ b/generator/internal/parser/protobuf.go
@@ -227,7 +227,8 @@ func makeAPIForProtobuf(serviceConfig *serviceconfig.Service, req *pluginpb.Code
 		if serviceConfig.Documentation != nil {
 			result.Description = serviceConfig.Documentation.Summary
 		}
-		enabledMixinMethods, mixinFileDesc = loadMixins(serviceConfig)
+		withLongrunning := requiresLongrunningMixin(req)
+		enabledMixinMethods, mixinFileDesc = loadMixins(serviceConfig, withLongrunning)
 		packageName := ""
 		for _, api := range serviceConfig.Apis {
 			packageName, _ = splitApiName(api.Name)
@@ -324,28 +325,6 @@ func makeAPIForProtobuf(serviceConfig *serviceconfig.Service, req *pluginpb.Code
 		result.Services = append(result.Services, fileServices...)
 	}
 
-	needsLongrunningMixin := func() bool {
-		for _, service := range result.Services {
-			for _, method := range service.Methods {
-				if method.OperationInfo == nil {
-					continue
-				}
-				// This service has an LRO, it needs to implement `GetOperation`.
-				target := service.ID + ".GetOperation"
-				if _, ok := result.State.MethodByID[target]; ok {
-					// The method already exists, nothing to do.
-					continue
-				}
-				return true
-			}
-		}
-		return false
-	}()
-
-	if needsLongrunningMixin {
-		enabledMixinMethods, mixinFileDesc = appendLongrunningMixin(enabledMixinMethods, mixinFileDesc)
-	}
-
 	// Add the mixing methods to the existing services.
 	for _, service := range result.Services {
 		for _, f := range mixinFileDesc {
@@ -377,6 +356,24 @@ func makeAPIForProtobuf(serviceConfig *serviceconfig.Service, req *pluginpb.Code
 	return result
 }
 
+// requiresLongrunningMixin finds out if any method returns a LRO. This is used
+// to forcibly load the longrunning mixin. It needs to happen before the proto
+// descriptors are converted to the `api.*`, as that conversion requires the
+// mixin.
+func requiresLongrunningMixin(req *pluginpb.CodeGeneratorRequest) bool {
+	for _, f := range req.GetSourceFileDescriptors() {
+		for _, s := range f.Service {
+			for _, m := range s.Method {
+				info := parseOperationInfo(f.GetPackage(), m)
+				if info != nil && m.GetOutputType() == ".google.protobuf.Operation" {
+					return true
+				}
+			}
+		}
+	}
+	return false
+}
+
 var descriptorpbToTypez = map[descriptorpb.FieldDescriptorProto_Type]api.Typez{
 	descriptorpb.FieldDescriptorProto_TYPE_DOUBLE:   api.DOUBLE_TYPE,
 	descriptorpb.FieldDescriptorProto_TYPE_FLOAT:    api.FLOAT_TYPE,