Generate Json Schemas for Kotlin

slinkydeveloper · slinkydeveloper · commit 0e6492d66f65 · 2025-04-14T16:08:22.000+02:00
diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
@@ -181,6 +181,24 @@
     [libraries.victools-jsonschema-module-jackson.version]
       ref = 'victools-json-schema'
 
+  [libraries.schema-kenerator-core]
+    module = 'io.github.smiley4:schema-kenerator-core'
+
+    [libraries.schema-kenerator-core.version]
+      ref = 'schema-kenerator'
+
+  [libraries.schema-kenerator-serialization]
+    module = 'io.github.smiley4:schema-kenerator-serialization'
+
+    [libraries.schema-kenerator-serialization.version]
+      ref = 'schema-kenerator'
+
+  [libraries.schema-kenerator-jsonschema]
+    module = 'io.github.smiley4:schema-kenerator-jsonschema'
+
+    [libraries.schema-kenerator-jsonschema.version]
+      ref = 'schema-kenerator'
+
 [plugins]
   aggregate-javadoc = 'io.freefair.aggregate-javadoc:8.6'
   dependency-license-report = 'com.github.jk1.dependency-license-report:2.0'
@@ -213,3 +231,4 @@
   spring-boot = '3.4.4'
   vertx = '4.5.11'
   victools-json-schema = '4.37.0'
