Add asynchronous connection primitives.

  - Extend connection_status and add polling_status.
  - Add ?startonly parameter to connection and connect_poll method.
  - Add reset_start and reset_poll methods.

Author: paurkedal

examples/async.ml

@@ -1,5 +1,8 @@
+open Printf
 open Postgresql
 
+let failwith_f fmt = ksprintf failwith fmt
+
 let _ =
   if Array.length Sys.argv <> 2 then (
     Printf.printf "\
@@ -22,29 +25,51 @@ let fetch_single_result c =
   | None -> assert false
   | Some r -> assert (fetch_result c = None); r
 
-let main () =
-  let c = new connection ~conninfo:Sys.argv.(1) () in
-  c#set_nonblocking true;
+(* See http://www.postgresql.org/docs/devel/static/libpq-connect.html *)
+let rec finish_conn socket_fd connect_poll = function
+  | Polling_failed ->
+    printf "f\n%!"
+  | Polling_reading ->
+    printf "r,%!";
+    ignore (Unix.select [socket_fd] [] [] (-1.0));
+    finish_conn socket_fd connect_poll (connect_poll ())
+  | Polling_writing ->
+    printf "w,%!";
+    ignore (Unix.select [] [socket_fd] [] (-1.0));
+    finish_conn socket_fd connect_poll (connect_poll ())
+  | Polling_ok ->
+    printf "c\n%!"
+
+let test (c : connection) =
+  (* Create a table using a non-prepared statement. *)
   c#send_query "\
     CREATE TEMPORARY TABLE postgresql_ocaml_async \
       (id SERIAL PRIMARY KEY, a INTEGER NOT NULL, b TEXT NOT NULL)";
   assert ((fetch_single_result c)#status = Command_ok);
+
+  (* Populate using a prepared statement. *)
   c#send_prepare "test_ins"
                  "INSERT INTO postgresql_ocaml_async (a, b) VALUES ($1, $2)";
   assert ((fetch_single_result c)#status = Command_ok);
   c#send_query_prepared ~params:[|"2"; "two"|] "test_ins";
   assert ((fetch_single_result c)#status = Command_ok);
   c#send_query_prepared ~params:[|"3"; "three"|] "test_ins";
   assert ((fetch_single_result c)#status = Command_ok);
+
+  (* Prepare a select statement. *)
   c#send_prepare "test_sel" "SELECT * FROM postgresql_ocaml_async";
   assert ((fetch_single_result c)#status = Command_ok);
+
+  (* Describe it. *)
   c#send_describe_prepared "test_sel";
   let r = fetch_single_result c in
   assert (r#status = Command_ok);
   assert (r#nfields = 3);
   assert (r#fname 0 = "id");
   assert (r#fname 1 = "a");
   assert (r#fname 2 = "b");
+
+  (* Run it. *)
   c#send_query_prepared "test_sel";
   let r = fetch_single_result c in
   assert (r#status = Tuples_ok);
@@ -55,6 +80,22 @@ let main () =
                   (r#getvalue i 0) (r#getvalue i 1) (r#getvalue i 2)
   done
 
+let main () =
+  (* Async connect and test. *)
+  let c = new connection ~conninfo:Sys.argv.(1) ~startonly:true () in
+  finish_conn (Obj.magic c#socket) (fun () -> c#connect_poll) Polling_writing;
+  if c#status = Bad then failwith_f "Connection failed: %s" c#error_message;
+  assert (c#status = Ok);
+  c#set_nonblocking true;
+  test c;
+
+  (* Async reset and test again. *)
+  if not c#reset_start then failwith_f "reset_start failed: %s" c#error_message;
+  finish_conn (Obj.magic c#socket) (fun () -> c#reset_poll) Polling_writing;
+  if c#status = Bad then failwith_f "Reset connection bad: %s" c#error_message;
+  assert (c#status = Ok);
+  test c
+
 let _ =
   try main () with
   | Error e -> prerr_endline (string_of_error e)

lib/postgresql.ml

@@ -238,7 +238,20 @@ let () =
   Callback.register "Postgresql.null" null;
   init ()
 
-type connection_status = Ok | Bad
+type connection_status =
+  | Ok | Bad
+  | Connection_started
+  | Connection_made
+  | Connection_awaiting_response
+  | Connection_auth_ok
+  | Connection_setenv
+  | Connection_ssl_startup
+
+type polling_status =
+  | Polling_failed
+  | Polling_reading
+  | Polling_writing
+  | Polling_ok
 
 type conninfo_option =
   {
@@ -307,7 +320,7 @@ module Stub = struct
   type result
 
   external conn_isnull : connection -> bool = "PQconn_isnull" "noalloc"
-  external connect : string -> connection = "PQconnectdb_stub"
+  external connect : string -> bool -> connection = "PQconnectdb_stub"
   external finish : connection -> unit = "PQfinish_stub"
   external reset : connection -> unit = "PQreset_stub"
 
@@ -383,6 +396,15 @@ module Stub = struct
 
   (* Asynchronous Query Processing *)
 
+  external connect_poll :
+    connection -> polling_status = "PQconnectPoll_stub" "noalloc"
+
+  external reset_start :
+    connection -> bool = "PQresetStart_stub" "noalloc"
+
+  external reset_poll :
+    connection -> polling_status = "PQresetPoll_stub" "noalloc"
+
   external set_nonblocking :
     connection -> bool -> int = "PQsetnonblocking_stub" "noalloc"
 
@@ -613,7 +635,7 @@ let protectx ~f x ~(finally : 'a -> unit) =
   res
 
 class connection ?host ?hostaddr ?port ?dbname ?user ?password ?options ?tty
-    ?requiressl ?conninfo =
+    ?requiressl ?conninfo ?(startonly = false) =
 
   let conn_info =
     match conninfo with
@@ -642,9 +664,9 @@ class connection ?host ?hostaddr ?port ?dbname ?user ?password ?options ?tty
         Buffer.contents b in
 
   fun () ->
-    let my_conn = Stub.connect conn_info in
+    let my_conn = Stub.connect conn_info startonly in
     let () =
-      if Stub.connection_status my_conn <> Ok then (
+      if Stub.connection_status my_conn = Bad then (
         let s = Stub.error_message my_conn in
         Stub.finish my_conn;
         raise (Error (Connection_failure s)))
@@ -903,6 +925,10 @@ object (self)
 
   (* Asynchronous operations and non blocking mode *)
 
+  method connect_poll = wrap_conn Stub.connect_poll
+  method reset_start = wrap_conn Stub.reset_start
+  method reset_poll = wrap_conn Stub.reset_poll
+
   method set_nonblocking b =
     wrap_conn (fun conn ->
       if Stub.set_nonblocking conn b <> 0 then signal_error conn)

lib/postgresql.mli

@@ -361,7 +361,22 @@ end
 (** {6 Handling database connections} *)
 
 (** Status of a connection *)
-type connection_status = Ok | Bad
+type connection_status =
+  | Ok | Bad
+  (* Non-blocking: *)
+  | Connection_started
+  | Connection_made
+  | Connection_awaiting_response
+  | Connection_auth_ok
+  | Connection_setenv
+  | Connection_ssl_startup
+
+(** Polling status used while establishing a connection asynchronously. *)
+type polling_status =
+  | Polling_failed
+  | Polling_reading
+  | Polling_writing
+  | Polling_ok
 
 (** Record of connection options *)
 type conninfo_option =
@@ -383,6 +398,10 @@ external conndefaults : unit -> conninfo_option array = "PQconndefaults_stub"
 
     When [conninfo] is given, it will be used instead of all other
     optional arguments.
+
+    @param startonly If true, initiate a non-blocking connect procedure, which
+      involves cooperative calls to {!connect_poll} before the connection is
+      usable.
 *)
 class connection :
   ?host : string ->  (* Default: none *)
@@ -395,6 +414,7 @@ class connection :
   ?tty : string ->  (* Default: none *)
   ?requiressl : string ->  (* Default: none *)
   ?conninfo : string ->  (* Default: none *)
+  ?startonly : bool -> (* Default: false *)
   unit ->
   (** @raise Error if there is a connection failure. *)
 object
@@ -723,6 +743,33 @@ object
 
   (** Asynchronous operations and non blocking mode *)
 
+  method connect_poll : polling_status
+  (** After creating a connection with [~startonly:true], {!connect_poll}
+      must be called a number of times before the connection can be used.  The
+      precise procedure is described in libpq manual, but the following code
+      should capture the idea, assuming monadic concurrency primitives
+      [return] and [>>=] along with polling functions [wait_for_read] and
+      [wait_for_write]:
+      {[
+        let my_async_connect () =
+          let c = new connection () in
+          let rec establish_connection = function
+            | Polling_failed | Polling_ok -> return c
+            | Polling_reading -> wait_for_read c#socket >>= fun () ->
+                                 establish_connection c#connect_poll
+            | Polling_writing -> wait_for_write c#socket >>= fun () ->
+                                 establish_connection c#connect_poll in
+          establish_connection Polling_writing
+      ]}
+      See also [examples/async.ml]. *)
+
+  method reset_start : bool
+  (** An asynchronous variant of {!reset}.  Use {!reset_poll} to
+      finish re-establishing the connection. *)
+
+  method reset_poll : polling_status
+  (** Used analogously to {!connect_poll} after calling {!reset_start} *)
+
   method set_nonblocking : bool -> unit
   (** [set_nonblocking b] sets state of the connection to nonblocking if
       [b] is true and to blocking otherwise.

lib/postgresql_stubs.c

@@ -209,8 +209,6 @@ static inline void np_decr_refcount(np_callback *c)
 /* Database Connection Functions */
 
 /* Missing:
-     PQconnectStart, PQconnectPoll, PQresetStart, PQresetPoll:
-       for non-blocking connection
      PQgetssl: the SSL structure used in the connection
 */
 
@@ -244,21 +242,26 @@ static inline void free_conn(value v_conn)
   }
 }
 
-CAMLprim value PQconnectdb_stub(value v_conn_info)
+CAMLprim value PQconnectdb_stub(value v_conn_info, value v_startonly)
 {
   PGconn *conn;
   value v_conn;
-
-  size_t len = caml_string_length(v_conn_info) + 1;
   PGcancel *cancel;
-  char *conn_info = caml_stat_alloc(len);
-  memcpy(conn_info, String_val(v_conn_info), len);
 
-  caml_enter_blocking_section();
-    conn = PQconnectdb(conn_info);
+  if (Bool_val(v_startonly)) {
+    conn = PQconnectStart(String_val(v_conn_info));
     cancel = PQgetCancel(conn);
-    free(conn_info);
-  caml_leave_blocking_section();
+  }
+  else {
+    size_t len = caml_string_length(v_conn_info) + 1;
+    char *conn_info = caml_stat_alloc(len);
+    memcpy(conn_info, String_val(v_conn_info), len);
+    caml_enter_blocking_section();
+      conn = PQconnectdb(conn_info);
+      cancel = PQgetCancel(conn);
+      free(conn_info);
+    caml_leave_blocking_section();
+  }
 
   /* One may raise this 30 to 500 for instance if the program takes
      responsibility of closing connections */
@@ -349,6 +352,9 @@ static inline value make_string(const char *s)
     return ret(fun(get_conn(v_conn))); \
   }
 
+conn_info(PQconnectPoll, Val_int)
+conn_info(PQresetStart, Val_int)
+conn_info(PQresetPoll, Val_int)
 conn_info(PQdb, make_string)
 conn_info(PQuser, make_string)
 conn_info(PQpass, make_string)