move to the suggested way of doing things

MahdiBM · MahdiBM · commit 6c85b3e16774 · 2025-02-20T03:48:15.000+03:30
diff --git a/.gitignore b/.gitignore
@@ -1,5 +1,6 @@
 .DS_Store
 /.build
+/.index-build
 /Packages
 /*.xcodeproj
 DerivedData
diff --git a/Sources/PostgresNIO/Connection/PostgresConnection.swift b/Sources/PostgresNIO/Connection/PostgresConnection.swift
@@ -438,6 +438,43 @@ 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.
+    public func query<Result>(
+        _ query: PostgresQuery,
+        logger: Logger,
+        file: String = #fileID,
+        line: Int = #line,
+        _ consume: (PostgresRowSequence) async throws -> Result
+    ) async throws -> (Result, PostgresQueryMetadata) {
+        var logger = logger
+        logger[postgresMetadataKey: .connectionID] = "\(self.id)"
+
+        guard query.binds.count <= Int(UInt16.max) else {
+            throw PSQLError(code: .tooManyParameters, query: query, file: file, line: line)
+        }
+        let promise = self.channel.eventLoop.makePromise(of: PSQLRowStream.self)
+        let context = ExtendedQueryContext(
+            query: query,
+            logger: logger,
+            promise: promise
+        )
+
+        self.channel.write(HandlerTask.extendedQuery(context), promise: nil)
+
+        do {
+            let rowSequence = try await promise.futureResult.map({ $0.asyncSequence() }).get()
+            let result = try await consume(rowSequence)
+            let metadata = try await rowSequence.drainAndCollectMetadata()
+            return (result, metadata)
+        } catch var error as PSQLError {
+            error.file = file
+            error.line = line
+            error.query = query
+            throw error // rethrow with more metadata
+        }
+    }
+
     /// Start listening for a channel
     public func listen(_ channel: String) async throws -> PostgresNotificationSequence {
         let id = self.internalListenID.loadThenWrappingIncrement(ordering: .relaxed)
@@ -531,6 +568,42 @@ 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.
+    @discardableResult
+    public func execute(
+        _ query: PostgresQuery,
+        logger: Logger,
+        file: String = #fileID,
+        line: Int = #line
+    ) async throws -> PostgresQueryMetadata {
+        var logger = logger
+        logger[postgresMetadataKey: .connectionID] = "\(self.id)"
+
+        guard query.binds.count <= Int(UInt16.max) else {
+            throw PSQLError(code: .tooManyParameters, query: query, file: file, line: line)
+        }
+        let promise = self.channel.eventLoop.makePromise(of: PSQLRowStream.self)
+        let context = ExtendedQueryContext(
+            query: query,
+            logger: logger,
+            promise: promise
+        )
+
+        self.channel.write(HandlerTask.extendedQuery(context), promise: nil)
+
+        do {
+            let rowSequence = try await promise.futureResult.map({ $0.asyncSequence() }).get()
+            let metadata = try await rowSequence.drainAndCollectMetadata()
+            return metadata
+        } catch var error as PSQLError {
+            error.file = file
+            error.line = line
+            error.query = query
+            throw error // rethrow with more metadata
+        }
+    }
+
     #if compiler(>=6.0)
     /// Puts the connection into an open transaction state, for the provided `closure`'s lifetime.
     ///
diff --git a/Sources/PostgresNIO/New/PSQLRowStream.swift b/Sources/PostgresNIO/New/PSQLRowStream.swift
@@ -276,7 +276,63 @@ final class PSQLRowStream: @unchecked Sendable {
             return self.eventLoop.makeFailedFuture(error)
         }
     }
-    
+
+    // MARK: Drain on EventLoop
+
+    func drain() -> EventLoopFuture<Void> {
+        if self.eventLoop.inEventLoop {
+            return self.drain0()
+        } else {
+            return self.eventLoop.flatSubmit {
+                self.drain0()
+            }
+        }
+    }
+
+    private func drain0() -> EventLoopFuture<Void> {
+        self.eventLoop.preconditionInEventLoop()
+
+        switch self.downstreamState {
+        case .waitingForConsumer(let bufferState):
+            switch bufferState {
+            case .streaming(var buffer, let dataSource):
+                let promise = self.eventLoop.makePromise(of: Void.self)
+
+                buffer.removeAll()
+                self.downstreamState = .iteratingRows(onRow: { _ in }, promise, dataSource)
+                // immediately request more
+                dataSource.request(for: self)
+
+                return promise.futureResult
+
+            case .finished(_, let summary):
+                self.downstreamState = .consumed(.success(summary))
+                return self.eventLoop.makeSucceededVoidFuture()
+
+            case .failure(let error):
+                self.downstreamState = .consumed(.failure(error))
+                return self.eventLoop.makeFailedFuture(error)
+            }
+        case .asyncSequence(let consumer, let dataSource, _):
+            consumer.finish()
+
+            let promise = self.eventLoop.makePromise(of: Void.self)
+
+            self.downstreamState = .iteratingRows(onRow: { _ in }, promise, dataSource)
+            // immediately request more
+            dataSource.request(for: self)
+
+            return promise.futureResult
+        case .consumed(.success):
+            // already drained
+            return self.eventLoop.makeSucceededVoidFuture()
+        case .consumed(let .failure(error)):
+            return self.eventLoop.makeFailedFuture(error)
+        default:
+            preconditionFailure("Invalid state: \(self.downstreamState)")
+        }
+    }
+
     internal func noticeReceived(_ notice: PostgresBackendMessage.NoticeResponse) {
         self.logger.debug("Notice Received", metadata: [
             .notice: "\(notice)"
diff --git a/Sources/PostgresNIO/New/PostgresRowSequence.swift b/Sources/PostgresNIO/New/PostgresRowSequence.swift
@@ -30,6 +30,10 @@ public struct PostgresRowSequence: AsyncSequence, Sendable {
             columns: self.columns
         )
     }
+
+    func drain() {
+        self.backing
+    }
 }
 
 extension PostgresRowSequence {
@@ -64,43 +68,11 @@ extension PostgresRowSequence {
 extension PostgresRowSequence.AsyncIterator: Sendable {}
 
 extension PostgresRowSequence {
-    /// Collect and return all rows.
-    /// - Returns: The rows.
-    public func collect() async throws -> [PostgresRow] {
-        var result = [PostgresRow]()
-        for try await row in self {
-            result.append(row)
-        }
-        return result
-    }
-
-    /// Collect and return all rows, alongside the query metadata.
-    /// - Returns: The query metadata and the rows.
-    public func collectWithMetadata() async throws -> (metadata: PostgresQueryMetadata, rows: [PostgresRow]) {
-        let rows = try await self.collect()
-        guard let metadata = PostgresQueryMetadata(string: self.rowStream.commandTag) else {
-            throw PSQLError.invalidCommandTag(self.rowStream.commandTag)
-        }
-        return (metadata, rows)
-    }
-
-    /// Consumes all rows and returns the query metadata.
-    ///
-    /// If you don't need the returned query metadata, just use the for-try-await-loop syntax:
-    /// ```swift
-    /// for try await row in myRowSequence {
-    ///     /// Process each row
-    /// }
-    /// ```
-    ///
-    /// - Parameter onRow: Processes each row.
-    /// - Returns: The query metadata.
-    public func consume(
-        onRow: @Sendable (PostgresRow) throws -> ()
-    ) async throws -> PostgresQueryMetadata {
-        for try await row in self {
-            try onRow(row)
-        }
+    /// Collects the query metadata.
+    /// Should be called after the sequence is consumed, otherwise throws `PSQLError`.
+    /// - Returns: The metadata.
+    func drainAndCollectMetadata() async throws -> PostgresQueryMetadata {
+        try await self.rowStream.drain().get()
         guard let metadata = PostgresQueryMetadata(string: self.rowStream.commandTag) else {
             throw PSQLError.invalidCommandTag(self.rowStream.commandTag)
         }
diff --git a/Tests/IntegrationTests/AsyncTests.swift b/Tests/IntegrationTests/AsyncTests.swift
@@ -47,7 +47,7 @@ final class AsyncPostgresConnectionTests: XCTestCase {
         }
     }
 
-    func testSelect10kRowsAndConsume() async throws {
+    func testSelect10kRowsWithMetadata() async throws {
         let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
         defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
         let eventLoop = eventLoopGroup.next()
@@ -56,24 +56,28 @@ final class AsyncPostgresConnectionTests: XCTestCase {
         let end = 10000
 
         try await withTestConnection(on: eventLoop) { connection in
-            let rows = try await connection.query("SELECT generate_series(\(start), \(end));", logger: .psqlTest)
-
-            let counter = ManagedAtomic(0)
-            let metadata = try await rows.consume { row in
-                let element = try row.decode(Int.self)
-                let newCounter = counter.wrappingIncrementThenLoad(ordering: .relaxed)
-                XCTAssertEqual(element, newCounter)
+            let (result, metadata) = try await connection.query(
+                "SELECT generate_series(\(start), \(end));",
+                logger: .psqlTest
+            ) { rows in
+                var counter = 0
+                for try await row in rows {
+                    let element = try row.decode(Int.self)
+                    XCTAssertEqual(element, counter + 1)
+                    counter += 1
+                }
+                return counter
             }
 
             XCTAssertEqual(metadata.command, "SELECT")
             XCTAssertEqual(metadata.oid, nil)
-            XCTAssertEqual(metadata.rows, 10000)
+            XCTAssertEqual(metadata.rows, end)
 
-            XCTAssertEqual(counter.load(ordering: .relaxed), end)
+            XCTAssertEqual(result, end)
         }
     }
 
-    func testSelect10kRowsAndCollect() async throws {
+    func testSelectRowsWithMetadataNotConsumedAtAll() async throws {
         let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
         defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
         let eventLoop = eventLoopGroup.next()
@@ -82,20 +86,56 @@ final class AsyncPostgresConnectionTests: XCTestCase {
         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
-            }
+            let (_, metadata) = try await connection.query(
+                "SELECT generate_series(\(start), \(end));",
+                logger: .psqlTest
+            ) { _ in }
 
             XCTAssertEqual(metadata.command, "SELECT")
             XCTAssertEqual(metadata.oid, nil)
-            XCTAssertEqual(metadata.rows, 10000)
+            XCTAssertEqual(metadata.rows, end)
+        }
+    }
 
-            XCTAssertEqual(counter, end)
+    func testSelectRowsWithMetadataNotFullyConsumed() async throws {
+        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
+        let eventLoop = eventLoopGroup.next()
+
+        try await withTestConnection(on: eventLoop) { connection in
+            do {
+                _ = try await connection.query(
+                    "SELECT generate_series(1, 10000);",
+                    logger: .psqlTest
+                ) { rows in
+                    for try await _ in rows { break }
+                }
+                XCTFail("Expected a failure")
+            } catch is CancellationError {
+                // Expected
+            } catch {
+                XCTFail("Expected 'CancellationError', got: \(String(reflecting: error))")
+            }
+        }
+    }
+
+    func testExecuteRowsWithMetadata() 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 metadata = try await connection.execute(
+                "SELECT generate_series(\(start), \(end));",
+                logger: .psqlTest
+            )
+
+            XCTAssertEqual(metadata.command, "SELECT")
+            XCTAssertEqual(metadata.oid, nil)
+            XCTAssertEqual(metadata.rows, end)
         }
     }
 
@@ -294,8 +334,7 @@ final class AsyncPostgresConnectionTests: XCTestCase {
                 unsafeSQL: "INSERT INTO table1 VALUES \(insertionValues)",
                 binds: binds
             )
-            let result = try await connection.query(insertionQuery, logger: .psqlTest)
-            let metadata = try await result.collectWithMetadata().metadata
+            let metadata = try await connection.execute(insertionQuery, logger: .psqlTest)
             XCTAssertEqual(metadata.rows, rowsCount)
 
             let dropQuery = PostgresQuery(unsafeSQL: "DROP TABLE table1")
diff --git a/docker-compose.yml b/docker-compose.yml
@@ -10,6 +10,9 @@ x-shared-config: &shared_config
     - 5432:5432
 
 services:
+  psql-17:
+    image: postgres:17
+    <<: *shared_config
   psql-16:
     image: postgres:16
     <<: *shared_config
