Support binary data transfer in COPY FROM

My benchmark of transferring the integers from 0 to 1,000,000 both as an integer and as a string was about the same speed as the old text-based transfer. I believe that the binary transfer will start to show significant benefits when transferring binary data, other fields that don't need to be represented as fields and also means that the user doesn't need to worry about escapping their data.

Author: ahoppen

Sources/PostgresNIO/Connection/PostgresConnection+CopyFrom.swift

@@ -1,3 +1,110 @@
+import NIO
+
+#if compiler(>=6.0)
+/// Handle to send binary data for a `COPY ... FROM STDIN` query to the backend.
+///
+/// It takes care of serializing `PostgresEncodable` column types into the binary format that Postgres expects.
+public struct PostgresBinaryCopyFromWriter: ~Copyable {
+    /// Handle to serialize columns into a row that is being written by `PostgresBinaryCopyFromWriter`.
+    public struct ColumnWriter: ~Copyable {
+        /// The `PostgresBinaryCopyFromWriter` that is gathering the serialized data.
+        ///
+        /// We need to model this as `UnsafeMutablePointer` because we can't express in the Swift type system that
+        /// `ColumnWriter` never exceeds the lifetime of `PostgresBinaryCopyFromWriter`.
+        @usableFromInline
+        let underlying: UnsafeMutablePointer<PostgresBinaryCopyFromWriter>
+
+        /// The number of columns that have been written by this `ColumnWriter`.
+        @usableFromInline
+        var columns: UInt16 = 0
+
+        @usableFromInline
+        init(underlying: UnsafeMutablePointer<PostgresBinaryCopyFromWriter>) {
+            self.underlying = underlying
+        }
+
+        /// Serialize a single column to a row.
+        ///
+        /// - Important: It is critical that that data type encoded here exactly matches the data type in the
+        ///   databasse. For example, if the database stores an a 4-bit integer the corresponding `writeColumn` must
+        ///   be called with an `Int32`. Serializing an integer of a different width will cause a deserialization
+        ///   failure in the backend.
+        @inlinable
+        public mutating func writeColumn(_ column: (some PostgresEncodable)?) throws {
+            columns += 1
+            try underlying.pointee.writeColumn(column)
+        }
+    }
+
+    /// The underlying `PostgresCopyFromWriter` that sends the serialized data to the backend.
+    @usableFromInline let underlying: PostgresCopyFromWriter
+
+    /// The buffer in which we accumulate binary data. Once this buffer exceeds `bufferSize`, we flush it to
+    /// the backend.
+    @usableFromInline var buffer = ByteBuffer()
+
+    /// Once `buffer` exceeds this size, it gets flushed to the backend.
+    @usableFromInline let bufferSize: Int
+
+    init(underlying: PostgresCopyFromWriter, bufferSize: Int) {
+        self.underlying = underlying
+        // Allocate 10% more than the buffer size because we only flush the buffer once it has exceeded `bufferSize`
+        buffer.reserveCapacity(bufferSize + bufferSize / 10)
+        self.bufferSize = bufferSize
+    }
+
+    /// Serialize a single row to the backend. Call `writeColumn` on `columnWriter` for every column that should be
+    /// included in the row.
+    @inlinable
+    public mutating func writeRow(_ body: (_ columnWriter: inout ColumnWriter) throws -> Void) async throws {
+        // Write a placeholder for the number of columns
+        let columnIndex = buffer.writerIndex
+        buffer.writeInteger(UInt16(0))
+
+        let columns = try withUnsafeMutablePointer(to: &self) { pointerToSelf in
+            // Important: We need to ensure that `pointerToSel` (and thus `ColumnWriter`) does not exceed the lifetime
+            // of `self` because it is holding an unsafe reference to it.
+            //
+            // We achieve this because `ColumnWriter` is non-Copyable and thus the client can't store a copy to it.
+            // Futhermore `columnWriter` is destroyed before the end of `withUnsafeMutablePointer`, which holds `self`
+            // alive.
+            var columnWriter = ColumnWriter(underlying: pointerToSelf)
+
+            try body(&columnWriter)
+
+            return columnWriter.columns
+        }
+
+        // Fill in the number of columns
+        buffer.setInteger(columns, at: columnIndex)
+
+        if buffer.readableBytes > bufferSize {
+            try await flush()
+        }
+    }
+
+    /// Serialize a single column to the buffer. Should only be called by `ColumnWriter`.
+    @inlinable
+    mutating func writeColumn(_ column: (some PostgresEncodable)?) throws {
+        if let column {
+            let sizeIndex = buffer.readableBytes
+            buffer.writeInteger(Int32(0))
+            try column.encode(into: &buffer, context: .default)
+            buffer.setInteger(Int32(buffer.readableBytes - sizeIndex - 4), at: sizeIndex)
+        } else {
+            buffer.writeInteger(Int32(-1))
+        }
+    }
+
+    /// Flush any pending data in the buffer to the backend.
+    @usableFromInline
+    mutating func flush(isolation: (any Actor)? = #isolation) async throws {
+        try await underlying.write(buffer)
+        buffer.clear()
+    }
+}
+#endif
+
 /// Handle to send data for a `COPY ... FROM STDIN` query to the backend.
 public struct PostgresCopyFromWriter: Sendable {
     private let channelHandler: NIOLoopBound<PostgresChannelHandler>
@@ -115,15 +222,25 @@ public struct PostgresCopyFromFormat: Sendable {
         public init() {}
     }
 
+    /// Options that can be used to modify the `binary` format of a COPY operation.
+    public struct BinaryOptions: Sendable {
+        public init() {}
+    }
+
     enum Format {
         case text(TextOptions)
+        case binary(BinaryOptions)
     }
 
     var format: Format
 
     public static func text(_ options: TextOptions) -> PostgresCopyFromFormat {
         return PostgresCopyFromFormat(format: .text(options))
     }
+
+    public static func binary(_ options: BinaryOptions) -> PostgresCopyFromFormat {
+        return PostgresCopyFromFormat(format: .binary(options))
+    }
 }
 
 #if compiler(>=6.0)
