From 33b8ab7d0e0937177b3b6c826bcb4c1e0da8dec9 Mon Sep 17 00:00:00 2001
From: Carlos O'Ryan <coryan@google.com>
Date: Fri, 24 Jan 2025 16:11:00 +0000
Subject: [PATCH] feat(generator): automatic `longrunning` mixin (#812)

Some services need the `longrunning` mixin, but their service config YAML file
does not provide it. Automatically add it when needed.
---
 generator/internal/parser/mixin.go    | 22 +++++++++++++++++-----
 generator/internal/parser/protobuf.go | 21 ++++++++++++++++++++-
 2 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/generator/internal/parser/mixin.go b/generator/internal/parser/mixin.go
index c03b411f9..eb588abe3 100644
--- a/generator/internal/parser/mixin.go
+++ b/generator/internal/parser/mixin.go
@@ -40,15 +40,27 @@ 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
+	hasLongrunning := false
+	for _, api := range serviceConfig.GetApis() {
+		// Only insert the service if needed. We want to preserve the order
+		// to make the generated code reproducible, so we cannot use a map.
+		if api.GetName() == longrunningService {
+			hasLongrunning = true
+		}
+		apiNames = append(apiNames, api.GetName())
+	}
+	if withLongrunning && !hasLongrunning {
+		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:
diff --git a/generator/internal/parser/protobuf.go b/generator/internal/parser/protobuf.go
index a60be717a..96b13bdaa 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)
@@ -355,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.longrunning.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,