googleapis · summer-ji-eng · Dec 10, 2020 · Dec 3, 2020 · Dec 9, 2020 · Nov 25, 2020
diff --git a/src/main/java/com/google/api/generator/gapic/composer/ServiceClientClassComposer.java b/src/main/java/com/google/api/generator/gapic/composer/ServiceClientClassComposer.java
@@ -566,13 +566,13 @@ private static List<MethodDefinition> createMethodVariants(
               .build();
 
       Optional<String> methodSampleCode = Optional.empty();
-      if (!method.isPaged() && !method.hasLro()) {
-        // TODO(summerji): Remove the condition check once finished the implementation on paged
-        // sample code and lro sample code.
+      if (!method.hasLro()) {
+        // TODO(summerji): Remove the condition check once finished the implementation on lro sample
+        // code.
         methodSampleCode =
             Optional.of(
                 ServiceClientSampleCodeComposer.composeRpcMethodHeaderSampleCode(
-                    method, types.get(clientName), signature, resourceNames));
+                    method, types.get(clientName), signature, resourceNames, messageTypes));
       }
       MethodDefinition.Builder methodVariantBuilder =
           MethodDefinition.builder()

diff --git a/src/main/java/com/google/api/generator/gapic/composer/ServiceClientSampleCodeComposer.java b/src/main/java/com/google/api/generator/gapic/composer/ServiceClientSampleCodeComposer.java
@@ -16,26 +16,32 @@
 
 import com.google.api.gax.core.FixedCredentialsProvider;
 import com.google.api.generator.engine.ast.AssignmentExpr;
+import com.google.api.generator.engine.ast.CommentStatement;
 import com.google.api.generator.engine.ast.ConcreteReference;
 import com.google.api.generator.engine.ast.Expr;
 import com.google.api.generator.engine.ast.ExprStatement;
+import com.google.api.generator.engine.ast.ForStatement;
+import com.google.api.generator.engine.ast.LineComment;
 import com.google.api.generator.engine.ast.MethodInvocationExpr;
+import com.google.api.generator.engine.ast.Statement;
 import com.google.api.generator.engine.ast.TryCatchStatement;
 import com.google.api.generator.engine.ast.TypeNode;
 import com.google.api.generator.engine.ast.VaporReference;
 import com.google.api.generator.engine.ast.Variable;
 import com.google.api.generator.engine.ast.VariableExpr;
 import com.google.api.generator.gapic.composer.samplecode.SampleCodeWriter;
+import com.google.api.generator.gapic.model.Field;
+import com.google.api.generator.gapic.model.Message;
 import com.google.api.generator.gapic.model.Method;
 import com.google.api.generator.gapic.model.MethodArgument;
 import com.google.api.generator.gapic.model.ResourceName;
 import com.google.api.generator.gapic.utils.JavaStyle;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
@@ -186,8 +192,24 @@ public static String composeRpcMethodHeaderSampleCode(
       Method method,
       TypeNode clientType,
       List<MethodArgument> arguments,
-      Map<String, ResourceName> resourceNames) {
+      Map<String, ResourceName> resourceNames,
+      Map<String, Message> messageTypes) {
     // TODO(summerji): Add other types RPC methods' sample code.
+    if (method.isPaged()) {
+      // Find the repeated field.
+      Message methodOutputMessage = messageTypes.get(method.outputType().reference().simpleName());
+      Field repeatedPagedResultsField = methodOutputMessage.findAndUnwrapFirstRepeatedField();
+      Preconditions.checkNotNull(
+          repeatedPagedResultsField,
+          String.format(
+              "No repeated field found on message %s for method %s",
+              methodOutputMessage.name(), method.name()));
+
+      TypeNode repeatedResponseType = repeatedPagedResultsField.type();
+      return SampleCodeWriter.write(
+          composeUnaryPagedRpcMethodSampleCode(
+              method, clientType, repeatedResponseType, arguments, resourceNames));
+    }
     return SampleCodeWriter.write(
         composeUnaryRpcMethodSampleCode(method, clientType, arguments, resourceNames));
   }