@@ -156,6 +273,8 @@ private func buildCopyFromQuery(
             // Set the delimiter as a Unicode code point. This avoids the possibility of SQL injection.
             queryOptions.append("DELIMITER U&'\\\(String(format: "%04x", delimiter.value))'")
         }
+    case .binary:
+        queryOptions.append("FORMAT binary")
     }
     precondition(!queryOptions.isEmpty)
     query += " WITH ("
@@ -165,6 +284,50 @@ private func buildCopyFromQuery(
 }
 
 extension PostgresConnection {
+    /// Copy data into a table using a `COPY <table name> FROM STDIN` query, transferring data in a binary format.
+    ///
+    /// - 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.
+    ///   - bufferSize: How many bytes to accumulate a local buffer before flushing it to the database. Can affect
+    ///     performance characteristics of the copy operation.
+    ///   - 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.
+    ///
+    /// - Important: The table and column names are inserted into the `COPY FROM` query as passed and might thus be
+    ///   susceptible to SQL injection. Ensure no untrusted data is contained in these strings.
+    public func copyFromBinary(
+        table: String,
+        columns: [String] = [],
+        options: PostgresCopyFromFormat.BinaryOptions = .init(),
+        bufferSize: Int = 100_000,
+        logger: Logger,
+        isolation: isolated (any Actor)? = #isolation,
+        file: String = #fileID,
+        line: Int = #line,
+        writeData: @escaping @Sendable (inout PostgresBinaryCopyFromWriter) async throws -> Void
+    )  async throws {
+        try await copyFrom(table: table, columns: columns, format: .binary(PostgresCopyFromFormat.BinaryOptions()), logger: logger) { writer in
+            var header = ByteBuffer()
+            header.writeString("PGCOPY\n")
+            header.writeInteger(UInt8(0xff))
+            header.writeString("\r\n\0")
+
+            // Flag fields
+            header.writeInteger(UInt32(0))
+
+            // Header extension area length
+            header.writeInteger(UInt32(0))
+            try await writer.write(header)
+
+            var binaryWriter = PostgresBinaryCopyFromWriter(underlying: writer, bufferSize: bufferSize)
+            try await writeData(&binaryWriter)
+            try await binaryWriter.flush()
+        }
+    }
+
     /// Copy data into a table using a `COPY <table name> FROM STDIN` query.
     ///
     /// - Parameters:

Tests/IntegrationTests/PSQLIntegrationTests.swift

@@ -488,5 +488,40 @@ final class IntegrationTests: XCTestCase {
             XCTAssertEqual((error as? PSQLError)?.serverInfo?[.sqlState], "42601") // scanner_yyerror
         }
     }
