Expose query metadata in PostgresRowSequence

Author: MahdiBM

Sources/PostgresNIO/New/PSQLRowStream.swift

@@ -44,7 +44,7 @@ final class PSQLRowStream: @unchecked Sendable {
     }
 
     internal let rowDescription: [RowDescription.Column]
-    private let lookupTable: [String: Int]
+    internal let lookupTable: [String: Int]
     private var downstreamState: DownstreamState
 
     init(
@@ -114,7 +114,7 @@ final class PSQLRowStream: @unchecked Sendable {
             self.downstreamState = .consumed(.failure(error))
         }
 
-        return PostgresRowSequence(producer.sequence, lookupTable: self.lookupTable, columns: self.rowDescription)
+        return PostgresRowSequence(producer.sequence, stream: self)
     }
 
     func demand() {

Sources/PostgresNIO/New/PostgresRowSequence.swift

@@ -9,14 +9,18 @@ public struct PostgresRowSequence: AsyncSequence, Sendable {
 
     typealias BackingSequence = NIOThrowingAsyncSequenceProducer<DataRow, Error, AdaptiveRowBuffer, PSQLRowStream>
 
-    let backing: BackingSequence
-    let lookupTable: [String: Int]
-    let columns: [RowDescription.Column]
+    private let backing: BackingSequence
+    private let stream: PSQLRowStream
+    var lookupTable: [String: Int] {
+        self.stream.lookupTable
+    }
+    var columns: [RowDescription.Column] {
+        self.stream.rowDescription
+    }
 
-    init(_ backing: BackingSequence, lookupTable: [String: Int], columns: [RowDescription.Column]) {
+    init(_ backing: BackingSequence, stream: PSQLRowStream) {
         self.backing = backing
-        self.lookupTable = lookupTable
-        self.columns = columns
+        self.stream = stream
     }
 
     public func makeAsyncIterator() -> AsyncIterator {
@@ -67,6 +71,12 @@ extension PostgresRowSequence {
         }
         return result
     }
+
+    public func collectWithMetadata() async throws -> (metadata: PostgresQueryMetadata?, rows: [PostgresRow]) {
+        let rows = try await self.collect()
+        let metadata = PostgresQueryMetadata(string: self.stream.commandTag)
+        return (metadata, rows)
+    }
 }
 
 struct AdaptiveRowBuffer: NIOAsyncSequenceProducerBackPressureStrategy {

Tests/IntegrationTests/AsyncTests.swift

@@ -46,6 +46,31 @@ final class AsyncPostgresConnectionTests: XCTestCase {
         }
     }
 
+    func testSelect10kRowsAndCollect() async throws {
+        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
+        let eventLoop = eventLoopGroup.next()
+
+        let start = 1
+        let end = 10000
+
+        try await withTestConnection(on: eventLoop) { connection in
+            let rows = try await connection.query("SELECT generate_series(\(start), \(end));", logger: .psqlTest)
+            let (metadata, elements) = try await rows.collectWithMetadata()
+            var counter = 0
+            for row in elements {
+                let element = try row.decode(Int.self)
+                XCTAssertEqual(element, counter + 1)
+                counter += 1
+            }
+            XCTAssertEqual(metadata?.command, "SELECT")
+            XCTAssertEqual(metadata?.oid, nil)
+            XCTAssertEqual(metadata?.rows, 10000)
+
+            XCTAssertEqual(counter, end)
+        }
+    }
+
     func testSelectActiveConnection() async throws {
         let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
         defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
@@ -207,7 +232,7 @@ final class AsyncPostgresConnectionTests: XCTestCase {
 
         try await withTestConnection(on: eventLoop) { connection in
             // Max binds limit is UInt16.max which is 65535 which is 3 * 5 * 17 * 257
-            // Max columns limit is 1664, so we will only make 5 * 257 columns which is less
+            // Max columns limit appears to be ~1600, so we will only make 5 * 257 columns which is less
             // Then we will insert 3 * 17 rows
             // In the insertion, there will be a total of 3 * 17 * 5 * 257 == UInt16.max bindings
             // If the test is successful, it means Postgres supports UInt16.max bindings
@@ -241,13 +266,9 @@ final class AsyncPostgresConnectionTests: XCTestCase {
                 unsafeSQL: "INSERT INTO table1 VALUES \(insertionValues)",
                 binds: binds
             )
-            try await connection.query(insertionQuery, logger: .psqlTest)
-
-            let countQuery = PostgresQuery(unsafeSQL: "SELECT COUNT(*) FROM table1")
-            let countRows = try await connection.query(countQuery, logger: .psqlTest)
-            var countIterator = countRows.makeAsyncIterator()
-            let insertedRowsCount = try await countIterator.next()?.decode(Int.self, context: .default)
-            XCTAssertEqual(rowsCount, insertedRowsCount)
+            let result = try await connection.query(insertionQuery, logger: .psqlTest)
+            let metadata = try await result.collectWithMetadata().metadata
+            XCTAssertEqual(metadata?.rows, rowsCount)
 
             let dropQuery = PostgresQuery(unsafeSQL: "DROP TABLE table1")
             try await connection.query(dropQuery, logger: .psqlTest)