Add message types to support COPY operations

ahoppen · ahoppen · commit e04e1c6c75d6 · 2025-07-07T12:36:32.000+02:00
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.
diff --git a/Sources/PostgresNIO/New/Connection State Machine/ConnectionStateMachine.swift b/Sources/PostgresNIO/New/Connection State Machine/ConnectionStateMachine.swift
@@ -752,6 +752,12 @@ struct ConnectionStateMachine {
         return self.modify(with: action)
     }
     
+    mutating func copyInResponseReceived(
+        _ copyInResponse: PostgresBackendMessage.CopyInResponseMessage
+    ) -> 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))
diff --git a/Sources/PostgresNIO/New/Connection State Machine/ExtendedQueryStateMachine.swift b/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.CopyInResponseMessage
+    ) -> Action {
+        return self.setAndFireError(.unexpectedBackendMessage(.copyInResponse(copyInResponse)))
+    }
+
     mutating func emptyQueryResponseReceived() -> Action {
         guard case .bindCompleteReceived(let queryContext) = self.state else {
             return self.setAndFireError(.unexpectedBackendMessage(.emptyQueryResponse))
diff --git a/Sources/PostgresNIO/New/Messages/CopyInMessage.swift b/Sources/PostgresNIO/New/Messages/CopyInMessage.swift
@@ -0,0 +1,44 @@
+extension PostgresBackendMessage {
+    struct CopyInResponseMessage: Hashable {
+        enum Format: Int {
+            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: Int(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 let columnFormat = Format(rawValue: Int(rawColumnFormat)) else {
+                    throw PSQLPartialDecodingError.unexpectedValue(value: rawColumnFormat)
+                }
+                columnFormatCodes.append(columnFormat)
+            }
+            
+            return CopyInResponseMessage(format: format, columnFormats: columnFormatCodes)
+        }
+    }
+}
+
+extension PostgresBackendMessage.CopyInResponseMessage: CustomDebugStringConvertible {
+    var debugDescription: String {
+        "format: \(format), columnFormats: \(columnFormats)"
+    }
+}
diff --git a/Sources/PostgresNIO/New/PostgresBackendMessage.swift b/Sources/PostgresNIO/New/PostgresBackendMessage.swift
@@ -29,6 +29,7 @@ enum PostgresBackendMessage: Hashable {
     case bindComplete
     case closeComplete
     case commandComplete(String)
+    case copyInResponse(CopyInResponseMessage)
     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:
+            preconditionFailure("Unknown message kind: \(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:
diff --git a/Sources/PostgresNIO/New/PostgresChannelHandler.swift b/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:
diff --git a/Sources/PostgresNIO/New/PostgresFrontendMessageEncoder.swift b/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
diff --git a/Tests/PostgresNIOTests/New/Connection State Machine/ExtendedQueryStateMachineTests.swift b/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
diff --git a/Tests/PostgresNIOTests/New/Extensions/PSQLBackendMessageEncoder.swift b/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.CopyInResponseMessage: 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)
diff --git a/Tests/PostgresNIOTests/New/Extensions/PSQLFrontendMessageDecoder.swift b/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")
 
diff --git a/Tests/PostgresNIOTests/New/Extensions/PostgresFrontendMessage.swift b/Tests/PostgresNIOTests/New/Extensions/PostgresFrontendMessage.swift
@@ -36,6 +36,14 @@ enum PostgresFrontendMessage: Equatable {
         let secretKey: Int32
     }
 
+    struct CopyData: Equatable {
+        let data: ByteBuffer
+    }
+
+    struct CopyFail: Equatable {
+        let message: String
+    }
+
     enum Close: Hashable {
         case preparedStatement(String)
         case portal(String)
@@ -170,6 +178,9 @@ enum PostgresFrontendMessage: Equatable {
 
     case bind(Bind)
     case cancel(Cancel)
+    case copyData(CopyData)
+    case copyDone
+    case copyFail(CopyFail)
     case close(Close)
     case describe(Describe)
     case execute(Execute)
@@ -186,6 +197,9 @@ enum PostgresFrontendMessage: Equatable {
     enum ID: UInt8, Equatable {
         
         case bind
+        case copyData
+        case copyDone
+        case copyFail
         case close
         case describe
         case execute
@@ -201,12 +215,18 @@ enum PostgresFrontendMessage: Equatable {
             switch rawValue {
             case UInt8(ascii: "B"):
                 self = .bind
+            case UInt8(ascii: "c"):
+                self = .copyDone
             case UInt8(ascii: "C"):
                 self = .close
+            case UInt8(ascii: "d"):
+                self = .copyData
             case UInt8(ascii: "D"):
                 self = .describe
             case UInt8(ascii: "E"):
                 self = .execute
+            case UInt8(ascii: "f"):
+                self = .copyFail
             case UInt8(ascii: "H"):
                 self = .flush
             case UInt8(ascii: "P"):
@@ -230,6 +250,12 @@ enum PostgresFrontendMessage: Equatable {
             switch self {
             case .bind:
                 return UInt8(ascii: "B")
+            case .copyData:
+                return UInt8(ascii: "d")
+            case .copyDone:
+                return UInt8(ascii: "c")
+            case .copyFail:
+                return UInt8(ascii: "f")
             case .close:
                 return UInt8(ascii: "C")
             case .describe:
@@ -263,6 +289,12 @@ extension PostgresFrontendMessage {
             return .bind
         case .cancel:
             preconditionFailure("Cancel messages don't have an identifier")
+        case .copyData:
+            return .copyData
+        case .copyDone:
+            return .copyDone
+        case .copyFail:
+            return .copyFail
         case .close:
             return .close
         case .describe:
diff --git a/Tests/PostgresNIOTests/New/Messages/CopyTests.swift b/Tests/PostgresNIOTests/New/Messages/CopyTests.swift

Original file line number	Diff line number	Diff line change
`@@ -114,7 +114,7 @@ class ExtendedQueryStateMachineTests: XCTestCase {`
`114`	`114`	`.failQuery(promise, with: psqlError, cleanupContext: .init(action: .close, tasks: [], error: psqlError, closePromise: nil)))`
`115`	`115`	`}`
`116`	`116`
`117`		`- func testExtendedQueryIsCancelledImmediatly() {`
	`117`	`+ func testExtendedQueryIsCancelledImmediately() {`
`118`	`118`	`var state = ConnectionStateMachine.readyForQuery()`
`119`	`119`
`120`	`120`	`let logger = Logger.psqlTest`