Add message types to support COPY operations (#569)

This adds the infrastrucutre to decode messages needed for COPY operations. It does not implement the handling support itself yet. That will be added in a follow-up PR.

Author: ahoppen

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

@@ -752,6 +752,12 @@ struct ConnectionStateMachine {
         return self.modify(with: action)
     }
 
+    mutating func copyInResponseReceived(
+        _ copyInResponse: PostgresBackendMessage.CopyInResponse
+    ) -> ConnectionAction {
+        return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.copyInResponse(copyInResponse)))
+    }
+
     mutating func emptyQueryResponseReceived() -> ConnectionAction {
         guard case .extendedQuery(var queryState, let connectionContext) = self.state, !queryState.isComplete else {
             return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.emptyQueryResponse))

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

@@ -91,7 +91,7 @@ struct ExtendedQueryStateMachine {
     mutating func cancel() -> Action {
         switch self.state {
         case .initialized:
-            preconditionFailure("Start must be called immediatly after the query was created")
+            preconditionFailure("Start must be called immediately after the query was created")
 
         case .messagesSent(let queryContext),
              .parseCompleteReceived(let queryContext),
@@ -322,6 +322,12 @@ struct ExtendedQueryStateMachine {
         }
     }
 
+    mutating func copyInResponseReceived(
+        _ copyInResponse: PostgresBackendMessage.CopyInResponse
+    ) -> Action {
+        return self.setAndFireError(.unexpectedBackendMessage(.copyInResponse(copyInResponse)))
+    }
+
     mutating func emptyQueryResponseReceived() -> Action {
         guard case .bindCompleteReceived(let queryContext) = self.state else {
             return self.setAndFireError(.unexpectedBackendMessage(.emptyQueryResponse))

Sources/PostgresNIO/New/Messages/CopyInMessage.swift

@@ -0,0 +1,44 @@
+extension PostgresBackendMessage {
+    struct CopyInResponse: Hashable {
+        enum Format: Int8 {
+            case textual = 0
+            case binary = 1
+        }
+
+        let format: Format
+        let columnFormats: [Format]
+
+        static func decode(from buffer: inout ByteBuffer) throws -> Self {
+            guard let rawFormat = buffer.readInteger(endianness: .big, as: Int8.self) else {
+                throw PSQLPartialDecodingError.expectedAtLeastNRemainingBytes(1, actual: buffer.readableBytes)
+            }
+            guard let format = Format(rawValue: rawFormat) else {
+                throw PSQLPartialDecodingError.unexpectedValue(value: rawFormat)
+            }
+            
+            guard let numColumns = buffer.readInteger(endianness: .big, as: Int16.self) else {
+                throw PSQLPartialDecodingError.expectedAtLeastNRemainingBytes(2, actual: buffer.readableBytes)
+            }
+            var columnFormatCodes: [Format] = []
+            columnFormatCodes.reserveCapacity(Int(numColumns))
+
+            for _ in 0..<numColumns {
+                guard let rawColumnFormat = buffer.readInteger(endianness: .big, as: Int16.self) else {
+                    throw PSQLPartialDecodingError.expectedAtLeastNRemainingBytes(2, actual: buffer.readableBytes)
+                }
+                guard Int8.min <= rawColumnFormat, rawColumnFormat <= Int8.max, let columnFormat = Format(rawValue: Int8(rawColumnFormat)) else {
+                    throw PSQLPartialDecodingError.unexpectedValue(value: rawColumnFormat)
+                }
+                columnFormatCodes.append(columnFormat)
+            }
+            
+            return CopyInResponse(format: format, columnFormats: columnFormatCodes)
+        }
+    }
+}
+
+extension PostgresBackendMessage.CopyInResponse: CustomDebugStringConvertible {
+    var debugDescription: String {
+        "format: \(format), columnFormats: \(columnFormats)"
+    }
+}

Sources/PostgresNIO/New/PostgresBackendMessage.swift

@@ -29,6 +29,7 @@ enum PostgresBackendMessage: Hashable {
     case bindComplete
     case closeComplete
     case commandComplete(String)
+    case copyInResponse(CopyInResponse)
     case dataRow(DataRow)
     case emptyQueryResponse
     case error(ErrorResponse)
@@ -96,6 +97,9 @@ extension PostgresBackendMessage {
             }
             return .commandComplete(commandTag)
 
+        case .copyInResponse:
+            return try .copyInResponse(.decode(from: &buffer))
+        
         case .dataRow:
             return try .dataRow(.decode(from: &buffer))
 
@@ -131,9 +135,9 @@ extension PostgresBackendMessage {
 
         case .rowDescription:
             return try .rowDescription(.decode(from: &buffer))
-            
-        case .copyData, .copyDone, .copyInResponse, .copyOutResponse, .copyBothResponse, .functionCallResponse, .negotiateProtocolVersion:
-            preconditionFailure()
+
+        case .copyData, .copyDone, .copyOutResponse, .copyBothResponse, .functionCallResponse, .negotiateProtocolVersion:
+            throw PSQLPartialDecodingError.unknownMessageKind(messageID)
         }
     }
 }
@@ -151,6 +155,8 @@ extension PostgresBackendMessage: CustomDebugStringConvertible {
             return ".closeComplete"
         case .commandComplete(let commandTag):
             return ".commandComplete(\(String(reflecting: commandTag)))"
+        case .copyInResponse(let copyInResponse):
+            return ".copyInResponse(\(String(reflecting: copyInResponse)))"
         case .dataRow(let dataRow):
             return ".dataRow(\(String(reflecting: dataRow)))"
         case .emptyQueryResponse:

Sources/PostgresNIO/New/PostgresBackendMessageDecoder.swift

@@ -189,6 +189,12 @@ struct PSQLPartialDecodingError: Error {
             description: "Expected the integer to be positive or null, but got \(actual).",
             file: file, line: line)
     }
+
+    static func unknownMessageKind(_ messageID: PostgresBackendMessage.ID, file: String = #fileID, line: Int = #line) -> Self {
+        return PSQLPartialDecodingError(
+            description: "Unknown message kind: \(messageID)",
+            file: file, line: line)
+    }
 }
 
 extension ByteBuffer {

Sources/PostgresNIO/New/PostgresChannelHandler.swift

@@ -136,6 +136,8 @@ final class PostgresChannelHandler: ChannelDuplexHandler {
             action = self.state.closeCompletedReceived()
         case .commandComplete(let commandTag):
             action = self.state.commandCompletedReceived(commandTag)
+        case .copyInResponse(let copyInResponse):
+            action = self.state.copyInResponseReceived(copyInResponse)
         case .dataRow(let dataRow):
             action = self.state.dataRowReceived(dataRow)
         case .emptyQueryResponse:

Sources/PostgresNIO/New/PostgresFrontendMessageEncoder.swift

@@ -167,6 +167,28 @@ struct PostgresFrontendMessageEncoder {
         self.buffer.writeMultipleIntegers(UInt32(8), Self.sslRequestCode)
     }
 
+    /// Adds the `CopyData` message ID and `dataLength` to the message buffer but not the actual data.
+    /// 
+    /// The caller of this function is expected to write the encoder's message buffer to the backend after calling this
+    /// function, followed by sending the actual data to the backend.
+    mutating func copyDataHeader(dataLength: UInt32) {
+        self.clearIfNeeded()
+        self.buffer.psqlWriteMultipleIntegers(id: .copyData, length: dataLength)
+    }
+
+    mutating func copyDone() {
+        self.clearIfNeeded()
+        self.buffer.psqlWriteMultipleIntegers(id: .copyDone, length: 0)
+    }
+
+    mutating func copyFail(message: String) {
+        self.clearIfNeeded()
+        var messageBuffer = ByteBuffer()
+        messageBuffer.writeNullTerminatedString(message)
+        self.buffer.psqlWriteMultipleIntegers(id: .copyFail, length: UInt32(messageBuffer.readableBytes))
+        self.buffer.writeImmutableBuffer(messageBuffer)
+    }
+
     mutating func sync() {
         self.clearIfNeeded()
         self.buffer.psqlWriteMultipleIntegers(id: .sync, length: 0)
@@ -197,6 +219,9 @@ struct PostgresFrontendMessageEncoder {
 private enum FrontendMessageID: UInt8, Hashable, Sendable {
     case bind = 66 // B
     case close = 67 // C
+    case copyData = 100 // d
+    case copyDone = 99 // c
+    case copyFail = 102 // f
     case describe = 68 // D
     case execute = 69 // E
     case flush = 72 // H

Tests/PostgresNIOTests/New/Connection State Machine/ExtendedQueryStateMachineTests.swift

@@ -114,7 +114,7 @@ class ExtendedQueryStateMachineTests: XCTestCase {
                        .failQuery(promise, with: psqlError, cleanupContext: .init(action: .close, tasks: [], error: psqlError, closePromise: nil)))
     }
 
-    func testExtendedQueryIsCancelledImmediatly() {
+    func testExtendedQueryIsCancelledImmediately() {
         var state = ConnectionStateMachine.readyForQuery()
 
         let logger = Logger.psqlTest

Tests/PostgresNIOTests/New/Extensions/PSQLBackendMessageEncoder.swift

@@ -28,6 +28,8 @@ struct PSQLBackendMessageEncoder: MessageToByteEncoder {
         case .commandComplete(let string):
             self.encode(messageID: message.id, payload: StringPayload(string), into: &buffer)
 
+        case .copyInResponse(let copyInResponse):
+            self.encode(messageID: message.id, payload: copyInResponse, into: &buffer)
         case .dataRow(let row):
             self.encode(messageID: message.id, payload: row, into: &buffer)
 
@@ -99,6 +101,8 @@ extension PostgresBackendMessage {
             return .closeComplete
         case .commandComplete:
             return .commandComplete
+        case .copyInResponse:
+            return .copyInResponse
         case .dataRow:
             return .dataRow
         case .emptyQueryResponse:
@@ -184,6 +188,16 @@ extension PostgresBackendMessage.BackendKeyData: PSQLMessagePayloadEncodable {
     }
 }
 
+extension PostgresBackendMessage.CopyInResponse: PSQLMessagePayloadEncodable {
+    public func encode(into buffer: inout ByteBuffer) {
+        buffer.writeInteger(Int8(self.format.rawValue))
+        buffer.writeInteger(Int16(self.columnFormats.count))
+        for columnFormat in columnFormats {
+            buffer.writeInteger(Int16(columnFormat.rawValue))
+        }
+    }
+}
+
 extension DataRow: PSQLMessagePayloadEncodable {
     public func encode(into buffer: inout ByteBuffer) {
         buffer.writeInteger(self.columnCount, as: Int16.self)

Tests/PostgresNIOTests/New/Extensions/PSQLFrontendMessageDecoder.swift

@@ -168,6 +168,18 @@ extension PostgresFrontendMessage {
                 )
             )
 
+        case .copyData:
+            return .copyData(CopyData(data: buffer))
+
+        case .copyDone:
+            return .copyDone
+
+        case .copyFail:
+            guard let message = buffer.readNullTerminatedString() else {
+                throw PSQLPartialDecodingError.fieldNotDecodable(type: String.self)
+            }
+            return .copyFail(CopyFail(message: message))
+
         case .close:
             preconditionFailure("TODO: Unimplemented")