Address review comments

Author: ahoppen

Sources/PostgresNIO/Connection/PostgresConnection+CopyFrom.swift

@@ -1,16 +1,5 @@
 /// Handle to send data for a `COPY ... FROM STDIN` query to the backend.
 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`.
-        public let underlyingError: PSQLError
-    }
-
     private let channelHandler: NIOLoopBound<PostgresChannelHandler>
     private let eventLoop: any EventLoop
 
@@ -42,9 +31,9 @@ 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 `CopyCancellationError`.
-    public func write(_ byteBuffer: ByteBuffer) async throws {
+    /// - Throws: If an error occurs during the write of if the backend sent an `ErrorResponse` during the copy 
+    ///   operation, eg. to indicate that a **previous** `write` call had an invalid format.
+    public func write(_ byteBuffer: ByteBuffer, isolation: isolated (any Actor)? = #isolation) async throws {
         // Check for cancellation. This is cheap and makes sure that we regularly check for cancellation in the
         // `writeData` closure. It is likely that the user would forget to do so.
         try Task.checkCancellation()
@@ -82,7 +71,7 @@ public struct PostgresCopyFromWriter: Sendable {
 
     /// Finalize the data transfer, putting the state machine out of the copy mode and sending a `CopyDone` message to
     /// the backend.
-    func done() async throws {
+    func done(isolation: isolated (any Actor)? = #isolation) async throws {
         try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in
             if eventLoop.inEventLoop {
                 self.channelHandler.value.sendCopyDone(continuation: continuation)
@@ -96,37 +85,43 @@ public struct PostgresCopyFromWriter: Sendable {
 
     /// Finalize the data transfer, putting the state machine out of the copy mode and sending a `CopyFail` message to
     /// the backend.
-    func failed(error: any Error) async throws {
+    func failed(error: any Error, isolation: isolated (any Actor)? = #isolation) async throws {
         try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in
-            // TODO: Is it OK to use string interpolation to construct an error description to be sent to the backend
-            // here? We could also use a generic description, it doesn't really matter since we throw the user's error
-            // in `copyFrom`.
             if eventLoop.inEventLoop {
-                self.channelHandler.value.sendCopyFail(message: "\(error)", continuation: continuation)
+                self.channelHandler.value.sendCopyFail(message: "Client failed copy", continuation: continuation)
             } else {
                 eventLoop.execute {
-                    self.channelHandler.value.sendCopyFail(message: "\(error)", continuation: continuation)
+                    self.channelHandler.value.sendCopyFail(message: "Client failed copy", continuation: continuation)
                 }
             }
         }
     }
 }
 
 /// Specifies the format in which data is transferred to the backend in a COPY operation.
-public enum PostgresCopyFromFormat: Sendable {
+/// 
+/// See the Postgres documentation at https://www.postgresql.org/docs/current/sql-copy.html for the option's meanings
+/// and their default values.
+public struct PostgresCopyFromFormat: Sendable {
     /// Options that can be used to modify the `text` format of a COPY operation.
     public struct TextOptions: Sendable {
         /// The delimiter that separates columns in the data.
         ///
         /// See the `DELIMITER` option in Postgres's `COPY` command.
-        ///
-        /// Uses the default delimiter of the format
         public var delimiter: UnicodeScalar? = nil
 
         public init() {}
     }
 
-    case text(TextOptions)
+    enum Format {
+        case text(TextOptions)
+    }
+
+    var format: Format
+
+    public static func text(_ options: TextOptions) -> PostgresCopyFromFormat {
+        return PostgresCopyFromFormat(format: .text(options))
+    }
 }
 
 /// Create a `COPY ... FROM STDIN` query based on the given parameters.
@@ -138,14 +133,17 @@ private func buildCopyFromQuery(
     columns: [StaticString] = [],
     format: PostgresCopyFromFormat
 ) -> PostgresQuery {
-    // TODO: Should we put the table and column names in quotes to make them case-sensitive?
-    var query = "COPY \(table)"
+    var query = """
+        COPY "\(table)"
+        """
     if !columns.isEmpty {
-        query += "(" + columns.map(\.description).joined(separator: ",") + ")"
+        query += "("
+        query += columns.map { #"""# + $0.description + #"""# }.joined(separator: ",")
+        query += ")"
     }
     query += " FROM STDIN"
     var queryOptions: [String] = []
-    switch format {
+    switch format.format {
     case .text(let options):
         queryOptions.append("FORMAT text")
         if let delimiter = options.delimiter {
@@ -179,6 +177,7 @@ extension PostgresConnection {
         columns: [StaticString] = [],
         format: PostgresCopyFromFormat = .text(.init()),
         logger: Logger,
+        isolation: isolated (any Actor)? = #isolation,
         file: String = #fileID,
         line: Int = #line,
         writeData: @escaping @Sendable (PostgresCopyFromWriter) async throws -> Void
@@ -205,22 +204,13 @@ extension PostgresConnection {
             //    threw instead of the one that got relayed back, so it's better to ignore the error here.
             //  - The backend sent us an `ErrorResponse` during the copy, eg. because of an invalid format. This puts
             //    the `ExtendedQueryStateMachine` in the error state. Trying to send a `CopyFail` will throw but trigger
-            //    a `Sync` that takes the backend out of copy mode. If `writeData` threw the `CopyCancellationError`
-            //    from the `PostgresCopyFromWriter.write` call, `writer.failed` will throw with the same error, so it
-            //    doesn't matter that we ignore the error here. If the user threw some other error, it's better to honor
-            //    the user's error.
+            //    a `Sync` that takes the backend out of copy mode. If `writeData` threw the error from from the 
+            //    `PostgresCopyFromWriter.write` call, `writer.failed` will throw with the same error, so it doesn't 
+            //    matter that we ignore the error here. If the user threw some other error, it's better to honor the 
+            //    user's error.
             try? await writer.failed(error: error)
 
-            if let error = error as? PostgresCopyFromWriter.CopyCancellationError {
-                // If we receive a `CopyCancellationError` that is with almost certain likelihood because
-                // `PostgresCopyFromWriter.write` threw it - otherwise the user must have saved a previous
-                // `PostgresCopyFromWriter` error, which is very unlikely.
-                // Throw the underlying error because that contains the error message that was sent by the backend and
-                // is most actionable by the user.
-                throw error.underlyingError
-            } else {
-                throw error
-            }
+            throw error
         }
 
         // `writer.done` may fail, eg. because the backend sends an error response after receiving `CopyDone` or during
@@ -230,5 +220,4 @@ extension PostgresConnection {
         // above.
         try await writer.done()
     }
-
 }

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

@@ -427,7 +427,7 @@ struct ExtendedQueryStateMachine {
         if case .error(let error) = self.state {
             // The backend sent us an ErrorResponse during the copy operation. Indicate to the client that it should
             // abort the data transfer.
-            promise.fail(PostgresCopyFromWriter.CopyCancellationError(underlyingError: error))
+            promise.fail(error)
             return
         }
         guard case .copyingData(.readyToSend) = self.state else {

Tests/PostgresNIOTests/New/PostgresConnectionTests.swift

@@ -660,7 +660,7 @@ import Synchronization
         try await expectCopyFrom { writer in
             try await writer.write(ByteBuffer(staticString: "1\tAlice\n"))
         } validateCopyRequest: { copyRequest in
-            #expect(copyRequest.parse.query == "COPY copy_table(id,name) FROM STDIN WITH (FORMAT text)")
+            #expect(copyRequest.parse.query == #"COPY "copy_table"("id","name") FROM STDIN WITH (FORMAT text)"#)
             #expect(copyRequest.bind.parameters == [])
         } mockBackend: { channel, _ in
             let data = try await channel.waitForCopyData()
@@ -677,7 +677,7 @@ import Synchronization
         try await expectCopyFrom(format: .text(options)) { writer in
             try await writer.write(ByteBuffer(staticString: "1,Alice\n"))
         } validateCopyRequest: { copyRequest in
-            #expect(copyRequest.parse.query == #"COPY copy_table(id,name) FROM STDIN WITH (FORMAT text,DELIMITER U&'\002c')"#)
+            #expect(copyRequest.parse.query == #"COPY "copy_table"("id","name") FROM STDIN WITH (FORMAT text,DELIMITER U&'\002c')"#)
             #expect(copyRequest.bind.parameters == [])
         } mockBackend: { channel, _ in
             let data = try await channel.waitForCopyData()
@@ -688,19 +688,17 @@ import Synchronization
     }
 
     @Test func testCopyFromWriterFails() async throws {
-        struct MyError: Error, CustomStringConvertible {
-            var description: String { "My error" }
-        }
+        struct MyError: Error {}
 
         try await expectCopyFrom { writer in
             throw MyError()
         } validateCopyFromError: { error in
             #expect(error is MyError, "Expected error of type MyError, got \(error)")
         } mockBackend: { channel, _ in
             let data = try await channel.waitForCopyData()
-            #expect(data.result == .failed(message: "My error"))
+            #expect(data.result == .failed(message: "Client failed copy"))
             try await channel.writeInbound(PostgresBackendMessage.error(.init(fields: [
-                .message: "COPY from stdin failed: My error",
+                .message: "COPY from stdin failed: Client failed copy",
                 .sqlState : "57014" // query_canceled
             ])))
         }
@@ -783,7 +781,7 @@ import Synchronization
                 try await writer.write(ByteBuffer(staticString: "2\tBob\n"))
                 Issue.record("Expected error to be thrown")
             } catch {
-                #expect(error is PostgresCopyFromWriter.CopyCancellationError, "Received unexpected error: \(error)")
+                #expect((error as? PSQLError)?.serverInfo?[.sqlState] == "22P02")
                 throw error
             }
         } validateCopyFromError: { error in
@@ -800,7 +798,7 @@ import Synchronization
         }
     }
 
-    @Test func testCopyFromCallerDoesNotRethrowCopyCancellationError() async throws {
+    @Test func testCopyFromCallerDoesNotRethrowFromWriteCall() async throws {
         struct MyError: Error, CustomStringConvertible {
             var description: String { "My error" }
         }
@@ -816,7 +814,7 @@ import Synchronization
                 try await writer.write(ByteBuffer(staticString: "2\tBob\n"))
                 Issue.record("Expected error to be thrown")
             } catch {
-                #expect(error is PostgresCopyFromWriter.CopyCancellationError, "Received unexpected error: \(error)")
+                #expect((error as? PSQLError)?.serverInfo?[.sqlState] == "22P02")
                 throw MyError()
             }
         } validateCopyFromError: { error in
@@ -903,10 +901,10 @@ import Synchronization
             cancelCopy()
 
             let data = try await channel.waitForCopyData()
-            #expect(data.result == .failed(message: "CancellationError()"))
+            #expect(data.result == .failed(message: "Client failed copy"))
 
             try await channel.writeInbound(PostgresBackendMessage.error(.init(fields: [
-                .message: "COPY from stdin failed: CancellationError()",
+                .message: "COPY from stdin failed: Client failed copy",
                 .sqlState : "57014" // query_canceled
             ])))
         }