Add support for multi-attribute keys in DDB Mapper

gradle/libs.versions.toml

@@ -19,7 +19,7 @@ smithy-kotlin-codegen-version = "0.35.22"
 smithy-version = "1.64.0"
 
 # testing
-ddb-local-version = "2.5.2"
+ddb-local-version = "2.6.1"
 junit-version = "5.13.2"
 kotest-version = "5.9.1"
 kotlinx-benchmark-version = "0.4.12"

hll/dynamodb-mapper/dynamodb-mapper-codegen/src/main/kotlin/aws/sdk/kotlin/hll/dynamodbmapper/codegen/model/MapperTypes.kt

@@ -48,16 +48,55 @@ public object MapperTypes {
         public fun itemSchema(typeVar: String): TypeRef =
             TypeRef(MapperPkg.Hl.Items, "ItemSchema", genericArgs = listOf(TypeVar(typeVar)))
 
-        public fun itemSchemaPartitionKey(objectType: TypeRef, pkType: TypeRef): TypeRef =
-            TypeRef(MapperPkg.Hl.Items, "ItemSchema.PartitionKey", genericArgs = listOf(objectType, pkType))
+        public fun itemSchemaPartitionKey(objectType: TypeRef, pkTypes: List<TypeRef>): TypeRef =
+            TypeRef(
+                MapperPkg.Hl.Items,
+                "ItemSchema.PartitionKey",
+                genericArgs = listOf(
+                    objectType,
+                    keyType(pkTypes),
+                ),
+            )
+
+        public fun itemSchemaCompositeKey(objectType: TypeRef, pkTypes: List<TypeRef>, skTypes: List<TypeRef>): TypeRef =
+            TypeRef(
+                MapperPkg.Hl.Items,
+                "ItemSchema.CompositeKey",
+                genericArgs = listOf(
+                    objectType,
+                    keyType(pkTypes),
+                    keyType(skTypes),
+                ),
+            )
+
+        public fun keySpec(keyTypes: List<TypeRef>): TypeRef {
+            val keySize = keyTypes.size
+            require(keySize in 1..4) { "KeySpec subtypes must have between 1 and 4 keys, found $keySize" }
+            return TypeRef(MapperPkg.Hl.Items, "KeySpec.Key$keySize", keyTypes)
+        }
+
+        public val KeySpecByteArray: TypeRef = TypeRef(MapperPkg.Hl.Items, "KeySpec.byteArray")
+        public fun keySpecNumber(numberTypeRef: TypeRef? = null): TypeRef = TypeRef(
+            MapperPkg.Hl.Items,
+            "KeySpec.number",
+            genericArgs = listOfNotNull(numberTypeRef),
+        )
+        public val KeySpecString: TypeRef = TypeRef(MapperPkg.Hl.Items, "KeySpec.string")
 
-        public fun itemSchemaCompositeKey(objectType: TypeRef, pkType: TypeRef, skType: TypeRef): TypeRef =
-            TypeRef(MapperPkg.Hl.Items, "ItemSchema.CompositeKey", genericArgs = listOf(objectType, pkType, skType))
+        public val KeySpecThenByteArray: TypeRef = TypeRef(MapperPkg.Hl.Items, "thenByteArray")
+        public fun keySpecThenNumber(numberTypeRef: TypeRef): TypeRef = TypeRef(
+            MapperPkg.Hl.Items,
+            "thenNumber",
+            genericArgs = listOfNotNull(numberTypeRef),
+        )
+        public val KeySpecThenString: TypeRef = TypeRef(MapperPkg.Hl.Items, "thenString")
+
+        public fun keyType(keyTypes: List<TypeRef>): TypeRef {
+            val keySize = keyTypes.size
+            require(keySize in 1..4) { "KeyType subtypes must have between 1 and 4 keys, found $keySize" }
+            return TypeRef(MapperPkg.Hl.Items, "KeyType.Key$keySize", keyTypes)
+        }
 
-        public fun keySpec(keyType: TypeRef): TypeRef = TypeRef(MapperPkg.Hl.Items, "KeySpec", genericArgs = listOf(keyType))
-        public val KeySpecByteArray: TypeRef = TypeRef(MapperPkg.Hl.Items, "KeySpec.ByteArray")
-        public val KeySpecNumber: TypeRef = TypeRef(MapperPkg.Hl.Items, "KeySpec.Number")
-        public val KeySpecString: TypeRef = TypeRef(MapperPkg.Hl.Items, "KeySpec.String")
         public val AttributeDescriptor: TypeRef = TypeRef(MapperPkg.Hl.Items, "AttributeDescriptor")
 
         public fun itemConverter(objectType: TypeRef): TypeRef =
@@ -68,16 +107,28 @@ public object MapperTypes {
 
     public object Model {
         public val intersectKeys: TypeRef = TypeRef(MapperPkg.Hl.Model, "intersectKeys")
-        public fun tablePartitionKey(objectType: TypeRef, pkType: TypeRef): TypeRef = TypeRef(
+
+        public fun tablePartitionKey(objectType: TypeRef, pkTypes: List<TypeRef>): TypeRef = TypeRef(
             MapperPkg.Hl.Model,
             "Table.PartitionKey",
-            genericArgs = listOf(objectType, pkType),
-        )
-        public fun tableCompositeKey(objectType: TypeRef, pkType: TypeRef, skType: TypeRef): TypeRef = TypeRef(
-            MapperPkg.Hl.Model,
-            "Table.CompositeKey",
-            genericArgs = listOf(objectType, pkType, skType),
+            genericArgs = listOf(objectType, Items.keyType(pkTypes)),
         )
+
+        public fun tableCompositeKey(
+            objectType: TypeRef,
+            pkTypes: List<TypeRef>,
+            skTypes: List<TypeRef>,
+        ): TypeRef {
+            require(pkTypes.size in 1..4) { "Partition keys must have between 1 and 4 attributes, found ${pkTypes.size}" }
+            require(skTypes.size in 1..4) { "Sort keys must have between 1 and 4 attributes, found ${skTypes.size}" }
+
+            return TypeRef(
+                MapperPkg.Hl.Model,
+                "Table.CompositeKey",
+                genericArgs = listOf(objectType, Items.keyType(pkTypes), Items.keyType(skTypes)),
+            )
+        }
+
         public val toItem: TypeRef = TypeRef(MapperPkg.Hl.Model, "toItem")
     }
 

hll/dynamodb-mapper/dynamodb-mapper-schema-codegen/src/main/kotlin/aws/sdk/kotlin/hll/dynamodbmapper/codegen/annotations/rendering/SchemaRenderer.kt

@@ -51,25 +51,22 @@ internal class SchemaRenderer(
     private val properties = classDeclaration
         .getAllProperties()
         .filterNot { it.modifiers.contains(Modifier.PRIVATE) || it.isAnnotationPresent(DynamoDbIgnore::class) }
+        .toList()
 
-    init {
-        check(properties.count { it.isPk } == 1) {
-            "Expected exactly one @DynamoDbPartitionKey annotation on a property"
-        }
-        check(properties.count { it.isSk } <= 1) {
-            "Expected at most one @DynamoDbSortKey annotation on a property"
+    private val partitionKeyProps = properties.filter { it.isPk }.also {
+        require(it.size in 1..4) {
+            "Expected between 1 and 4 properties annotated with @DynamoDbPartitionKey, found ${it.size}"
         }
     }
 
-    private val partitionKeyProp = properties.single { it.isPk }
-    private val partitionKeyName = partitionKeyProp
-        .getAnnotationsByType(DynamoDbAttribute::class)
-        .singleOrNull()?.name ?: partitionKeyProp.name
+    private val sortKeyProps = properties.filter { it.isSk }.also {
+        require(it.size in 0..4) {
+            "Expected between 0 and 4 properties annotated with @DynamoDbSortKey, found ${it.size}"
+        }
+    }
 
-    private val sortKeyProp = properties.singleOrNull { it.isSk }
-    private val sortKeyName = sortKeyProp
-        ?.getAnnotationsByType(DynamoDbAttribute::class)
-        ?.singleOrNull()?.name ?: sortKeyProp?.name
+    private val partitionKeyTypeRefs = partitionKeyProps.map { it.typeRef }
+    private val sortKeyTypeRefs = sortKeyProps.map { it.typeRef }
 
     /**
      * Skip rendering a class builder if:
@@ -283,52 +280,84 @@ internal class SchemaRenderer(
         }
 
     private fun renderSchema() {
-        val schemaType = if (sortKeyProp != null) {
-            MapperTypes.Items.itemSchemaCompositeKey(classType, partitionKeyProp.typeRef, sortKeyProp.typeRef)
+        val schemaType = if (sortKeyProps.isEmpty()) {
+            MapperTypes.Items.itemSchemaPartitionKey(classType, partitionKeyTypeRefs)
         } else {
-            MapperTypes.Items.itemSchemaPartitionKey(classType, partitionKeyProp.typeRef)
+            MapperTypes.Items.itemSchemaCompositeKey(classType, partitionKeyTypeRefs, sortKeyTypeRefs)
         }
 
         withBlock("#Lobject #L : #T {", "}", ctx.attributes.visibility, schemaName, schemaType) {
             write("override val converter: #1T = #1T", itemConverter)
-            write("override val partitionKey: #T = #T(#S)", MapperTypes.Items.keySpec(partitionKeyProp.keySpec), partitionKeyProp.keySpecType, partitionKeyName)
-            if (sortKeyProp != null) {
-                write("override val sortKey: #T = #T(#S)", MapperTypes.Items.keySpec(sortKeyProp.keySpec), sortKeyProp.keySpecType, sortKeyName!!)
+
+            writeInline("override val partitionKey: #T = ", MapperTypes.Items.keySpec(partitionKeyTypeRefs))
+            keySpecInstantiation(partitionKeyProps)
+            write()
+
+            if (sortKeyProps.isNotEmpty()) {
+                writeInline("override val sortKey: #T = ", MapperTypes.Items.keySpec(sortKeyTypeRefs))
+                keySpecInstantiation(sortKeyProps)
+                write()
             }
         }
+
         blankLine()
     }
 
-    private val KSPropertyDeclaration.keySpec: TypeRef
-        get() = when (typeName) {
-            "kotlin.ByteArray" -> Types.Kotlin.ByteArray
-            "kotlin.Int" -> Types.Kotlin.Number
-            "kotlin.String" -> Types.Kotlin.String
-            else -> error("Unsupported key type $typeName, expected ByteArray, Int, or String")
+    private fun keySpecInstantiation(keyProps: List<KSPropertyDeclaration>) {
+        val first = keyProps.first()
+        val rest = keyProps.drop(1)
+
+        val firstTypeRef = when (first.typeRef) {
+            Types.Kotlin.ByteArray -> MapperTypes.Items.KeySpecByteArray
+
+            Types.Kotlin.Byte,
+            Types.Kotlin.Int,
+            Types.Kotlin.Long,
+            Types.Kotlin.Short,
+            -> MapperTypes.Items.keySpecNumber(first.typeRef)
+
+            Types.Kotlin.String -> MapperTypes.Items.KeySpecString
+
+            else -> error("Unsupported key attribute type ${first.typeRef}")
         }
 
-    private val KSPropertyDeclaration.keySpecType: TypeRef
-        get() = when (typeName) {
-            "kotlin.ByteArray" -> MapperTypes.Items.KeySpecByteArray
-            "kotlin.Int" -> MapperTypes.Items.KeySpecNumber
-            "kotlin.String" -> MapperTypes.Items.KeySpecString
-            else -> error("Unsupported key type $typeName, expected ByteArray, Int, or String")
+        writeInline("#T(#S)", firstTypeRef, first.ddbName)
+
+        rest.forEach { prop ->
+            val typeRef = when (prop.typeRef) {
+                Types.Kotlin.ByteArray -> MapperTypes.Items.KeySpecThenByteArray
+
+                Types.Kotlin.Byte,
+                Types.Kotlin.Int,
+                Types.Kotlin.Long,
+                Types.Kotlin.Short,
+                -> MapperTypes.Items.keySpecThenNumber(prop.typeRef)
+
+                Types.Kotlin.String -> MapperTypes.Items.KeySpecThenString
+
+                else -> error("Unsupported key attribute type ${prop.typeRef}")
+            }
+
+            writeInline(".#T(#S)", typeRef, prop.ddbName)
         }
+    }
 
     private fun renderGetTable() {
         docs("Returns a reference to a table named [name] containing items representing [#T]", classType)
 
+        val tableType = if (sortKeyProps.isEmpty()) {
+            MapperTypes.Model.tablePartitionKey(classType, partitionKeyTypeRefs)
+        } else {
+            MapperTypes.Model.tableCompositeKey(classType, partitionKeyTypeRefs, sortKeyTypeRefs)
+        }
+
         val fnName = "get${className}Table"
         write(
             "#Lfun #T.#L(name: String): #T = #L(name, #L)",
             ctx.attributes.visibility,
             MapperTypes.DynamoDbMapper,
             fnName,
-            if (sortKeyProp != null) {
-                MapperTypes.Model.tableCompositeKey(classType, partitionKeyProp.typeRef, sortKeyProp.typeRef)
-            } else {
-                MapperTypes.Model.tablePartitionKey(classType, partitionKeyProp.typeRef)
-            },
+            tableType,
             "getTable",
             schemaName,
         )