Implement Copy Support (#18)

Co-authored-by: hfhbd <[email protected]>

Author: hfhbd

build.gradle.kts

@@ -48,6 +48,7 @@ kotlin {
         commonTest {
             dependencies {
                 implementation(kotlin("test"))
+                implementation("org.jetbrains.kotlinx:kotlinx-coroutines-test:1.6.2")
             }
         }
     }

src/commonMain/kotlin/app/softwork/sqldelight/postgresdriver/PostgresNativeDriver.kt

@@ -83,10 +83,14 @@ class PostgresNativeDriver(private var conn: CPointer<PGconn>) : SqlDriver {
                 )
             }
         }.check(conn)
-        val rows = PQcmdTuples(result)!!.toKString()
-        result.clear()
-        val resultRows = rows.toLongOrNull() ?: 0
-        return QueryResult.Value(value = resultRows)
+
+        return QueryResult.Value(value = result.rows)
+    }
+
+    private val CPointer<PGresult>.rows: Long get() {
+        val rows = PQcmdTuples(this)!!.toKString()
+        clear()
+        return rows.toLongOrNull() ?: 0
     }
 
     private fun preparedStatementExists(identifier: Int): Boolean {
@@ -194,6 +198,19 @@ class PostgresNativeDriver(private var conn: CPointer<PGconn>) : SqlDriver {
             transaction = enclosingTransaction
         }
     }
+
+    fun copy(stdin: String): Long {
+        val status = PQputCopyData(conn, stdin, stdin.encodeToByteArray().size)
+        check(status == 1) {
+            conn.error()
+        }
+        val end = PQputCopyEnd(conn, null)
+        check(end == 1) {
+            conn.error()
+        }
+        val result = PQgetResult(conn).check(conn)
+        return result.rows
+    }
 }
 
 private fun CPointer<PGconn>?.error(): String {
@@ -214,7 +231,7 @@ private fun CPointer<PGconn>.exec(sql: String) {
 
 private fun CPointer<PGresult>?.check(conn: CPointer<PGconn>): CPointer<PGresult> {
     val status = PQresultStatus(this)
-    check(status == PGRES_TUPLES_OK || status == PGRES_COMMAND_OK) {
+    check(status == PGRES_TUPLES_OK || status == PGRES_COMMAND_OK || status == PGRES_COPY_IN) {
         conn.error()
     }
     return this!!

src/commonTest/kotlin/app/softwork/sqldelight/postgresdriver/PostgresNativeDriverTest.kt

@@ -119,4 +119,20 @@ class PostgresNativeDriverTest {
             )
         }
     }
+
+    @Test
+    fun copyTest() {
+        val driver = PostgresNativeDriver(
+            host = "localhost",
+            port = 5432,
+            user = "postgres",
+            database = "postgres",
+            password = "password"
+        )
+        assertEquals(0, driver.execute(null, "DROP TABLE IF EXISTS copying;", parameters = 0).value)
+        assertEquals(0, driver.execute(null, "CREATE TABLE copying(a int primary key);", parameters = 0).value)
+        driver.execute(42, "COPY copying FROM STDIN (FORMAT CSV);", 0)
+        val results = driver.copy("1\n2\n")
+        assertEquals(2, results)
+    }
 }

testing/src/commonMain/sqldelight/app/softwork/sqldelight/postgresdriver/Foo.sq

@@ -3,3 +3,6 @@ INSERT INTO foo VALUES ?;
 
 get:
 SELECT * FROM foo;
+
+copy:
+COPY foo FROM STDIN (FORMAT CSV);

testing/src/commonTest/kotlin/app/softwork/sqldelight/postgresdriver/PostgresNativeDriverTest.kt

@@ -25,10 +25,36 @@ class PostgresNativeDriverTest {
             date = LocalDate(2020, Month.DECEMBER, 12),
             timestamp = LocalDateTime(2014, Month.AUGUST, 1, 12, 1, 2, 0),
             instant = Instant.fromEpochMilliseconds(10L),
-            uuid = UUID.NIL,
-            interval = 42.seconds
+            interval = 42.seconds,
+            uuid = UUID.NIL
         )
         queries.create(foo)
         assertEquals(foo, queries.get().executeAsOne())
     }
+
+    @Test
+    fun copyTest() {
+        val driver = PostgresNativeDriver(
+            host = "localhost",
+            port = 5432,
+            user = "postgres",
+            database = "postgres",
+            password = "password"
+        )
+        val queries = NativePostgres(driver).fooQueries
+        NativePostgres.Schema.migrate(driver, 0, NativePostgres.Schema.version)
+        queries.copy()
+        val result = driver.copy("42,answer,2020-12-12,2014-08-01T12:01:02.0000,1970-01-01T00:00:00.010Z,PT42S,00000000-0000-0000-0000-000000000000")
+        assertEquals(1, result)
+        val foo = Foo(
+            a = 42,
+            b = "answer",
+            date = LocalDate(2020, Month.DECEMBER, 12),
+            timestamp = LocalDateTime(2014, Month.AUGUST, 1, 12, 1, 2, 0),
+            instant = Instant.fromEpochMilliseconds(10L),
+            interval = 42.seconds,
+            uuid = UUID.NIL,
+        )
+        assertEquals(foo, queries.get().executeAsOne())
+    }
 }