@@ -204,52 +226,12 @@ static TryCatchStatement composeUnaryRpcMethodSampleCode(
                 .setName(JavaStyle.toLowerCamelCase(clientType.reference().name()))
                 .setType(clientType)
                 .build());
-    // List of rpc method arguments' variable expressions.
-    List<VariableExpr> rpcMethodArgVarExprs =
-        arguments.stream()
-            .map(
-                arg ->
-                    VariableExpr.withVariable(
-                        Variable.builder()
-                            .setName(JavaStyle.toLowerCamelCase(arg.name()))
-                            .setType(arg.type())
-                            .build()))
-            .collect(Collectors.toList());
-    // List of rpc method arguments' default value expression.
-    List<ResourceName> resourceNameList =
-        resourceNames.values().stream().collect(Collectors.toList());
+    List<VariableExpr> rpcMethodArgVarExprs = createRpcMethodArgumentVariableExprs(arguments);
     List<Expr> rpcMethodArgDefaultValueExprs =
-        arguments.stream()
-            .map(
-                arg ->
-                    !isStringTypedResourceName(arg, resourceNames)
-                        ? DefaultValueComposer.createDefaultValue(arg, resourceNames)
-                        : MethodInvocationExpr.builder()
-                            .setExprReferenceExpr(
-                                DefaultValueComposer.createDefaultValue(
-                                    resourceNames.get(
-                                        arg.field().resourceReference().resourceTypeString()),
-                                    resourceNameList,
-                                    arg.field().name()))
-                            .setMethodName("toString")
-                            .setReturnType(TypeNode.STRING)
-                            .build())
-            .collect(Collectors.toList());
-
-    List<Expr> bodyExprs = new ArrayList<>();
-    Preconditions.checkState(
-        rpcMethodArgVarExprs.size() == rpcMethodArgDefaultValueExprs.size(),
-        "Expected the number of method arguments to match the number of default values.");
-    bodyExprs.addAll(
-        IntStream.range(0, rpcMethodArgVarExprs.size())
-            .mapToObj(
-                i ->
-                    AssignmentExpr.builder()
-                        .setVariableExpr(
-                            (rpcMethodArgVarExprs.get(i)).toBuilder().setIsDecl(true).build())
-                        .setValueExpr(rpcMethodArgDefaultValueExprs.get(i))
-                        .build())
-            .collect(Collectors.toList()));
+        createRpcMethodArgumentDefaultValueExprs(arguments, resourceNames);
+    List<Expr> bodyExprs =
+        createAssignmentsForVarExprsWithValueExprs(
+            rpcMethodArgVarExprs, rpcMethodArgDefaultValueExprs);
     // Invoke current method based on return type.
     // e.g. if return void, echoClient.echo(..); or,
     // e.g. if return other type, EchoResponse response = echoClient.echo(...);
@@ -290,8 +272,125 @@ static TryCatchStatement composeUnaryRpcMethodSampleCode(
         .build();
   }
 
