Implement COPY … FROM STDIN queries

This implements support for COPY operations using `COPY … FROM STDIN` queries for fast data transfer from the client to the backend.

Author: ahoppen

Sources/PostgresNIO/Connection/PostgresConnection+CopyFrom.swift

@@ -0,0 +1,211 @@
+/// 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
+
+    init(handler: PostgresChannelHandler, eventLoop: any EventLoop) {
+        self.channelHandler = NIOLoopBound(handler, eventLoop: eventLoop)
+        self.eventLoop = eventLoop
+    }
+
+    private func writeAssumingInEventLoop(_ byteBuffer: ByteBuffer, _ continuation: CheckedContinuation<Void, any Error>) {
+        precondition(eventLoop.inEventLoop)
+        let promise = eventLoop.makePromise(of: Void.self)
+        self.channelHandler.value.checkBackendCanReceiveCopyData(promise: promise)
+        promise.futureResult.map {
+            if eventLoop.inEventLoop {
+                self.channelHandler.value.sendCopyData(byteBuffer)
+            } else {
+                eventLoop.execute {
+                    self.channelHandler.value.sendCopyData(byteBuffer)
+                }
+            }
+        }.whenComplete { result in
+            continuation.resume(with: result)
+        }
+    }
+
+    /// 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 {
+        // 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()
+
+        // TODO: Listen for task cancellation while we are waiting for backpressure to clear.
+        try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in
+            if eventLoop.inEventLoop {
+                writeAssumingInEventLoop(byteBuffer, continuation)
+            } else {
+                eventLoop.execute {
+                    writeAssumingInEventLoop(byteBuffer, continuation)
+                }
+            }
+        }
+    }
+
+    /// 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 {
+        try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in
+            if eventLoop.inEventLoop {
+                self.channelHandler.value.sendCopyDone(continuation: continuation)
+            } else {
+                eventLoop.execute {
+                    self.channelHandler.value.sendCopyDone(continuation: continuation)
+                }
+            }
+        }
+    }
+
+    /// 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 {
+        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)
+            } else {
+                eventLoop.execute {
+                    self.channelHandler.value.sendCopyFail(message: "\(error)", continuation: continuation)
+                }
+            }
+        }
+    }
+}
+
+/// Specifies the format in which data is transferred to the backend in a COPY operation.
+public enum 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)
+}
+
+/// Create a `COPY ... FROM STDIN` query based on the given parameters.
+///
+/// An empty `columns` array signifies that no columns should be specified in the query and that all columns will be
+/// copied by the caller.
+private func buildCopyFromQuery(
+    table: StaticString,
+    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)"
+    if !columns.isEmpty {
+        query += "(" + columns.map(\.description).joined(separator: ",") + ")"
+    }
+    query += " FROM STDIN"
+    var queryOptions: [String] = []
+    switch format {
+    case .text(let options):
+        queryOptions.append("FORMAT text")
+        if let delimiter = options.delimiter {
+            // Set the delimiter as a Unicode code point. This avoids the possibility of SQL injection.
+            queryOptions.append("DELIMITER U&'\\\(String(format: "%04x", delimiter.value))'")
+        }
+    }
+    precondition(!queryOptions.isEmpty)
+    query += " WITH ("
+    query += queryOptions.map { "\($0)" }.joined(separator: ",")
+    query += ")"
+    return "\(unescaped: query)"
+}
+
+extension PostgresConnection {
+    /// Copy data into a table using a `COPY <table name> FROM STDIN` query.
+    ///
+    /// - Parameters:
+    ///   - table: The name of the table into which to copy the data.
+    ///   - columns: The name of the columns to copy. If an empty array is passed, all columns are assumed to be copied.
+    ///   - format: Options that specify the format of the data that is produced by `writeData`.
+    ///   - writeData: Closure that produces the data for the table, to be streamed to the backend. Call `write` on the
+    ///     writer provided by the closure to send data to the backend and return from the closure once all data is sent.
+    ///     Throw an error from the closure to fail the data transfer. The error thrown by the closure will be rethrown
+    ///     by the `copyFrom` function.
+    ///
+    /// - Note: The table and column names are inserted into the SQL query verbatim. They are forced to be compile-time
+    ///   specified to avoid runtime SQL injection attacks.
+    public func copyFrom(
+        table: StaticString,
+        columns: [StaticString] = [],
+        format: PostgresCopyFromFormat = .text(.init()),
+        logger: Logger,
+        file: String = #fileID,
+        line: Int = #line,
+        writeData: @escaping @Sendable (PostgresCopyFromWriter) async throws -> Void
+    )  async throws {
+        var logger = logger
+        logger[postgresMetadataKey: .connectionID] = "\(self.id)"
+        let writer: PostgresCopyFromWriter = try await withCheckedThrowingContinuation { continuation in
+            let context = ExtendedQueryContext(
+                copyFromQuery: buildCopyFromQuery(table: table, columns: columns, format: format),
+                triggerCopy: continuation,
+                logger: logger
+            )
+            self.channel.write(HandlerTask.extendedQuery(context), promise: nil)
+        }
+
+        do {
+            try await writeData(writer)
+        } catch {
+            // We need to send a `CopyFail` to the backend to put it out of copy mode. This will most likely throw, most
+            // notably for the following two reasons. In both of them, it's better to ignore the error thrown by
+            // `writer.failed` and instead throw the error from `writeData`:
+            //  - We send `CopyFail` and the backend replies with an `ErrorResponse` that relays the `CopyFail` message.
+            //    This took the backend out of copy mode but it's more informative to the user to see the error they
+            //    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.
+            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
+            }
+        }
+
+        // `writer.done` may fail, eg. because the backend sends an error response after receiving `CopyDone` or during
+        // the transfer of the last bit of data so that the user didn't call `PostgresCopyFromWriter.write` again, which
+        // would have checked the error state. In either of these cases, calling `writer.done` puts the backend out of
+        // copy mode, so we don't need to send another `CopyFail`. Thus, this must not be handled in the `do` block
+        // above.
+        try await writer.done()
+    }
+
+}

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

