Adds prepared statement support to client (#459)

Author: fabianfett

Sources/PostgresNIO/Pool/PostgresClient.swift

@@ -342,6 +342,48 @@ public final class PostgresClient: Sendable {
         }
     }
 
+    /// Execute a prepared statement, taking care of the preparation when necessary
+    public func execute<Statement: PostgresPreparedStatement, Row>(
+        _ preparedStatement: Statement,
+        logger: Logger,
+        file: String = #fileID,
+        line: Int = #line
+    ) async throws -> AsyncThrowingMapSequence<PostgresRowSequence, Row> where Row == Statement.Row {
+        let bindings = try preparedStatement.makeBindings()
+
+        do {
+            let connection = try await self.leaseConnection()
+
+            let promise = connection.channel.eventLoop.makePromise(of: PSQLRowStream.self)
+            let task = HandlerTask.executePreparedStatement(.init(
+                name: String(reflecting: Statement.self),
+                sql: Statement.sql,
+                bindings: bindings,
+                bindingDataTypes: Statement.bindingDataTypes,
+                logger: logger,
+                promise: promise
+            ))
+            connection.channel.write(task, promise: nil)
+
+            promise.futureResult.whenFailure { _ in
+                self.pool.releaseConnection(connection)
+            }
+
+            return try await promise.futureResult
+                .map { $0.asyncSequence(onFinish: { self.pool.releaseConnection(connection) }) }
+                .get()
+                .map { try preparedStatement.decodeRow($0) }
+        } catch var error as PSQLError {
+            error.file = file
+            error.line = line
+            error.query = .init(
+                unsafeSQL: Statement.sql,
+                binds: bindings
+            )
+            throw error // rethrow with more metadata
+        }
+    }
+
     /// The client's run method. Users must call this function in order to start the client's background task processing
     /// like creating and destroying connections and running timers. 
     ///

Tests/IntegrationTests/PostgresClientTests.swift

@@ -25,16 +25,17 @@ final class PostgresClientTests: XCTestCase {
                 await client.run()
             }
 
-            for i in 0..<10000 {
+            let iterations = 1000
+
+            for i in 0..<iterations {
                 taskGroup.addTask {
                     try await client.withConnection() { connection in
                         _ = try await connection.query("SELECT 1", logger: logger)
                     }
-                    print("done: \(i)")
                 }
             }
 
-            for _ in 0..<10000 {
+            for _ in 0..<iterations {
                 _ = await taskGroup.nextResult()!
             }
 
@@ -78,6 +79,80 @@ final class PostgresClientTests: XCTestCase {
         }
     }
 
+    func testQueryTable() async throws {
+        let tableName = "test_client_prepared_statement"
+
+        var mlogger = Logger(label: "test")
+        mlogger.logLevel = .debug
+        let logger = mlogger
+        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 8)
+        self.addTeardownBlock {
+            try await eventLoopGroup.shutdownGracefully()
+        }
+
+        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 SERIAL PRIMARY KEY,
+                        uuid UUID NOT NULL
+                    );
+                    """,
+                    logger: logger
+                )
+
+                for _ in 0..<1000 {
+                    try await client.query(
+                        """
+                        INSERT INTO "\(unescaped: tableName)" (uuid) VALUES (\(UUID()));
+                        """,
+                        logger: logger
+                    )
+                }
+
+                let rows = try await client.query(#"SELECT id, uuid FROM "\#(unescaped: tableName)";"#, logger: logger).decode((Int, UUID).self)
+                for try await (id, uuid) in rows {
+                    logger.info("id: \(id), uuid: \(uuid.uuidString)")
+                }
+
+                struct Example: PostgresPreparedStatement {
+                    static let sql = "SELECT id, uuid FROM test_client_prepared_statement WHERE id < $1"
+                    typealias Row = (Int, UUID)
+                    var id: Int
+                    func makeBindings() -> PostgresBindings {
+                        var bindings = PostgresBindings()
+                        bindings.append(self.id)
+                        return bindings
+                    }
+                    func decodeRow(_ row: PostgresNIO.PostgresRow) throws -> Row {
+                        try row.decode(Row.self)
+                    }
+                }
+
+                for try await (id, uuid) in try await client.execute(Example(id: 200), logger: logger) {
+                    logger.info("id: \(id), uuid: \(uuid.uuidString)")
+                }
+
+                try await client.query(
+                    """
+                    DROP TABLE "\(unescaped: tableName)";
+                    """,
+                    logger: logger
+                )
+
+                taskGroup.cancelAll()
+            }
+        } catch {
+            XCTFail("Unexpected error: \(String(reflecting: error))")
+        }
+    }
 }
 
 @available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *)

Tests/PostgresNIOTests/New/PostgresConnectionTests.swift

@@ -155,7 +155,7 @@ class PostgresConnectionTests: XCTestCase {
                     _ = try await iterator.next()
                     XCTFail("Did not expect to not throw")
                 } catch {
-                    print(error)
+                    self.logger.error("error", metadata: ["error": "\(error)"])
                 }
             }