Update Persistence with attributeModifier Member

Sources/Persistence/Persistence.swift

@@ -16,23 +16,53 @@ public typealias AttributeValue = DynamoDBClientTypes.AttributeValue
 
 /// A type that persists collections of attributes.
 public struct Persistence {
+    /// A closure to modify the attributes of persisted values.
+    ///
+    /// Use this to add additional attributes like a timestamp or to perform validation of all
+    /// persisted values.
+    var attributeModifier: @Sendable ([String: AttributeValue]) throws -> [String: AttributeValue]
+
     /// A closure to create a new item or replace an old item with a new one.
     var put: @Sendable ([String: AttributeValue]) async throws -> Void
 
     /// Creates an instance.
     ///
-    /// - Parameter put: A closure to persist item attributes.
+    /// - Parameters:
+    ///   - put: A closure to persist item attributes.
+    ///   - attributeModifier: A closure to modify the attributes of all persisted values.
     public init(
-        put: @escaping @Sendable ([String: AttributeValue]) async throws -> Void
+        put: @escaping @Sendable ([String: AttributeValue]) async throws -> Void,
+        attributeModifier: @escaping @Sendable ([String: AttributeValue]) throws -> [String: AttributeValue] = { $0 }
     ) {
         self.put = put
+        self.attributeModifier = attributeModifier
     }
 
     /// Persists the given value.
     ///
     /// - Parameter contact: The value to persist.
     public func put<T: AttributeValueConvertible>(_ value: T) async throws {
-        try await put(value.attributes)
+        try await put(try attributeModifier(value.attributes))
+    }
+}
+
+public extension Persistence {
+    /// Returns an instance that adds a `CreatedAt` timestamp from the given `TimeStampProvider` to
+    /// all persisted entities.
+    ///
+    /// - Parameters:
+    ///   - timestampName: The name of the timestamp attribute.
+    ///   - timestampProvider: A timestamp provider.
+    ///   - put: A closure to persist item attributes.
+    /// - Returns: A `Persistence` instance.
+    static func addingTimestamp(
+        named timestampName: String,
+        from timestampProvider: TimestampProvider,
+        put: @escaping @Sendable ([String: AttributeValue]) async throws -> Void
+    ) -> Self {
+        .init(
+            put: put,
+            attributeModifier: { $0.merging([timestampName: .s(timestampProvider.timestamp())]) { _, new in new } })
     }
 }
 
@@ -60,27 +90,15 @@ public struct PersistenceFactory {
 public extension PersistenceFactory {
     /// Returns a live implementation.
     ///
-    /// - Parameter attributeProvider: An optional closure returning additional attributes to any
-    /// values that are persisted.
+    /// - Parameter attributeModifier: A closure to modify the attributes of persisted values.
     static func live(
-        additionalAttributes attributeProvider: (() throws -> [String: AttributeValue])? = nil
+        attributeModifier: @escaping @Sendable ([String: AttributeValue]) throws -> [String: AttributeValue] = { $0 }
     ) -> Self {
         .init(make: { region, tableName in
             let dbClient = try DynamoDBClient(region: region)
-
-            let inputProvider: ([String: AttributeValue]) throws -> PutItemInput
-            if let attributeProvider {
-                inputProvider = { attributes in
-                    PutItemInput(
-                        item: try attributes.merging(attributeProvider()) { _, new in new },
-                        tableName: tableName
-                    )
-                }
-            } else {
-                inputProvider = { PutItemInput(item: $0, tableName: tableName) }
-            }
-
-            return Persistence(put: { _ = try await dbClient.putItem(input: inputProvider($0)) })
+            return Persistence(
+                put: { _ = try await dbClient.putItem(input: PutItemInput(item: $0, tableName: tableName)) },
+                attributeModifier: attributeModifier)
         })
     }
 }

Tests/PersistenceTests/PersistenceTests.swift

@@ -10,7 +10,40 @@ import Foundation
 import XCTest
 
 final class PersistenceTests: XCTestCase {
+
+    struct Model: Codable, Equatable, AttributeValueConvertible {
+        let bool: Bool
+        let string: String
+
+        var attributes: [String: AttributeValue] {
+            let values: [CodingKeys: AttributeValue] = [
+                .bool: .bool(bool),
+                .string: .s(string)
+            ]
+            return values.attributeValues()
+        }
+    }
+
+    let timeoutInterval: TimeInterval =  0.1
+
     func testPersistence() async throws {
-        // ...
+        let expected: [String: AttributeValue] = [
+            "Bool": .bool(true),
+            "CreatedAt": .s("1970-01-01T00:00:00Z"),
+            "String": .s("string")
+        ]
+
+        let timestampProvider = TimestampProvider(
+            dateProvider: { Date(timeIntervalSince1970: 0) },
+            formatter: ISO8601DateFormatter())
+
+        let expectation = expectation(description: "Model persisted")
+        let sut = Persistence.addingTimestamp(named: "CreatedAt", from: timestampProvider) { actual in
+            XCTAssertEqual(actual, expected)
+            expectation.fulfill()
+        }
+
+        try await sut.put(Model(bool: true, string: "string"))
+        await fulfillment(of: [expectation], timeout: timeoutInterval)
     }
 }