Raw handlers (#303)

Author: slinkydeveloper

sdk-api-gen-common/src/main/java/dev/restate/sdk/gen/model/Handler.java

@@ -9,18 +9,25 @@
 package dev.restate.sdk.gen.model;
 
 import java.util.Objects;
+import org.jspecify.annotations.Nullable;
 
 public class Handler {
 
   private final CharSequence name;
   private final HandlerType handlerType;
+  private final @Nullable String inputAccept;
   private final PayloadType inputType;
   private final PayloadType outputType;
 
   public Handler(
-      CharSequence name, HandlerType handlerType, PayloadType inputType, PayloadType outputType) {
+      CharSequence name,
+      HandlerType handlerType,
+      @Nullable String inputAccept,
+      PayloadType inputType,
+      PayloadType outputType) {
     this.name = name;
     this.handlerType = handlerType;
+    this.inputAccept = inputAccept;
     this.inputType = inputType;
     this.outputType = outputType;
   }
@@ -33,6 +40,10 @@ public HandlerType getHandlerType() {
     return handlerType;
   }
 
+  public String getInputAccept() {
+    return inputAccept;
+  }
+
   public PayloadType getInputType() {
     return inputType;
   }
@@ -48,6 +59,7 @@ public static Builder builder() {
   public static class Builder {
     private CharSequence name;
     private HandlerType handlerType;
+    private String inputAccept;
     private PayloadType inputType;
     private PayloadType outputType;
 
@@ -61,6 +73,11 @@ public Builder withHandlerType(HandlerType handlerType) {
       return this;
     }
 
+    public Builder withInputAccept(String inputAccept) {
+      this.inputAccept = inputAccept;
+      return this;
+    }
+
     public Builder withInputType(PayloadType inputType) {
       this.inputType = inputType;
       return this;
@@ -96,7 +113,11 @@ public Handler validateAndBuild() {
       }
 
       return new Handler(
-          Objects.requireNonNull(name), Objects.requireNonNull(handlerType), inputType, outputType);
+          Objects.requireNonNull(name),
+          Objects.requireNonNull(handlerType),
+          inputAccept,
+          inputType,
+          outputType);
     }
   }
 }

sdk-api-gen-common/src/main/java/dev/restate/sdk/gen/template/HandlebarsTemplateEngine.java

@@ -107,7 +107,12 @@ private ServiceTemplateModel(
 
       this.handlers =
           inner.getMethods().stream()
-              .map(h -> new HandlerTemplateModel(h, handlerNamesToPrefix))
+              .map(
+                  h ->
+                      new HandlerTemplateModel(
+                          h,
+                          this.generatedClassSimpleNamePrefix + "Definitions.Serde",
+                          handlerNamesToPrefix))
               .collect(Collectors.toList());
     }
   }
