fix: Allow custom HttpRules for REST LROs

gapic-generator-java/src/main/java/com/google/api/generator/gapic/composer/common/AbstractTransportServiceStubClassComposer.java

@@ -188,6 +188,7 @@ public GapicClass generate(GapicContext context, Service service) {
     Map<String, Message> messageTypes = context.messages();
     List<Statement> classStatements =
         createClassStatements(
+            context,
             service,
             protoMethodNameToDescriptorVarExprs,
             callableClassMemberVarExprs,
@@ -278,6 +279,7 @@ protected List<Statement> createTypeRegistry(Service service) {
   }
 
   protected List<Statement> createClassStatements(
+      GapicContext context,
       Service service,
       Map<String, VariableExpr> protoMethodNameToDescriptorVarExprs,
       Map<String, VariableExpr> callableClassMemberVarExprs,
@@ -649,6 +651,7 @@ protected List<MethodDefinition> createConstructorMethods(
     if (generateOperationsStubLogic(service)) {
       secondCtorExprs.addAll(
           createOperationsStubInitExpr(
+              context,
               service,
               thisExpr,
               operationsStubClassVarExpr,
@@ -758,6 +761,7 @@ protected List<MethodDefinition> createConstructorMethods(
   }
 
   protected List<Expr> createOperationsStubInitExpr(
+      GapicContext context,
       Service service,
       Expr thisExpr,
       VariableExpr operationsStubClassVarExpr,

gapic-generator-java/src/main/java/com/google/api/generator/gapic/composer/rest/HttpJsonServiceStubClassComposer.java

@@ -14,6 +14,7 @@
 
 package com.google.api.generator.gapic.composer.rest;
 
+import com.google.api.HttpRule;
 import com.google.api.core.InternalApi;
 import com.google.api.gax.httpjson.ApiMethodDescriptor;
 import com.google.api.gax.httpjson.ApiMethodDescriptor.MethodType;
@@ -63,6 +64,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.BiMap;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.protobuf.TypeRegistry;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -73,6 +75,7 @@
 import java.util.Set;
 import java.util.function.BiFunction;
 import java.util.function.Function;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
 public class HttpJsonServiceStubClassComposer extends AbstractTransportServiceStubClassComposer {
@@ -89,6 +92,7 @@ public class HttpJsonServiceStubClassComposer extends AbstractTransportServiceSt
                   .setType(FIXED_REST_TYPESTORE.get(TypeRegistry.class.getSimpleName()))
                   .build())
           .build();
+  private static final String LRO_NAME_PREFIX = "google.longrunning.Operations";
 
   protected HttpJsonServiceStubClassComposer() {
     super(RestContext.instance());
@@ -110,6 +114,7 @@ private static TypeStore createStaticTypes() {
             HttpJsonOperationSnapshot.class,
             HttpJsonStubCallableFactory.class,
             Map.class,
+            ImmutableMap.class,
             ProtoMessageRequestFormatter.class,
             ProtoMessageResponseParser.class,
             ProtoRestSerializer.class,
@@ -1075,6 +1080,7 @@ private List<Expr> getMethodTypeExpr(Method protoMethod) {
 
   @Override
   protected List<Expr> createOperationsStubInitExpr(
+      GapicContext context,
       Service service,
       Expr thisExpr,
       VariableExpr operationsStubClassVarExpr,
@@ -1088,6 +1094,40 @@ protected List<Expr> createOperationsStubInitExpr(
     if (standardOpStub.equals(operationsStubType.reference().fullName())) {
       arguments.add(TYPE_REGISTRY_VAR_EXPR);
     }
+    Map<String, String> operationCustomHttpBindings =
+        filterCustomHttpBindingsMap(
+            parseCustomHttpBindings(context), x -> x.getKey().contains(LRO_NAME_PREFIX));
+    if (operationCustomHttpBindings.size() > 0) {
+      Expr operationCustomHttpBindingsBuilderExpr =
+          MethodInvocationExpr.builder()
+              .setStaticReferenceType(FIXED_REST_TYPESTORE.get(ImmutableMap.class.getSimpleName()))
+              .setMethodName("builder")
+              .setGenerics(Arrays.asList(TypeNode.STRING.reference(), TypeNode.STRING.reference()))
+              .build();
+
+      for (Map.Entry<String, String> entrySet : operationCustomHttpBindings.entrySet()) {
+        String selector = entrySet.getKey();
+        String path = entrySet.getValue();
+        operationCustomHttpBindingsBuilderExpr =
+            MethodInvocationExpr.builder()
+                .setExprReferenceExpr(operationCustomHttpBindingsBuilderExpr)
+                .setMethodName("put")
+                .setArguments(
+                    Arrays.asList(
+                        ValueExpr.withValue(StringObjectValue.withValue(selector)),
+                        ValueExpr.withValue(StringObjectValue.withValue(path))))
+                .build();
+      }
+
+      operationCustomHttpBindingsBuilderExpr =
+          MethodInvocationExpr.builder()
+              .setExprReferenceExpr(operationCustomHttpBindingsBuilderExpr)
+              .setMethodName("build")
+              .setReturnType(FIXED_REST_TYPESTORE.get(ImmutableMap.class.getSimpleName()))
+              .build();
+
+      arguments.add(operationCustomHttpBindingsBuilderExpr);
+    }
 
     return Collections.singletonList(
         AssignmentExpr.builder()
@@ -1103,6 +1143,43 @@ protected List<Expr> createOperationsStubInitExpr(
             .build());
   }
 
+  private Map<String, String> parseCustomHttpBindings(GapicContext context) {
+    Map<String, String> customHttpBindings = new HashMap<>();
+    com.google.api.Service service = context.serviceYamlProto();
+    if (service != null && service.getHttp() != null) {
+      for (HttpRule httpRule : service.getHttp().getRulesList()) {
+        customHttpBindings.put(httpRule.getSelector(), getValueBasedOnPatternCase(httpRule));
+      }
+    }
+    return customHttpBindings;
+  }
+
+  private Map<String, String> filterCustomHttpBindingsMap(
+      Map<String, String> customHttpBindings, Predicate<Map.Entry<String, String>> predicate) {
+    return customHttpBindings.entrySet().stream()
+        .filter(predicate)
+        .collect(Collectors.toMap(x -> x.getKey(), x -> x.getValue()));
+  }
+
+  private String getValueBasedOnPatternCase(HttpRule httpRule) {
+    switch (httpRule.getPatternCase().getNumber()) {
+      case 2:
+        return httpRule.getGet();
+      case 3:
+        return httpRule.getPut();
+      case 4:
+        return httpRule.getPost();
+      case 5:
+        return httpRule.getDelete();
+      case 6:
+        return httpRule.getPatch();
+      case 8:
+        return httpRule.getCustom().getPath();
+      default:
+        return null;
+    }
+  }
+
   @Override
   protected List<Statement> createLongRunningClient(Service service, TypeStore typeStore) {
     Method pollingMethod = service.operationPollingMethod();

gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/ProtoMessageRequestFormatter.java

@@ -147,6 +147,12 @@ public Builder<RequestT> setAdditionalPaths(String... rawAdditionalPaths) {
       return this;
     }
 
+    @InternalApi
+    public Builder<RequestT> updateRawPath(String rawPath) {
+      this.rawPath = rawPath;
+      return this;
+    }
+
     @InternalApi
     public Builder<RequestT> updateRawPath(String target, String replacement) {
       this.rawPath = this.rawPath.replace(target, replacement);