Fix #17

oharaandrew314 · Mar 1, 2024 · f6c4a01 · f6c4a01
1 parent d950a14
commit f6c4a01
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 3 deletions.
diff --git a/src/main/kotlin/io/andrewohara/dynamokt/ImmutableDataClassAttribute.kt b/src/main/kotlin/io/andrewohara/dynamokt/ImmutableDataClassAttribute.kt
@@ -79,19 +79,19 @@ internal fun <Table: Any, Attr: Any?> KProperty1<Table, Attr>.toImmutableDataCla
     dataClass: KClass<Table>,
     schemaCache: MetaTableSchemaCache
 ): ImmutableAttribute<Table, ImmutableDataClassBuilder, Attr> {
-    val converter = findAnnotation<DynamoKtConverted>()
+    val customConverter = findAnnotation<DynamoKtConverted>()
         ?.converter
         ?.let { it as KClass<AttributeConverter<Attr>> }
         ?.let { initConverter(it) }
-        ?: AttributeConverterProvider.defaultProvider().converterFor(returnType.toEnhancedType(schemaCache))
+    val converter = customConverter ?: 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(schemaCache) as EnhancedType<Attr>
+            customConverter?.type() ?: returnType.toEnhancedType(schemaCache) as EnhancedType<Attr>
         )
         .name(dynamoName)
         .getter(::get)

diff --git a/src/test/kotlin/io/andrewohara/dynamokt/DataClassTableSchemaTest.kt b/src/test/kotlin/io/andrewohara/dynamokt/DataClassTableSchemaTest.kt
@@ -1,10 +1,15 @@
 package io.andrewohara.dynamokt
 
+import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
+import com.fasterxml.jackson.module.kotlin.readValue
 import io.kotest.assertions.throwables.shouldThrow
 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.AttributeConverter
+import software.amazon.awssdk.enhanced.dynamodb.AttributeValueType
+import software.amazon.awssdk.enhanced.dynamodb.EnhancedType
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue
 import java.lang.IllegalArgumentException
 import java.lang.IllegalStateException
@@ -166,4 +171,44 @@ class DataClassTableSchemaTest {
             ))
         )
     }
+
+    data class Subtype(
+        val map: Map<String, Any>
+    )
+
+    class SubtypeConverter: AttributeConverter<Subtype> {
+        private val jackson = jacksonObjectMapper()
+
+        override fun transformFrom(input: Subtype): AttributeValue = AttributeValue.fromS(jackson.writeValueAsString(input.map))
+        override fun transformTo(input: AttributeValue) = Subtype(
+            map = jackson.readValue(input.s())
+        )
+        override fun type(): EnhancedType<Subtype> = EnhancedType.of(Subtype::class.java)
+        override fun attributeValueType() = AttributeValueType.S
+    }
+
+    @Test
+    fun `custom converter for data class`() {
+        data class CustomDataClass(
+            @DynamoKtPartitionKey
+            val id: String,
+            @DynamoKtConverted(SubtypeConverter::class)
+            val subtype: Subtype
+        )
+
+        val schema = DataClassTableSchema(CustomDataClass::class)
+
+        val item = CustomDataClass(
+            id = "foo",
+            subtype = Subtype(mapOf("foo" to "bar", "num" to 1))
+        )
+
+        val map = mapOf(
+            "id" to AttributeValue.fromS("foo"),
+            "subtype" to AttributeValue.fromS("""{"foo":"bar","num":1}""")
+        )
+
+        schema.itemToMap(item, true) shouldBe map
+        schema.mapToItem(map) shouldBe item
+    }
 }