diff --git a/.gitignore b/.gitignore index d0d66c24..bb8bc4ba 100644 --- 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 index e267d8f9..afb9a1a2 100644 --- a/Sources/PostgresNIO/Connection/PostgresConnection.swift +++ b/Sources/PostgresNIO/Connection/PostgresConnection.swift @@ -438,6 +438,56 @@ extension PostgresConnection { } } + /// 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``. + /// 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( + _ 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 (rowStream, rowSequence) = try await promise.futureResult.map { rowStream in + (rowStream, rowStream.asyncSequence()) + }.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 + } + } + /// Start listening for a channel public func listen(_ channel: String) async throws -> PostgresNotificationSequence { let id = self.internalListenID.loadThenWrappingIncrement(ordering: .relaxed) @@ -531,6 +581,52 @@ extension PostgresConnection { } } + /// 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, + 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 rowStream = try await promise.futureResult.get() + try await rowStream.drain().get() + guard let metadata = PostgresQueryMetadata(string: rowStream.commandTag) else { + throw PSQLError.invalidCommandTag(rowStream.commandTag) + } + 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 index ee925d0e..fe7adbc0 100644 --- a/Sources/PostgresNIO/New/PSQLRowStream.swift +++ b/Sources/PostgresNIO/New/PSQLRowStream.swift @@ -276,7 +276,70 @@ final class PSQLRowStream: @unchecked Sendable { return self.eventLoop.makeFailedFuture(error) } } - + + // MARK: Drain on EventLoop + + func drain() -> EventLoopFuture { + if self.eventLoop.inEventLoop { + return self.drain0() + } else { + return self.eventLoop.flatSubmit { + self.drain0() + } + } + } + + private func drain0() -> EventLoopFuture { + self.eventLoop.preconditionInEventLoop() + + switch self.downstreamState { + case .waitingForConsumer(let bufferState): + switch bufferState { + case .streaming(var buffer, let dataSource): + let promise = self.eventLoop.makePromise(of: [PostgresRow].self) + + buffer.removeAll() + self.downstreamState = .waitingForAll([], promise, dataSource) + // immediately request more + dataSource.request(for: self) + + return promise.futureResult.map { _ in } + + 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, let onFinish): + consumer.finish() + onFinish() + + let promise = self.eventLoop.makePromise(of: [PostgresRow].self) + + self.downstreamState = .waitingForAll([], promise, dataSource) + // immediately request more + dataSource.request(for: self) + + return promise.futureResult.map { _ in } + case .consumed(.success): + // already drained + 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)") + } + } + 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 index 3936b51e..0077a837 100644 --- a/Sources/PostgresNIO/New/PostgresRowSequence.swift +++ b/Sources/PostgresNIO/New/PostgresRowSequence.swift @@ -60,6 +60,8 @@ extension PostgresRowSequence { extension PostgresRowSequence.AsyncIterator: Sendable {} extension PostgresRowSequence { + /// Collects all rows into an array. + /// - Returns: The rows. public func collect() async throws -> [PostgresRow] { var result = [PostgresRow]() for try await row in self { diff --git a/Sources/PostgresNIO/Pool/PostgresClient.swift b/Sources/PostgresNIO/Pool/PostgresClient.swift index d54e34eb..14038476 100644 --- a/Sources/PostgresNIO/Pool/PostgresClient.swift +++ b/Sources/PostgresNIO/Pool/PostgresClient.swift @@ -435,6 +435,61 @@ 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``. + /// 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( + _ 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( _ preparedStatement: Statement, diff --git a/Tests/IntegrationTests/AsyncTests.swift b/Tests/IntegrationTests/AsyncTests.swift index b4c8e93f..07be3412 100644 --- a/Tests/IntegrationTests/AsyncTests.swift +++ b/Tests/IntegrationTests/AsyncTests.swift @@ -46,6 +46,98 @@ final class AsyncPostgresConnectionTests: XCTestCase { } } + func testSelect10kRowsWithMetadata() 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 (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, end) + + XCTAssertEqual(result, end) + } + } + + func testSelectRowsWithMetadataNotConsumedAtAll() 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.query( + "SELECT generate_series(\(start), \(end));", + logger: .psqlTest + ) { _ in } + + XCTAssertEqual(metadata.command, "SELECT") + XCTAssertEqual(metadata.oid, nil) + XCTAssertEqual(metadata.rows, 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 } + } + // This path is also fine + } 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) + } + } + func testSelectActiveConnection() async throws { let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) } @@ -207,7 +299,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 +333,8 @@ 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 metadata = try await connection.execute(insertionQuery, logger: .psqlTest) + XCTAssertEqual(metadata.rows, rowsCount) let dropQuery = PostgresQuery(unsafeSQL: "DROP TABLE table1") try await connection.query(dropQuery, logger: .psqlTest) diff --git a/Tests/IntegrationTests/PostgresClientTests.swift b/Tests/IntegrationTests/PostgresClientTests.swift index 34a8ad2a..a7b26e05 100644 --- a/Tests/IntegrationTests/PostgresClientTests.swift +++ b/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" diff --git a/Tests/PostgresNIOTests/New/PSQLRowStreamTests.swift b/Tests/PostgresNIOTests/New/PSQLRowStreamTests.swift index 65ca26c3..b2619db3 100644 --- a/Tests/PostgresNIOTests/New/PSQLRowStreamTests.swift +++ b/Tests/PostgresNIOTests/New/PSQLRowStreamTests.swift @@ -10,7 +10,7 @@ final class PSQLRowStreamTests: XCTestCase { let logger = Logger(label: "PSQLRowStreamTests") let eventLoop = EmbeddedEventLoop() - func testEmptyStream() { + func testEmptyStreamAndDrainDoesNotThrowErrorAfterConsumption() { let stream = PSQLRowStream( source: .noRows(.success(.tag("INSERT 0 1"))), eventLoop: self.eventLoop, @@ -19,20 +19,29 @@ final class PSQLRowStreamTests: XCTestCase { XCTAssertEqual(try stream.all().wait(), []) XCTAssertEqual(stream.commandTag, "INSERT 0 1") + + XCTAssertNoThrow(try stream.drain().wait()) } - + func testFailedStream() { let stream = PSQLRowStream( source: .noRows(.failure(PSQLError.serverClosedConnection(underlying: nil))), eventLoop: self.eventLoop, logger: self.logger ) - + + let expectedError = PSQLError.serverClosedConnection(underlying: nil) + XCTAssertThrowsError(try stream.all().wait()) { - XCTAssertEqual($0 as? PSQLError, .serverClosedConnection(underlying: nil)) + XCTAssertEqual($0 as? PSQLError, expectedError) + } + + // Drain should work + XCTAssertThrowsError(try stream.drain().wait()) { + XCTAssertEqual($0 as? PSQLError, expectedError) } } - + func testGetArrayAfterStreamHasFinished() { let dataSource = CountingDataSource() let stream = PSQLRowStream( @@ -75,37 +84,37 @@ final class PSQLRowStreamTests: XCTestCase { ) 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 consumer let future = stream.all() 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")) - + var rows: [PostgresRow]? XCTAssertNoThrow(rows = try future.wait()) XCTAssertEqual(rows?.count, 6) } - + func testOnRowAfterStreamHasFinished() { let dataSource = CountingDataSource() let stream = PSQLRowStream( @@ -231,6 +240,196 @@ final class PSQLRowStreamTests: XCTestCase { XCTAssertEqual(stream.commandTag, "SELECT 6") } + func testEmptyStreamDrainsSuccessfully() { + let stream = PSQLRowStream( + source: .noRows(.success(.tag("INSERT 0 1"))), + eventLoop: self.eventLoop, + logger: self.logger + ) + + XCTAssertNoThrow(try stream.drain().wait()) + XCTAssertEqual(stream.commandTag, "INSERT 0 1") + } + + func testDrainFailedStream() { + let stream = PSQLRowStream( + source: .noRows(.failure(PSQLError.serverClosedConnection(underlying: nil))), + eventLoop: self.eventLoop, + logger: self.logger + ) + + let expectedError = PSQLError.serverClosedConnection(underlying: nil) + + XCTAssertThrowsError(try stream.drain().wait()) { + XCTAssertEqual($0 as? PSQLError, expectedError) + } + } + + func testDrainAfterStreamHasFinished() { + 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") + stream.receive(completion: .success("SELECT 2")) + + // attach consumer + XCTAssertNoThrow(try stream.drain().wait()) + XCTAssertEqual(dataSource.hitDemand, 0) // TODO: Is this right? + } + + func testDrainBeforeStreamHasFinished() { + 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 consumer + let future = 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 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", diff --git a/docker-compose.yml b/docker-compose.yml index 3eff4249..cb9c9404 100644 --- a/docker-compose.yml +++ b/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}" @@ -10,6 +8,9 @@ x-shared-config: &shared_config - 5432:5432 services: + psql-17: + image: postgres:17 + <<: *shared_config psql-16: image: postgres:16 <<: *shared_config