+
+    func testCopyFromBinary() async throws {
+        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 2)
+        defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
+        let eventLoop = eventLoopGroup.next()
+
+        let conn = try await PostgresConnection.test(on: eventLoop).get()
+        defer { XCTAssertNoThrow(try conn.close().wait()) }
+
+        _ = try? await conn.query("DROP TABLE copy_table", logger: .psqlTest).get()
+        _ = try await conn.query("CREATE TABLE copy_table (id INT, name VARCHAR(100))", logger: .psqlTest).get()
+
+        try await conn.copyFromBinary(table: "copy_table", columns: ["id", "name"], logger: .psqlTest) { writer in
+            let records: [(id: Int, name: String)] = [
+                (1, "Alice"),
+                (42, "Bob")
+            ]
+            for record in records {
+                try await writer.writeRow { columnWriter in
+                    try columnWriter.writeColumn(Int32(record.id))
+                    try columnWriter.writeColumn(record.name)
+                }
+            }
+        }
+        let rows = try await conn.query("SELECT id, name FROM copy_table").get().rows.map { try $0.decode((Int, String).self) }
+        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)
+        XCTAssertEqual(rows[1].1, "Bob")
+    }
+
     #endif
 }

Tests/PostgresNIOTests/New/PostgresConnectionTests.swift

@@ -912,6 +912,68 @@ class PostgresConnectionTests: XCTestCase {
             try await connection.closeFuture.get()
         }
     }
+
+    func testCopyFromBinary() async throws {
+        let (connection, channel) = try await self.makeTestConnectionWithAsyncTestingChannel()
+
+        try await withThrowingTaskGroup(of: Void.self) { taskGroup async throws -> () in
+            taskGroup.addTask {
+                try await connection.copyFromBinary(table: "copy_table", logger: .psqlTest) { writer in
+                    try await writer.writeRow { columnWriter in
+                        try columnWriter.writeColumn(Int32(1))
+                        try columnWriter.writeColumn("Alice")
+                    }
+                    try await writer.writeRow { columnWriter in
+                        try columnWriter.writeColumn(Int32(2))
+                        try columnWriter.writeColumn("Bob")
+                    }
+                }
+            }
+
+            let copyRequest = try await channel.waitForUnpreparedRequest()
+            XCTAssertEqual(copyRequest.parse.query, #"COPY "copy_table" FROM STDIN WITH (FORMAT binary)"#)
+
+            try await channel.sendUnpreparedRequestWithNoParametersBindResponse()
+            try await channel.writeInbound(PostgresBackendMessage.copyInResponse(.init(format: .binary, columnFormats: [.binary, .binary])))
+
+            let copyData = try await channel.waitForCopyData()
+            XCTAssertEqual(copyData.result, .done)
+            var data = copyData.data
+            // Signature
+            XCTAssertEqual(data.readString(length: 7), "PGCOPY\n")
+            XCTAssertEqual(data.readInteger(as: UInt8.self), 0xff)
+            XCTAssertEqual(data.readString(length: 3), "\r\n\0")
+            // Flags
+            XCTAssertEqual(data.readInteger(as: UInt32.self), 0)
+            // Header extension area length
+            XCTAssertEqual(data.readInteger(as: UInt32.self), 0)
+
+            struct Row: Equatable {
+                let id: Int32
+                let name: String
+            }
+            var rows: [Row] = []
+            while data.readableBytes > 0 {
+                // Number of columns
+                XCTAssertEqual(data.readInteger(as: UInt16.self), 2)
+                // 'id' column
+                XCTAssertEqual(data.readInteger(as: UInt32.self), 4)
+                let id = try XCTUnwrap(data.readInteger(as: Int32.self))
+                // 'name' column length
+                let nameLength = try XCTUnwrap(data.readInteger(as: UInt32.self))
+                let name = try XCTUnwrap(data.readString(length: Int(nameLength)))
+                rows.append(Row(id: id, name: name))
+            }
+            XCTAssertEqual(rows, [
+                Row(id: 1, name: "Alice"),
+                Row(id: 2, name: "Bob")
+            ])
+            try await channel.writeInbound(PostgresBackendMessage.commandComplete("COPY 1"))
+
+            try await channel.waitForPostgresFrontendMessage(\.sync)
+            try await channel.writeInbound(PostgresBackendMessage.readyForQuery(.idle))
+        }
+    }
     #endif
 
     func makeTestConnectionWithAsyncTestingChannel() async throws -> (PostgresConnection, NIOAsyncTestingChannel) {