Add an ability to pick a container for streaming a method

SuduIDE · Sep 5, 2023 · 17d96f4 · 17d96f4
1 parent 52a482c
commit 17d96f4
Show file tree

Hide file tree

Showing 11 changed files with 61 additions and 54 deletions.
diff --git a/DOCS.md b/DOCS.md
@@ -244,4 +244,4 @@ option (protogen.gen_method) = false;
   parameters likewise it isn't a domain object.
 * `(protogen.make_nullable)` annotates method with `@Nullable` and wraps response into `nullifyIfNotFound` method
   of `BaseGrpcClient`
-* `(protogen.stream_to_list)` if the output is streaming collects it into java.util.List
+* `(protogen.stream_to_container)` if the output is streaming collects it into a specified container
diff --git a/generator/src/main/java/org/sudu/protogen/generator/client/MethodGeneratorBase.java b/generator/src/main/java/org/sudu/protogen/generator/client/MethodGeneratorBase.java
@@ -68,7 +68,7 @@ protected TypeModel getReturnType() {
         } else if (method.getOutputType().isDomain()) {
             type = context.typeProcessor().processType(method.getOutputType(), context);
         } else {
-            if (method.getOutputType().getFields().size() == 0) {
+            if (method.getOutputType().getFields().isEmpty()) {
                 type = new TypeModel(TypeName.VOID);
             } else {
                 throw new ProtogenException(("Unable to create a method returning %s because request consist of more than " +
@@ -79,8 +79,9 @@ protected TypeModel getReturnType() {
             if (type.getTypeName().toString().equalsIgnoreCase("void")) {
                 return type;
             }
-            if (method.isStreamingToList()) {
-                return new RepeatedType(type, Field.RepeatedContainer.LIST);
+            var containerO = method.getStreamToContainer();
+            if (containerO.isPresent()) {
+                return new RepeatedType(type, containerO.get());
             } else {
                 return new IteratorType(type);
             }

diff --git a/...or/src/main/java/org/sudu/protogen/generator/client/StreamingGrpcCallMethodGenerator.java b/...or/src/main/java/org/sudu/protogen/generator/client/StreamingGrpcCallMethodGenerator.java
@@ -7,10 +7,10 @@
 import org.sudu.protogen.generator.type.RepeatedType;
 import org.sudu.protogen.generator.type.TypeModel;
 import org.sudu.protogen.protobuf.Method;
+import org.sudu.protogen.protobuf.RepeatedContainer;
 import org.sudu.protogen.utils.Name;
 
 import javax.lang.model.element.Modifier;
-import java.lang.reflect.ParameterizedType;
 import java.util.Iterator;
 import java.util.List;
 import java.util.stream.StreamSupport;
@@ -36,7 +36,7 @@ protected CodeBlock body(TypeModel returnType, List<ParameterSpec> params) {
             return CodeBlock.of("");
         }
         CodeBlock call = commonGen.onlyCallReturnExpression(returnType, params);
-        if (method.isStreamingToList()) {
+        if (method.getStreamToContainer().isPresent()) {
             Validate.validState(returnType instanceof RepeatedType);
             RepeatedType repType = (RepeatedType) returnType;
             return CodeBlock.builder()
@@ -55,8 +55,11 @@ protected CodeBlock returnExpression(TypeModel returnType, List<ParameterSpec> p
         if (returnType.getTypeName().toString().equalsIgnoreCase("void")) {
             return commonGen.onlyCallReturnExpression(returnType, params);
         }
-        if (method.isStreamingToList()) {
-            return CodeBlock.of("$T.stream(iterable.spliterator(), false).toList()", StreamSupport.class);
+        var containerO = method.getStreamToContainer();
+        if (containerO.isPresent()) {
+            RepeatedContainer container = containerO.get();
+            return CodeBlock.of("$T.stream(iterable.spliterator(), false).collect($L)",
+                    StreamSupport.class, container.getCollectorExpr());
         } else {
             Validate.validState(returnType instanceof IteratorType);
             IteratorType itType = (IteratorType) returnType;

diff --git a/generator/src/main/java/org/sudu/protogen/generator/type/RepeatedType.java b/generator/src/main/java/org/sudu/protogen/generator/type/RepeatedType.java
@@ -3,16 +3,16 @@
 import com.squareup.javapoet.CodeBlock;
 import com.squareup.javapoet.ParameterizedTypeName;
 import org.jetbrains.annotations.NotNull;
-import org.sudu.protogen.protobuf.Field;
+import org.sudu.protogen.protobuf.RepeatedContainer;
 import org.sudu.protogen.utils.Name;
 
 public class RepeatedType extends TypeModel {
 
     private final TypeModel elementModel;
 
-    private final Field.RepeatedContainer repeatedType;
+    private final RepeatedContainer repeatedType;
 
-    public RepeatedType(TypeModel typeName, Field.RepeatedContainer repeatedType) {
+    public RepeatedType(TypeModel typeName, RepeatedContainer repeatedType) {
         super(ParameterizedTypeName.get(
                 repeatedType.getTypeName(),
                 typeName.getTypeName().box()
@@ -39,7 +39,7 @@ public CodeBlock toGrpcTransformer(CodeBlock expr) {
 
     @Override
     public CodeBlock fromGrpcTransformer(CodeBlock expr) {
-        if (elementModel instanceof DomainType || elementModel instanceof UnfoldedType || repeatedType != Field.RepeatedContainer.LIST) {
+        if (elementModel instanceof DomainType || elementModel instanceof UnfoldedType || repeatedType != RepeatedContainer.LIST) {
             CodeBlock lambdaParameter = CodeBlock.builder().add("i").build();
             CodeBlock mappingLambda = CodeBlock.builder()
                     .add("$L -> $L", lambdaParameter, elementModel.fromGrpcTransformer(lambdaParameter))

diff --git a/generator/src/main/java/org/sudu/protogen/protobuf/Field.java b/generator/src/main/java/org/sudu/protogen/protobuf/Field.java
@@ -1,13 +1,9 @@
 package org.sudu.protogen.protobuf;
 
-import com.squareup.javapoet.ClassName;
-import com.squareup.javapoet.CodeBlock;
 import org.apache.commons.lang3.Validate;
 import org.jetbrains.annotations.Nullable;
-import protogen.Options;
 
 import java.util.Optional;
-import java.util.stream.Collectors;
 
 public abstract class Field {
 
@@ -93,32 +89,4 @@ public enum Type {
         MESSAGE
     }
 
-    public enum RepeatedContainer {
-        LIST(ClassName.get("java.util", "List"), CodeBlock.of("$T.toList()", Collectors.class)),
-        SET(ClassName.get("java.util", "Set"), CodeBlock.of("$T.toSet()", Collectors.class));
-
-        private final ClassName typeName;
-        private final CodeBlock collectorExpr;
-
-        RepeatedContainer(ClassName typeName, CodeBlock collectorExpr) {
-            this.typeName = typeName;
-            this.collectorExpr = collectorExpr;
-        }
-
-        public static RepeatedContainer fromGrpc(Options.RepeatedContainer proto) {
-            return switch (proto) {
-                case UNRECOGNIZED -> throw new IllegalArgumentException();
-                case LIST -> RepeatedContainer.LIST;
-                case SET -> RepeatedContainer.SET;
-            };
-        }
-
-        public ClassName getTypeName() {
-            return typeName;
-        }
-
-        public CodeBlock getCollectorExpr() {
-            return collectorExpr;
-        }
-    }
 }
diff --git a/generator/src/main/java/org/sudu/protogen/protobuf/Method.java b/generator/src/main/java/org/sudu/protogen/protobuf/Method.java
@@ -24,10 +24,6 @@ public final String generatedName() {
         return getNameOption().orElseGet(this::getName);
     }
 
-    public final boolean isStreamingToList() {
-        return getStreamToListOption().orElse(false);
-    }
-
     public final boolean doGenerate() {
         if (isInputStreaming()) return false;
         return getGenerateOption()
@@ -54,7 +50,7 @@ public final Field unfoldedResponseField() {
 
     protected abstract Optional<Boolean> getNullableOption();
 
-    protected abstract Optional<Boolean> getStreamToListOption();
+    public abstract Optional<RepeatedContainer> getStreamToContainer();
 
     protected abstract Optional<String> getNameOption();
 }
diff --git a/generator/src/main/java/org/sudu/protogen/protobuf/RepeatedContainer.java b/generator/src/main/java/org/sudu/protogen/protobuf/RepeatedContainer.java
@@ -0,0 +1,36 @@
+package org.sudu.protogen.protobuf;
+
+import com.squareup.javapoet.ClassName;
+import com.squareup.javapoet.CodeBlock;
+import protogen.Options;
+
+import java.util.stream.Collectors;
+
+public enum RepeatedContainer {
+    LIST(ClassName.get("java.util", "List"), CodeBlock.of("$T.toList()", Collectors.class)),
+    SET(ClassName.get("java.util", "Set"), CodeBlock.of("$T.toSet()", Collectors.class));
+
+    private final ClassName typeName;
+    private final CodeBlock collectorExpr;
+
+    RepeatedContainer(ClassName typeName, CodeBlock collectorExpr) {
+        this.typeName = typeName;
+        this.collectorExpr = collectorExpr;
+    }
+
+    public static RepeatedContainer fromGrpc(Options.RepeatedContainer proto) {
+        return switch (proto) {
+            case UNRECOGNIZED -> throw new IllegalArgumentException();
+            case LIST -> RepeatedContainer.LIST;
+            case SET -> RepeatedContainer.SET;
+        };
+    }
+
+    public ClassName getTypeName() {
+        return typeName;
+    }
+
+    public CodeBlock getCollectorExpr() {
+        return collectorExpr;
+    }
+}
diff --git a/generator/src/main/java/org/sudu/protogen/protoc/adaptor/Field.java b/generator/src/main/java/org/sudu/protogen/protoc/adaptor/Field.java
@@ -2,6 +2,7 @@
 
 import com.google.protobuf.Descriptors;
 import org.jetbrains.annotations.Nullable;
+import org.sudu.protogen.protobuf.RepeatedContainer;
 import org.sudu.protogen.protoc.Options;
 
 import java.util.Objects;

diff --git a/generator/src/main/java/org/sudu/protogen/protoc/adaptor/Method.java b/generator/src/main/java/org/sudu/protogen/protoc/adaptor/Method.java
@@ -1,6 +1,7 @@
 package org.sudu.protogen.protoc.adaptor;
 
 import com.google.protobuf.Descriptors;
+import org.sudu.protogen.protobuf.RepeatedContainer;
 import org.sudu.protogen.protoc.Options;
 
 import java.util.Objects;
@@ -68,8 +69,9 @@ protected Optional<Boolean> getNullableOption() {
     }
 
     @Override
-    protected Optional<Boolean> getStreamToListOption() {
-        return Options.wrapExtension(descriptor.getOptions(), protogen.Options.streamToList);
+    public Optional<RepeatedContainer> getStreamToContainer() {
+        return Options.wrapExtension(descriptor.getOptions(), protogen.Options.streamToContainer)
+                .map(RepeatedContainer::fromGrpc);
     }
 
     @Override

diff --git a/options/src/main/proto/protogen/options.proto b/options/src/main/proto/protogen/options.proto
@@ -140,8 +140,8 @@ extend google.protobuf.MethodOptions {
   bool unfold_request = 5502;
   // Annotates method with @Nullable and wraps response into nullifyIfNotFound method of BaseGrpcClient
   bool nullable = 5503;
-  // If method is output-streaming, collects its output into java.util.List using spliterator -> toList
-  bool stream_to_list = 5504;
+  // If method is output-streaming, collects its output into a specified container using spliterator
+  RepeatedContainer stream_to_container = 5504;
   /*
    * Sets generated Method name.
    *

diff --git a/tests/src/test/proto/client.proto b/tests/src/test/proto/client.proto
@@ -71,7 +71,7 @@ service SomeService {
   rpc streamOneField(google.protobuf.Empty) returns (stream GrpcOneFieldResponse);
   rpc streamDomain(google.protobuf.Empty) returns (stream GrpcDomainResponse);
   rpc streamDomainList(google.protobuf.Empty) returns (stream GrpcDomainResponse) {
-    option (.protogen.stream_to_list) = true;
+    option (.protogen.stream_to_container) = LIST;
   }
   rpc streamMulti(google.protobuf.Empty) returns (stream GrpcMultiFieldResponse);
 }