feat: add support for DDB Mapper getItem overloads that specify primary key(s) (#1583)

Author: ianbotsf

.changes/af8e4342-f541-4d83-88a1-2570461d7e7e.json

@@ -0,0 +1,8 @@
+{
+    "id": "af8e4342-f541-4d83-88a1-2570461d7e7e",
+    "type": "feature",
+    "description": "Add support for DynamoDB Mapper `getItem` overloads that specify primary key(s)",
+    "issues": [
+        "awslabs/aws-sdk-kotlin#1577"
+    ]
+}

hll/dynamodb-mapper/dynamodb-mapper/common/src/aws/sdk/kotlin/hll/dynamodbmapper/model/internal/TableImpl.kt

@@ -9,8 +9,13 @@ import aws.sdk.kotlin.hll.dynamodbmapper.items.ItemSchema
 import aws.sdk.kotlin.hll.dynamodbmapper.model.Index
 import aws.sdk.kotlin.hll.dynamodbmapper.model.Table
 import aws.sdk.kotlin.hll.dynamodbmapper.model.TableSpec
-import aws.sdk.kotlin.hll.dynamodbmapper.operations.TableOperations
-import aws.sdk.kotlin.hll.dynamodbmapper.operations.TableOperationsImpl
+import aws.sdk.kotlin.hll.dynamodbmapper.model.itemOf
+import aws.sdk.kotlin.hll.dynamodbmapper.operations.*
+import aws.sdk.kotlin.hll.dynamodbmapper.pipeline.Interceptor
+import aws.sdk.kotlin.hll.dynamodbmapper.pipeline.LReqContext
+import aws.sdk.kotlin.services.dynamodb.model.AttributeValue
+import aws.sdk.kotlin.services.dynamodb.model.GetItemRequest as LowLevelGetItemRequest
+import aws.sdk.kotlin.services.dynamodb.model.GetItemResponse as LowLevelGetItemResponse
 
 internal fun <T, PK> tableImpl(
     mapper: DynamoDbMapper,
@@ -35,7 +40,17 @@ internal fun <T, PK> tableImpl(
             schema: ItemSchema.CompositeKey<T, PK, SK>,
         ): Index.CompositeKey<T, PK, SK> = indexImpl(mapper, tableName, name, schema)
 
-        override suspend fun getItem(partitionKey: PK) = TODO("not yet implemented")
+        override suspend fun getItem(partitionKey: PK): T? {
+            val keyItem = itemOf(schema.partitionKey.toField(partitionKey))
+            val interceptor = KeyInsertionInterceptor<T>(keyItem)
+            val op = getItemOperation(specImpl).let {
+                it.copy(
+                    interceptors = it.interceptors.prepend(interceptor),
+                )
+            }
+            val hRes = op.execute(GetItemRequest { })
+            return hRes.item
+        }
     }
 }
 
@@ -61,6 +76,36 @@ internal fun <T, PK, SK> tableImpl(
             schema: ItemSchema.CompositeKey<T, PK, SK>,
         ): Index.CompositeKey<T, PK, SK> = indexImpl(mapper, tableName, name, schema)
 
-        override suspend fun getItem(partitionKey: PK, sortKey: SK) = TODO("Not yet implemented")
+        override suspend fun getItem(partitionKey: PK, sortKey: SK): T? {
+            val keyItem = itemOf(
+                schema.partitionKey.toField(partitionKey),
+                schema.sortKey.toField(sortKey),
+            )
+            val interceptor = KeyInsertionInterceptor<T>(keyItem)
+            val op = getItemOperation(specImpl).let {
+                it.copy(
+                    interceptors = it.interceptors.prepend(interceptor),
+                )
+            }
+            val hRes = op.execute(GetItemRequest { })
+            return hRes.item
+        }
     }
 }
