Handle error cases more exhaustively

Author: ahoppen

Sources/PostgresNIO/Connection/PostgresConnection.swift

@@ -698,6 +698,17 @@ extension PostgresConnection {
 
 /// A handle to send 
 public struct PostgresCopyFromWriter: Sendable {
+    /// The backend failed the copy data transfer, which means that no more data sent by the frontend would be processed.
+    /// 
+    /// The `PostgresCopyFromWriter` should cancel the data transfer.
+    public struct CopyCancellationError: Error {
+        /// The error that the backend sent us which cancelled the data transfer.
+        /// 
+        /// Note that this error is related to previous `write` calls since a `CopyCancellationError` is thrown before
+        /// new data is written by `write`.
+        let underlyingError: PSQLError
+    }
+
     private let channelHandler: NIOLoopBound<PostgresChannelHandler>
     private let context: NIOLoopBound<ChannelHandlerContext>
     private let eventLoop: any EventLoop
@@ -709,16 +720,30 @@ public struct PostgresCopyFromWriter: Sendable {
     }
 
     /// Send data for a `COPY ... FROM STDIN` operation to the backend.
+    /// 
+    /// If the backend encountered an error during the data transfer and thus cannot process any more data, this throws
+    /// a `CopyFailedError`.
     public func write(_ byteBuffer: ByteBuffer) async throws {
-        await withCheckedContinuation { (continuation: CheckedContinuation<Void, Never>) in 
+        // First, wait that we have a writable buffer. This also throws a `CopyFailedError` in case the backend sent an 
+        // error during the data transfer and thus cannot process any more data.
+        try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in 
             if eventLoop.inEventLoop {
-                self.channelHandler.value.copyData(byteBuffer, context: self.context.value, readyForMoreWriteContinuation: continuation)
+                self.channelHandler.value.waitForWritableBuffer(context: self.context.value, continuation)
             } else {
                 eventLoop.execute {
-                    self.channelHandler.value.copyData(byteBuffer, context: self.context.value, readyForMoreWriteContinuation: continuation)
+                    self.channelHandler.value.waitForWritableBuffer(context: self.context.value, continuation)
                 }
             }
         }
+
+        // Run the actual data transfer
+        if eventLoop.inEventLoop {
+            self.channelHandler.value.copyData(byteBuffer, context: self.context.value)
+        } else {
+            eventLoop.execute {
+                self.channelHandler.value.copyData(byteBuffer, context: self.context.value)
+            }
+        }
     }
 
     /// Finalize the data transfer, putting the state machine out of the copy mode and sending a `CopyDone` message to 
@@ -748,6 +773,17 @@ public struct PostgresCopyFromWriter: Sendable {
             }
         }
     }