@@ -126,14 +131,18 @@ static class HandlerTemplateModel {
     public final String inputSerdeDecl;
     public final String boxedInputFqcn;
     public final String inputSerdeFieldName;
+    public final String inputAcceptContentType;
+    public final String inputSerdeRef;
 
     public final boolean outputEmpty;
     public final String outputFqcn;
     public final String outputSerdeDecl;
     public final String boxedOutputFqcn;
     public final String outputSerdeFieldName;
+    public final String outputSerdeRef;
 
-    private HandlerTemplateModel(Handler inner, Set<String> handlerNamesToPrefix) {
+    private HandlerTemplateModel(
+        Handler inner, String definitionsClass, Set<String> handlerNamesToPrefix) {
       this.name = inner.getName().toString();
       this.methodName = (handlerNamesToPrefix.contains(this.name) ? "_" : "") + this.name;
       this.handlerType = inner.getHandlerType().toString();
@@ -146,13 +155,16 @@ private HandlerTemplateModel(Handler inner, Set<String> handlerNamesToPrefix) {
       this.inputFqcn = inner.getInputType().getName();
       this.inputSerdeDecl = inner.getInputType().getSerdeDecl();
       this.boxedInputFqcn = inner.getInputType().getBoxed();
-      this.inputSerdeFieldName = "SERDE_" + this.name.toUpperCase() + "_INPUT";
+      this.inputSerdeFieldName = this.name.toUpperCase() + "_INPUT";
+      this.inputAcceptContentType = inner.getInputAccept();
+      this.inputSerdeRef = definitionsClass + "." + this.inputSerdeFieldName;
 
       this.outputEmpty = inner.getOutputType().isEmpty();
       this.outputFqcn = inner.getOutputType().getName();
       this.outputSerdeDecl = inner.getOutputType().getSerdeDecl();
       this.boxedOutputFqcn = inner.getOutputType().getBoxed();
-      this.outputSerdeFieldName = "SERDE_" + this.name.toUpperCase() + "_OUTPUT";
+      this.outputSerdeFieldName = this.name.toUpperCase() + "_OUTPUT";
+      this.outputSerdeRef = definitionsClass + "." + this.outputSerdeFieldName;
     }
   }
 }

sdk-api-gen-common/src/main/java/dev/restate/sdk/gen/utils/AnnotationUtils.java

@@ -0,0 +1,23 @@
+// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
+//
+// This file is part of the Restate Java SDK,
+// which is released under the MIT license.
+//
+// You can find a copy of the license in file LICENSE in the root
+// directory of this repository or package, or at
+// https://github.com/restatedev/sdk-java/blob/main/LICENSE
+package dev.restate.sdk.gen.utils;
+
+import java.lang.annotation.Annotation;
+import java.util.Objects;
+
+public class AnnotationUtils {
+  public static Object getAnnotationDefaultValue(
+      Class<? extends Annotation> annotation, String name) {
+    try {
+      return Objects.requireNonNull(annotation.getMethod(name).getDefaultValue());
+    } catch (NoSuchMethodException e) {
+      throw new RuntimeException(e);
+    }
+  }
+}

sdk-api-gen/src/main/java/dev/restate/sdk/gen/ElementConverter.java

@@ -11,32 +11,31 @@
 import dev.restate.sdk.Context;
 import dev.restate.sdk.ObjectContext;
 import dev.restate.sdk.SharedObjectContext;
-import dev.restate.sdk.annotation.Exclusive;
-import dev.restate.sdk.annotation.Shared;
-import dev.restate.sdk.annotation.Workflow;
+import dev.restate.sdk.annotation.*;
 import dev.restate.sdk.common.ServiceType;
 import dev.restate.sdk.gen.model.*;
+import dev.restate.sdk.gen.model.Handler;
+import dev.restate.sdk.gen.model.Service;
+import dev.restate.sdk.gen.utils.AnnotationUtils;
 import dev.restate.sdk.workflow.WorkflowContext;
 import dev.restate.sdk.workflow.WorkflowSharedContext;
 import java.util.List;
-import java.util.Objects;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import javax.annotation.processing.Messager;
-import javax.lang.model.element.ElementKind;
-import javax.lang.model.element.ExecutableElement;
-import javax.lang.model.element.Modifier;
-import javax.lang.model.element.TypeElement;
+import javax.lang.model.element.*;
 import javax.lang.model.type.TypeKind;
 import javax.lang.model.type.TypeMirror;
 import javax.lang.model.util.Elements;
 import javax.lang.model.util.Types;
 import javax.tools.Diagnostic;
+import org.jspecify.annotations.Nullable;
 
 public class ElementConverter {
 
   private static final PayloadType EMPTY_PAYLOAD =
       new PayloadType(true, "", "Void", "dev.restate.sdk.common.CoreSerdes.VOID");
+  private static final String RAW_SERDE = "dev.restate.sdk.common.CoreSerdes.RAW";
 
   private final Messager messager;
   private final Elements elements;
@@ -197,18 +196,12 @@ private Handler fromExecutableElement(ServiceType serviceType, ExecutableElement
       return new Handler.Builder()
           .withName(element.getSimpleName())
           .withHandlerType(handlerType)
-          .withInputType(
-              element.getParameters().size() > 1
-                  ? payloadFromType(element.getParameters().get(1).asType())
-                  : EMPTY_PAYLOAD)
-          .withOutputType(
-              !element.getReturnType().getKind().equals(TypeKind.VOID)
-                  ? payloadFromType(element.getReturnType())
-                  : EMPTY_PAYLOAD)
+          .withInputAccept(inputAcceptFromParameterList(element.getParameters()))
+          .withInputType(inputPayloadFromParameterList(element.getParameters()))
+          .withOutputType(outputPayloadFromExecutableElement(element))
           .validateAndBuild();
     } catch (Exception e) {
-      messager.printMessage(
-          Diagnostic.Kind.ERROR, "Error when building handler: " + e.getMessage(), element);
+      messager.printMessage(Diagnostic.Kind.ERROR, "Error when building handler: " + e, element);
       return null;
     }
   }
@@ -280,13 +273,89 @@ private void validateFirstParameterType(Class<?> clazz, ExecutableElement elemen
     }
   }
 