+
+private typealias GetItemInterceptor<T> =
+    Interceptor<T, GetItemRequest<T>, LowLevelGetItemRequest, LowLevelGetItemResponse, GetItemResponse<T>>
+
+private class KeyInsertionInterceptor<T>(private val newKey: Map<String, AttributeValue>) : GetItemInterceptor<T> {
+    override fun modifyBeforeInvocation(ctx: LReqContext<T, GetItemRequest<T>, LowLevelGetItemRequest>) =
+        ctx.lowLevelRequest.copy {
+            if (key == null) {
+                key = newKey
+            }
+        }
+}
+
+private fun <T> List<T>.prepend(element: T): List<T> = buildList(size + 1) {
+    add(element)
+    addAll(this)
+}

hll/dynamodb-mapper/dynamodb-mapper/common/src/aws/sdk/kotlin/hll/dynamodbmapper/pipeline/internal/Operation.kt

@@ -5,20 +5,33 @@
 package aws.sdk.kotlin.hll.dynamodbmapper.pipeline.internal
 
 import aws.sdk.kotlin.hll.dynamodbmapper.items.ItemSchema
-import aws.sdk.kotlin.hll.dynamodbmapper.pipeline.*
-
-internal class Operation<T, HReq, LReq, LRes, HRes>(
-    private val initialize: (HReq) -> HReqContextImpl<T, HReq>,
-    private val serialize: (HReq, ItemSchema<T>) -> LReq,
-    private val lowLevelInvoke: suspend (LReq) -> LRes,
-    private val deserialize: (LRes, ItemSchema<T>) -> HRes,
-    interceptors: Collection<InterceptorAny>,
+import aws.sdk.kotlin.hll.dynamodbmapper.pipeline.Interceptor
+import aws.sdk.kotlin.hll.dynamodbmapper.pipeline.InterceptorAny
+
+internal data class Operation<T, HReq, LReq, LRes, HRes>(
+    val initialize: (HReq) -> HReqContextImpl<T, HReq>,
+    val serialize: (HReq, ItemSchema<T>) -> LReq,
+    val lowLevelInvoke: suspend (LReq) -> LRes,
+    val deserialize: (LRes, ItemSchema<T>) -> HRes,
+    val interceptors: List<Interceptor<T, HReq, LReq, LRes, HRes>>,
 ) {
-    private val interceptors = interceptors.map {
-        // Will cause runtime ClassCastExceptions during interceptor invocation if the types don't match. Is that ok?
-        @Suppress("UNCHECKED_CAST")
-        it as Interceptor<T, HReq, LReq, LRes, HRes>
-    }
+    constructor(
+        initialize: (HReq) -> HReqContextImpl<T, HReq>,
+        serialize: (HReq, ItemSchema<T>) -> LReq,
+        lowLevelInvoke: suspend (LReq) -> LRes,
+        deserialize: (LRes, ItemSchema<T>) -> HRes,
+        interceptors: Collection<InterceptorAny>,
+    ) : this(
+        initialize,
+        serialize,
+        lowLevelInvoke,
+        deserialize,
+        interceptors.map {
+            // Will cause runtime ClassCastExceptions during interceptor invocation if the types don't match. Is that ok?
+            @Suppress("UNCHECKED_CAST")
+            it as Interceptor<T, HReq, LReq, LRes, HRes>
+        },
+    )
 
     suspend fun execute(hReq: HReq): HRes {
         val hReqContext = doInitialize(hReq)

hll/dynamodb-mapper/dynamodb-mapper/jvm/test/aws/sdk/kotlin/hll/dynamodbmapper/operations/GetItemTest.kt

@@ -125,4 +125,24 @@ class GetItemTest : DdbLocalTest() {
         val tableCapacity = assertNotNull(cc.table)
         assertEquals(0.5, tableCapacity.capacityUnits)
     }
+
+    @Test
+    fun testPkGetItemByScalarKey() = runTest {
+        val table = mapper().getTable(PK_TABLE_NAME, pkSchema)
+
+        val item = assertNotNull(table.getItem(1))
+        assertEquals("foo", item.value)
+
+        assertNull(table.getItem(2))
+    }
+
+    @Test
+    fun testCkGetItemByScalarKeys() = runTest {
+        val table = mapper().getTable(CK_TABLE_NAME, ckSchema)
+
+        val item = assertNotNull(table.getItem("abcd", 42))
+        assertEquals("foo", item.value)
+
+        assertNull(table.getItem("abcd", 43))
+    }
 }