Add put_copy_data, put_copy_end, and get_copy_data.

These were introduced in libpq 7.4, obsoleting the other functions
dealing with COPY data.

Author: paurkedal

examples/copy.ml

@@ -0,0 +1,62 @@
+open Postgresql
+open Printf
+
+let _ =
+  if Array.length Sys.argv <> 2 then (
+    Printf.printf "\
+      Usage:  copy conninfo\n\
+      Connect to PostgreSQL with [conninfo] (e.g. \"host=localhost\"),\n";
+    exit 1)
+
+let create_sql = "\
+  CREATE TEMPORARY TABLE pgo_copy_test (\
+    id SERIAL PRIMARY KEY, \
+    name text UNIQUE NOT NULL, \
+    x integer NOT NULL\
+  )"
+
+let populate ?error_msg (c : connection) =
+  let _ = c#exec ~expect:[Copy_in] "COPY pgo_copy_test (name, x) FROM STDIN" in
+  for i = 0 to 9999 do
+    match c#put_copy_data (sprintf "c%d\t%d\n" i i) with
+    | Put_copy_error | Put_copy_not_queued -> assert false
+    | Put_copy_queued -> ()
+  done;
+  begin match c#put_copy_end ?error_msg () with
+  | Put_copy_error | Put_copy_not_queued -> assert false
+  | Put_copy_queued -> ()
+  end;
+  match c#get_result with
+  | None -> assert false
+  | Some result ->
+    begin match error_msg, result#status with
+    | None, Command_ok -> ()
+    | Some _msg, Fatal_error -> ()
+    | _ -> assert false
+    end
+
+let verify (c : connection) =
+  let _ =
+    c#exec ~expect:[Copy_out] "COPY pgo_copy_test (id, name, x) TO STDOUT" in
+  for i = 0 to 9999 do
+    match c#get_copy_data () with
+    | Get_copy_data data ->
+      assert (data = sprintf "%d\tc%d\t%d\n" (i + 1) i i)
+    | Get_copy_wait | Get_copy_end | Get_copy_error -> assert false
+  done;
+  match c#get_copy_data () with
+  | Get_copy_end -> ()
+  | _ -> assert false
+
+let main () =
+  let c = new connection ~conninfo:Sys.argv.(1) () in
+  let _ = c#exec ~expect:[Command_ok] create_sql in
+  populate c;
+  populate c ~error_msg:"test failure";
+  verify c;
+  c#finish
+
+let () =
+  try main () with
+  | Error e -> prerr_endline (string_of_error e)
+  | e -> prerr_endline (Printexc.to_string e)

examples/dune

@@ -1,4 +1,4 @@
 (executables
-  (names async binary cursor dump populate prompt test_lo)
+  (names async binary copy cursor dump populate prompt test_lo)
   (libraries postgresql)
 )

src/postgresql.ml

@@ -294,6 +294,17 @@ type result_status =
 
 external result_status : result_status -> string = "PQresStatus_stub"
 
+type put_copy_result =
+  | Put_copy_queued
+  | Put_copy_not_queued
+  | Put_copy_error
+
+type get_copy_result =
+  | Get_copy_data of string
+  | Get_copy_wait
+  | Get_copy_end
+  | Get_copy_error
+
 type getline_result = EOF | LineRead | BufFull
 
 type getline_async_result =
@@ -505,6 +516,19 @@ module Stub = struct
 
   (* Functions Associated with the COPY Command *)
 
+  external put_copy_data :
+    connection -> string ->
+    (int [@untagged]) -> (int [@untagged]) -> (int [@untagged])
+    = "PQputCopyData_bc" "PQputCopyData_stub"
+
+  external put_copy_end :
+    connection -> string option -> (int [@untagged])
+    = "PQputCopyEnd_bc" "PQputCopyEnd_stub"
+
+  external get_copy_data :
+    connection -> (int [@untagged]) -> get_copy_result
+    = "PQgetCopyData_bc" "PQgetCopyData_stub"
+
   external getline :
     connection -> Bytes.t ->
     (int [@untagged]) -> (int [@untagged]) -> (int [@untagged])
@@ -976,6 +1000,29 @@ object (self)
 
   (* Low level *)
 
+  method put_copy_data ?(pos = 0) ?len buf =
+    let buf_len = String.length buf in
+    let len = match len with Some len -> len | None -> buf_len - pos in
+    if len < 0 || pos < 0 || pos + len > buf_len then
+      invalid_arg "Postgresql.connection#put_copy_data";
+    wrap_conn (fun conn ->
+      match Stub.put_copy_data conn buf pos len with
+      | -1 -> Put_copy_error
+      | 0 -> Put_copy_not_queued
+      | 1 -> Put_copy_queued
+      | _ -> assert false)
+
+  method put_copy_end ?error_msg () =
+    wrap_conn (fun conn ->
+      match Stub.put_copy_end conn error_msg with
+      | -1 -> Put_copy_error
+      | 0 -> Put_copy_not_queued
+      | 1 -> Put_copy_queued
+      | _ -> assert false)
+
+  method get_copy_data ?(async = false) () =
+    wrap_conn (fun conn -> Stub.get_copy_data conn (if async then 1 else 0))
+
   method getline ?(pos = 0) ?len buf =
     let buf_len = Bytes.length buf in
     let len = match len with Some len -> len | None -> buf_len - pos in

src/postgresql.mli

@@ -124,6 +124,20 @@ type result_status =
   | Copy_both
   | Single_tuple    (** One tuple of a result set ({!set_single_row_mode}) *)
 