+  schema-kenerator = '2.1.2'
diff --git a/sdk-serde-jackson/src/test/java/dev/restate/serde/jackson/JacksonSerdesTest.java b/sdk-serde-jackson/src/test/java/dev/restate/serde/jackson/JacksonSerdesTest.java
@@ -13,17 +13,21 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.node.ObjectNode;
 import dev.restate.serde.Serde;
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
 import java.util.stream.Stream;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
 
 class JacksonSerdesTest {
 
+  record Recursive(String value, Recursive rec) {}
+
   public static class Person {
 
     private final String name;
@@ -75,4 +79,11 @@ private static Stream<Arguments> roundtripTestCases() {
   <T> void roundtrip(T value, Serde<T> serde) {
     assertThat(serde.deserialize(serde.serialize(value))).isEqualTo(value);
   }
+
+  @Test
+  void schemaGenWorksWithRecursion() {
+    ObjectNode node =
+        (ObjectNode) ((Serde.JsonSchema) JacksonSerdes.of(Recursive.class).jsonSchema()).schema();
+    assertThat(node.at("/properties/rec/$ref").textValue()).isEqualTo("#");
+  }
 }
diff --git a/sdk-serde-kotlinx/build.gradle.kts b/sdk-serde-kotlinx/build.gradle.kts
@@ -8,6 +8,12 @@ description = "Restate SDK Kotlinx Serialization integration"
 dependencies {
   api(libs.kotlinx.serialization.json)
   implementation(libs.kotlinx.serialization.core)
+  implementation(libs.schema.kenerator.core)
+  implementation(libs.schema.kenerator.serialization)
+  implementation(libs.schema.kenerator.jsonschema)
 
   implementation(project(":common"))
+
+  testImplementation(libs.junit.jupiter)
+  testImplementation(libs.assertj)
 }
diff --git a/sdk-serde-kotlinx/src/main/kotlin/dev/restate/serde/kotlinx/KotlinSerializationSerdeFactory.kt b/sdk-serde-kotlinx/src/main/kotlin/dev/restate/serde/kotlinx/KotlinSerializationSerdeFactory.kt
@@ -13,20 +13,25 @@ import dev.restate.serde.Serde
 import dev.restate.serde.SerdeFactory
 import dev.restate.serde.TypeRef
 import dev.restate.serde.TypeTag
+import io.github.smiley4.schemakenerator.jsonschema.JsonSchemaSteps.compileReferencing
+import io.github.smiley4.schemakenerator.jsonschema.JsonSchemaSteps.generateJsonSchema
+import io.github.smiley4.schemakenerator.jsonschema.JsonSchemaSteps.withTitle
+import io.github.smiley4.schemakenerator.jsonschema.TitleBuilder
+import io.github.smiley4.schemakenerator.jsonschema.data.RefType
+import io.github.smiley4.schemakenerator.jsonschema.data.TitleType
+import io.github.smiley4.schemakenerator.jsonschema.jsonDsl.JsonArray
+import io.github.smiley4.schemakenerator.jsonschema.jsonDsl.JsonNode
+import io.github.smiley4.schemakenerator.jsonschema.jsonDsl.JsonObject
+import io.github.smiley4.schemakenerator.jsonschema.jsonDsl.JsonTextValue
+import io.github.smiley4.schemakenerator.serialization.SerializationSteps.analyzeTypeUsingKotlinxSerialization
+import io.github.smiley4.schemakenerator.serialization.SerializationSteps.initial
 import java.nio.charset.StandardCharsets
 import kotlin.reflect.KClass
 import kotlin.reflect.KType
 import kotlinx.serialization.*
-import kotlinx.serialization.builtins.*
-import kotlinx.serialization.descriptors.PrimitiveKind
-import kotlinx.serialization.descriptors.SerialDescriptor
-import kotlinx.serialization.descriptors.StructureKind
-import kotlinx.serialization.encodeToString
+import kotlinx.serialization.builtins.nullable
 import kotlinx.serialization.json.Json
-import kotlinx.serialization.json.JsonArray
-import kotlinx.serialization.json.JsonElement
 import kotlinx.serialization.json.JsonNull
-import kotlinx.serialization.json.JsonTransformingSerializer
 import kotlinx.serialization.modules.SerializersModule
 
 /**
@@ -104,6 +109,8 @@ constructor(private val json: Json = Json.Default) : SerdeFactory {
 
     /** Creates a [Serde] implementation using the `kotlinx.serialization` json module. */
     fun <T : Any?> jsonSerde(json: Json = Json.Default, serializer: KSerializer<T>): Serde<T> {
+      val schema = jsonSchema(json, serializer).prettyPrint()
+
       return object : Serde<T> {
         @Suppress("WRONG_NULLABILITY_FOR_JAVA_OVERRIDE")
         override fun serialize(value: T?): Slice {
@@ -124,75 +131,72 @@ constructor(private val json: Json = Json.Default) : SerdeFactory {
         }
 
         override fun jsonSchema(): Serde.Schema {
-          val schema: JsonSchema = serializer.descriptor.jsonSchema()
-          return Serde.StringifiedJsonSchema(Json.encodeToString(schema))
+          return Serde.StringifiedJsonSchema(schema)
         }
       }
     }
 
-    @Serializable
-    @PublishedApi
-    internal data class JsonSchema(
-        @Serializable(with = StringListSerializer::class) val type: List<String>? = null,
-        val format: String? = null,
-    ) {
-      companion object {
-        val INT = JsonSchema(type = listOf("number"), format = "int32")
-
-        val LONG = JsonSchema(type = listOf("number"), format = "int64")
-
-        val DOUBLE = JsonSchema(type = listOf("number"), format = "double")
-
-        val FLOAT = JsonSchema(type = listOf("number"), format = "float")
-
-        val STRING = JsonSchema(type = listOf("string"))
-
-        val BOOLEAN = JsonSchema(type = listOf("boolean"))
-
-        val OBJECT = JsonSchema(type = listOf("object"))
-
-        val LIST = JsonSchema(type = listOf("array"))
-
-        val ANY = JsonSchema()
-      }
-    }
-
-    object StringListSerializer :
-        JsonTransformingSerializer<List<String>>(ListSerializer(String.Companion.serializer())) {
-      override fun transformSerialize(element: JsonElement): JsonElement {
-        require(element is JsonArray)
-        return element.singleOrNull() ?: element
+    private fun <T : Any?> jsonSchema(json: Json, serializer: KSerializer<T>): JsonNode =
+        runCatching {
+              val intermediateStep =
+                  initial(serializer.descriptor)
+                      .analyzeTypeUsingKotlinxSerialization {
+                        serializersModule = json.serializersModule
+                      }
+                      .generateJsonSchema()
+                      .withTitle(type = TitleType.SIMPLE)
+              val compiledSchema = intermediateStep.compileReferencing(RefType.SIMPLE)
+
+              // In case of nested schemas, compileReferencing also contains self schema...
+              val rootSchemaName =
+                  TitleBuilder.BUILDER_SIMPLE(
+                      compiledSchema.typeData, intermediateStep.typeDataById)
+
+              // If schema is not json object, then it's boolean, so we're good no need for
+              // additional manipulation
+              if (compiledSchema.json !is JsonObject) {
+                return compiledSchema.json
+              }
+
+              // Assemble the final schema now
+              val rootNode = compiledSchema.json as JsonObject
+              // Add $schema
+              rootNode.properties.put(
+                  "\$schema", JsonTextValue("https://json-schema.org/draft/2020-12/schema"))
+              // Add $defs
+              val definitions =
+                  compiledSchema.definitions.filter { it.key != rootSchemaName }.toMutableMap()
+              if (definitions.isNotEmpty()) {
+                rootNode.properties.put("\$defs", JsonObject(definitions))
+              }
+              // Replace all $refs
+              rootNode.fixRefsPrefix("#/definitions/$rootSchemaName")
+
+              return rootNode
+            }
+            .getOrDefault(JsonObject(mutableMapOf()))
+
+    private fun JsonNode.fixRefsPrefix(rootDefinition: String) {
+      when (this) {
+        is JsonArray -> this.items.forEach { it.fixRefsPrefix(rootDefinition) }
+        is JsonObject -> this.fixRefsPrefix(rootDefinition)
+        else -> {}
       }
     }
 
-    /**
-     * Super simplistic json schema generation. We should replace this with an appropriate library.
-     */
-    @OptIn(ExperimentalSerializationApi::class)
-    @PublishedApi
-    internal fun SerialDescriptor.jsonSchema(): JsonSchema {
-      var schema =
-          when (this.kind) {
-            PrimitiveKind.BOOLEAN -> JsonSchema.BOOLEAN
-            PrimitiveKind.BYTE -> JsonSchema.INT
-            PrimitiveKind.CHAR -> JsonSchema.STRING
-            PrimitiveKind.DOUBLE -> JsonSchema.DOUBLE
-            PrimitiveKind.FLOAT -> JsonSchema.FLOAT
-            PrimitiveKind.INT -> JsonSchema.INT
-            PrimitiveKind.LONG -> JsonSchema.LONG
-            PrimitiveKind.SHORT -> JsonSchema.INT
-            PrimitiveKind.STRING -> JsonSchema.STRING
-            StructureKind.LIST -> JsonSchema.LIST
-            StructureKind.MAP -> JsonSchema.OBJECT
-            else -> JsonSchema.ANY
+    private fun JsonObject.fixRefsPrefix(rootDefinition: String) {
+      this.properties.computeIfPresent("\$ref") { key, node ->
+        if (node is JsonTextValue) {
+          if (node.value.startsWith(rootDefinition)) {
+            JsonTextValue("#/" + node.value.removePrefix(rootDefinition))
+          } else {
+            JsonTextValue("#/\$defs/" + node.value.removePrefix("#/definitions/"))
           }
-
-      // Add nullability constraint
-      if (this.isNullable && schema.type != null) {
-        schema = schema.copy(type = schema.type.plus("null"))
+        } else {
+          node
+        }
       }
-
-      return schema
+      this.properties.values.forEach { it.fixRefsPrefix(rootDefinition) }
     }
   }
 
diff --git a/sdk-serde-kotlinx/src/test/kotlin/dev/restate/serde/kotlinx/KotlinxSerdeTest.kt b/sdk-serde-kotlinx/src/test/kotlin/dev/restate/serde/kotlinx/KotlinxSerdeTest.kt
