#11 ✨ support recursive table structure

Author: oharaandrew314

src/main/kotlin/io/andrewohara/dynamokt/DataClassTableSchema.kt

@@ -1,25 +1,45 @@
 package io.andrewohara.dynamokt
 
 import software.amazon.awssdk.enhanced.dynamodb.TableSchema
+import software.amazon.awssdk.enhanced.dynamodb.internal.mapper.MetaTableSchemaCache
 import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticImmutableTableSchema
 import kotlin.reflect.KClass
 import kotlin.reflect.KVisibility
 import kotlin.reflect.full.declaredMemberProperties
 import kotlin.reflect.full.primaryConstructor
 
-
 fun <Item: Any> DataClassTableSchema(dataClass: KClass<Item>): TableSchema<Item> {
-    val props = dataClass.declaredMemberProperties.sortedBy { it.name }
-    val params = dataClass.primaryConstructor!!.parameters.sortedBy { it.name }
+    require(dataClass.isData) { "$dataClass must be a data class" }
+    return dataClassTableSchema(dataClass, MetaTableSchemaCache())
+}
+
+internal fun <Item: Any> dataClassTableSchema(dataClass: KClass<Item>, schemaCache: MetaTableSchemaCache): TableSchema<Item> {
+    val props = dataClass.declaredMemberProperties
+    require(props.size == dataClass.primaryConstructor!!.parameters.size) {
+        "${dataClass.simpleName} properties MUST all be declared in the constructor"
+    }
+
     val constructor = dataClass.primaryConstructor
         ?.takeIf { it.visibility == KVisibility.PUBLIC }
         ?: error("${dataClass.simpleName} must have a public primary constructor")
 
-    require(dataClass.isData) { "$dataClass must be a data class" }
-    require(props.size == params.size) { "${dataClass.simpleName} properties MUST all be declared in the constructor" }
+    val metaSchema = schemaCache.getOrCreate(dataClass.java)
 
     return StaticImmutableTableSchema.builder(dataClass.java, ImmutableDataClassBuilder::class.java)
         .newItemBuilder({ ImmutableDataClassBuilder(constructor) }, { it.build() as Item })
-        .attributes(dataClass.declaredMemberProperties.map { it.toImmutableDataClassAttribute(dataClass) })
+        .attributes(props.map { it.toImmutableDataClassAttribute(dataClass, schemaCache) })
         .build()
+        .also { metaSchema.initialize(it) }
+}
+
+internal fun <Item: Any> recursiveDataClassTableSchema(dataClass: KClass<Item>, schemaCache: MetaTableSchemaCache): TableSchema<Item> {
+    val metaTableSchema = schemaCache.get(dataClass.java)
+
+    if (metaTableSchema.isPresent) {
+        return if (metaTableSchema.get().isInitialized) {
+            metaTableSchema.get().concreteTableSchema()
+        } else metaTableSchema.get()
+    }
+
+    return dataClassTableSchema(dataClass, schemaCache)
 }

src/main/kotlin/io/andrewohara/dynamokt/ImmutableDataClassAttribute.kt

@@ -3,6 +3,7 @@ package io.andrewohara.dynamokt
 import software.amazon.awssdk.enhanced.dynamodb.AttributeConverter
 import software.amazon.awssdk.enhanced.dynamodb.AttributeConverterProvider
 import software.amazon.awssdk.enhanced.dynamodb.EnhancedType
+import software.amazon.awssdk.enhanced.dynamodb.internal.mapper.MetaTableSchemaCache
 import software.amazon.awssdk.enhanced.dynamodb.mapper.ImmutableAttribute
 import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag
 import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags
@@ -15,23 +16,23 @@ import kotlin.reflect.full.findAnnotation
 import kotlin.reflect.full.staticFunctions
 import kotlin.reflect.jvm.javaType
 