+(** Result of put_copy_data and put_copy_end *)
+type put_copy_result =
+  | Put_copy_queued     (** Data queued *)
+  | Put_copy_not_queued (** Data not queued due to full bufffers (async only) *)
+  | Put_copy_error      (** Copying failed, see [#error_message] for details *)
+
+(** Result of get_copy_data *)
+type get_copy_result =
+  | Get_copy_data of string (** Data corresponding to one row is retured *)
+  | Get_copy_wait   (** The next row is still being received (async only); wait
+                        for read-only, call [consume_input], and try again *)
+  | Get_copy_end    (** All data has been successfully retrieved *)
+  | Get_copy_error  (** Copying failed, see [#error_message] for details *)
+
 (** Result of getline *)
 type getline_result =
   | EOF       (** End of input reached *)
@@ -817,6 +831,40 @@ object
 
   (** Low level *)
 
+  method put_copy_data : ?pos : int -> ?len : int -> string -> put_copy_result
+  (** [put_copy_data ?pos ?len buf] sends [buf] of length [len] starting at
+      [pos] to the backend server, which must be in copy-in mode.  In
+      non-blocking mode, returns {!Put_copy_not_queued} if the data was not
+      queued due to full buffers.
+
+      @param pos default = 0
+      @param len default = String.length - pos
+
+      @raise Invalid_argument if the buffer parameters are invalid.
+   *)
+
+  method put_copy_end : ?error_msg : string -> unit -> put_copy_result
+  (** [put_copy_end ?error_msg ()] terminates the copy-in mode, leaving the
+      connection in [Command_ok] or failed state.  In non-blocking mode, returns
+      {!Put_copy_not_queued} if the termination message was not queued due to
+      full buffers.  Also, to ensure delivery of data in non-blocking mode,
+      repeatedly wait for write-ready an call {!#flush}.
+
+      @param error_msg if set, force the copy operation to fail with the given
+        message.
+   *)
+
+  method get_copy_data : ?async : bool -> unit -> get_copy_result
+  (** [get_copy_data ?async ()] retrieves the next row of data if available.
+      Only single complete rows are returned.  In synchronous mode, the call
+      will wait for completion of the next row.  In asynchronous mode it will
+      return immediately with [Get_copy_wait] if the row transfer is incomplete.
+      In that case, wait for read-ready and call {!#consume_input} before
+      retrying.
+
+      @param async default = false
+   *)
+
   method getline : ?pos : int -> ?len : int -> Bytes.t -> getline_result
   (** [getline ?pos ?len buf] reads a newline-terminated line of at most
       [len] characters into [buf] starting at position [pos].

src/postgresql_stubs.c

@@ -1211,6 +1211,86 @@ CAMLprim value PQnotifies_stub(value v_conn)
 
 /* Functions Associated with the COPY Command */
 
+CAMLprim intnat PQputCopyData_stub(
+  value v_conn, value v_buf, intnat pos, intnat len)
+{
+  CAMLparam2(v_conn, v_buf);
+  PGconn *conn = get_conn(v_conn);
+  intnat res;
+  char *buf = caml_stat_alloc(len);
+  memcpy(buf, String_val(v_buf) + pos, len);
+  caml_enter_blocking_section();
+    res = PQputCopyData(conn, buf, len);
+  caml_leave_blocking_section();
+  caml_stat_free(buf);
+  CAMLreturn(res);
+}
+
+CAMLprim value PQputCopyData_bc(
+  value v_conn, value v_buf, value v_pos, value v_len)
+{
+  return
+    Val_int(PQputCopyData_stub(v_conn, v_buf, Int_val(v_pos), Int_val(v_len)));
+}
+
+CAMLprim intnat PQputCopyEnd_stub(value v_conn, value v_msg_opt)
+{
+  CAMLparam2(v_conn, v_msg_opt);
+  intnat res;
+  PGconn *conn = get_conn(v_conn);
+  char *msg = NULL;
+  if (Is_block(v_msg_opt)) {
+    value v_msg = Field(v_msg_opt, 0);
+    size_t msg_len = caml_string_length(v_msg);
+    msg = caml_stat_alloc(msg_len + 1);
+    memcpy(msg, String_val(v_msg), msg_len);
+    msg[msg_len] = '\0';
+  }
+  caml_enter_blocking_section();
+    res = PQputCopyEnd(conn, msg);
+  caml_leave_blocking_section();
+  if (msg)
+    caml_stat_free(msg);
+  CAMLreturn(res);
+}
+
+CAMLprim value PQputCopyEnd_bc(value v_conn, value v_msg)
+{
+  return Val_int(PQputCopyEnd_stub(v_conn, v_msg));
+}
+
+CAMLprim value PQgetCopyData_stub(value v_conn, intnat async)
+{
+  CAMLparam1(v_conn);
+  CAMLlocal2(v_buf, v_result);
+  PGconn *conn = get_conn(v_conn);
+  char *buf;
+  intnat res;
+  caml_enter_blocking_section();
+    res = PQgetCopyData(conn, &buf, async);
+  caml_leave_blocking_section();
+  switch (res) {
+    case 0:
+      CAMLreturn(Val_int(0)); /* Get_copy_wait */
+    case -1:
+      CAMLreturn(Val_int(1)); /* Get_copy_end */
+    case -2:
+      CAMLreturn(Val_int(2)); /* Get_copy_error */
+    default:
+      v_buf = caml_alloc_string(res);
+      memcpy(String_val(v_buf), buf, res);
+      PQfreemem(buf);
+      v_result = caml_alloc(1, 0); /* Get_copy_data */
+      Store_field(v_result, 0, v_buf);
+      CAMLreturn(v_result);
+  }
+}
+
+CAMLprim value PQgetCopyData_bc(value v_conn, value v_async)
+{
+  return PQgetCopyData_stub(v_conn, Int_val(v_async));
+}
+
 CAMLprim intnat PQgetline_stub(
   value v_conn, value v_buf, intnat pos, intnat len)
 {