If-not-found logic refactoring

SuduIDE · Sep 29, 2023 · 012adb6 · 012adb6
1 parent 40dcf41
commit 012adb6
Show file tree

Hide file tree

Showing 9 changed files with 159 additions and 72 deletions.
diff --git a/generator/src/main/java/org/sudu/protogen/descriptors/RepeatedContainer.java b/generator/src/main/java/org/sudu/protogen/descriptors/RepeatedContainer.java
@@ -4,10 +4,7 @@
 import com.squareup.javapoet.CodeBlock;
 import protogen.Options;
 
-import java.util.Iterator;
-import java.util.List;
-import java.util.Set;
-import java.util.Spliterators;
+import java.util.*;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
@@ -24,6 +21,11 @@ public CodeBlock getCollectorExpr() {
         public CodeBlock getToStreamExpr(CodeBlock thisRef) {
             return CodeBlock.of("$L.stream()", thisRef);
         }
+
+        @Override
+        public CodeBlock getEmptyOne() {
+            return CodeBlock.of("new $T<>()", ArrayList.class);
+        }
     },
 
     SET(ClassName.get(Set.class)) {
@@ -36,6 +38,11 @@ public CodeBlock getCollectorExpr() {
         public CodeBlock getToStreamExpr(CodeBlock thisRef) {
             return CodeBlock.of("$L.stream()", thisRef);
         }
+
+        @Override
+        public CodeBlock getEmptyOne() {
+            return CodeBlock.of("new $T<>()", HashSet.class);
+        }
     },
 
     ITERATOR(ClassName.get(Iterator.class)) {
@@ -63,6 +70,11 @@ public CodeBlock convertListToInstance(CodeBlock thisRef) {
         public CodeBlock convertInstanceToIterable(CodeBlock thisRef) {
             return CodeBlock.of("() -> $L", thisRef);
         }
+
+        @Override
+        public CodeBlock getEmptyOne() {
+            return CodeBlock.of("$T.emptyIterator()", Collections.class);
+        }
     },
 
     STREAM(ClassName.get(Stream.class)) {
@@ -85,6 +97,11 @@ public CodeBlock convertListToInstance(CodeBlock thisRef) {
         public CodeBlock convertInstanceToIterable(CodeBlock thisRef) {
             return CodeBlock.of("$L.toList()", thisRef);
         }
+
+        @Override
+        public CodeBlock getEmptyOne() {
+            return CodeBlock.of("$T.empty()", Stream.class);
+        }
     };
 
     private final ClassName typeName;
@@ -113,6 +130,8 @@ public static RepeatedContainer fromGrpc(Options.RepeatedContainer proto) {
      */
     public abstract CodeBlock getCollectorExpr();
 
+    public abstract CodeBlock getEmptyOne();
+
     /**
      * Builds an expression converting a java.Util.List instance into the container
      * Used to build .getSomeList().... expressions

diff --git a/generator/src/main/java/org/sudu/protogen/generator/client/ApiMethodGeneratorBase.java b/generator/src/main/java/org/sudu/protogen/generator/client/ApiMethodGeneratorBase.java
@@ -2,6 +2,8 @@
 
 import com.squareup.javapoet.*;
 import org.jetbrains.annotations.Nullable;
+import org.sudu.protogen.EmptyIfNotFound;
+import org.sudu.protogen.NullifyIfNotFound;
 import org.sudu.protogen.descriptors.Method;
 import org.sudu.protogen.generator.GenerationContext;
 import org.sudu.protogen.generator.field.FieldGenerator;
@@ -32,7 +34,7 @@ public ApiMethodGeneratorBase(GenerationContext context, Method method, TypeMode
 
     public MethodSpec generate() {
         List<ParameterSpec> params = params();
-        return MethodSpec.methodBuilder(method.generatedName())
+        MethodSpec.Builder builder = MethodSpec.methodBuilder(method.generatedName())
                 .addModifiers(method.getAccessModifier())
                 .returns(returnType.getTypeName())
                 .addParameters(params)
@@ -41,8 +43,12 @@ public MethodSpec generate() {
                         method.ifNotFoundBehavior() == Options.IfNotFound.NULLIFY
                                 ? context.configuration().nullableAnnotationClass()
                                 : context.configuration().nonnullAnnotationClass()
-                )
-                .build();
+                );
+        switch (method.ifNotFoundBehavior()) {
+            case NULLIFY -> builder.addAnnotation(NullifyIfNotFound.class);
+            case EMPTY -> builder.addAnnotation(EmptyIfNotFound.class);
+        }
+        return builder.build();
     }
 
     private List<ParameterSpec> params() {

diff --git a/generator/src/main/java/org/sudu/protogen/generator/client/StubCallMethodGenerator.java b/generator/src/main/java/org/sudu/protogen/generator/client/StubCallMethodGenerator.java
@@ -1,16 +1,19 @@
 package org.sudu.protogen.generator.client;
 
 import com.squareup.javapoet.*;
-import org.sudu.protogen.BaseClientUtils;
+import org.jetbrains.annotations.Nullable;
 import org.sudu.protogen.descriptors.Method;
 import org.sudu.protogen.descriptors.RepeatedContainer;
 import org.sudu.protogen.generator.GenerationContext;
 import org.sudu.protogen.generator.type.RepeatedType;
 import org.sudu.protogen.generator.type.TypeModel;
-import protogen.Options;
+import org.sudu.protogen.generator.type.UnfoldedType;
 
 import javax.lang.model.element.Modifier;
 import java.util.List;
+import java.util.function.Supplier;
+
+import static protogen.Options.IfNotFound.*;
 
 public class StubCallMethodGenerator {
 
@@ -33,46 +36,95 @@ public MethodSpec generate() {
         return MethodSpec.methodBuilder(method.generatedName() + "StubCall")
                 .addModifiers(Modifier.PRIVATE)
                 .addParameters(parameters())
-                .addCode(body())
+                .addCode(new BodyGenerator().withIfNotFound().get())
                 .returns(returnType.getTypeName())
                 .addAnnotation(
-                        method.ifNotFoundBehavior() == Options.IfNotFound.NULLIFY
+                        method.ifNotFoundBehavior() == NULLIFY
                                 ? context.configuration().nullableAnnotationClass()
                                 : context.configuration().nonnullAnnotationClass()
                 )
                 .build();
     }
 
-    private CodeBlock body() {
-        CodeBlock returnExpr = CodeBlock.of("$N.$L(request)", stubField, method.getName());
-        CodeBlock body = CodeBlock.of("");
-
-        if (returnType instanceof RepeatedType repType) { // i.e. method is streaming non-void
-            body = CodeBlock.builder()
-                    .addStatement(
-                            "var iterator = $L",
-                            returnExpr
-                    )
-                    .build();
-            // I write mapping here manually because input is always an Iterator<Grpc..> and output is specified by the RepeatedContainer option
-            // So RepeatedType.fromGrpcTransformer is not suitable because it does only T<U> <--> T<V> mappings
-            CodeBlock mappingExpr = CodeBlock.of("i -> $L", repType.getElementModel().fromGrpcTransformer(CodeBlock.of("i")));
-            returnExpr = CodeBlock.of("$L\n.map($L)$L", RepeatedContainer.ITERATOR.getToStreamExpr(CodeBlock.of("iterator")), mappingExpr, repType.getRepeatedType().getCollectorExpr());
-        } else {
-            // Don't transform if returnType is streaming (RepeatedType), see the comment above
-            returnExpr = returnType.fromGrpcTransformer(returnExpr);
+    private class BodyGenerator implements Supplier<CodeBlock> {
+
+        @Override
+        public CodeBlock get() {
+            CodeBlock returnExpr = CodeBlock.of("$N.$L(request)", stubField, method.getName());
+            CodeBlock body = CodeBlock.of("");
+
+            if (returnType instanceof RepeatedType repType) { // i.e. method is streaming non-void
+                body = CodeBlock.builder()
+                        .addStatement(
+                                "var iterator = $L",
+                                returnExpr
+                        )
+                        .build();
+                // I write mapping here manually because input is always an Iterator<Grpc..> and output is specified by the RepeatedContainer option
+                // So RepeatedType.fromGrpcTransformer is not suitable because it does only T<U> <--> T<V> mappings
+                CodeBlock mappingExpr = CodeBlock.of("i -> $L", repType.getElementModel().fromGrpcTransformer(CodeBlock.of("i")));
+                returnExpr = CodeBlock.of("$L\n.map($L)$L", RepeatedContainer.ITERATOR.getToStreamExpr(CodeBlock.of("iterator")), mappingExpr, repType.getRepeatedType().getCollectorExpr());
+            } else {
+                // Don't transform if returnType is streaming (RepeatedType), see the comment above
+                returnExpr = returnType.fromGrpcTransformer(returnExpr);
+            }
+            if (returnType.getTypeName() != TypeName.VOID) {
+                returnExpr = CodeBlock.of("return $L", returnExpr);
+            }
+            return CodeBlock.builder().add(body).addStatement(returnExpr).build();
         }
 
-        switch (method.ifNotFoundBehavior()) {
-            case NULLIFY ->
-                    returnExpr = CodeBlock.of("$T.nullifyIfNotFound(() -> $L)", BaseClientUtils.class, returnExpr);
-            case EMPTY -> returnExpr = CodeBlock.of("$T.emptyIfNotFound(() -> $L)", BaseClientUtils.class, returnExpr);
+        private BodyGenerator withIfNotFound() {
+            return new IfNotFoundDecorator(this);
         }
 
-        if (returnType.getTypeName() != TypeName.VOID) {
-            returnExpr = CodeBlock.of("return $L", returnExpr);
+        private class IfNotFoundDecorator extends BodyGenerator {
+
+            private final BodyGenerator generator;
+
+            public IfNotFoundDecorator(BodyGenerator generator) {
+                this.generator = generator;
+            }
+
+            @Override
+            public CodeBlock get() {
+                if (method.ifNotFoundBehavior() == IGNORE) return generator.get();
+                return CodeBlock.of("""
+                                try {$>
+                                $L
+                                $<} catch ($T ex) {$>
+                                if (ex.getStatus().getCode() == $T.NOT_FOUND.getCode()) {$>
+                                $L
+                                $<}
+                                throw ex;
+                                $<}
+                                """,
+                        generator.get(),
+                        ClassName.get("io.grpc", "StatusRuntimeException"),
+                        ClassName.get("io.grpc", "Status"),
+                        ifNotFoundBehaviour()
+                );
+            }
+
+            private CodeBlock ifNotFoundBehaviour() {
+                if (method.ifNotFoundBehavior() == EMPTY) {
+                    RepeatedContainer container = returnTypeContainer();
+                    if (container == null) return CodeBlock.of("");
+                    return CodeBlock.of("return $L;", container.getEmptyOne());
+                }
+                if (method.ifNotFoundBehavior() == NULLIFY) return CodeBlock.of("return null;");
+                return CodeBlock.of("");
+            }
+
+            @Nullable
+            private RepeatedContainer returnTypeContainer() {
+                if (returnType instanceof RepeatedType rt)
+                    return rt.getRepeatedType();
+                if (returnType instanceof UnfoldedType ut && ut.getType() instanceof RepeatedType rt)
+                    return rt.getRepeatedType();
+                return null;
+            }
         }
-        return CodeBlock.builder().add(body).addStatement(returnExpr).build();
     }
 
     private List<ParameterSpec> parameters() {

diff --git a/generator/src/main/java/org/sudu/protogen/generator/type/UnfoldedType.java b/generator/src/main/java/org/sudu/protogen/generator/type/UnfoldedType.java
@@ -21,6 +21,10 @@ public UnfoldedType(TypeModel type, Message originalMessage) {
         this.unfoldedTypeName = originalMessage.getProtobufTypeName();
     }
 
+    public TypeModel getType() {
+        return type;
+    }
+
     @Override
     public CodeBlock fromGrpcTransformer(CodeBlock expr) {
         return type.fromGrpcTransformer(CodeBlock.of("$L.$L()", expr, type.getterMethod(unfoldedFieldName)));

diff --git a/options/src/main/java/org/sudu/protogen/BaseClientUtils.java b/options/src/main/java/org/sudu/protogen/BaseClientUtils.java
diff --git a/options/src/main/java/org/sudu/protogen/EmptyIfNotFound.java b/options/src/main/java/org/sudu/protogen/EmptyIfNotFound.java
@@ -0,0 +1,9 @@
+package org.sudu.protogen;
+
+import java.lang.annotation.*;
+
+@Documented
+@Retention(RetentionPolicy.CLASS)
+@Target(ElementType.METHOD)
+public @interface EmptyIfNotFound {
+}
diff --git a/options/src/main/java/org/sudu/protogen/NullifyIfNotFound.java b/options/src/main/java/org/sudu/protogen/NullifyIfNotFound.java
@@ -0,0 +1,9 @@
+package org.sudu.protogen;
+
+import java.lang.annotation.*;
+
+@Documented
+@Retention(RetentionPolicy.CLASS)
+@Target(ElementType.METHOD)
+public @interface NullifyIfNotFound {
+}
diff --git a/tests/src/test/java/org/sudu/protogen/test/A.java b/tests/src/test/java/org/sudu/protogen/test/A.java
@@ -0,0 +1,18 @@
+package org.sudu.protogen.test;
+
+import io.grpc.Channel;
+import org.jetbrains.annotations.Nullable;
+import org.sudu.protogen.test.client.DefaultSomeClient;
+import org.sudu.protogen.test.client.Domain;
+
+public class A extends DefaultSomeClient {
+
+    public A(Channel channel) {
+        super(channel);
+    }
+
+    @Override
+    public @Nullable Domain commonMultiNullable(String r1, String r2) {
+        return super.commonMultiNullable(r1, r2);
+    }
+}
diff --git a/tests/src/test/proto/client.proto b/tests/src/test/proto/client.proto
@@ -62,7 +62,9 @@ service SomeService {
   rpc commonVoid(GrpcMultiFieldRequest) returns (GrpcVoidResponse);
   rpc commonInt(GrpcMultiFieldRequest) returns (google.protobuf.Int32Value);
   rpc commonEmpty(google.protobuf.Empty) returns (google.protobuf.Empty);
-  rpc commonRetList(google.protobuf.Empty) returns (GrpcListResponse);
+  rpc commonRetList(google.protobuf.Empty) returns (GrpcListResponse) {
+    option (.protogen.if_not_found) = EMPTY;
+  }
   rpc commonMulti(GrpcMultiFieldRequest) returns (GrpcDomain);
   rpc commonMultiNullable(GrpcMultiFieldRequest) returns (GrpcDomain) {
     option (.protogen.if_not_found) = NULLIFY;
@@ -81,7 +83,9 @@ service SomeService {
   rpc streamVoidById(GrpcDomainId) returns (stream GrpcVoidResponse);
   rpc streamOneField(google.protobuf.Empty) returns (stream GrpcOneFieldResponse);
   rpc streamOneFieldFromMulti(GrpcMultiFieldRequest) returns (stream GrpcOneFieldResponse);
-  rpc streamDomain(google.protobuf.Empty) returns (stream GrpcDomainResponse);
+  rpc streamDomain(google.protobuf.Empty) returns (stream GrpcDomainResponse) {
+    option (.protogen.if_not_found) = EMPTY;
+  }
   rpc streamDomainList(google.protobuf.Empty) returns (stream GrpcDomainResponse) {
     option (.protogen.stream_to_container) = LIST;
   }