add a test, add a func in PostgresClient + better docs

Author: MahdiBM

Sources/PostgresNIO/Connection/PostgresConnection.swift

@@ -438,8 +438,15 @@ extension PostgresConnection {
         }
     }
 
-    // use this for queries where you want to consume the rows.
-    // we can use the `consume` scope to better ensure structured concurrency when consuming the rows.
+    /// Run a query on the Postgres server the connection is connected to, returning the metadata.
+    ///
+    /// - Parameters:
+    ///   - query: The ``PostgresQuery`` to run
+    ///   - logger: The `Logger` to log into for the query
+    ///   - file: The file, the query was started in. Used for better error reporting.
+    ///   - line: The line, the query was started in. Used for better error reporting.
+    ///   - consume: The closure to consume the ``PostgresRowSequence``.
+    /// - Returns: The result of the `consume` closure as well as the query metadata.
     public func query<Result>(
         _ query: PostgresQuery,
         logger: Logger,
@@ -573,8 +580,15 @@ extension PostgresConnection {
         }
     }
 
-    // use this for queries where you want to consume the rows.
-    // we can use the `consume` scope to better ensure structured concurrency when consuming the rows.
+    /// Execute a statement on the Postgres server the connection is connected to,
+    /// returning the metadata.
+    ///
+    /// - Parameters:
+    ///   - query: The ``PostgresQuery`` to run
+    ///   - logger: The `Logger` to log into for the query
+    ///   - file: The file, the query was started in. Used for better error reporting.
+    ///   - line: The line, the query was started in. Used for better error reporting.
+    /// - Returns: The query metadata.
     @discardableResult
     public func execute(
         _ query: PostgresQuery,

Sources/PostgresNIO/Pool/PostgresClient.swift

@@ -435,6 +435,60 @@ public final class PostgresClient: Sendable, ServiceLifecycle.Service {
         }
     }
 
+    /// Run a query on the Postgres server the connection is connected to, returning the metadata.
+    ///
+    /// - Parameters:
+    ///   - query: The ``PostgresQuery`` to run
+    ///   - logger: The `Logger` to log into for the query
+    ///   - file: The file, the query was started in. Used for better error reporting.
+    ///   - line: The line, the query was started in. Used for better error reporting.
+    ///   - consume: The closure to consume the ``PostgresRowSequence``.
+    /// - Returns: The result of the `consume` closure as well as the query metadata.
+    public func query<Result>(
+        _ query: PostgresQuery,
+        logger: Logger? = nil,
+        file: String = #fileID,
+        line: Int = #line,
+        _ consume: (PostgresRowSequence) async throws -> Result
+    ) async throws -> (Result, PostgresQueryMetadata) {
+        let logger = logger ?? Self.loggingDisabled
+
+        do {
+            guard query.binds.count <= Int(UInt16.max) else {
+                throw PSQLError(code: .tooManyParameters, query: query, file: file, line: line)
+            }
+
+            let connection = try await self.leaseConnection()
+
+            var logger = logger
+            logger[postgresMetadataKey: .connectionID] = "\(connection.id)"
+
+            let promise = connection.channel.eventLoop.makePromise(of: PSQLRowStream.self)
+            let context = ExtendedQueryContext(
+                query: query,
+                logger: logger,
+                promise: promise
+            )
+
+            connection.channel.write(HandlerTask.extendedQuery(context), promise: nil)
+
+            let (rowStream, rowSequence) = try await promise.futureResult.map { rowStream in
+                (rowStream, rowStream.asyncSequence(onFinish: { self.pool.releaseConnection(connection) }))
+            }.get()
+            let result = try await consume(rowSequence)
+            try await rowStream.drain().get()
+            guard let metadata = PostgresQueryMetadata(string: rowStream.commandTag) else {
+                throw PSQLError.invalidCommandTag(rowStream.commandTag)
+            }
+            return (result, metadata)
+        } catch var error as PSQLError {
+            error.file = file
+            error.line = line
+            error.query = query
+            throw error // rethrow with more metadata
+        }
+    }
+
     /// Execute a prepared statement, taking care of the preparation when necessary
     public func execute<Statement: PostgresPreparedStatement, Row>(
         _ preparedStatement: Statement,

Tests/IntegrationTests/PostgresClientTests.swift

@@ -181,7 +181,6 @@ final class PostgresClientTests: XCTestCase {
         }
     }
 
-
     func testQueryDirectly() async throws {
         var mlogger = Logger(label: "test")
         mlogger.logLevel = .debug
@@ -218,6 +217,47 @@ final class PostgresClientTests: XCTestCase {
         }
     }
 
+    func testQueryMetadataDirectly() 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 clientConfig = PostgresClient.Configuration.makeTestConfiguration()
+        let client = PostgresClient(configuration: clientConfig, eventLoopGroup: eventLoopGroup, backgroundLogger: logger)
+
+        await withThrowingTaskGroup(of: Void.self) { taskGroup in
+            taskGroup.addTask {
+                await client.run()
+            }
+
+            for i in 0..<10000 {
+                taskGroup.addTask {
+                    do {
+                        let (_, metadata) = try await client.query("SELECT 1", logger: logger) { _ in
+                            // Don't consume the row, the function itself should drain the row
+                        }
+                        XCTAssertEqual(metadata.command, "SELECT")
+                        XCTAssertNil(metadata.oid)
+                        XCTAssertEqual(metadata.rows, 1)
+                        logger.info("Success", metadata: ["run": "\(i)"])
+                    } catch {
+                        XCTFail("Unexpected error: \(error)")
+                    }
+                }
+            }
+
+            for _ in 0..<10000 {
+                _ = await taskGroup.nextResult()!
+            }
+
+            taskGroup.cancelAll()
+        }
+    }
+
     func testQueryTable() async throws {
         let tableName = "test_client_prepared_statement"
 

docker-compose.yml

@@ -1,5 +1,3 @@
-version: '3.7'
-
 x-shared-config: &shared_config
   environment:
     POSTGRES_HOST_AUTH_METHOD: "${POSTGRES_HOST_AUTH_METHOD:-scram-sha-256}"