unikernel/duniverse/ocaml-caqti/caqti-driver-pgx/lib/caqti_driver_pgx.ml

@@ -0,0 +1,791 @@
+(* Copyright (C) 2021--2025  Petter A. Urkedal <paurkedal@gmail.com>
+ *
+ * This library is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU Lesser General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or (at your
+ * option) any later version, with the LGPL-3.0 Linking Exception.
+ *
+ * This library is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
+ * License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * and the LGPL-3.0 Linking Exception along with this library.  If not, see
+ * <http://www.gnu.org/licenses/> and <https://spdx.org>, respectively.
+ *)
+
+[@@@alert "-caqti_private"]
+
+open Caqti_template
+open Caqti_platform
+open Postgresql_conv
+open Printf
+
+let rec find_map_list f = function
+ | [] -> None
+ | x :: xs -> (match f x with | None -> find_map_list f xs | Some _ as y -> y)
+
+type Caqti_error.msg += Pgx_msg of string * Pgx.Error_response.t
+let () =
+  let pp ppf = function
+   | Pgx_msg (msg, _) -> Format.pp_print_string ppf msg
+   | _ -> assert false
+  in
+  let cause = function
+   | Pgx_msg (_, error_response) ->
+      cause_of_sqlstate error_response.Pgx.Error_response.code
+   | _ ->
+      assert false
+  in
+  Caqti_error.define_msg ~pp ~cause [%extension_constructor Pgx_msg]
+
+exception Failed_with_msg of Caqti_error.msg
+
+let host_of_string str =
+  (match Domain_name.of_string str with
+   | Ok dom ->
+      (match Domain_name.host dom with
+       | Ok dom -> Some dom
+       | Error _ -> None)
+   | Error _ -> None)
+
+let pg_type_name : type a. a Field_type.t -> string = function
+ | Bool -> "bool"
+ | Int -> "int8"
+ | Int16 -> "int2"
+ | Int32 -> "int4"
+ | Int64 -> "int8"
+ | Float -> "float8"
+ | String -> "text"
+ | Octets -> "bytea"
+ | Pdate -> "date"
+ | Ptime -> "timestamptz"
+ | Ptime_span -> "interval"
+ | Enum name -> name
+
+let encode_field
+    : type a. a Field_type.t -> a -> Pgx.Value.t
+    = fun field_type x ->
+  (match field_type with
+   | Bool -> Pgx.Value.of_bool x
+   | Int -> Pgx.Value.of_int x
+   | Int16 -> Pgx.Value.of_int x
+   | Int32 -> Pgx.Value.of_int32 x
+   | Int64 -> Pgx.Value.of_int64 x
+   | Float -> Pgx.Value.of_float x
+   | String -> Pgx.Value.of_string x
+   | Enum _ -> Pgx.Value.of_string x
+   | Octets -> Pgx.Value.of_binary x
+   | Pdate -> Pgx.Value.of_string (Conv.iso8601_of_pdate x)
+   | Ptime -> Pgx.Value.of_string (pgstring_of_ptime x)
+   | Ptime_span -> Pgx.Value.of_string (pgstring_of_ptime_span x))
+
+let query_string ~subst templ =
+  let templ = Query.expand ~final:true subst templ in
+
+  let rec extract_quotes : Query.t -> _ = function
+   | V (ft, v) -> fun (n, acc) -> (n + 1, encode_field ft v :: acc)
+   | Q s -> fun (n, acc) -> (n + 1, Pgx.Value.of_string s :: acc)
+   | L _ | P _ -> Fun.id
+   | E _ -> fun _ -> assert false
+   | S qs -> List_ext.fold extract_quotes qs
+  in
+  let nQ, rev_quotes = extract_quotes templ (0, []) in
+
+  let buf = Buffer.create 64 in
+  let rec write_query_string : Query.t -> _ = function
+   | L s -> fun jQ -> Buffer.add_string buf s; jQ
+   | V _ -> fun jQ -> bprintf buf "$%d" jQ; jQ + 1
+   | Q _ -> fun jQ -> bprintf buf "$%d" jQ; jQ + 1
+   | P j -> fun jQ -> bprintf buf "$%d" (nQ + 1 + j); jQ
+   | E _ -> assert false
+   | S qs -> List_ext.fold write_query_string qs
+  in
+  let _jQ = write_query_string templ 1 in
+  (Buffer.contents buf, rev_quotes)
+
+let encode_param ~uri t param =
+  let write_value ~uri:_ ft fv acc = encode_field ft fv :: acc in
+  let write_null ~uri:_ _ acc = Pgx.Value.null :: acc in
+  try
+    Request_utils.encode_param ~uri {write_value; write_null} t param []
+      |> List.rev |> Result.ok
+  with Caqti_error.Exn (#Caqti_error.call as err) ->
+    Error err
+
+let decode_field : type a. uri: Uri.t -> a Field_type.t -> Pgx.Value.t -> a =
+  fun ~uri field_type v ->
+  let wrap_conv_exn f s =
+    (match f s with
+     | y -> y
+     | exception Pgx.Value.Conversion_failure msg_str ->
+        let msg = Caqti_error.Msg msg_str in
+        let typ = Row_type.field field_type in
+        Request_utils.raise_decode_rejected ~uri ~typ msg)
+  in
+  let wrap_conv_res f s =
+    (match f s with
+     | Ok y -> y
+     | Error msg_str ->
+        let msg = Caqti_error.Msg msg_str in
+        let typ = Row_type.field field_type in
+        Request_utils.raise_decode_rejected ~uri ~typ msg)
+  in
+  (match field_type with
+   | Bool -> wrap_conv_exn Pgx.Value.to_bool_exn v
+   | Int -> wrap_conv_exn Pgx.Value.to_int_exn v
+   | Int16 -> wrap_conv_exn Pgx.Value.to_int_exn v
+   | Int32 -> wrap_conv_exn Pgx.Value.to_int32_exn v
+   | Int64 -> wrap_conv_exn Pgx.Value.to_int64_exn v
+   | Float -> wrap_conv_exn Pgx.Value.to_float_exn v
+   | String -> wrap_conv_exn Pgx.Value.to_string_exn v
+   | Enum _ -> wrap_conv_exn Pgx.Value.to_string_exn v
+   | Octets -> wrap_conv_exn Pgx.Value.to_binary_exn v
+   | Pdate ->
+      v |> wrap_conv_exn Pgx.Value.to_string_exn
+        |> wrap_conv_res Conv.pdate_of_iso8601
+   | Ptime ->
+      v |> wrap_conv_exn Pgx.Value.to_string_exn
+        |> wrap_conv_res Conv.ptime_of_rfc3339_utc
+   | Ptime_span ->
+      v |> wrap_conv_exn Pgx.Value.to_string_exn
+        |> wrap_conv_res ptime_span_of_pgstring)
+
+let decode_row ~uri row_type =
+  let read_value ~uri ft = function
+   | [] -> assert false
+   | field :: fields ->
+      let y = decode_field ~uri ft field in
+      (y, fields)
+  in
+  let rec skip_null n xs =
+    if n = 0 then Some xs else
+    (match xs with
+     | [] -> assert false
+     | x :: xs' when Pgx.Value.(compare null) x = 0 -> skip_null (n - 1) xs'
+     | _ :: _ -> None)
+  in
+  let decode = Request_utils.decode_row ~uri {read_value; skip_null} row_type in
+  fun fields ->
+    try
+      let (y, fields) = decode fields in
+      assert (fields = []);
+      Ok y
+    with
+     | Caqti_error.Exn (#Caqti_error.retrieve as err) -> Error err
+
+module Q = struct
+  let select_type_oid = "SELECT oid FROM pg_catalog.pg_type WHERE typname = $1"
+end
+
+type connect_arg = {
+  host: string option;
+  port: int option;
+  user: string option;
+  password: string option;
+  database: string option;
+  unix_domain_socket_dir: string option;
+}
+
+let (let/?) m f = match m with Ok x -> f x | Error _ as r -> r
+
+let parse_uri uri =
+  let reject msg =
+    let msg = Caqti_error.Msg msg in
+    Error (Caqti_error.connect_rejected ~uri msg)
+  in
+  let/? host_or_unix_domain_socket_dir =
+    (match Uri.host uri, Uri.get_query_param uri "host" with
+     | None, None -> Ok None
+     | Some s, None | None, Some s -> Ok (Some s)
+     | Some _, Some _ -> reject "Conflicting host in URI.")
+  in
+  let host, unix_domain_socket_dir =
+    (match host_or_unix_domain_socket_dir with
+     | None | Some "" -> (None, None)
+     | Some s -> if s.[0] = '/' then (None, Some s) else (Some s, None))
+  in
+  let/? port =
+    (match Uri.port uri, Uri.get_query_param uri "port" with
+     | None, None -> Ok None
+     | Some i, None -> Ok (Some i)
+     | None, Some s ->
+        (try Ok (Some (int_of_string s)) with
+         | Failure _ -> reject "Non-integer port number in URI.")
+     | Some _, Some _ -> reject "Conflicting host in URI.")
+  in
+  let/? user =
+    (match Uri.user uri, Uri.get_query_param uri "user" with
+     | None, None -> Ok None
+     | Some s, None | None, Some s -> Ok (Some s)
+     | Some _, Some _ -> reject "Conflicting user in URI.")
+  in
+  let/? password =
+    (match Uri.password uri, Uri.get_query_param uri "password" with
+     | None, None -> Ok None
+     | Some s, None | None, Some s -> Ok (Some s)
+     | Some _, Some _ -> reject "Conflicting password in URI.")
+  in
+  let/? database =
+    (match String.split_on_char '/' (Uri.path uri),
+           Uri.get_query_param uri "dbname" with
+     | [], _ -> assert false
+     | [""], dbname -> Ok dbname
+     | [""; dbname], None -> Ok (Some dbname)
+     | [""; _], Some _ -> reject "Conflicting dbname in URI."
+     | _ -> reject "Too many path components in URI.")
+  in
+  Ok {host; port; user; password; database; unix_domain_socket_dir}
+
+let dialect = Dialect.create_pgsql
+  ~server_version:(Version.of_string_unsafe "")
+  ~client_library:`pgx
+  ()
+
+module Connect_functor (System : Caqti_platform.System_sig.S) = struct
+  open System
+  open System.Fiber.Infix
+
+  let ( let*? ) m f = m >>= function Ok x -> f x | Error _ as r -> Fiber.return r
+  let ( let+? ) m f = m >|= function Ok x -> Ok (f x) | Error _ as r -> r
+  let ( >>=? ) = ( let*? )
+  let ( >|=? ) = ( let+? )
+
+  let intercept h f =
+    Fiber.catch
+      (fun () -> f () >|= fun y -> Ok y)
+      (function
+       | Failed_with_msg msg ->
+          Fiber.return (Error (h msg))
+       | Pgx.PostgreSQL_Error (msg, err) ->
+          Fiber.return (Error (h (Pgx_msg (msg, err))))
+       | End_of_file ->
+          Fiber.return (Error (h (Caqti_error.Msg "Unexpected EOF from server.")))
+       | Failure msg -> (* Raised by our Pgx.Io implementation. *)
+          Fiber.return (Error (h (Caqti_error.Msg msg)))
+       | exn ->
+          (match Net.convert_io_exception exn with
+           | Some msg -> Fiber.return (Error (h msg))
+           | None -> raise exn))
+
+  let intercept_request_failed ~uri ~query =
+    intercept (Caqti_error.request_failed ~uri ~query)
+  let intercept_connect_failed ~uri =
+    intercept (Caqti_error.connect_failed ~uri)
+
+  type ssl_config =
+    Ssl_config : {
+      impl: (module Net.TLS_PROVIDER with type tls_config = 'a);
+      config: 'a;
+      host: [`host] Domain_name.t option;
+    } -> ssl_config
+
+  (* We need to pass stdenv into open_connection below.  This means that
+   * PGX will be instantiated for each connection. *)
+  module Pass_stdenv
+    (Connect_env : sig val sw : Switch.t val stdenv : stdenv end) =
+  struct
+    open Connect_env
+
+    module Pgx_with_io = Pgx.Make (struct
+      type 'a t = 'a Fiber.t
+      let return = Fiber.return
+      let ( >>= ) = ( >>= )
+      let catch = Fiber.catch
+
+      include Net
+
+      type in_channel = Socket.t
+      type out_channel = Socket.t
+
+      type sockaddr = Unix of string | Inet of string * int
+
+      let open_connection sockaddr =
+        let connect sockaddr =
+          Net.connect_tcp ~sw ~stdenv sockaddr
+            >|= Result.map (fun socket -> (socket, socket))
+        in
+        (match sockaddr with
+         | Unix path ->
+            connect (Sockaddr.unix path)
+         | Inet (host_or_ipaddr, port) ->
+            (match Ipaddr.of_string host_or_ipaddr with
+             | Ok ipaddr ->
+                connect (Sockaddr.tcp (ipaddr, port))
+             | Error _ ->
+                (match host_of_string host_or_ipaddr with
+                 | None ->
+                    failwith
+                      ("Cannot resolve invalid host name " ^ host_or_ipaddr)
+                 | Some host ->
+                    getaddrinfo ~stdenv host port >>= (function
+                     | Ok [] ->
+                        failwith "The host name does not resolve."
+                     | Ok (sockaddr :: _) ->
+                        connect sockaddr
+                     | Error (`Msg msg) ->
+                        failwith msg))))
+        >|= function Ok conn -> conn | Error msg -> raise (Failed_with_msg msg)
+
+      let output_char = Socket.output_char
+      let output_string = Socket.output_string
+      let flush = Socket.flush
+
+      let output_binary_int oc x =
+        let buf = Bytes.create 4 in
+        Bytes.set_int32_be buf 0 (Int32.of_int x);
+        Socket.output_string oc (Bytes.to_string buf)
+
+      let input_char = Socket.input_char
+      let really_input = Socket.really_input
+
+      (* This closes the output channel instead of the input channel; cf.
+       * Unix.open_connection. *)
+      let close_in = Socket.close
+
+      let input_binary_int ic =
+        let buf = Bytes.create 4 in
+        Socket.really_input ic buf 0 4 >|= fun () ->
+        Int32.to_int (Bytes.get_int32_be buf 0)
+
+      type nonrec ssl_config = ssl_config
+
+      let upgrade_ssl =
+        let upgrade ?ssl_config socket _ =
+          (match ssl_config, tcp_flow_of_socket socket with
+           | None, _ -> assert false (* we don't use `Auto *)
+           | _, None -> assert false (* we only upgrade once *)
+           | Some (Ssl_config {impl; config; host}), Some tcp_flow ->
+              let module Impl = (val impl) in
+              Impl.start_tls ~config ?host tcp_flow >|= function
+               | Ok tls_flow ->
+                  let socket = socket_of_tls_flow ~sw tls_flow in
+                  (socket, socket)
+               | Error msg -> raise (Failed_with_msg msg))
+        in
+        `Supported upgrade
+
+      let getlogin () = failwith "The DB user must be provided."
+
+      let debug msg = Log.debug (fun f -> f "%s" msg)
+
+      let protect f ~finally = Fiber.finally f finally
+
+      module Sequencer = struct
+        type 'a monad = 'a Fiber.t
+        include Sequencer
+      end
+    end)
+  end
+
+  module Make_connection_base
+    (Pgx_with_io : Pgx.S with type 'a Io.t = 'a Fiber.t
+                          and type Io.ssl_config = ssl_config)
+    (Connection_arg : sig
+      val subst : Query.subst
+      val uri : Uri.t
+      val db_arg : Pgx_with_io.t
+      val select_type_oid : Pgx_with_io.Prepared.s
+      val dynamic_capacity : int
+    end) =
+  struct
+    open Connection_arg
+
+    let db_txn = ref None
+
+    module Response = struct
+
+      type ('b, 'm) t = {
+        query: string;
+        row_type: 'b Row_type.t;
+        prepared: Pgx_with_io.Prepared.s;
+        params: Pgx.Value.t list;
+      }
+
+      let returned_count _ = Fiber.return (Error `Unsupported)
+      let affected_count _ = Fiber.return (Error `Unsupported)
+
+      let reject ~query msg =
+        Error (Caqti_error.response_rejected ~uri ~query (Caqti_error.Msg msg))
+
+      let exec {query; prepared; params; _} =
+        intercept_request_failed ~uri ~query (fun () ->
+          Pgx_with_io.Prepared.execute prepared ~params) >|=
+        (function
+         | Ok [] -> Ok ()
+         | Ok _ ->
+            reject ~query "Received multiple rows where none were expected."
+         | Error _ as r -> r)
+
+      let find {query; row_type; prepared; params} =
+        intercept_request_failed ~uri ~query (fun () ->
+          Pgx_with_io.Prepared.execute prepared ~params) >|=
+        (function
+         | Ok [row] -> decode_row ~uri row_type row
+         | Ok [] ->
+            reject ~query "Received no rows where one was expected."
+         | Ok _ ->
+            reject ~query "Received more than one row where one was expected."
+         | Error _ as r -> r)
+
+      let find_opt {query; row_type; prepared; params} =
+        intercept_request_failed ~uri ~query (fun () ->
+          Pgx_with_io.Prepared.execute prepared ~params) >|=
+        (function
+         | Ok [] -> Ok None
+         | Ok [row] ->
+            decode_row ~uri row_type row |> Result.map (fun x -> Some x)
+         | Ok _ ->
+            reject ~query
+              "Received two or more rows where at most one was expected."
+         | Error _ as r -> r)
+
+      let fold f {query; row_type; prepared; params} =
+        let decode = decode_row ~uri row_type in
+        let f acc row =
+          Fiber.return @@ match acc with
+           | Ok acc ->
+              (match decode row with
+               | Ok row -> Ok (f row acc)
+               | Error _ as r -> r)
+           | Error _ as r -> r
+        in
+        fun acc ->
+          intercept_request_failed ~uri ~query begin fun () ->
+            Pgx_with_io.Prepared.execute_fold ~f prepared ~params ~init:(Ok acc)
+          end >|= Stdlib.Result.join
+
+      let fold_s f {query; row_type; prepared; params} =
+        let decode = decode_row ~uri row_type in
+        let f acc row =
+          (match acc with
+           | Ok acc ->
+              (match decode row with
+               | Ok row -> f row acc
+               | Error _ as r -> Fiber.return r)
+           | Error _ as r -> Fiber.return r)
+        in
+        fun acc ->
+          intercept_request_failed ~uri ~query begin fun () ->
+            Pgx_with_io.Prepared.execute_fold ~f prepared ~params ~init:(Ok acc)
+          end >|= Stdlib.Result.join
+
+      let iter_s f {query; row_type; prepared; params} =
+        let decode = decode_row ~uri row_type in
+        let f acc row =
+          (match acc with
+           | Ok () ->
+              (match decode row with
+               | Ok row -> f row
+               | Error _ as r -> Fiber.return r)
+           | Error _ as r -> Fiber.return r)
+        in
+        intercept_request_failed ~uri ~query begin fun () ->
+          Pgx_with_io.Prepared.execute_fold ~f prepared ~params ~init:(Ok ())
+        end >|= Stdlib.Result.join
+
+      let to_stream resp () =
+        fold List.cons resp [] >|= Result.map List.rev >|= function
+         | Ok [] -> Stream.Nil
+         | Ok (row :: rows) -> Stream.Cons (row, Stream.of_list rows)
+         | Error err -> (Stream.Error err)
+    end
+
+    type prepared = {
+      query: string;
+      pgx_prepared: Pgx_with_io.Prepared.s;
+      rev_quotes: Pgx.Value.t list;
+    }
+
+    module Pcache =
+      Request_cache.Make (struct type t = prepared let weight _ = 1 end)
+
+    let in_use = ref false
+    let pcache : Pcache.t = Pcache.create ~dynamic_capacity dialect
+
+    let reset _ = Fiber.return () (* FIXME *)
+
+    let using_db f =
+      if !in_use then
+        failwith "Invalid concurrent usage of PostgreSQL connection detected.";
+      in_use := true;
+      let db = match !db_txn with None -> db_arg | Some db -> db in
+      Fiber.cleanup
+        (fun () -> f db >|= fun res -> in_use := false; res)
+        (fun () -> reset db >|= fun _ -> in_use := false)
+
+    let type_oid_cache = Hashtbl.create 11
+
+    let field_type_oid ft =
+      let name = pg_type_name ft in
+      (match Hashtbl.find_opt type_oid_cache name with
+       | Some oid -> Fiber.return (Ok oid)
+       | None ->
+          let params = [Pgx.Value.of_string name] in
+          let*? row =
+            intercept_request_failed ~uri ~query:Q.select_type_oid (fun () ->
+              Pgx_with_io.Prepared.execute select_type_oid ~params)
+          in
+          let fail s =
+            let msg = Caqti_error.Msg s in
+            Fiber.return (Error
+              (Caqti_error.request_failed ~uri ~query:Q.select_type_oid msg))
+          in
+          let failf fmt = Format.kasprintf fail fmt in
+          (match row with
+           | [[v]] ->
+              (match Pgx.Value.to_int32 v with
+               | Some oid ->
+                  Hashtbl.add type_oid_cache name oid;
+                  Fiber.return (Ok oid)
+               | None ->
+                  failf "Expected an int32 in response from OID request.")
+           | [] ->
+              failf "OID for type %a not found." Field_type.pp ft
+           | [_] ->
+              failf "Expected single field result from OID request."
+           | _ ->
+              failf "Expected at most one row from OID request."))
+
+    let type_oids param_type =
+      let rec loop :
+          type a. a Row_type.t -> Pgx.oid list ->
+          (Pgx.oid list, _) result Fiber.t =
+        (function
+         | Field ft -> fun acc -> field_type_oid ft >|=? fun ft -> ft :: acc
+         | Option t -> loop t
+         | Product (_, prod) ->
+            let rec loop_prod : type i. (i, a) Row_type.product -> _ =
+              (function
+               | Proj_end -> fun acc -> Fiber.return (Ok acc)
+               | Proj (t, _, prod) ->
+                  let loop_t = loop t in
+                  let loop_prod = loop_prod prod in
+                  fun acc -> loop_prod acc >>=? loop_t)
+            in
+            loop_prod prod
+         | Annot (_, t) -> loop t)
+      in
+      loop (Row_type.option param_type) []
+
+    let pp_request_with_param ppf =
+      Request.make_pp_with_param ~subst ~dialect () ppf
+
+    let free_prepared prepared =
+      intercept_request_failed ~uri ~query:"DEALLOCATE"
+        (fun () -> Pgx_with_io.Prepared.close prepared.pgx_prepared)
+
+    let deallocate req =
+      (match Request.prepare_policy req with
+       | Dynamic | Static ->
+          (match Pcache.deallocate pcache req with
+           | None ->
+              Fiber.return (Ok ())
+           | Some (prepared, commit) ->
+              free_prepared prepared >|=? commit)
+       | Direct ->
+          failwith "deallocate called on oneshot request")
+
+    let deallocate_some () =
+      let rec loop = function
+       | [] -> Fiber.return (Ok ())
+       | prepared :: orphans ->
+          let*? () = free_prepared prepared in
+          loop orphans
+      in
+      let orphans, commit = Pcache.trim pcache in
+      loop orphans >|=? commit
+
+    let fresh_name = Request_utils.fresh_name_generator "caq"
+
+    let call ~f req param =
+      using_db @@ fun db ->
+      deallocate_some () >>=? fun () ->
+      Log.debug ~src:Logging.request_log_src (fun f ->
+        f "Sending %a" pp_request_with_param (req, param)) >>= fun () ->
+      let pre_prepare () =
+        let templ = Request.query req dialect in
+        let query, rev_quotes = query_string ~subst templ in
+        let*? param_types = type_oids (Request.param_type req) in
+        let+? string_oid = field_type_oid Field_type.String in
+        let quote_types = List.rev_map (fun _ -> string_oid) rev_quotes in
+        let types = List.rev_append quote_types param_types in
+        (query, types, rev_quotes)
+      in
+      let post_prepare pq =
+        (match encode_param ~uri (Request.param_type req) param with
+         | Error _ as r -> Fiber.return r
+         | Ok regular_params ->
+            let params = List.rev_append pq.rev_quotes regular_params in
+            f {
+              Response.row_type = Request.row_type req;
+              query = pq.query;
+              prepared = pq.pgx_prepared;
+              params;
+            })
+      in
+      (match Request.prepare_policy req with
+       | Dynamic | Static ->
+          (match Pcache.find_and_promote pcache req with
+           | Some pq -> Fiber.return (Ok pq)
+           | None ->
+              let*? query, types, rev_quotes = pre_prepare () in
+              let name = fresh_name () in
+              let+? pgx_prepared =
+                intercept_request_failed ~uri ~query (fun () ->
+                  Pgx_with_io.Prepared.prepare ~name ~query ~types db)
+              in
+              let pq = {query; pgx_prepared; rev_quotes} in
+              Pcache.add pcache req pq;
+              pq)
+          >>=? post_prepare
+       | Direct ->
+          let*? query, types, rev_quotes = pre_prepare () in
+          Pgx_with_io.Prepared.with_prepare db ~types ~query
+            ~f:(fun pgx_prepared ->
+                  post_prepare {query; pgx_prepared; rev_quotes}))
+
+    let disconnect () =
+      using_db @@ fun _ ->
+      Fiber.catch
+        (fun () ->
+          (match !db_txn with
+           | None -> Fiber.return ()
+           | Some db -> Pgx_with_io.close db) >>= fun () ->
+          Pgx_with_io.close db_arg)
+        (function
+         | Pgx.PostgreSQL_Error (msg, _) ->
+            Log.err (fun f -> f "Failed to disconnect %a: %s" Uri.pp uri msg)
+         | exn ->
+            raise exn)
+
+    let validate () =
+      (* No need to handle exceptions here, since alive is a catch-all wrapper
+       * around ping. *)
+      using_db Pgx_with_io.alive
+
+    let check f = f true
+
+    let start () =
+      (match !db_txn with
+       | None ->
+          intercept_request_failed ~uri ~query:"BEGIN" begin fun () ->
+            Pgx_with_io.begin_work db_arg >|= fun db ->
+            db_txn := Some db
+          end
+       | Some _ ->
+          failwith "A transaction already in progress.")
+
+    let commit () =
+      (match !db_txn with
+       | Some db ->
+          intercept_request_failed ~uri ~query:"COMMIT" begin fun () ->
+            db_txn := None;
+            Pgx_with_io.commit db
+          end
+       | None ->
+          failwith "No transaction to commit.")
+
+    let rollback () =
+      (match !db_txn with
+       | Some db ->
+          intercept_request_failed ~uri ~query:"ROLLBACK" begin fun () ->
+            db_txn := None;
+            Pgx_with_io.rollback db
+          end
+       | None ->
+          failwith "No transaction to roll back.")
+
+    let set_statement_timeout t =
+      let t_arg =
+        (match t with
+         | None -> 0
+         | Some t -> max 1 (int_of_float (t *. 1000.0 +. 500.0)))
+      in
+      let query = sprintf "SET statement_timeout TO %d" t_arg in
+      using_db @@ fun db ->
+        intercept_request_failed ~uri ~query @@ fun () ->
+          Pgx_with_io.execute_unit db query
+  end
+
+  let driver_info = Caqti_driver_info.of_dialect dialect
+
+  module type CONNECTION = Caqti_connection_sig.S
+    with type 'a fiber := 'a Fiber.t
+     and type ('a, 'err) stream := ('a, 'err) System.Stream.t
+
+  let find_tls_provider ~config ?host () =
+    let with_config (module Tls_provider : Net.TLS_PROVIDER) =
+      (match Caqti_connect_config.get Tls_provider.tls_config_key config with
+       | None -> None
+       | Some config ->
+          Some (Ssl_config {impl = (module Tls_provider); config; host}))
+    in
+    (match find_map_list with_config (Net.tls_providers config) with
+     | None -> `No
+     | Some ssl_config -> `Always ssl_config)
+
+  let connect ~sw ~stdenv ~subst ~config uri =
+
+    (* Create PGX connection and helper functions. *)
+    let*? {host; port; user; password; database; unix_domain_socket_dir} =
+      Fiber.return (parse_uri uri)
+    in
+    let ssl =
+      let host = Option.bind host host_of_string in
+      find_tls_provider ~config ?host ()
+    in
+    let open Pass_stdenv (struct let sw = sw let stdenv = stdenv end) in
+    let*? db =
+      intercept_connect_failed ~uri
+        (Pgx_with_io.connect
+          ~ssl ?host ?port ?user ?password ?database ?unix_domain_socket_dir)
+    in
+    let prepare_post_connect query =
+      intercept_request_failed ~uri ~query
+        (fun () -> Pgx_with_io.Prepared.prepare db ~query)
+      >|= Result.map_error (fun err -> `Post_connect err)
+    in
+    let execute_post_connect ?params query =
+      intercept_request_failed ~uri ~query
+        (fun () -> Pgx_with_io.execute ?params db query)
+    in
+
+    (* Run setup. *)
+    let*? () =
+      let query = "SET TimeZone TO 'UTC'" in
+      let reject msg =
+        let msg = Caqti_error.Msg msg in
+        let err = Caqti_error.response_rejected ~uri ~query msg in
+        Error (`Post_connect err)
+      in
+      execute_post_connect query >|= function
+       | Ok [] -> Ok ()
+       | Ok _ -> reject "Invalid response from setup request."
+       | Error err -> Error (`Post_connect err)
+    in
+
+    (* Return Caqti connection module. *)
+    let+? select_type_oid = prepare_post_connect Q.select_type_oid in
+    let module B = Make_connection_base (Pgx_with_io) (struct
+      let subst = subst dialect
+      let uri = uri
+      let db_arg = db
+      let select_type_oid = select_type_oid
+      let dynamic_capacity =
+        Caqti_connect_config.(get dynamic_prepare_capacity) config
+    end) in
+    let module Connection = struct
+      let driver_info = driver_info
+      let dialect = dialect
+      let driver_connection = None
+      include B
+      include Connection_utils.Make_convenience (System) (B)
+      include Connection_utils.Make_populate (System) (B)
+    end in
+    (module Connection : CONNECTION)
+end
+
+let () =
+  Caqti_platform.Driver_loader.register "pgx" (module Connect_functor)

unikernel/duniverse/ocaml-caqti/caqti-driver-pgx/lib/caqti_driver_pgx.mli

@@ -0,0 +1,39 @@
+(* Copyright (C) 2022  Petter A. Urkedal <paurkedal@gmail.com>
+ *
+ * This library is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU Lesser General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or (at your
+ * option) any later version, with the LGPL-3.0 Linking Exception.
+ *
+ * This library is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
+ * License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * and the LGPL-3.0 Linking Exception along with this library.  If not, see
+ * <http://www.gnu.org/licenses/> and <https://spdx.org>, respectively.
+ *)
+
+(** PostgreSQL driver for Caqti based on pgx
+
+    This driver is implemented in terms of the pgx library.
+
+    It handles URIs of the form
+    {[
+      pgx://<user>:<password>@<host-or-directory>:<port>/<database>
+    ]}
+    where [<host-or-directory>] is either a host name, an IP number, or a
+    directory containing the Unix domain socket of a locally running PostgreSQL
+    server.  In the latter case, the directory must be percent-encoded, e.g.
+    {[
+      pgx://jdoe@%2fvar%2frun%2fpostgresql
+    ]}
+
+    The interface provided by this module {e should normally not be used by
+    applications}, but provides access to some pgx specifics in case they are
+    needed. *)
+
+(** {1 Error Details} *)
+
+type Caqti_error.msg += Pgx_msg of string * Pgx.Error_response.t

unikernel/duniverse/ocaml-caqti/caqti-driver-pgx/lib/dune

@@ -0,0 +1,14 @@
+(rule (copy# ../../shared/postgresql_conv.ml postgresql_conv.ml))
+
+(library
+ (name caqti_driver_pgx)
+ (public_name caqti-driver-pgx)
+ (flags (:standard -alert -caqti_unstable))
+ (library_flags (:standard -linkall))
+ (libraries caqti caqti.platform domain-name ipaddr pgx))
+
+(plugin
+ (package caqti-driver-pgx)
+ (name caqti-driver-pgx)
+ (libraries caqti-driver-pgx)
+ (site (caqti plugins)))