Add withTransaction API (#519)

Co-authored-by: Fabian Fett <[email protected]>

Author: thoven87

Diff for: Sources/PostgresNIO/Pool/PostgresClient.swift

@@ -307,6 +307,28 @@ public final class PostgresClient: Sendable, ServiceLifecycle.Service {
 
         return try await closure(connection)
     }
+    
+    /// Lease a connection for the provided `closure`'s lifetime.
+    /// A transation starts with call to withConnection
+    /// A transaction should end with a call to COMMIT or ROLLBACK
+    /// COMMIT is called upon successful completion and ROLLBACK is called should any steps fail
+    ///
+    /// - Parameter closure: A closure that uses the passed `PostgresConnection`. The closure **must not** capture
+    ///                      the provided `PostgresConnection`.
+    /// - Returns: The closure's return value.
+    public func withTransaction<Result>(_ process: (PostgresConnection) async throws -> Result) async throws -> Result {
+        try await withConnection { connection in
+            try await connection.query("BEGIN;", logger: self.backgroundLogger)
+            do {
+                let value = try await process(connection)
+                try await connection.query("COMMIT;", logger: self.backgroundLogger)
+                return value
+            } catch {
+                try await connection.query("ROLLBACK;", logger: self.backgroundLogger)
+                throw error
+            }
+        }
+    }
 
     /// Run a query on the Postgres server the client is connected to.
     ///

Diff for: Tests/IntegrationTests/PostgresClientTests.swift

@@ -42,6 +42,110 @@ final class PostgresClientTests: XCTestCase {
             taskGroup.cancelAll()
         }
     }
+    
+    func testTransaction() async throws {
+        var mlogger = Logger(label: "test")
+        mlogger.logLevel = .debug
+        let logger = mlogger
+        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 8)
+        self.addTeardownBlock {
+            try await eventLoopGroup.shutdownGracefully()
+        }
+        
+        let tableName = "test_client_transactions"
+        
+        let clientConfig = PostgresClient.Configuration.makeTestConfiguration()
+        let client = PostgresClient(configuration: clientConfig, eventLoopGroup: eventLoopGroup, backgroundLogger: logger)
+        
+        do {
+            try await withThrowingTaskGroup(of: Void.self) { taskGroup in
+                taskGroup.addTask {
+                    await client.run()
+                }
+                
+                try await client.query(
+                    """
+                    CREATE TABLE IF NOT EXISTS "\(unescaped: tableName)" (
+                        id INT PRIMARY KEY GENERATED ALWAYS AS IDENTITY,
+                        uuid UUID NOT NULL
+                    );
+                    """,
+                    logger: logger
+                )
+                
+                let iterations = 1000
+                
+                for _ in 0..<iterations {
+                    taskGroup.addTask {
+                        let _ = try await client.withTransaction { transaction in
+                            try await transaction.query(
+                            """
+                            INSERT INTO "\(unescaped: tableName)" (uuid) VALUES (\(UUID()));
+                            """,
+                            logger: logger
+                            )
+                        }
+                    }
+                }
+                
+                for _ in 0..<iterations {
+                    _ = await taskGroup.nextResult()!
+                }
+                
+                let rows = try await client.query(#"SELECT COUNT(1)::INT AS table_size FROM "\#(unescaped: tableName)";"#, logger: logger).decode(Int.self)
+                for try await (count) in rows {
+                    XCTAssertEqual(count, iterations)
+                }
+                
+                /// Test roll back
+                taskGroup.addTask {
+                    
+                    do {
+                        let _ = try await client.withTransaction { transaction in
+                            /// insert valid data
+                            try await transaction.query(
+                                """
+                                INSERT INTO "\(unescaped: tableName)" (uuid) VALUES (\(UUID()));
+                                """,
+                                logger: logger
+                            )
+                            
+                            /// insert invalid data
+                            try await transaction.query(
+                                """
+                                INSERT INTO "\(unescaped: tableName)" (uuid) VALUES (\(iterations));
+                                """,
+                                logger: logger
+                            )
+                        }
+                    } catch {
+                        XCTAssertNotNil(error)
+                        guard let error = error as? PSQLError else { return XCTFail("Unexpected error type") }
+                    
+                        XCTAssertEqual(error.code, .server)
+                        XCTAssertEqual(error.serverInfo?[.severity], "ERROR")
+                    }
+                }
+                
+                let row = try await client.query(#"SELECT COUNT(1)::INT AS table_size FROM "\#(unescaped: tableName)";"#, logger: logger).decode(Int.self)
+                
+                for try await (count) in row {
+                    XCTAssertEqual(count, iterations)
+                }
+                
+                try await client.query(
+                    """
+                    DROP TABLE "\(unescaped: tableName)";
+                    """,
+                    logger: logger
+                )
+                
+                taskGroup.cancelAll()
+            }
+        } catch {
+            XCTFail("Unexpected error: \(String(reflecting: error))")
+        }
+    }
 
     func testApplicationNameIsForwardedCorrectly() async throws {
         var mlogger = Logger(label: "test")