+
+    /// Send a `Sync` message to the backend.
+    func sync() {
+        if eventLoop.inEventLoop {
+            self.channelHandler.value.sendSync(context: self.context.value)
+        } else {
+            eventLoop.execute {
+                self.channelHandler.value.sendSync(context: self.context.value)
+            }
+        }
+    }
 }
 
 extension PostgresConnection {
@@ -777,8 +813,15 @@ extension PostgresConnection {
 
         do {
             try await writeData(writer)
+        } catch let error as PostgresCopyFromWriter.CopyCancellationError {
+            // If the copy was cancelled because the backend sent us an error, we need to send a `Sync` message to put 
+            // the backend out of the copy mode.
+            writer.sync()
+            throw error.underlyingError
         } catch {
-            try await writer.failed(error: error)
+            // Throw the error from the `writeData` closure instead of the one that Postgres gives us upon receiving the 
+            // `CopyFail` message.
+            try? await writer.failed(error: error)
             throw error
         }
         try await writer.done()

Sources/PostgresNIO/New/Connection State Machine/ConnectionStateMachine.swift

@@ -88,8 +88,9 @@ struct ConnectionStateMachine {
         case sendParseDescribeBindExecuteSync(PostgresQuery)
         case sendBindExecuteSync(PSQLExecuteStatement)
         case failQuery(EventLoopPromise<PSQLRowStream>, with: PSQLError, cleanupContext: CleanUpContext?)
-        /// Fail a query's execution by throwing an error on the given continuation.
-        case failQueryContinuation(any AnyErrorContinuation, with: PSQLError, cleanupContext: CleanUpContext?)
+        /// Fail a query's execution by throwing an error on the given continuation. When `sync` is `true`, send a 
+        /// `sync` message to the backend.
+        case failQueryContinuation(any AnyErrorContinuation, with: PSQLError, cleanupContext: CleanUpContext?, sync: Bool)
         case succeedQuery(EventLoopPromise<PSQLRowStream>, with: QueryResult)
         case succeedQueryContinuation(CheckedContinuation<Void, any Error>)
 
@@ -124,8 +125,24 @@ struct ConnectionStateMachine {
     }
 
     enum ChannelWritabilityChangedAction {
+        /// No action needs to be taken based on the writability change.
         case none
-        case resumeContinuation(CheckedContinuation<Void, Never>)
+
+        /// Resume the given continuation successfully.
+        case resumeContinuation(CheckedContinuation<Void, any Error>)
+    }
+
+    enum WaitForWritableBufferAction {
+        /// The channel has backpressure and cannot handle any data right now. We should flush the channel to help  
+        /// relieve backpressure. Once the channel is writable again, this will be communicated via 
+        /// `channelWritabilityChanged`
+        case waitForBackpressureRelieve
+
+        /// Resume the given continuation successfully.
+        case resumeContinuation(CheckedContinuation<Void, any Error>)
+
+        /// Fail the continuation with the given error.
+        case failContinuation(CheckedContinuation<Void, any Error>, error: any Error)
     }
 
     private var state: State
@@ -609,7 +626,7 @@ struct ConnectionStateMachine {
             case .executeStatement(_, let promise), .unnamed(_, let promise):
                 return .failQuery(promise, with: psqlErrror, cleanupContext: nil)
             case .copyFrom(_, let triggerCopy):
-                return .failQueryContinuation(triggerCopy, with: psqlErrror, cleanupContext: nil)
+                return .failQueryContinuation(triggerCopy, with: psqlErrror, cleanupContext: nil, sync: false)
             case .prepareStatement(_, _, _, let promise):
                 return .failPreparedStatementCreation(promise, with: psqlErrror, cleanupContext: nil)
             }
@@ -798,16 +815,20 @@ struct ConnectionStateMachine {
         return self.modify(with: action)
     }
 
-    /// Assuming that the channel to the backend is not writable, wait for the write buffer to become writable again and 
-    /// then resume `continuation`.
-    mutating func waitForWritableBuffer(continuation: CheckedContinuation<Void, Never>) {
+    /// Wait fo `channel` to be writable and be able to handle more `CopyData` messages. Resume the given continuation
+    /// when the channel is able handle more data. 
+    /// 
+    /// This fails the continuation with a `PostgresCopyFromWriter.CopyCancellationError` when the server has cancelled
+    /// the data transfer to indicate that the frontend should not send any more data.
+    mutating func waitForWritableBuffer(channel: any Channel, continuation: CheckedContinuation<Void, any Error>) -> WaitForWritableBufferAction {
         guard case .extendedQuery(var queryState, let connectionContext) = self.state else {
             preconditionFailure("Copy mode is only supported for extended queries")
         }
 
         self.state = .modifying // avoid CoW
-        queryState.waitForWritableBuffer(continuation: continuation)
+        let action = queryState.waitForWritableBuffer(channel: channel, continuation: continuation)
         self.state = .extendedQuery(queryState, connectionContext)
+        return action
     }
 
     /// Put the state machine out of the copying mode and send a `CopyDone` message to the backend.
@@ -955,8 +976,8 @@ struct ConnectionStateMachine {
             case .failQuery(let promise, with: let error):
                 return .failQuery(promise, with: error, cleanupContext: cleanupContext)
 
-            case .failQueryContinuation(let continuation, with: let error):
-                return .failQueryContinuation(continuation, with: error, cleanupContext: cleanupContext)
+            case .failQueryContinuation(let continuation, with: let error, let sync):
+                return .failQueryContinuation(continuation, with: error, cleanupContext: cleanupContext, sync: sync)
 
             case .forwardStreamError(let error, let read):
                 return .forwardStreamError(error, read: read, cleanupContext: cleanupContext)
@@ -1127,9 +1148,9 @@ extension ConnectionStateMachine {
         case .failQuery(let requestContext, with: let error):
             let cleanupContext = self.setErrorAndCreateCleanupContextIfNeeded(error)
             return .failQuery(requestContext, with: error, cleanupContext: cleanupContext)
-        case .failQueryContinuation(let continuation, with: let error):
+        case .failQueryContinuation(let continuation, with: let error, let sync):
             let cleanupContext = self.setErrorAndCreateCleanupContextIfNeeded(error)
-            return .failQueryContinuation(continuation, with: error, cleanupContext: cleanupContext)
+            return .failQueryContinuation(continuation, with: error, cleanupContext: cleanupContext, sync: sync)
         case .succeedQuery(let requestContext, with: let result):
             return .succeedQuery(requestContext, with: result)
         case .succeedQueryContinuation(let continuation):

Sources/PostgresNIO/New/Connection State Machine/ExtendedQueryStateMachine.swift

@@ -8,7 +8,7 @@ struct ExtendedQueryStateMachine {
 
         /// The write channel has backpressure. Once that is relieved, we should resume the given continuation to allow more
         /// data to be sent by the client.
-        case pendingBackpressureRelieve(CheckedContinuation<Void, Never>)
+        case pendingBackpressureRelieve(CheckedContinuation<Void, any Error>)
     }
 
     private enum State {
@@ -50,8 +50,9 @@ struct ExtendedQueryStateMachine {
 
         // --- general actions
         case failQuery(EventLoopPromise<PSQLRowStream>, with: PSQLError)
-        /// Fail a query's execution by throwing an error on the given continuation.
-        case failQueryContinuation(AnyErrorContinuation, with: PSQLError)
+        /// Fail a query's execution by throwing an error on the given continuation. If `sync` is `true`, send a `sync`
+        /// message to the backend to put it out of the copy mode.
+        case failQueryContinuation(AnyErrorContinuation, with: PSQLError, sync: Bool)
         case succeedQuery(EventLoopPromise<PSQLRowStream>, with: QueryResult)
         case succeedQueryContinuation(CheckedContinuation<Void, any Error>)
 
@@ -142,14 +143,14 @@ struct ExtendedQueryStateMachine {
             case .unnamed(_, let eventLoopPromise), .executeStatement(_, let eventLoopPromise):
                 return .failQuery(eventLoopPromise, with: .queryCancelled)
             case .copyFrom(_, let triggerCopy):
-                return .failQueryContinuation(triggerCopy, with: .queryCancelled)
+                return .failQueryContinuation(triggerCopy, with: .queryCancelled, sync: false)
             case .prepareStatement(_, _, _, let eventLoopPromise):
                 return .failPreparedStatementCreation(eventLoopPromise, with: .queryCancelled)
             }
         case .copyingData:
             return .sendCopyFailed(message: "Query cancelled")
         case .copyingFinished(let continuation):
-            return .failQueryContinuation(continuation, with: .queryCancelled)
+            return .failQueryContinuation(continuation, with: .queryCancelled, sync: true)
 
         case .streaming(let columns, var streamStateMachine):
             precondition(!self.isCancelled)
@@ -392,24 +393,40 @@ struct ExtendedQueryStateMachine {
         }
     }
 
-    /// Assuming that the channel to the backend is not writable, wait for the write buffer to become writable again and 
-    /// then resume `continuation`.
-    mutating func waitForWritableBuffer(continuation: CheckedContinuation<Void, Never>) {
+    /// Wait fo `channel` to be writable and be able to handle more `CopyData` messages. Resume the given continuation
+    /// when the channel is able handle more data. 
+    /// 
+    /// This fails the continuation with a `PostgresCopyFromWriter.CopyCancellationError` when the server has cancelled
+    /// the data transfer to indicate that the frontend should not send any more data.
+    mutating func waitForWritableBuffer(
+        channel: any Channel, 
+        continuation: CheckedContinuation<Void, any Error>
+    ) -> ConnectionStateMachine.WaitForWritableBufferAction {
+        if case .error(let error) = self.state {
+            return .failContinuation(continuation, error: PostgresCopyFromWriter.CopyCancellationError(underlyingError: error))
+        }
         guard case .copyingData(let copyingSubstate) = self.state else {
             preconditionFailure("Must be in copy mode to copy data")
         }
         guard case .readyToSend = copyingSubstate else {
             preconditionFailure("Not ready to send data")
         }
+        if channel.isWritable {
+            return .resumeContinuation(continuation)
+        }
         return avoidingStateMachineCoW { state in
             // Even if the buffer isn't writable, we write the current chunk of data to it. We just don't resume
             // the continuation. This will prevent more writes from happening to build up more write backpressure.
             state = .copyingData(.pendingBackpressureRelieve(continuation))
+            return .waitForBackpressureRelieve
         }
     }
 
     /// Put the state machine out of the copying mode and send a `CopyDone` message to the backend.
     mutating func sendCopyDone(continuation: CheckedContinuation<Void, any Error>) -> Action {
+        if case .error(let error) = self.state {
+            return .failQueryContinuation(continuation, with: error, sync: true)
+        }
         guard case .copyingData = self.state else {
             preconditionFailure("Must be in copy mode to send CopyDone")
         }
@@ -421,6 +438,9 @@ struct ExtendedQueryStateMachine {
 
     /// Put the state machine out of the copying mode and send a `CopyFail` message to the backend.
     mutating func sendCopyFailed(message: String, continuation: CheckedContinuation<Void, any Error>) -> Action {
+        if case .error(let error) = self.state {
+            return .failQueryContinuation(continuation, with: error, sync: true)
+        }
         guard case .copyingData = self.state else {
             preconditionFailure("Must be in copy mode to send CopyFail")
         }
@@ -627,7 +647,7 @@ struct ExtendedQueryStateMachine {
                 case .unnamed(_, let eventLoopPromise), .executeStatement(_, let eventLoopPromise):
                     return .failQuery(eventLoopPromise, with: error)
                 case .copyFrom(_, let triggerCopy):
-                    return .failQueryContinuation(triggerCopy, with: error)
+                    return .failQueryContinuation(triggerCopy, with: error, sync: false)
                 case .prepareStatement(_, _, _, let eventLoopPromise):
                     return .failPreparedStatementCreation(eventLoopPromise, with: error)
                 }
@@ -637,7 +657,7 @@ struct ExtendedQueryStateMachine {
             return .evaluateErrorAtConnectionLevel(error)
         case .copyingFinished(let continuation):
             self.state = .error(error)
-            return .failQueryContinuation(continuation, with: error)
+            return .failQueryContinuation(continuation, with: error, sync: true)
         case .drain:
             self.state = .error(error)
             return .evaluateErrorAtConnectionLevel(error)

Sources/PostgresNIO/New/PostgresChannelHandler.swift

@@ -171,17 +171,26 @@ final class PostgresChannelHandler: ChannelDuplexHandler {
         self.run(action, with: context)
     }
 
-    /// Send a `CopyData` message to the backend using the given data.
+    /// Wait for the channel to be writable and be able to handle more `CopyData` messages. Resume the given 
+    /// continuation when the channel is able handle more data. 
     /// 
-    /// `readyForMoreWriteContinuation` is resumed when the channel is able to handle more data to be written to it.
-    func copyData(_ data: ByteBuffer, context: ChannelHandlerContext, readyForMoreWriteContinuation: CheckedContinuation<Void, Never>) {
-        self.encoder.copyData(data: data)
-        if context.channel.isWritable {
-            readyForMoreWriteContinuation.resume()
-        } else {
-            self.state.waitForWritableBuffer(continuation: readyForMoreWriteContinuation)
-            context.flush()
+    /// This fails the continuation with a `PostgresCopyFromWriter.CopyCancellationError` when the server has cancelled
+    /// the data transfer to indicate that the frontend should not send any more data.
+    func waitForWritableBuffer(context: ChannelHandlerContext, _ continuation: CheckedContinuation<Void, any Error>) {
+        let action = self.state.waitForWritableBuffer(channel: context.channel, continuation: continuation)
+        switch action {
+        case .waitForBackpressureRelieve:
+            context.channel.flush()
+        case .resumeContinuation(let continuation):
+            continuation.resume()
+        case .failContinuation(_, error: let error):
+            continuation.resume(throwing: error)
         }
+    }
+
+    /// Send a `CopyData` message to the backend using the given data.
+    func copyData(_ data: ByteBuffer, context: ChannelHandlerContext) {
+        self.encoder.copyData(data: data)
         context.write(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil)
     }
 
@@ -197,6 +206,12 @@ final class PostgresChannelHandler: ChannelDuplexHandler {
         self.run(action, with: context)
     }
 
+    /// Send a `Sync` message to the backend.
+    func sendSync(context: ChannelHandlerContext) {
+        self.encoder.sync()
+        context.writeAndFlush(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil)
+    }
+
     func channelReadComplete(context: ChannelHandlerContext) {
         let action = self.state.channelReadComplete()
         self.run(action, with: context)
@@ -398,11 +413,15 @@ final class PostgresChannelHandler: ChannelDuplexHandler {
             if let cleanupContext = cleanupContext {
                 self.closeConnectionAndCleanup(cleanupContext, context: context)
             }
-        case .failQueryContinuation(let continuation, with: let error, let cleanupContext):
-            continuation.resume(throwing: error)
+        case .failQueryContinuation(let continuation, with: let error, let cleanupContext, let sync):
+            if sync {
+                self.encoder.sync()
+                context.writeAndFlush(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil)
+            }
             if let cleanupContext = cleanupContext {
                 self.closeConnectionAndCleanup(cleanupContext, context: context)
             }
+            continuation.resume(throwing: error)
         case .triggerCopyData(let triggerCopy):
             let writer = PostgresCopyFromWriter(handler: self, context: context, eventLoop: eventLoop)
             triggerCopy.resume(returning: writer)

Tests/IntegrationTests/PSQLIntegrationTests.swift

@@ -401,7 +401,10 @@ final class IntegrationTests: XCTestCase {
             }
         }, logger: .psqlTest)
         let rows = try await conn.query("SELECT id, name FROM copy_table").get().rows.map { try $0.decode((Int, String).self) }
-        XCTAssertEqual(rows.count, 2)
+        guard rows.count == 2 else {
+            XCTFail("Expected 2 columns, received \(rows.count)")
+            return
+        }
         XCTAssertEqual(rows[0].0, 1)
         XCTAssertEqual(rows[0].1, "Alice")
         XCTAssertEqual(rows[1].0, 42)