+  @VisibleForTesting
+  static TryCatchStatement composeUnaryPagedRpcMethodSampleCode(
+      Method method,
+      TypeNode clientType,
+      TypeNode repeatedResponseType,
+      List<MethodArgument> arguments,
+      Map<String, ResourceName> resourceNames) {
+    VariableExpr clientVarExpr =
+        VariableExpr.withVariable(
+            Variable.builder()
+                .setName(JavaStyle.toLowerCamelCase(clientType.reference().name()))
+                .setType(clientType)
+                .build());
+    List<VariableExpr> rpcMethodArgVarExprs = createRpcMethodArgumentVariableExprs(arguments);
+    List<Expr> rpcMethodArgDefaultValueExprs =
+        createRpcMethodArgumentDefaultValueExprs(arguments, resourceNames);
+    List<Expr> bodyExprs =
+        createAssignmentsForVarExprsWithValueExprs(
+            rpcMethodArgVarExprs, rpcMethodArgDefaultValueExprs);
+    // For loop paged response item on iterateAll method.
+    // e.g. for (LogEntry element : loggingServiceV2Client.ListLogs(parent).iterateAll()) {
+    //          //doThingsWith(element);
+    //      }
+    MethodInvocationExpr clientMethodExpr =
+        MethodInvocationExpr.builder()
+            .setExprReferenceExpr(clientVarExpr)
+            .setMethodName(JavaStyle.toLowerCamelCase(method.name()))
+            .setArguments(
+                rpcMethodArgVarExprs.stream().map(e -> (Expr) e).collect(Collectors.toList()))
+            .build();
+    Expr clientMethodIteratorAllExpr =
+        MethodInvocationExpr.builder()
+            .setExprReferenceExpr(clientMethodExpr)
+            .setMethodName("iterateAll")
+            .setReturnType(repeatedResponseType)
+            .build();
+    ForStatement loopIteratorStatement =
+        ForStatement.builder()
+            .setLocalVariableExpr(
+                VariableExpr.builder()
+                    .setVariable(
+                        Variable.builder().setName("element").setType(repeatedResponseType).build())
+                    .setIsDecl(true)
+                    .build())
+            .setCollectionExpr(clientMethodIteratorAllExpr)
+            .setBody(
+                Arrays.asList(
+                    CommentStatement.withComment(
+                        LineComment.withComment("doThingsWith(element);"))))
+            .build();
+
+    List<Statement> bodyStatements =
+        bodyExprs.stream().map(e -> ExprStatement.withExpr(e)).collect(Collectors.toList());
+    bodyStatements.add(loopIteratorStatement);
+
+    return TryCatchStatement.builder()
+        .setTryResourceExpr(assignClientVariableWithCreateMethodExpr(clientVarExpr))
+        .setTryBody(bodyStatements)
+        .setIsSampleCode(true)
+        .build();
+  }
+
   // ==================================Helpers===================================================//
 
+  // Create a list of RPC method arguments' variable expressions.
+  private static List<VariableExpr> createRpcMethodArgumentVariableExprs(
+      List<MethodArgument> arguments) {
+    return arguments.stream()
+        .map(
+            arg ->
+                VariableExpr.withVariable(
+                    Variable.builder()
+                        .setName(JavaStyle.toLowerCamelCase(arg.name()))
+                        .setType(arg.type())
+                        .build()))
+        .collect(Collectors.toList());
+  }
+
+  // Create a list of RPC method arguments' default value expression.
+  private static List<Expr> createRpcMethodArgumentDefaultValueExprs(
+      List<MethodArgument> arguments, Map<String, ResourceName> resourceNames) {
+    List<ResourceName> resourceNameList =
+        resourceNames.values().stream().collect(Collectors.toList());
+    Function<MethodArgument, MethodInvocationExpr> stringResourceNameDefaultValueExpr =
+        arg ->
+            MethodInvocationExpr.builder()
+                .setExprReferenceExpr(
+                    DefaultValueComposer.createDefaultValue(
+                        resourceNames.get(arg.field().resourceReference().resourceTypeString()),
+                        resourceNameList,
+                        arg.field().name()))
+                .setMethodName("toString")
+                .setReturnType(TypeNode.STRING)
+                .build();
+    return arguments.stream()
+        .map(
+            arg ->
+                !isStringTypedResourceName(arg, resourceNames)
+                    ? DefaultValueComposer.createDefaultValue(arg, resourceNames)
+                    : stringResourceNameDefaultValueExpr.apply(arg))
+        .collect(Collectors.toList());
+  }
+
+  // Create a list of assignment expressions for variable expr with value expr.
+  private static List<Expr> createAssignmentsForVarExprsWithValueExprs(
+      List<VariableExpr> variableExprs, List<Expr> valueExprs) {
+    Preconditions.checkState(
+        variableExprs.size() == valueExprs.size(),
+        "Expected the number of method arguments to match the number of default values.");
+    return IntStream.range(0, variableExprs.size())
+        .mapToObj(
+            i ->
+                AssignmentExpr.builder()
+                    .setVariableExpr(variableExprs.get(i).toBuilder().setIsDecl(true).build())
+                    .setValueExpr(valueExprs.get(i))
+                    .build())
+        .collect(Collectors.toList());
+  }
+
   // Assign client variable expr with create client.
   // e.g EchoClient echoClient = EchoClient.create()
   private static AssignmentExpr assignClientVariableWithCreateMethodExpr(