Merge pull request #100 from amzn/haotzhu

Support onTransactionWrite with GSI key in InMemoryDynamoDBCompositePrimaryKeyTableWithIndex

Sources/SmokeDynamoDB/DynamoDBCompositePrimaryKeyGSILogic.swift

@@ -20,12 +20,19 @@ import SmokeHTTPClient
 import DynamoDBModel
 import NIO
 
+// Provide a default `PolymorphicWriteEntry` for the `DynamoDBCompositePrimaryKeyGSILogic` for backwards compatibility
+public struct NoOpPolymorphicWriteEntry: PolymorphicWriteEntry {
+    public func handle<Context>(context: Context) throws -> Context.WriteEntryTransformType where Context : PolymorphicWriteEntryContext {
+        fatalError("Unimplemented")
+    }
+}
+
 /**
   A protocol that simulates the logic of a GSI reacting to events on the main table.
  */
 public protocol DynamoDBCompositePrimaryKeyGSILogic {
     associatedtype GSIAttributesType: PrimaryKeyAttributes
-
+    associatedtype WriteEntryType: PolymorphicWriteEntry = NoOpPolymorphicWriteEntry
     /**
      * Called when an item is inserted on the main table. Can be used to transform the provided item to the item that would be made available on the GSI.
      */
@@ -78,6 +85,13 @@ public protocol DynamoDBCompositePrimaryKeyGSILogic {
      */
     func onDeleteItem<AttributesType>(forKey key: CompositePrimaryKey<AttributesType>,
                                       gsiDataStore: InMemoryDynamoDBCompositePrimaryKeyTable) async throws
+
+    /**
+     * Called when an transact write in the main table. Can be used to also transact write the corresponding item on the GSI.
+
+     */
+    func onTransactWrite(_ entries: [WriteEntryType],
+                         gsiDataStore: InMemoryDynamoDBCompositePrimaryKeyTable) async throws
 #endif
 }
 
@@ -114,6 +128,11 @@ public extension DynamoDBCompositePrimaryKeyGSILogic {
                                       gsiDataStore: InMemoryDynamoDBCompositePrimaryKeyTable) async throws {
         try await onDeleteItem(forKey: key, gsiDataStore: gsiDataStore).get()
     }
+
+    // provide default for backwards compatibility
+    func onTransactWrite(_ entries: [WriteEntryType],
+                         gsiDataStore: InMemoryDynamoDBCompositePrimaryKeyTable) async throws {
+        // do nothing
+    }
 #endif
 }
-

Sources/SmokeDynamoDB/InMemoryDynamoDBCompositePrimaryKeyTableWithIndex.swift

@@ -1,4 +1,4 @@
-// swiftlint:disable cyclomatic_complexity
+// swiftlint:disable cyclomatic_complexity type_body_length
 // Copyright 2018-2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License").
@@ -76,8 +76,17 @@ public struct InMemoryDynamoDBCompositePrimaryKeyTableWithIndex<GSILogic: Dynamo
         }
     }
 
-    public func transactWrite<WriteEntryType: PolymorphicWriteEntry>(_ entries: [WriteEntryType]) async throws {
-        return try await self.primaryTable.transactWrite(entries)
+    public func transactWrite<WriteEntryType>(_ entries: [WriteEntryType]) async throws where WriteEntryType : PolymorphicWriteEntry {
+        let transformedEntries = entries.map { entry -> GSILogic.WriteEntryType in
+            guard let transformedEntry = entry as? GSILogic.WriteEntryType else {
+                fatalError("Unable to transform transactWrite of type \(type(of: entry)) to \(GSILogic.WriteEntryType.self)")
+            }
+
+            return transformedEntry
+        }
+
+        try await self.primaryTable.transactWrite(transformedEntries)
+        try await self.gsiLogic.onTransactWrite(transformedEntries, gsiDataStore: self.gsiDataStore)
     }
 
     public func transactWrite<WriteEntryType: PolymorphicWriteEntry,