-  private PayloadType payloadFromType(TypeMirror typeMirror) {
-    Objects.requireNonNull(typeMirror);
-    return new PayloadType(
-        false, typeMirror.toString(), boxedType(typeMirror), serdeDecl(typeMirror));
+  private String inputAcceptFromParameterList(List<? extends VariableElement> element) {
+    if (element.size() <= 1) {
+      return null;
+    }
+
+    Accept accept = element.get(1).getAnnotation(Accept.class);
+    if (accept == null) {
+      return null;
+    }
+    return accept.value();
+  }
+
+  private PayloadType inputPayloadFromParameterList(List<? extends VariableElement> element) {
+    if (element.size() <= 1) {
+      return EMPTY_PAYLOAD;
+    }
+
+    Element parameterElement = element.get(1);
+    return payloadFromTypeMirrorAndAnnotations(
+        parameterElement.asType(),
+        parameterElement.getAnnotation(Json.class),
+        parameterElement.getAnnotation(Raw.class),
+        parameterElement);
+  }
+
+  private PayloadType outputPayloadFromExecutableElement(ExecutableElement element) {
+    return payloadFromTypeMirrorAndAnnotations(
+        element.getReturnType(),
+        element.getAnnotation(Json.class),
+        element.getAnnotation(Raw.class),
+        element);
+  }
+
+  private PayloadType payloadFromTypeMirrorAndAnnotations(
+      TypeMirror ty, @Nullable Json jsonAnnotation, @Nullable Raw rawAnnotation, Element element) {
+    if (ty.getKind().equals(TypeKind.VOID)) {
+      if (rawAnnotation != null || jsonAnnotation != null) {
+        messager.printMessage(
+            Diagnostic.Kind.ERROR, "Unexpected annotation for void type.", element);
+      }
+      return EMPTY_PAYLOAD;
+    }
+    // Some validation
+    if (rawAnnotation != null && jsonAnnotation != null) {
+      messager.printMessage(
+          Diagnostic.Kind.ERROR,
+          "A parameter cannot be annotated both with @Raw and @Json.",
+          element);
+    }
+    if (rawAnnotation != null
+        && !types.isSameType(ty, types.getArrayType(types.getPrimitiveType(TypeKind.BYTE)))) {
+      messager.printMessage(
+          Diagnostic.Kind.ERROR,
+          "A parameter annotated with @Raw MUST be of type byte[], was " + ty,
+          element);
+    }
+
+    String serdeDecl = rawAnnotation != null ? RAW_SERDE : jsonSerdeDecl(ty);
+    if (rawAnnotation != null
+        && !rawAnnotation
+            .contentType()
+            .equals(AnnotationUtils.getAnnotationDefaultValue(Raw.class, "contentType"))) {
+      serdeDecl = contentTypeDecoratedSerdeDecl(serdeDecl, rawAnnotation.contentType());
+    }
+    if (jsonAnnotation != null
+        && !jsonAnnotation
+            .contentType()
+            .equals(AnnotationUtils.getAnnotationDefaultValue(Json.class, "contentType"))) {
+      serdeDecl = contentTypeDecoratedSerdeDecl(serdeDecl, jsonAnnotation.contentType());
+    }
+
+    return new PayloadType(false, ty.toString(), boxedType(ty), serdeDecl);
+  }
+
+  private static String contentTypeDecoratedSerdeDecl(String serdeDecl, String contentType) {
+    return "dev.restate.sdk.common.Serde.withContentType(\""
+        + contentType
+        + "\", "
+        + serdeDecl
+        + ")";
   }
 
-  private static String serdeDecl(TypeMirror ty) {
+  private static String jsonSerdeDecl(TypeMirror ty) {
     switch (ty.getKind()) {
       case BOOLEAN:
         return "dev.restate.sdk.common.CoreSerdes.JSON_BOOLEAN";

sdk-api-gen/src/main/java/dev/restate/sdk/gen/ServiceProcessor.java

@@ -35,6 +35,7 @@
 @SupportedSourceVersion(SourceVersion.RELEASE_11)
 public class ServiceProcessor extends AbstractProcessor {
 
+  private HandlebarsTemplateEngine definitionsCodegen;
   private HandlebarsTemplateEngine bindableServiceFactoryCodegen;
   private HandlebarsTemplateEngine bindableServiceCodegen;
   private HandlebarsTemplateEngine clientCodegen;
@@ -47,6 +48,18 @@ public synchronized void init(ProcessingEnvironment processingEnv) {
 
     FilerTemplateLoader filerTemplateLoader = new FilerTemplateLoader(processingEnv.getFiler());
 
+    this.definitionsCodegen =
+        new HandlebarsTemplateEngine(
+            "Definitions",
+            filerTemplateLoader,
+            Map.of(
+                ServiceType.WORKFLOW,
+                "templates/Definitions.hbs",
+                ServiceType.SERVICE,
+                "templates/Definitions.hbs",
+                ServiceType.VIRTUAL_OBJECT,
+                "templates/Definitions.hbs"),
+            RESERVED_METHOD_NAMES);
     this.bindableServiceFactoryCodegen =
         new HandlebarsTemplateEngine(
             "BindableServiceFactory",
@@ -108,6 +121,7 @@ public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment
       try {
         ThrowingFunction<String, Writer> fileCreator =
             name -> filer.createSourceFile(name, e.getKey()).openWriter();
+        this.definitionsCodegen.generate(fileCreator, e.getValue());
         this.bindableServiceFactoryCodegen.generate(fileCreator, e.getValue());
         this.bindableServiceCodegen.generate(fileCreator, e.getValue());
         this.clientCodegen.generate(fileCreator, e.getValue());