@@ -88,7 +88,27 @@ struct ConnectionStateMachine {
         case sendParseDescribeBindExecuteSync(PostgresQuery)
         case sendBindExecuteSync(PSQLExecuteStatement)
         case failQuery(EventLoopPromise<PSQLRowStream>, with: PSQLError, cleanupContext: CleanUpContext?)
+        /// Fail a query's execution by resuming the continuation with the given error. When `sync` is `true`, send a
+        /// `Sync` message to the backend.
+        case failQueryContinuation(AnyErrorContinuation, with: PSQLError, sync: Bool, cleanupContext: CleanUpContext?)
+        /// Fail a query's execution by resuming the continuation with the given error and send a `Sync` message to the
+        /// backend.
         case succeedQuery(EventLoopPromise<PSQLRowStream>, with: QueryResult)
+        /// Succeed the continuation with a void result. When `sync` is `true`, send a `Sync` message to the backend.
+        case succeedQueryContinuation(CheckedContinuation<Void, any Error>, sync: Bool)
+
+        /// Trigger a data transfer returning a `PostgresCopyFromWriter` to the given continuation.
+        ///
+        /// Once the data transfer is triggered, it will send `CopyData` messages to the backend. After that the state
+        /// machine needs to be prodded again to send a `CopyDone` or `CopyFail` by calling
+        /// `PostgresChannelHandler.sendCopyDone` or `PostgresChannelHandler.sendCopyFail`.
+        case triggerCopyData(CheckedContinuation<PostgresCopyFromWriter, any Error>)
+
+        /// Send a `CopyDone` and `Sync` message to the backend.
+        case sendCopyDoneAndSync
+
+        /// Send a `CopyFail` message to the backend with the given error message.
+        case sendCopyFail(message: String)
 
         // --- streaming actions
         // actions if query has requested next row but we are waiting for backend
@@ -107,6 +127,14 @@ struct ConnectionStateMachine {
         case failClose(CloseCommandContext, with: PSQLError, cleanupContext: CleanUpContext?)
     }
 
+    enum ChannelWritabilityChangedAction {
+        /// No action needs to be taken based on the writability change.
+        case none
+
+        /// Resume the given continuation successfully.
+        case succeedPromise(EventLoopPromise<Void>)
+    }
+
     private var state: State
     private let requireBackendKeyData: Bool
     private var taskQueue = CircularBuffer<PSQLTask>()
@@ -587,6 +615,8 @@ struct ConnectionStateMachine {
             switch queryContext.query {
             case .executeStatement(_, let promise), .unnamed(_, let promise):
                 return .failQuery(promise, with: psqlErrror, cleanupContext: nil)
+            case .copyFrom(_, let triggerCopy):
+                return .failQueryContinuation(.copyFromWriter(triggerCopy), with: psqlErrror, sync: false, cleanupContext: nil)
             case .prepareStatement(_, _, _, let promise):
                 return .failPreparedStatementCreation(promise, with: psqlErrror, cleanupContext: nil)
             }
@@ -660,6 +690,16 @@ struct ConnectionStateMachine {
             preconditionFailure("Invalid state: \(self.state)")
         }
     }
+
+    mutating func channelWritabilityChanged(isWritable: Bool) -> ChannelWritabilityChangedAction {
+        guard case .extendedQuery(var queryState, let connectionContext) = state else {
+            return .none
+        }
+        self.state = .modifying // avoid CoW
+        let action = queryState.channelWritabilityChanged(isWritable: isWritable)
+        self.state = .extendedQuery(queryState, connectionContext)
+        return action
+    }
 
     // MARK: - Running Queries -
 
@@ -752,10 +792,55 @@ struct ConnectionStateMachine {
         return self.modify(with: action)
     }
 
-    mutating func copyInResponseReceived(
-        _ copyInResponse: PostgresBackendMessage.CopyInResponse
-    ) -> ConnectionAction {
-        return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.copyInResponse(copyInResponse)))
+    mutating func copyInResponseReceived(_ copyInResponse: PostgresBackendMessage.CopyInResponse) -> ConnectionAction {
+        guard case .extendedQuery(var queryState, let connectionContext) = self.state, !queryState.isComplete else {
+            return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.copyInResponse(copyInResponse)))
+        }
+
+        self.state = .modifying // avoid CoW
+        let action = queryState.copyInResponseReceived(copyInResponse)
+        self.state = .extendedQuery(queryState, connectionContext)
+        return self.modify(with: action)
+    }
+
+
+    /// Succeed the promise when the channel to the backend is writable and the backend is ready to receive more data.
+    ///
+    /// The promise may be failed if the backend indicated that it can't handle any more data by sending an
+    /// `ErrorResponse`. This is mostly the case when malformed data is sent to it. In that case, the data transfer
+    /// should be aborted to avoid unnecessary work.
+    mutating func checkBackendCanReceiveCopyData(channelIsWritable: Bool, promise: EventLoopPromise<Void>) {
+        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.checkBackendCanReceiveCopyData(channelIsWritable: channelIsWritable, promise: promise)
+        self.state = .extendedQuery(queryState, connectionContext)
+    }
+
+    /// 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>) -> ConnectionAction {
+        guard case .extendedQuery(var queryState, let connectionContext) = self.state else {
+            preconditionFailure("Copy mode is only supported for extended queries")
+        }
+
+        self.state = .modifying // avoid CoW
+        let action = queryState.sendCopyDone(continuation: continuation)
+        self.state = .extendedQuery(queryState, connectionContext)
+        return self.modify(with: action)
+    }
+
+    /// Put the state machine out of the copying mode and send a `CopyFail` message to the backend.
+    mutating func sendCopyFail(message: String, continuation: CheckedContinuation<Void, any Error>) -> ConnectionAction {
+        guard case .extendedQuery(var queryState, let connectionContext) = self.state else {
+            preconditionFailure("Copy mode is only supported for extended queries")
+        }
+
+        self.state = .modifying // avoid CoW
+        let action = queryState.sendCopyFail(message: message, continuation: continuation)
+        self.state = .extendedQuery(queryState, connectionContext)
+        return self.modify(with: action)
     }
 
     mutating func emptyQueryResponseReceived() -> ConnectionAction {
@@ -866,14 +951,21 @@ struct ConnectionStateMachine {
                  .forwardRows,
                  .forwardStreamComplete,
                  .wait,
-                 .read:
+                 .read,
+                 .triggerCopyData,
+                 .sendCopyDoneAndSync,
+                 .sendCopyFail,
+                 .succeedQueryContinuation:
                 preconditionFailure("Invalid query state machine action in state: \(self.state), action: \(action)")
 
             case .evaluateErrorAtConnectionLevel:
                 return .closeConnectionAndCleanup(cleanupContext)
 
-            case .failQuery(let queryContext, with: let error):
-                return .failQuery(queryContext, with: error, cleanupContext: cleanupContext)
+            case .failQuery(let promise, with: let error):
+                return .failQuery(promise, with: error, cleanupContext: cleanupContext)
+
+            case .failQueryContinuation(let continuation, with: let error, let sync):
+                return .failQueryContinuation(continuation, with: error, sync: sync, cleanupContext: cleanupContext)
 
             case .forwardStreamError(let error, let read):
                 return .forwardStreamError(error, read: read, cleanupContext: cleanupContext)
@@ -1044,8 +1136,19 @@ 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, let sync):
+            let cleanupContext = self.setErrorAndCreateCleanupContextIfNeeded(error)
+            return .failQueryContinuation(continuation, with: error, sync: sync, cleanupContext: cleanupContext)
         case .succeedQuery(let requestContext, with: let result):
             return .succeedQuery(requestContext, with: result)
+        case .succeedQueryContinuation(let continuation, let sync):
+            return .succeedQueryContinuation(continuation, sync: sync)
+        case .triggerCopyData(let triggerCopy):
+            return .triggerCopyData(triggerCopy)
+        case .sendCopyDoneAndSync:
+            return .sendCopyDoneAndSync
+        case .sendCopyFail(message: let message):
+            return .sendCopyFail(message: message)
         case .forwardRows(let buffer):
             return .forwardRows(buffer)
         case .forwardStreamComplete(let buffer, let commandTag):