more tests, more fixes

Author: MahdiBM

Sources/PostgresNIO/Connection/PostgresConnection.swift

@@ -446,6 +446,7 @@ extension PostgresConnection {
     ///   - 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``.
+    ///     DO NOT escape the row-sequence out of the closure.
     /// - Returns: The result of the `consume` closure as well as the query metadata.
     public func query<Result>(
         _ query: PostgresQuery,

Sources/PostgresNIO/New/PSQLRowStream.swift

@@ -328,6 +328,12 @@ final class PSQLRowStream: @unchecked Sendable {
             return self.eventLoop.makeSucceededVoidFuture()
         case .consumed(let .failure(error)):
             return self.eventLoop.makeFailedFuture(error)
+        case .waitingForAll(let rows, let promise, let dataSource):
+            self.downstreamState = .waitingForAll(rows, promise, dataSource)
+            // immediately request more
+            dataSource.request(for: self)
+
+            return promise.futureResult.map { _ in }
         default:
             preconditionFailure("Invalid state: \(self.downstreamState)")
         }

Sources/PostgresNIO/Pool/PostgresClient.swift

@@ -443,6 +443,7 @@ public final class PostgresClient: Sendable, ServiceLifecycle.Service {
     ///   - 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``.
+    ///     DO NOT escape the row-sequence out of the closure.
     /// - Returns: The result of the `consume` closure as well as the query metadata.
     public func query<Result>(
         _ query: PostgresQuery,

Tests/PostgresNIOTests/New/PSQLRowStreamTests.swift

@@ -317,6 +317,104 @@ final class PSQLRowStreamTests: XCTestCase {
         XCTAssertNoThrow(try future.wait())
     }
 
+    func testDrainBeforeStreamHasFinishedWhenThereIsAlreadyAConsumer() {
+        let dataSource = CountingDataSource()
+        let stream = PSQLRowStream(
+            source: .stream(
+                [self.makeColumnDescription(name: "foo", dataType: .text, format: .binary)],
+                dataSource
+            ),
+            eventLoop: self.eventLoop,
+            logger: self.logger
+        )
+        XCTAssertEqual(dataSource.hitDemand, 0)
+        XCTAssertEqual(dataSource.hitCancel, 0)
+
+        stream.receive([
+            [ByteBuffer(string: "0")],
+            [ByteBuffer(string: "1")]
+        ])
+
+        XCTAssertEqual(dataSource.hitDemand, 0, "Before we have a consumer demand is not signaled")
+
+        // attach consumers
+        let allFuture = stream.all()
+        XCTAssertEqual(dataSource.hitDemand, 1)
+        let drainFuture = stream.drain()
+        XCTAssertEqual(dataSource.hitDemand, 2)
+
+        stream.receive([
+            [ByteBuffer(string: "2")],
+            [ByteBuffer(string: "3")]
+        ])
+        XCTAssertEqual(dataSource.hitDemand, 3)
+
+        stream.receive([
+            [ByteBuffer(string: "4")],
+            [ByteBuffer(string: "5")]
+        ])
+        XCTAssertEqual(dataSource.hitDemand, 4)
+
+        stream.receive(completion: .success("SELECT 2"))
+
+        XCTAssertNoThrow(try drainFuture.wait())
+
+        var rows: [PostgresRow]?
+        XCTAssertNoThrow(rows = try allFuture.wait())
+        XCTAssertEqual(rows?.count, 6)
+    }
+
+    func testDrainBeforeStreamHasFinishedWhenThereIsAlreadyAnAsyncConsumer() {
+        let dataSource = CountingDataSource()
+        let stream = PSQLRowStream(
+            source: .stream(
+                [self.makeColumnDescription(name: "foo", dataType: .text, format: .binary)],
+                dataSource
+            ),
+            eventLoop: self.eventLoop,
+            logger: self.logger
+        )
+        XCTAssertEqual(dataSource.hitDemand, 0)
+        XCTAssertEqual(dataSource.hitCancel, 0)
+
+        stream.receive([
+            [ByteBuffer(string: "0")],
+            [ByteBuffer(string: "1")]
+        ])
+
+        XCTAssertEqual(dataSource.hitDemand, 0, "Before we have a consumer demand is not signaled")
+
+        // attach consumers
+        let rowSequence = stream.asyncSequence()
+        XCTAssertEqual(dataSource.hitDemand, 0)
+        let drainFuture = stream.drain()
+        XCTAssertEqual(dataSource.hitDemand, 1)
+
+        stream.receive([
+            [ByteBuffer(string: "2")],
+            [ByteBuffer(string: "3")]
+        ])
+        XCTAssertEqual(dataSource.hitDemand, 2)
+
+        stream.receive([
+            [ByteBuffer(string: "4")],
+            [ByteBuffer(string: "5")]
+        ])
+        XCTAssertEqual(dataSource.hitDemand, 3)
+
+        stream.receive(completion: .success("SELECT 2"))
+
+        XCTAssertNoThrow(try drainFuture.wait())
+
+        XCTAssertNoThrow {
+            let rows = try stream.eventLoop.makeFutureWithTask {
+                try? await rowSequence.collect()
+            }.wait()
+            XCTAssertEqual(dataSource.hitDemand, 4)
+            XCTAssertEqual(rows?.count, 6)
+        }
+    }
+
     func makeColumnDescription(name: String, dataType: PostgresDataType, format: PostgresFormat) -> RowDescription.Column {
         RowDescription.Column(
             name: "test",