Restructure the copyFrom function to take structured input instead of a query

Author: ahoppen

Sources/PostgresNIO/Connection/PostgresConnection.swift

@@ -786,25 +786,54 @@ public struct PostgresCopyFromWriter: Sendable {
     }
 }
 
+public struct CopyFromOptions {
+    let delimiter: StaticString?
+
+    public init(delimiter: StaticString? = nil) {
+        self.delimiter = delimiter
+    }
+}
+
+private func buildCopyFromQuery(
+    table: String,
+    columns: [StaticString]?,
+    options: CopyFromOptions
+) -> PostgresQuery {
+    var query = "COPY \(table)"
+    if let columns {
+        // TODO: Is using `StaticString` sufficient here to prevent against SQL injection attacks or should we try to 
+        // escape the identifiers, essentially re-implementing `PQescapeIdentifier`?
+        query += "(" + columns.map(\.description).joined(separator: ",") + ")"
+    }
+    query += " FROM STDIN"
+    var queryOptions: [String] = []
+    if let delimiter = options.delimiter {
+        queryOptions.append("DELIMITER '\(delimiter)'")
+    }
+    if !queryOptions.isEmpty {
+        query += " WITH "
+        query += queryOptions.map { "(\($0))" }.joined(separator: ",")
+    }
+    return "\(unescaped: query)"
+}
+
 extension PostgresConnection {
     // TODO: Instead of an arbitrary query, make this a structured data structure.
     // TODO: Write doc comment
     public func copyFrom(
-        _ query: PostgresQuery,
-        writeData: @escaping @Sendable (PostgresCopyFromWriter) async throws -> Void,
+        table: String,
+        columns: [StaticString]? = nil,
+        options: CopyFromOptions = CopyFromOptions(),
         logger: Logger,
+        writeData: @escaping @Sendable (PostgresCopyFromWriter) async throws -> Void,
         file: String = #fileID,
         line: Int = #line
     )  async throws {
         var logger = logger
         logger[postgresMetadataKey: .connectionID] = "\(self.id)"
-        guard query.binds.count <= Int(UInt16.max) else {
-            throw PSQLError(code: .tooManyParameters, query: query)
-        }
-
         let writer = try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<PostgresCopyFromWriter, any Error>) in
             let context = ExtendedQueryContext(
-                copyFromQuery: query,
+                copyFromQuery: buildCopyFromQuery(table: table, columns: columns, options: options),
                 triggerCopy: continuation,
                 logger: logger
             )

Tests/IntegrationTests/PSQLIntegrationTests.swift

@@ -389,17 +389,17 @@ final class IntegrationTests: XCTestCase {
 
         _ = 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.copyFrom("COPY copy_table FROM STDIN", writeData: { writer in
+        try await conn.copyFrom(table: "copy_table", columns: ["id", "name"], options: CopyFromOptions(delimiter: ","), logger: .psqlTest) { writer in
             let records: [(id: Int, name: String)] = [
                 (1, "Alice"),
                 (42, "Bob")
             ]
             for record in records {
                 var buffer = ByteBuffer()
-                buffer.writeString("\(record.id)\t\(record.name)\n")
+                buffer.writeString("\(record.id),\(record.name)\n")
                 try await writer.write(buffer)
             }
-        }, logger: .psqlTest)
+        }
         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)")

Tests/PostgresNIOTests/New/PostgresConnectionTests.swift

@@ -628,9 +628,9 @@ class PostgresConnectionTests: XCTestCase {
 
         try await withThrowingTaskGroup(of: Void.self) { taskGroup async throws -> () in
             taskGroup.addTask {
-                try await connection.copyFrom("COPY copy_table FROM STDIN", writeData: { writer in
+                try await connection.copyFrom(table: "copy_table", logger: .psqlTest) { writer in
                     try await writer.write(ByteBuffer(staticString: "1\tAlice\n"))
-                }, logger: .psqlTest)
+                }
             }
 
             let copyMessage = try await channel.waitForUnpreparedRequest()
@@ -656,9 +656,9 @@ class PostgresConnectionTests: XCTestCase {
         try await withThrowingTaskGroup(of: Void.self) { taskGroup async throws -> () in
             taskGroup.addTask {
                 await assertThrowsError(
-                    try await connection.copyFrom("COPY copy_table FROM STDIN", writeData: { writer in
+                    try await connection.copyFrom(table: "copy_table", logger: .psqlTest) { writer in
                         throw MyError()
-                    }, logger: .psqlTest)
+                    }
                 ) { error in 
                     XCTAssert(error is MyError, "Expected error of type MyError, got \(error)")
                 }
@@ -691,9 +691,9 @@ class PostgresConnectionTests: XCTestCase {
         try await withThrowingTaskGroup(of: Void.self) { taskGroup async throws -> () in
             taskGroup.addTask {
                 await assertThrowsError(
-                    try await connection.copyFrom("COPY copy_table FROM STDIN", writeData: { writer in
+                    try await connection.copyFrom(table: "copy_table", logger: .psqlTest) { writer in
                         try await writer.write(ByteBuffer(staticString: "1Alice\n"))
-                    }, logger: .psqlTest)
+                    }
                 ) { error in
                     XCTAssertEqual((error as? PSQLError)?.serverInfo?.underlying.fields[.sqlState], "22P02")
                 }
@@ -724,11 +724,11 @@ class PostgresConnectionTests: XCTestCase {
         try await withThrowingTaskGroup(of: Void.self) { taskGroup async throws -> () in
             taskGroup.addTask {
                 await assertThrowsError(
-                    try await connection.copyFrom("COPY copy_table FROM STDIN", writeData: { writer in
+                    try await connection.copyFrom(table: "copy_table", logger: .psqlTest) { writer in
                         try await writer.write(ByteBuffer(staticString: "1Alice\n"))
                         channel.flush()
                         _ = await XCTWaiter.fulfillment(of: [backendDidSendErrorExpectation])
-                    }, logger: .psqlTest)
+                    }
                 ) { error in
                     XCTAssertEqual((error as? PSQLError)?.serverInfo?.underlying.fields[.sqlState], "22P02")
                 }
@@ -764,7 +764,7 @@ class PostgresConnectionTests: XCTestCase {
         try await withThrowingTaskGroup(of: Void.self) { taskGroup async throws -> () in
             taskGroup.addTask {
                 await assertThrowsError(
-                    try await connection.copyFrom("COPY copy_table FROM STDIN", writeData: { writer in
+                    try await connection.copyFrom(table: "copy_table", logger: .psqlTest) { writer in
                         try await writer.write(ByteBuffer(staticString: "1Alice\n"))
                         channel.flush()
                         _ = await XCTWaiter.fulfillment(of: [expectation])
@@ -775,7 +775,7 @@ class PostgresConnectionTests: XCTestCase {
                             XCTAssert(error is PostgresCopyFromWriter.CopyCancellationError, "Received unexpected error: \(error)")
                             throw error
                         }
-                    }, logger: .psqlTest)
+                    }
                 ) { error in
                     XCTAssertEqual((error as? PSQLError)?.serverInfo?.underlying.fields[.sqlState], "22P02")
                 }
@@ -803,6 +803,29 @@ class PostgresConnectionTests: XCTestCase {
         }
     }
 
+    func testCopyDataWithOptions() async throws {
+        let (connection, channel) = try await self.makeTestConnectionWithAsyncTestingChannel()
+
+        try await withThrowingTaskGroup(of: Void.self) { taskGroup async throws -> () in
+            taskGroup.addTask {
+                try await connection.copyFrom(table: "copy_table", columns: ["id", "name"], options: CopyFromOptions(delimiter: ","), logger: .psqlTest) { writer in
+                    try await writer.write(ByteBuffer(staticString: "1,Alice\n"))
+                }
+            }
+
+            let copyMessage = try await channel.waitForUnpreparedRequest()
+            XCTAssertEqual(copyMessage.parse.query, "COPY copy_table(id,name) FROM STDIN WITH (DELIMITER ',')")
+            XCTAssertEqual(copyMessage.bind.parameters, [])
+            try await channel.sendUnpreparedRequestWithNoParametersBindResponse()
+            try await channel.sendCopyInResponseForTwoTextualColumns()
+            let data = try await channel.waitForCopyData()
+            XCTAssertEqual(String(buffer: data.data), "1,Alice\n")
+            XCTAssertEqual(data.result, .done)
+            try await channel.writeInbound(PostgresBackendMessage.commandComplete("COPY 1"))
+            try await channel.writeInbound(PostgresBackendMessage.readyForQuery(.idle))
+        }
+    }
+
     func makeTestConnectionWithAsyncTestingChannel() async throws -> (PostgresConnection, NIOAsyncTestingChannel) {
         let eventLoop = NIOAsyncTestingEventLoop()
         let channel = try await NIOAsyncTestingChannel(loop: eventLoop) { channel in