Add CopyOut! API

Author: cmey

src/copy.jl

@@ -81,3 +81,87 @@ function execute(
         Result(copy_end_result, jl_conn, kwargs...); throw_error=throw_error
     )
 end
+
+"""
+    CopyOut!(data, query) -> CopyOut!
+
+Create a `CopyOut!` query instance which can be executed to receive data from PostgreSQL via a
+`COPY <table_name> TO STDIN` query.
+
+`query` must be a `COPY TO STDIN` query as described in the [PostgreSQL documentation](https://www.postgresql.org/docs/10/sql-copy.html).
+`COPY TO` queries which use a file or `PROGRAM` source can instead use the standard
+[`execute`](@ref) query interface.
+
+`data` is an IOBuffer where strings of data received from PostgreSQL are written to.
+The data is received as text in CSV format.
+"""
+struct CopyOut!
+    data::IOBuffer
+    query::String
+end
+
+"""
+    execute(jl_conn::Connection, copyout::CopyOut!, args...;
+        throw_error::Bool=true, kwargs...
+    ) -> Result
+
+Runs [`execute`](@ref execute(::Connection, ::String)) on `copyout`'s query, then fills
+`copyout`'s data from the server.
+
+All other arguments are passed through to the `execute` call for the initial query.
+"""
+function execute(
+    jl_conn::Connection,
+    copy::CopyOut!,
+    parameters=nothing;
+    throw_error=true,
+    kwargs...,
+)
+    level = throw_error ? error : warn
+    if parameters !== nothing
+        string_params = string_parameters(parameters)
+        pointer_params = parameter_pointers(string_params)
+    end
+
+    copy_end_result = lock(jl_conn) do
+        if parameters === nothing
+            result = _execute(jl_conn.conn, copy.query)
+        else
+            result = _execute(jl_conn.conn, copy.query, pointer_params)
+        end
+        result_status = libpq_c.PQresultStatus(result)
+
+        if result_status != libpq_c.PGRES_COPY_OUT
+            if !(result_status in (libpq_c.PGRES_BAD_RESPONSE, libpq_c.PGRES_FATAL_ERROR))
+                level(LOGGER, Errors.JLResultError(
+                    "Expected PGRES_COPY_OUT after COPY query, got $result_status"
+                ))
+            end
+            return result
+        end
+
+        io = copy.data  # store csv string
+        async::Cint = 0  # blocking call
+        rowRef = Ref{Cstring}()
+        status_code = Cint(0)
+        while (status_code = libpq_c.PQgetCopyData(jl_conn.conn, rowRef, async)) > 0
+            rowPtr = rowRef[]
+            write(io, unsafe_string(rowPtr))
+            if rowPtr != C_NULL
+                libpq_c.PQfreemem(convert(Ptr{Cvoid}, rowPtr))
+            end
+        end
+        seekstart(io)  # rewind iobuffer so future user read will begin from start
+        if -2 == status_code
+            level(LOGGER, Errors.JLResultError(
+                    "PQgetCopyData error: $(error_message(jl_conn))"
+                ))
+        end
+
+        libpq_c.PQgetResult(jl_conn.conn)
+    end
+
+    return handle_result(
+        Result(copy_end_result, jl_conn, kwargs...); throw_error=throw_error
+    )
+end

test/runtests.jl

@@ -1,5 +1,6 @@
 using LibPQ
 using Test
+using CSV
 using Dates
 using DataFrames
 using DataFrames: eachrow
@@ -444,6 +445,59 @@ end
         end
     end
 
+    @testset "COPY TO" begin
+        @testset "Example COPY TO" begin
+            conn = LibPQ.Connection("dbname=postgres user=$DATABASE_USER")
+
+            # make data
+            no_nulls = map(string, 'a':'z')
+            yes_nulls = Union{String, Missing}[isodd(Int(c)) ? string(c) : missing for c in 'a':'z']
+            data = DataFrame(no_nulls=no_nulls, yes_nulls=yes_nulls)
+
+            row_strings = imap(eachrow(data)) do row
+                if ismissing(row[:yes_nulls])
+                    "$(row[:no_nulls]),\n"
+                else
+                    "$(row[:no_nulls]),$(row[:yes_nulls])\n"
+                end
+            end
+
+            # setup db table
+            result = execute(conn, """
+                CREATE TEMPORARY TABLE libpqjl_test (
+                    no_nulls    varchar(10) PRIMARY KEY,
+                    yes_nulls   varchar(10)
+                );
+            """)
+            @test status(result) == LibPQ.libpq_c.PGRES_COMMAND_OK
+            close(result)
+
+            # populate db table
+            copyin = LibPQ.CopyIn("COPY libpqjl_test FROM STDIN (FORMAT CSV);", row_strings)
+
+            result = execute(conn, copyin)
+            @test isopen(result)
+            @test status(result) == LibPQ.libpq_c.PGRES_COMMAND_OK
+            @test isempty(LibPQ.error_message(result))
+            close(result)
+
+            # test CopyOut!
+            databuf = IOBuffer()
+            copyout = LibPQ.CopyOut!(databuf, "COPY (SELECT * FROM libpqjl_test) TO STDOUT (FORMAT CSV, HEADER, ENCODING 'UTF8');")
+
+            result = execute(conn, copyout)
+            @test isopen(result)
+            @test status(result) == LibPQ.libpq_c.PGRES_COMMAND_OK
+            @test isempty(LibPQ.error_message(result))
+            close(result)
+
+            df = DataFrame(CSV.File(databuf, stringtype=String))
+            @test isequal(data, df)
+
+            close(conn)
+        end
+    end
+
     @testset "LibPQ.Connection" begin
         @testset "do" begin
             local saved_conn