-private fun KType.toEnhancedType(): EnhancedType<out Any> {
+private fun KType.toEnhancedType(schemaCache: MetaTableSchemaCache): EnhancedType<out Any> {
     return when(val clazz = classifier as KClass<Any>) {
         List::class -> {
-            val listType = arguments.first().type!!.toEnhancedType()
+            val listType = arguments.first().type!!.toEnhancedType(schemaCache)
             EnhancedType.listOf(listType)
         }
         Set::class -> {
-            val setType = arguments.first().type!!.toEnhancedType()
+            val setType = arguments.first().type!!.toEnhancedType(schemaCache)
             EnhancedType.setOf(setType)
         }
         Map::class -> {
-            val (key, value) = arguments.map { it.type!!.toEnhancedType() }
+            val (key, value) = arguments.map { it.type!!.toEnhancedType(schemaCache) }
             EnhancedType.mapOf(key, value)
         }
         else -> {
             if (clazz.isData) {
-                EnhancedType.documentOf(clazz.java, DataClassTableSchema(clazz))
+                EnhancedType.documentOf(clazz.java, recursiveDataClassTableSchema(clazz, schemaCache))
             } else {
                 EnhancedType.of(javaType)
             }
@@ -70,17 +71,24 @@ private fun KProperty1<out Any, *>.tags() = buildList {
     }
 }
 
-internal fun <Table: Any, Attr: Any?> KProperty1<Table, Attr>.toImmutableDataClassAttribute(dataClass: KClass<Table>): ImmutableAttribute<Table, ImmutableDataClassBuilder, Attr> {
+internal fun <Table: Any, Attr: Any?> KProperty1<Table, Attr>.toImmutableDataClassAttribute(
+    dataClass: KClass<Table>,
+    schemaCache: MetaTableSchemaCache
+): ImmutableAttribute<Table, ImmutableDataClassBuilder, Attr> {
     val converter = findAnnotation<DynamoKtConverted>()
         ?.converter
         ?.let { it as KClass<AttributeConverter<Attr>> }
         ?.let { initConverter(it) }
-        ?: AttributeConverterProvider.defaultProvider().converterFor(returnType.toEnhancedType())
+        ?: AttributeConverterProvider.defaultProvider().converterFor(returnType.toEnhancedType(schemaCache))
 
     val dynamoName = findAnnotation<DynamoKtAttribute>()?.name?: name
 
     return ImmutableAttribute
-        .builder(EnhancedType.of(dataClass.java), EnhancedType.of(ImmutableDataClassBuilder::class.java), returnType.toEnhancedType() as EnhancedType<Attr>)
+        .builder(
+            EnhancedType.of(dataClass.java),
+            EnhancedType.of(ImmutableDataClassBuilder::class.java),
+            returnType.toEnhancedType(schemaCache) as EnhancedType<Attr>
+        )
         .name(dynamoName)
         .getter(::get)
         .setter { builder, value -> builder[name] = value }

src/test/kotlin/io/andrewohara/dynamokt/DataClassTableSchemaTest.kt

@@ -5,6 +5,7 @@ import io.kotest.matchers.collections.shouldContainExactlyInAnyOrder
 import io.kotest.matchers.shouldBe
 import org.junit.jupiter.api.Test
 import software.amazon.awssdk.core.SdkBytes
+import software.amazon.awssdk.enhanced.dynamodb.TableSchema
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue
 import java.lang.IllegalArgumentException
 import java.lang.IllegalStateException
@@ -116,4 +117,19 @@ class DataClassTableSchemaTest {
             DataClassTableSchema(Person::class)
         }.message shouldBe "Person must have a public primary constructor"
     }
+
+    @Test
+    fun `recursive model`() {
+        data class Item(val id: Int, val inner: Item?)
+
+        val item = Item(1, Item(2, null))
+
+        DataClassTableSchema(Item::class).itemToMap(item, true) shouldBe mapOf(
+            "id" to AttributeValue.fromN("1"),
+            "inner" to AttributeValue.fromM(mapOf(
+                "id" to AttributeValue.fromN("2"),
+                "inner" to AttributeValue.fromNul(true)
+            ))
+        )
+    }
 }