Add title for classes (#398)

restatedev · Nov 15, 2024 · 12184e9 · 12184e9
1 parent 7954ab5
commit 12184e9
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 0 deletions.
diff --git a/examples/src/main/java/my/restate/sdk/examples/Counter.java b/examples/src/main/java/my/restate/sdk/examples/Counter.java
@@ -48,6 +48,13 @@ public long get(SharedObjectContext ctx) {
     return ctx.get(TOTAL).orElse(0L);
   }
 
+  /** Get the current counter value. */
+  @Shared
+  @Handler
+  public Node getNode(SharedObjectContext ctx) {
+    return null;
+  }
+
   /** Add a value, and get both the previous value and the new value. */
   @Handler
   public CounterUpdateResult getAndAdd(ObjectContext ctx, long request) {
@@ -64,6 +71,24 @@ public static void main(String[] args) {
     RestateHttpEndpointBuilder.builder().bind(new Counter()).buildAndListen();
   }
 
+  public static class Node {
+    private final String data;
+    private final Node inner;
+
+    public Node(String data, Node inner) {
+      this.data = data;
+      this.inner = inner;
+    }
+
+    public String getData() {
+      return data;
+    }
+
+    public Node getInner() {
+      return inner;
+    }
+  }
+
   public static class CounterUpdateResult {
     private final long newValue;
     private final long oldValue;

diff --git a/sdk-serde-jackson/src/main/java/dev/restate/sdk/serde/jackson/JacksonSerdes.java b/sdk-serde-jackson/src/main/java/dev/restate/sdk/serde/jackson/JacksonSerdes.java
@@ -10,13 +10,15 @@
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.github.victools.jsonschema.generator.*;
 import com.github.victools.jsonschema.module.jackson.JacksonModule;
 import com.github.victools.jsonschema.module.jackson.JacksonOption;
 import dev.restate.sdk.common.RichSerde;
 import dev.restate.sdk.common.Serde;
 import java.io.IOException;
+import java.util.stream.StreamSupport;
 import org.jspecify.annotations.Nullable;
 
 /**
@@ -59,6 +61,37 @@ private JacksonSerdes() {}
         new SchemaGeneratorConfigBuilder(
                 defaultMapper, SchemaVersion.DRAFT_2020_12, OptionPreset.PLAIN_JSON)
             .with(module);
+
+    // Make sure we use `title` for types
+    configBuilder
+        .forTypesInGeneral()
+        .withTypeAttributeOverride(
+            (schema, scope, context) -> {
+              if (schema.isObject()
+                  && !schema.hasNonNull(
+                      SchemaKeyword.TAG_TITLE.forVersion(
+                          context.getGeneratorConfig().getSchemaVersion()))) {
+                JsonNode typeKeyword =
+                    schema.get(
+                        SchemaKeyword.TAG_TYPE.forVersion(
+                            context.getGeneratorConfig().getSchemaVersion()));
+                boolean isObjectSchema =
+                    typeKeyword != null
+                        && ((typeKeyword.isTextual() && "object".equals(typeKeyword.textValue()))
+                            || (typeKeyword.isArray()
+                                && StreamSupport.stream(typeKeyword.spliterator(), false)
+                                    .anyMatch(
+                                        el -> el.isTextual() && "object".equals(el.textValue()))));
+                if (isObjectSchema) {
+                  schema.put(
+                      SchemaKeyword.TAG_TITLE.forVersion(
+                          context.getGeneratorConfig().getSchemaVersion()),
+                      scope.getSimpleTypeDescription());
+                }
+              }
+            });
+    ;
+
     schemaGenerator = new SchemaGenerator(configBuilder.build());
   }