1059 lines
37 KiB
OCaml
1059 lines
37 KiB
OCaml
(* Copyright (C) 2017--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 Printf
|
|
module Pg = Postgresql
|
|
|
|
let ( |>? ) = Result.bind
|
|
let ( %>? ) f g x = match f x with Ok y -> g y | Error _ as r -> r
|
|
|
|
let pct_encoder =
|
|
Uri.pct_encoder ~query_value:(`Custom (`Query_value, "", "=")) ()
|
|
|
|
module Q = struct
|
|
open Caqti_template.Create
|
|
|
|
let start = static T.(unit -->. unit) "BEGIN"
|
|
let commit = static T.(unit -->. unit) "COMMIT"
|
|
let rollback = static T.(unit -->. unit) "ROLLBACK"
|
|
|
|
let type_oid =
|
|
static T.(string -->? int)
|
|
"SELECT oid FROM pg_catalog.pg_type WHERE typname = ?"
|
|
|
|
let set_timezone_to_utc =
|
|
direct T.(unit -->. unit) "SET TimeZone TO 'UTC'"
|
|
|
|
let set_statement_timeout t =
|
|
direct_gen T.(unit -->. unit) @@ fun _ ->
|
|
(* Parameters are not supported for SET. *)
|
|
Q.lit "SET statement_timeout TO " @++ Q.int t
|
|
end
|
|
|
|
type Caqti_error.msg +=
|
|
| Connect_error_msg of {
|
|
error: Pg.error;
|
|
}
|
|
| Connection_error_msg of {
|
|
error: Pg.error;
|
|
connection_status: Pg.connection_status;
|
|
}
|
|
| Result_error_msg of {
|
|
error_message: string;
|
|
sqlstate: string;
|
|
}
|
|
|
|
let extract_connect_error error = Connect_error_msg {error}
|
|
|
|
let extract_communication_error connection error =
|
|
Connection_error_msg {
|
|
error;
|
|
connection_status = connection#status;
|
|
}
|
|
|
|
let extract_result_error result =
|
|
Result_error_msg {
|
|
error_message = result#error;
|
|
sqlstate = result#error_field Pg.Error_field.SQLSTATE;
|
|
}
|
|
|
|
let () =
|
|
let pp ppf = function
|
|
| Connect_error_msg {error; _} | Connection_error_msg {error; _} ->
|
|
Format.pp_print_string ppf (Pg.string_of_error error)
|
|
| Result_error_msg {error_message; _} ->
|
|
Format.pp_print_string ppf error_message
|
|
| _ ->
|
|
assert false
|
|
in
|
|
let cause = function
|
|
| Result_error_msg {sqlstate; _} ->
|
|
Postgresql_conv.cause_of_sqlstate sqlstate
|
|
| _ ->
|
|
assert false
|
|
in
|
|
Caqti_error.define_msg ~pp [%extension_constructor Connect_error_msg];
|
|
Caqti_error.define_msg ~pp [%extension_constructor Connection_error_msg];
|
|
Caqti_error.define_msg ~pp ~cause [%extension_constructor Result_error_msg]
|
|
|
|
let driver_info =
|
|
let dummy_dialect =
|
|
Dialect.create_pgsql
|
|
~server_version:(Version.of_string_unsafe "")
|
|
~client_library:`postgresql
|
|
()
|
|
in
|
|
Caqti_driver_info.of_dialect dummy_dialect
|
|
|
|
module Pg_ext = struct
|
|
include Postgresql_conv
|
|
|
|
(* Turns a constant into [(must_quote, encode)]. Note the slight difference
|
|
* from encoded parameters like booleans and quoting. *)
|
|
let query_string_of_value
|
|
: type a. Pg.connection -> a Field_type.t -> bool * (a -> string) =
|
|
fun db ->
|
|
let escape_string s = db#escape_string s in
|
|
(function
|
|
| Bool -> (false, string_of_bool)
|
|
| Int -> (false, string_of_int)
|
|
| Int16 -> (false, string_of_int)
|
|
| Int32 -> (false, Int32.to_string)
|
|
| Int64 -> (false, Int64.to_string)
|
|
| Float -> (false, Float.to_string)
|
|
| String -> (true, escape_string)
|
|
| Octets -> (true, escape_string)
|
|
| Pdate -> (true, Conv.iso8601_of_pdate)
|
|
| Ptime -> (true, pgstring_of_ptime)
|
|
| Ptime_span -> (true, pgstring_of_ptime_span)
|
|
| Enum _ -> (true, escape_string))
|
|
|
|
let query_string ~subst (db : Pg.connection) templ =
|
|
let buf = Buffer.create 64 in
|
|
let rec loop : Query.t -> _ = function
|
|
| L s -> Buffer.add_string buf s
|
|
| Q s ->
|
|
Buffer.add_char buf '\'';
|
|
Buffer.add_string buf (db#escape_string s);
|
|
Buffer.add_char buf '\''
|
|
| V (ft, v) ->
|
|
let quote, conv = query_string_of_value db ft in
|
|
if quote then Buffer.add_char buf '\'';
|
|
Buffer.add_string buf (conv v);
|
|
if quote then Buffer.add_char buf '\''
|
|
| P i -> bprintf buf "$%d" (i + 1)
|
|
| E _ -> assert false
|
|
| S frags -> List.iter loop frags
|
|
in
|
|
loop (Query.expand ~final:true subst templ);
|
|
Buffer.contents buf
|
|
|
|
let escaped_connvalue s =
|
|
let buf = Buffer.create (String.length s) in
|
|
let aux = function
|
|
| '\\' -> Buffer.add_string buf {|\\|}
|
|
| '\'' -> Buffer.add_string buf {|\'|}
|
|
| ch -> Buffer.add_char buf ch in
|
|
String.iter aux s;
|
|
Buffer.contents buf
|
|
|
|
let pop_uri_param present absent param uri =
|
|
(match Uri.get_query_param uri param with
|
|
| None ->
|
|
Ok (absent, uri)
|
|
| Some value_str ->
|
|
(match present value_str with
|
|
| value -> Ok (value, Uri.remove_query_param uri param)
|
|
| exception Failure msg ->
|
|
let msg = Caqti_error.Msg msg in
|
|
Error (Caqti_error.connect_rejected ~uri msg)))
|
|
|
|
let parse_notice_processing = function
|
|
| "quiet" -> `Quiet
|
|
| "stderr" -> `Stderr
|
|
| _ -> failwith "Invalid argument for notice_processing."
|
|
|
|
let parse_uri uri =
|
|
pop_uri_param parse_notice_processing `Quiet "notice_processing" uri
|
|
|>? fun (notice_processing, uri) ->
|
|
pop_uri_param bool_of_string false "use_single_row_mode" uri
|
|
|>? fun (use_single_row_mode, uri) ->
|
|
let conninfo =
|
|
if Uri.host uri <> None then Uri.to_string ~pct_encoder uri else
|
|
let mkparam k v = k ^ " = '" ^ escaped_connvalue v ^ "'" in
|
|
let mkparams (k, vs) = List.map (mkparam k) vs in
|
|
String.concat " " (List.flatten (List.map mkparams (Uri.query uri)))
|
|
in
|
|
Ok (conninfo, notice_processing, use_single_row_mode)
|
|
end
|
|
|
|
let bool_oid = Pg.oid_of_ftype Pg.BOOL
|
|
let int2_oid = Pg.oid_of_ftype Pg.INT2
|
|
let int4_oid = Pg.oid_of_ftype Pg.INT4
|
|
let int8_oid = Pg.oid_of_ftype Pg.INT8
|
|
let float8_oid = Pg.oid_of_ftype Pg.FLOAT8
|
|
let bytea_oid = Pg.oid_of_ftype Pg.BYTEA
|
|
let date_oid = Pg.oid_of_ftype Pg.DATE
|
|
let timestamp_oid = Pg.oid_of_ftype Pg.TIMESTAMPTZ
|
|
let interval_oid = Pg.oid_of_ftype Pg.INTERVAL
|
|
let unknown_oid = Pg.oid_of_ftype Pg.UNKNOWN
|
|
|
|
let init_param_types ~type_oid_cache =
|
|
let oid_of_field_type : type a. a Field_type.t -> _ = function
|
|
| Bool -> Ok bool_oid
|
|
| Int -> Ok int8_oid
|
|
| Int16 -> Ok int2_oid
|
|
| Int32 -> Ok int4_oid
|
|
| Int64 -> Ok int8_oid
|
|
| Float -> Ok float8_oid
|
|
| String -> Ok unknown_oid
|
|
| Octets -> Ok bytea_oid
|
|
| Pdate -> Ok date_oid
|
|
| Ptime -> Ok timestamp_oid
|
|
| Ptime_span -> Ok interval_oid
|
|
| Enum name -> Ok (Hashtbl.find type_oid_cache name)
|
|
in
|
|
let rec recurse : type a. _ -> _ -> a Row_type.t -> _ -> _
|
|
= fun pt bp -> function
|
|
| Field ft -> fun i ->
|
|
oid_of_field_type ft |>? fun oid ->
|
|
pt.(i) <- oid;
|
|
bp.(i) <- oid = bytea_oid;
|
|
Ok (i + 1)
|
|
| Option t ->
|
|
recurse pt bp t
|
|
| Product (_, prod) ->
|
|
let rec loop : type i. (i, a) Row_type.product -> _ = function
|
|
| Proj_end -> Result.ok
|
|
| Proj (t, _, prod) -> recurse pt bp t %>? loop prod
|
|
in
|
|
loop prod
|
|
| Annot (_, t0) ->
|
|
recurse pt bp t0
|
|
in
|
|
fun pt bp t ->
|
|
recurse pt bp t 0 |>? fun np ->
|
|
assert (np = Array.length pt);
|
|
assert (np = Array.length bp);
|
|
Ok ()
|
|
|
|
module type STRING_ENCODER = sig
|
|
val encode_string : string -> string
|
|
val encode_octets : string -> string
|
|
end
|
|
|
|
module Make_encoder (String_encoder : STRING_ENCODER) = struct
|
|
open String_encoder
|
|
|
|
let encode_field : type a. a Field_type.t -> a -> string =
|
|
fun field_type x ->
|
|
(match field_type with
|
|
| Bool -> Pg_ext.pgstring_of_bool x
|
|
| Int -> string_of_int x
|
|
| Int16 -> string_of_int x
|
|
| Int32 -> Int32.to_string x
|
|
| Int64 -> Int64.to_string x
|
|
| Float -> sprintf "%.17g" x
|
|
| String -> encode_string x
|
|
| Enum _ -> encode_string x
|
|
| Octets -> encode_octets x
|
|
| Pdate -> Conv.iso8601_of_pdate x
|
|
| Ptime -> Pg_ext.pgstring_of_ptime x
|
|
| Ptime_span -> Pg_ext.pgstring_of_ptime_span x)
|
|
|
|
let encode ~uri params t x =
|
|
let write_value ~uri:_ ft fv i =
|
|
let s = encode_field ft fv in
|
|
params.(i) <- s; i + 1
|
|
in
|
|
let write_null ~uri:_ _ i = i + 1 in
|
|
try
|
|
let n = Request_utils.encode_param ~uri {write_value; write_null} t x 0 in
|
|
assert (n = Array.length params);
|
|
Ok ()
|
|
with Caqti_error.Exn (#Caqti_error.call as err) -> Error err
|
|
end
|
|
|
|
module Param_encoder = Make_encoder (struct
|
|
let encode_string s = s
|
|
let encode_octets s = s
|
|
end)
|
|
|
|
let decode_field : type a. uri: Uri.t -> a Field_type.t -> string -> a =
|
|
fun ~uri field_type s ->
|
|
let wrap_conv_exn f s =
|
|
(try (f s) with
|
|
| _ ->
|
|
let msg = Caqti_error.Msg (sprintf "Invalid value %S." s) 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 ->
|
|
let msg = Caqti_error.Msg msg 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 Pg_ext.bool_of_pgstring s
|
|
| Int -> wrap_conv_exn int_of_string s
|
|
| Int16 -> wrap_conv_exn int_of_string s
|
|
| Int32 -> wrap_conv_exn Int32.of_string s
|
|
| Int64 -> wrap_conv_exn Int64.of_string s
|
|
| Float -> wrap_conv_exn float_of_string s
|
|
| String -> s
|
|
| Enum _ -> s
|
|
| Octets -> Postgresql.unescape_bytea s
|
|
| Pdate -> wrap_conv_res Conv.pdate_of_iso8601 s
|
|
| Ptime -> wrap_conv_res Conv.ptime_of_rfc3339_utc s
|
|
| Ptime_span -> wrap_conv_res Pg_ext.ptime_span_of_pgstring s)
|
|
|
|
let decode_row ~uri row_type =
|
|
let read_value ~uri ft (resp, i, j) =
|
|
let y = decode_field ~uri ft (resp#getvalue i j) in
|
|
(y, (resp, i, j + 1))
|
|
in
|
|
let skip_null n (resp, i, j) =
|
|
let j' = j + n in
|
|
let rec check k = k = j' || resp#getisnull i k && check (k + 1) in
|
|
if check j then Some (resp, i, j') else None
|
|
in
|
|
let decode = Request_utils.decode_row ~uri {read_value; skip_null} row_type in
|
|
fun (resp, i) ->
|
|
(match decode (resp, i, 0) with
|
|
| (y, (_, _, j)) -> assert (j = Row_type.length row_type); Ok y
|
|
| exception Caqti_error.Exn (`Decode_rejected _ as err) -> Error err)
|
|
|
|
type request_info = {
|
|
query_name: string;
|
|
query: string;
|
|
param_length: int;
|
|
param_types: Pg.oid array;
|
|
binary_params: bool array;
|
|
}
|
|
|
|
module Pcache =
|
|
Request_cache.Make (struct type t = request_info let weight _ = 1 end)
|
|
|
|
module Connect_functor
|
|
(System : Caqti_platform.System_sig.S)
|
|
(System_unix : Caqti_platform_unix.System_sig.S
|
|
with type 'a fiber := 'a System.Fiber.t
|
|
and type stdenv := System.stdenv) =
|
|
struct
|
|
open System
|
|
open System.Fiber.Infix
|
|
open System_utils.Monad_syntax (System.Fiber)
|
|
open System_unix
|
|
module H = Connection_utils.Make_helpers (System)
|
|
|
|
let ( let/? ) m f = match m with Ok x -> f x | Error _ as r -> Fiber.return r
|
|
let ( >|>=? ) m f = m >|= function Ok x -> f x | Error _ as r -> r
|
|
let ( let+*? ) = ( >|>=? )
|
|
|
|
let driver_info = driver_info
|
|
|
|
module Pg_io = struct
|
|
|
|
let communicate ~stdenv db step =
|
|
let aux fd =
|
|
let rec loop = function
|
|
| Pg.Polling_reading ->
|
|
let* _ = Unix.poll ~stdenv ~read:true fd in
|
|
(match step () with
|
|
| exception Pg.Error msg -> Fiber.return (Error msg)
|
|
| ps -> loop ps)
|
|
| Pg.Polling_writing ->
|
|
let* _ = Unix.poll ~stdenv ~write:true fd in
|
|
(match step () with
|
|
| exception Pg.Error msg -> Fiber.return (Error msg)
|
|
| ps -> loop ps)
|
|
| Pg.Polling_failed | Pg.Polling_ok ->
|
|
Fiber.return (Ok ())
|
|
in
|
|
loop Pg.Polling_writing
|
|
in
|
|
(match db#socket with
|
|
| exception Pg.Error msg -> Fiber.return (Error msg)
|
|
| socket -> Unix.wrap_fd aux (Obj.magic socket))
|
|
|
|
let get_next_result ~stdenv ~uri ~query db =
|
|
let rec retry fd =
|
|
db#consume_input;
|
|
if db#is_busy then
|
|
Unix.poll ~stdenv ~read:true fd >>= (fun _ -> retry fd)
|
|
else
|
|
Fiber.return (Ok db#get_result)
|
|
in
|
|
try Unix.wrap_fd retry (Obj.magic db#socket)
|
|
with Pg.Error err ->
|
|
let msg = extract_communication_error db err in
|
|
Fiber.return (Error (Caqti_error.request_failed ~uri ~query msg))
|
|
|
|
let get_one_result ~stdenv ~uri ~query db =
|
|
get_next_result ~stdenv ~uri ~query db >>=? function
|
|
| None ->
|
|
let msg = Caqti_error.Msg "No response received after send." in
|
|
Fiber.return (Error (Caqti_error.request_failed ~uri ~query msg))
|
|
| Some result ->
|
|
Fiber.return (Ok result)
|
|
|
|
let get_final_result ~stdenv ~uri ~query db =
|
|
get_one_result ~stdenv ~uri ~query db >>=? fun result ->
|
|
get_next_result ~stdenv ~uri ~query db >>=? function
|
|
| None ->
|
|
Fiber.return (Ok result)
|
|
| Some _ ->
|
|
let msg = Caqti_error.Msg "More than one response received." in
|
|
Fiber.return (Error (Caqti_error.response_rejected ~uri ~query msg))
|
|
|
|
let check_query_result ~uri ~query ~row_mult ~single_row_mode result =
|
|
let reject msg =
|
|
let msg = Caqti_error.Msg msg in
|
|
Error (Caqti_error.response_rejected ~uri ~query msg)
|
|
in
|
|
let fail msg =
|
|
let msg = Caqti_error.Msg msg in
|
|
Error (Caqti_error.request_failed ~uri ~query msg)
|
|
in
|
|
(match result#status with
|
|
| Pg.Command_ok ->
|
|
(match Row_mult.expose row_mult with
|
|
| `Zero -> Ok ()
|
|
| (`One | `Zero_or_one | `Zero_or_more) ->
|
|
reject "Tuples expected for this query.")
|
|
| Pg.Tuples_ok ->
|
|
if single_row_mode then
|
|
if result#ntuples = 0 then Ok () else
|
|
reject "Tuples returned in single-row-mode."
|
|
else
|
|
(match Row_mult.expose row_mult with
|
|
| `Zero ->
|
|
if result#ntuples = 0 then Ok () else
|
|
reject "No tuples expected for this query."
|
|
| `One ->
|
|
if result#ntuples = 1 then Ok () else
|
|
ksprintf reject "Received %d tuples, expected one."
|
|
result#ntuples
|
|
| `Zero_or_one ->
|
|
if result#ntuples <= 1 then Ok () else
|
|
ksprintf reject "Received %d tuples, expected at most one."
|
|
result#ntuples
|
|
| `Zero_or_more -> Ok ())
|
|
| Pg.Empty_query -> fail "The query was empty."
|
|
| Pg.Bad_response ->
|
|
let msg = extract_result_error result in
|
|
Error (Caqti_error.response_rejected ~uri ~query msg)
|
|
| Pg.Fatal_error ->
|
|
let msg = extract_result_error result in
|
|
Error (Caqti_error.request_failed ~uri ~query msg)
|
|
| Pg.Nonfatal_error -> Ok () (* TODO: Log *)
|
|
| Pg.Copy_out | Pg.Copy_in | Pg.Copy_both ->
|
|
reject "Received unexpected copy response."
|
|
| Pg.Single_tuple ->
|
|
if not single_row_mode then
|
|
reject "Received unexpected single tuple response." else
|
|
if result#ntuples <> 1 then
|
|
reject "Expected a single row in single-row mode." else
|
|
Ok ())
|
|
|
|
let check_command_result ~uri ~query result =
|
|
check_query_result
|
|
~uri ~query ~row_mult:Row_mult.zero ~single_row_mode:false result
|
|
end
|
|
|
|
(* Driver Interface *)
|
|
|
|
module type CONNECTION = Caqti_connection_sig.S
|
|
with type 'a fiber := 'a Fiber.t
|
|
and type ('a, 'err) stream := ('a, 'err) Stream.t
|
|
|
|
module Make_connection_base
|
|
(Connection_arg : sig
|
|
val stdenv : stdenv
|
|
val dialect : Dialect.t
|
|
val subst : Query.subst
|
|
val uri : Uri.t
|
|
val db : Pg.connection
|
|
val use_single_row_mode : bool
|
|
val dynamic_capacity : int
|
|
end) =
|
|
struct
|
|
open Connection_arg
|
|
|
|
let dialect = dialect
|
|
|
|
module Copy_encoder = Make_encoder (struct
|
|
|
|
let encode_string s =
|
|
let buf = Buffer.create (String.length s) in
|
|
for i = 0 to String.length s - 1 do
|
|
(match s.[i] with
|
|
| '\\' -> Buffer.add_string buf "\\\\"
|
|
| '\n' -> Buffer.add_string buf "\\n"
|
|
| '\r' -> Buffer.add_string buf "\\r"
|
|
| '\t' -> Buffer.add_string buf "\\t"
|
|
| c -> Buffer.add_char buf c)
|
|
done;
|
|
Buffer.contents buf
|
|
|
|
let encode_octets s = encode_string (db#escape_bytea s)
|
|
end)
|
|
|
|
let in_use = ref false
|
|
let in_transaction = ref false
|
|
let pcache : Pcache.t = Pcache.create ~dynamic_capacity dialect
|
|
|
|
let wrap_pg ~query f =
|
|
try Ok (f ()) with
|
|
| Postgresql.Error err ->
|
|
let msg = extract_communication_error db err in
|
|
Error (Caqti_error.request_failed ~uri ~query msg)
|
|
|
|
let reset () =
|
|
Log.warn (fun p ->
|
|
p "Lost connection to <%a>, reconnecting." Caqti_error.pp_uri uri)
|
|
>>= fun () ->
|
|
in_transaction := false;
|
|
(match db#reset_start with
|
|
| exception Pg.Error _ -> Fiber.return false
|
|
| true ->
|
|
Pcache.clear_and_discard pcache;
|
|
Pg_io.communicate ~stdenv db (fun () -> db#reset_poll) >|=
|
|
(function
|
|
| Error _ -> false
|
|
| Ok () -> (try db#status = Pg.Ok with Pg.Error _ -> false))
|
|
| false ->
|
|
Fiber.return false)
|
|
|
|
let rec retry_on_connection_error ?(n = 1) f =
|
|
if !in_transaction then f () else
|
|
(f () : (_, [> Caqti_error.call]) result Fiber.t) >>=
|
|
(function
|
|
| Ok _ as r -> Fiber.return r
|
|
| Error (`Request_failed
|
|
{Caqti_error.msg = Connection_error_msg
|
|
{error = Postgresql.Connection_failure _; _}; _})
|
|
as r when n > 0 ->
|
|
let* reset_ok = reset () in
|
|
if reset_ok then
|
|
retry_on_connection_error ~n:(n - 1) f
|
|
else
|
|
Fiber.return r
|
|
| Error _ as r -> Fiber.return r)
|
|
|
|
let send_simple_query query =
|
|
retry_on_connection_error begin fun () ->
|
|
Fiber.return @@ wrap_pg ~query begin fun () ->
|
|
db#send_query query;
|
|
db#consume_input
|
|
end
|
|
end
|
|
|
|
let send_direct_query ~single_row_mode request_info params =
|
|
let {query; param_types; binary_params; _} = request_info in
|
|
retry_on_connection_error begin fun () ->
|
|
Fiber.return @@ wrap_pg ~query begin fun () ->
|
|
db#send_query ~params ~param_types ~binary_params query;
|
|
if single_row_mode then db#set_single_row_mode;
|
|
db#consume_input
|
|
end
|
|
end
|
|
|
|
let send_prepared_query ~single_row_mode request_info params =
|
|
let {query_name; query; binary_params; _} = request_info in
|
|
assert (query_name <> "");
|
|
retry_on_connection_error begin fun () ->
|
|
Fiber.return @@ wrap_pg ~query begin fun () ->
|
|
db#send_query_prepared ~params ~binary_params query_name;
|
|
if single_row_mode then db#set_single_row_mode;
|
|
db#consume_input
|
|
end
|
|
end
|
|
|
|
let fetch_one_result ~query () =
|
|
Pg_io.get_one_result ~stdenv ~uri ~query db
|
|
|
|
let fetch_final_result ~query () =
|
|
Pg_io.get_final_result ~stdenv ~uri ~query db
|
|
|
|
let fetch_single_row ~query () =
|
|
Pg_io.get_one_result ~stdenv ~uri ~query db >>=? fun result ->
|
|
(match result#status with
|
|
| Pg.Single_tuple ->
|
|
assert (result#ntuples = 1);
|
|
Fiber.return (Ok (Some result))
|
|
| Pg.Tuples_ok ->
|
|
assert (result#ntuples = 0);
|
|
Pg_io.get_next_result ~stdenv ~uri ~query db >|>=?
|
|
(function
|
|
| None -> Ok None
|
|
| Some _ ->
|
|
let msg =
|
|
Caqti_error.Msg "Extra result after final single-row result." in
|
|
Error (Caqti_error.response_rejected ~uri ~query msg))
|
|
| _ ->
|
|
Fiber.return @@ Result.map (fun () -> None) @@
|
|
Pg_io.check_query_result
|
|
~uri ~query ~row_mult:Row_mult.zero_or_more ~single_row_mode:true
|
|
result)
|
|
|
|
let prepare {query_name; query; param_types; _} =
|
|
assert (query_name <> "");
|
|
retry_on_connection_error begin fun () ->
|
|
let*? () =
|
|
Fiber.return @@ wrap_pg ~query @@ fun () ->
|
|
db#send_prepare ~param_types query_name query;
|
|
db#consume_input
|
|
in
|
|
let+*? result = Pg_io.get_final_result ~stdenv ~uri ~query db in
|
|
Pg_io.check_command_result ~uri ~query result
|
|
end
|
|
|
|
let free_prepared request_info =
|
|
let query = sprintf "DEALLOCATE %s" request_info.query_name in
|
|
let*? () = send_simple_query query in
|
|
let+*? result = fetch_final_result ~query () in
|
|
Pg_io.check_query_result
|
|
~uri ~query ~row_mult:Row_mult.zero ~single_row_mode:false
|
|
result
|
|
|
|
module Response = struct
|
|
|
|
type source =
|
|
| Complete of Pg.result
|
|
| Single_row
|
|
|
|
type ('b, 'm) t = {
|
|
row_type: 'b Row_type.t;
|
|
source: source;
|
|
query: string;
|
|
}
|
|
|
|
let returned_count = function
|
|
| {source = Complete result; _} ->
|
|
Fiber.return (Ok result#ntuples)
|
|
| {source = Single_row; _} ->
|
|
Fiber.return (Error `Unsupported)
|
|
|
|
let affected_count = function
|
|
| {source = Complete result; _} ->
|
|
Fiber.return (Ok (int_of_string result#cmd_tuples))
|
|
| {source = Single_row; _} ->
|
|
Fiber.return (Error `Unsupported)
|
|
|
|
let exec _ = Fiber.return (Ok ())
|
|
|
|
let find = function
|
|
| {row_type; source = Complete result; _} ->
|
|
Fiber.return (decode_row ~uri row_type (result, 0))
|
|
| {source = Single_row; _} ->
|
|
assert false
|
|
|
|
let find_opt = function
|
|
| {row_type; source = Complete result; _} ->
|
|
Fiber.return begin
|
|
if result#ntuples = 0 then Ok None else
|
|
(match decode_row ~uri row_type (result, 0) with
|
|
| Ok y -> Ok (Some y)
|
|
| Error _ as r -> r)
|
|
end
|
|
| {source = Single_row; _} ->
|
|
assert false
|
|
|
|
let fold f {row_type; query; source} =
|
|
let decode = decode_row ~uri row_type in
|
|
(match source with
|
|
| Complete result ->
|
|
let n = result#ntuples in
|
|
let rec loop i acc =
|
|
if i = n then Ok acc else
|
|
(match decode (result, i) with
|
|
| Ok y -> loop (i + 1) (f y acc)
|
|
| Error _ as r -> r)
|
|
in
|
|
fun acc -> Fiber.return (loop 0 acc)
|
|
| Single_row ->
|
|
let rec loop acc =
|
|
fetch_single_row ~query () >>=? function
|
|
| None -> Fiber.return (Ok acc)
|
|
| Some result ->
|
|
(match decode (result, 0) with
|
|
| Ok y -> loop (f y acc)
|
|
| Error _ as r -> Fiber.return r)
|
|
in
|
|
loop)
|
|
|
|
let fold_s f {row_type; query; source} =
|
|
let decode = decode_row ~uri row_type in
|
|
(match source with
|
|
| Complete result ->
|
|
let n = result#ntuples in
|
|
let rec loop i acc =
|
|
if i = n then Fiber.return (Ok acc) else
|
|
(match decode (result, i) with
|
|
| Ok y -> f y acc >>=? loop (i + 1)
|
|
| Error _ as r -> Fiber.return r)
|
|
in
|
|
loop 0
|
|
| Single_row ->
|
|
let rec loop acc =
|
|
fetch_single_row ~query () >>=? function
|
|
| None -> Fiber.return (Ok acc)
|
|
| Some result ->
|
|
(match decode (result, 0) with
|
|
| Ok y -> f y acc >>=? loop
|
|
| Error _ as r -> Fiber.return r)
|
|
in
|
|
loop)
|
|
|
|
let iter_s f {row_type; query; source} =
|
|
let decode = decode_row ~uri row_type in
|
|
(match source with
|
|
| Complete result ->
|
|
let n = result#ntuples in
|
|
let rec loop i =
|
|
if i = n then Fiber.return (Ok ()) else
|
|
(match decode (result, i) with
|
|
| Ok y -> f y >>=? fun () -> loop (i + 1)
|
|
| Error _ as r -> Fiber.return r)
|
|
in
|
|
loop 0
|
|
| Single_row ->
|
|
let rec loop () =
|
|
fetch_single_row ~query () >>=? function
|
|
| None -> Fiber.return (Ok ())
|
|
| Some result ->
|
|
(match decode (result, 0) with
|
|
| Ok y -> f y >>=? fun () -> loop ()
|
|
| Error _ as r -> Fiber.return r)
|
|
in
|
|
loop ())
|
|
|
|
let to_stream {row_type; query; source} =
|
|
let decode = decode_row ~uri row_type in
|
|
(match source with
|
|
| Complete result ->
|
|
let n = result#ntuples in
|
|
let rec seq i () =
|
|
if i = n then Fiber.return Stream.Nil else
|
|
(match decode (result, i) with
|
|
| Ok y -> Fiber.return (Stream.Cons (y, seq (i + 1)))
|
|
| Error err -> Fiber.return (Stream.Error err))
|
|
in
|
|
seq 0
|
|
| Single_row ->
|
|
let rec seq () =
|
|
fetch_single_row ~query () >|= function
|
|
| Ok None -> Stream.Nil
|
|
| Ok (Some result) ->
|
|
(match decode (result, 0) with
|
|
| Ok y -> Stream.Cons (y, seq)
|
|
| Error err -> Stream.Error err)
|
|
| Error err -> Stream.Error err
|
|
in
|
|
seq)
|
|
end
|
|
|
|
let type_oid_cache = Hashtbl.create 19
|
|
|
|
let pp_request_with_param ppf =
|
|
Request.make_pp_with_param ~subst ~dialect () ppf
|
|
|
|
let fresh_static_name = Request_utils.fresh_name_generator "caqs"
|
|
let fresh_dynamic_name = Request_utils.fresh_name_generator "caqd"
|
|
|
|
let build_request_info request =
|
|
let templ = Request.query request dialect in
|
|
let query_name =
|
|
(match Request.prepare_policy request with
|
|
| Direct -> ""
|
|
| Static -> fresh_static_name ()
|
|
| Dynamic -> fresh_dynamic_name ())
|
|
in
|
|
let query = Pg_ext.query_string ~subst db templ in
|
|
let param_type = Request.param_type request in
|
|
let param_length = Row_type.length param_type in
|
|
let param_types = Array.make param_length 0 in
|
|
let binary_params = Array.make param_length false in
|
|
init_param_types ~type_oid_cache param_types binary_params param_type
|
|
|> Result.map @@ fun () ->
|
|
{query_name; query; param_length; param_types; binary_params}
|
|
|
|
let build_params request request_info param =
|
|
let param_type = Request.param_type request in
|
|
let params = Array.make request_info.param_length Pg.null in
|
|
Param_encoder.encode ~uri params param_type param
|
|
|> Result.map (fun () -> params)
|
|
|
|
let send_request ~single_row_mode request param =
|
|
(match Request.prepare_policy request with
|
|
| Direct ->
|
|
let/? request_info = build_request_info request in
|
|
let/? params = build_params request request_info param in
|
|
let+? () = send_direct_query ~single_row_mode request_info params in
|
|
request_info.query
|
|
| Dynamic | Static ->
|
|
let*? request_info =
|
|
(match Pcache.find_and_promote pcache request with
|
|
| Some request_info ->
|
|
Fiber.return (Ok request_info)
|
|
| None ->
|
|
let/? request_info = build_request_info request in
|
|
let+? () = prepare request_info in
|
|
Pcache.add pcache request request_info;
|
|
request_info)
|
|
in
|
|
let/? params = build_params request request_info param in
|
|
let+? () = send_prepared_query ~single_row_mode request_info params in
|
|
request_info.query)
|
|
|
|
let call_without_oids ~f request param =
|
|
Log.debug ~src:Logging.request_log_src (fun f ->
|
|
f "Sending %a" pp_request_with_param (request, param)) >>= fun () ->
|
|
|
|
let single_row_mode =
|
|
use_single_row_mode
|
|
&& Row_mult.can_be_many (Request.row_mult request)
|
|
in
|
|
|
|
(* Prepare, if requested, and send the query. *)
|
|
let*? query = send_request ~single_row_mode request param in
|
|
|
|
(* Fetch and process the result. *)
|
|
let row_type = Request.row_type request in
|
|
if single_row_mode then
|
|
f Response.{row_type; query; source = Single_row}
|
|
else begin
|
|
let row_mult = Request.row_mult request in
|
|
let*? result = fetch_final_result ~query () in
|
|
(match Pg_io.check_query_result
|
|
~uri ~query ~row_mult ~single_row_mode result with
|
|
| Ok () -> f Response.{row_type; query; source = Complete result}
|
|
| Error _ as r -> Fiber.return r)
|
|
end
|
|
|
|
let rec fetch_type_oids : type a. a Row_type.t -> _ = function
|
|
| Field (Enum name as field_type)
|
|
when not (Hashtbl.mem type_oid_cache name) ->
|
|
call_without_oids ~f:Response.find_opt Q.type_oid name >>=
|
|
(function
|
|
| Ok (Some oid) ->
|
|
Fiber.return (Ok (Hashtbl.add type_oid_cache name oid))
|
|
| Ok None ->
|
|
Log.warn (fun p ->
|
|
p "Failed to query OID for enum %s." name) >|= fun () ->
|
|
Error (Caqti_error.encode_missing ~uri ~field_type ())
|
|
| Error (`Encode_rejected _ | `Decode_rejected _ |
|
|
`Response_failed _ as err) ->
|
|
Log.err (fun p ->
|
|
p "Failed to fetch obtain OID for enum %s due to: %a"
|
|
name Caqti_error.pp err) >|= fun () ->
|
|
Error (Caqti_error.encode_missing ~uri ~field_type ())
|
|
| Error #Caqti_error.call as r ->
|
|
Fiber.return r)
|
|
| Field _ -> Fiber.return (Ok ())
|
|
| Option t -> fetch_type_oids t
|
|
| Product (_, prod) ->
|
|
let rec loop : type i. (i, a) Row_type.product -> _ = function
|
|
| Proj_end -> Fiber.return (Ok ())
|
|
| Proj (t, _, prod) -> fetch_type_oids t >>=? fun () -> loop prod
|
|
in
|
|
loop prod
|
|
| Annot (_, t0) -> fetch_type_oids t0
|
|
|
|
let using_db f =
|
|
if !in_use then
|
|
failwith "Invalid concurrent usage of PostgreSQL connection detected.";
|
|
in_use := true;
|
|
Fiber.cleanup
|
|
(fun () -> f () >|= fun res -> in_use := false; res)
|
|
(fun () -> reset () >|= fun _ -> in_use := false)
|
|
|
|
let deallocate request = using_db @@ fun () ->
|
|
(match Request.prepare_policy request with
|
|
| Direct -> failwith "deallocate called on direct request"
|
|
| Dynamic | Static ->
|
|
(match Pcache.deallocate pcache request with
|
|
| None -> Fiber.return (Ok ())
|
|
| Some (request_info, commit_remove) ->
|
|
free_prepared request_info >|= Result.map commit_remove))
|
|
|
|
let deallocate_some () =
|
|
let rec loop = function
|
|
| [] -> Fiber.return (Ok ())
|
|
| request_info :: orphans ->
|
|
let*? () = free_prepared request_info in
|
|
loop orphans
|
|
in
|
|
let orphans, commit = Pcache.trim pcache in
|
|
loop orphans >|=? commit
|
|
|
|
let call ~f req param = using_db @@ fun () ->
|
|
deallocate_some () >>=? fun () ->
|
|
fetch_type_oids (Request.param_type req) >>=? fun () ->
|
|
call_without_oids ~f req param
|
|
|
|
let disconnect () = using_db @@ fun () ->
|
|
try db#finish; Fiber.return () with Pg.Error err ->
|
|
Log.warn (fun p ->
|
|
p "While disconnecting from <%a>: %s"
|
|
Caqti_error.pp_uri uri (Pg.string_of_error err))
|
|
|
|
let validate () = using_db @@ fun () ->
|
|
if (try db#consume_input; db#status = Pg.Ok with Pg.Error _ -> false) then
|
|
Fiber.return true
|
|
else
|
|
reset ()
|
|
|
|
let check f = f (try db#status = Pg.Ok with Pg.Error _ -> false)
|
|
|
|
let exec q p = call ~f:Response.exec q p
|
|
let start () = exec Q.start () >|=? fun () -> in_transaction := true
|
|
let commit () = in_transaction := false; exec Q.commit ()
|
|
let rollback () = in_transaction := false; exec Q.rollback ()
|
|
|
|
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
|
|
call ~f:Response.exec (Q.set_statement_timeout t_arg) ()
|
|
|
|
let populate ~table ~columns row_type data =
|
|
let query =
|
|
sprintf "COPY %s (%s) FROM STDIN" table (String.concat "," columns)
|
|
in
|
|
let param_length = Row_type.length row_type in
|
|
let fail msg =
|
|
Fiber.return
|
|
(Error (Caqti_error.request_failed ~uri ~query (Caqti_error.Msg msg)))
|
|
in
|
|
let pg_error err =
|
|
let msg = extract_communication_error db err in
|
|
Fiber.return (Error (Caqti_error.request_failed ~uri ~query msg))
|
|
in
|
|
let put_copy_data data =
|
|
let rec loop fd =
|
|
match db#put_copy_data data with
|
|
| Pg.Put_copy_error ->
|
|
fail "Unable to put copy data"
|
|
| Pg.Put_copy_queued ->
|
|
Fiber.return (Ok ())
|
|
| Pg.Put_copy_not_queued ->
|
|
Unix.poll ~stdenv ~write:true fd >>= fun _ -> loop fd
|
|
in
|
|
(match db#socket with
|
|
| exception Pg.Error msg -> pg_error msg
|
|
| socket -> Unix.wrap_fd loop (Obj.magic socket))
|
|
in
|
|
let copy_row row =
|
|
let params = Array.make param_length "\\N" in
|
|
(match Copy_encoder.encode ~uri params row_type row with
|
|
| Ok () ->
|
|
Fiber.return (Ok (String.concat "\t" (Array.to_list params)))
|
|
| Error _ as r ->
|
|
Fiber.return r)
|
|
>>=? fun param_string -> put_copy_data (param_string ^ "\n")
|
|
in
|
|
begin
|
|
(* Send the copy command to start the transfer.
|
|
* Skip checking that there is only a single result: while in copy mode
|
|
* we can repeatedly get the latest result and it will always be
|
|
* Copy_in, so checking for a single result would trigger an error.
|
|
*)
|
|
send_simple_query query >>=? fun () ->
|
|
fetch_one_result ~query ()
|
|
>>=? fun result ->
|
|
(* We expect the Copy_in response only - turn other success responses
|
|
* into errors, and delegate error handling.
|
|
*)
|
|
(match result#status with
|
|
| Pg.Copy_in -> Fiber.return (Ok ())
|
|
| Pg.Command_ok -> fail "Received Command_ok when expecting Copy_in"
|
|
| _ -> Fiber.return (Pg_io.check_command_result ~uri ~query result))
|
|
>>=? fun () -> System.Stream.iter_s ~f:copy_row data
|
|
>>=? fun () ->
|
|
(* End the copy *)
|
|
let rec copy_end_loop fd =
|
|
match db#put_copy_end () with
|
|
| Pg.Put_copy_error ->
|
|
fail "Unable to finalize copy"
|
|
| Pg.Put_copy_not_queued ->
|
|
Unix.poll ~stdenv ~write:true fd >>= fun _ ->
|
|
copy_end_loop fd
|
|
| Pg.Put_copy_queued ->
|
|
Fiber.return (Ok ())
|
|
in
|
|
(match db#socket with
|
|
| exception Pg.Error msg -> pg_error msg
|
|
| socket -> Unix.wrap_fd copy_end_loop (Obj.magic socket))
|
|
>>=? fun () ->
|
|
(* After ending the copy, there will be a new result for the initial
|
|
* query.
|
|
*)
|
|
fetch_final_result ~query () >|>=? Pg_io.check_command_result ~uri ~query
|
|
end
|
|
end
|
|
|
|
let connect ~sw:_ ~stdenv ~subst ~config uri =
|
|
Fiber.return (Pg_ext.parse_uri uri)
|
|
>>=? fun (conninfo, notice_processing, use_single_row_mode) ->
|
|
(match new Pg.connection ~conninfo () with
|
|
| exception Pg.Error err ->
|
|
let msg = extract_connect_error err in
|
|
Fiber.return (Error (Caqti_error.connect_failed ~uri msg))
|
|
| db ->
|
|
Pg_io.communicate ~stdenv db (fun () -> db#connect_poll) >>=
|
|
(function
|
|
| Error err ->
|
|
let msg = extract_communication_error db err in
|
|
Fiber.return (Error (Caqti_error.connect_failed ~uri msg))
|
|
| Ok () ->
|
|
(match db#status <> Pg.Ok with
|
|
| exception Pg.Error err ->
|
|
let msg = extract_communication_error db err in
|
|
Fiber.return (Error (Caqti_error.connect_failed ~uri msg))
|
|
| true ->
|
|
let msg = Caqti_error.Msg db#error_message in
|
|
Fiber.return (Error (Caqti_error.connect_failed ~uri msg))
|
|
| false ->
|
|
db#set_notice_processing notice_processing;
|
|
let server_version =
|
|
let v0, v1, v2 = db#server_version in
|
|
Version.of_string_unsafe
|
|
(if v0 < 10 then
|
|
Printf.sprintf "%d.%d.%d" v0 v1 v2
|
|
else
|
|
Printf.sprintf "%d.%d" v0 (v1 * 100 + v2))
|
|
in
|
|
let module B = Make_connection_base
|
|
(struct
|
|
let dialect =
|
|
Dialect.create_pgsql
|
|
~server_version ~client_library:`postgresql ()
|
|
let subst = subst dialect
|
|
let stdenv = stdenv
|
|
let uri = uri
|
|
let db = db
|
|
let use_single_row_mode = use_single_row_mode
|
|
let dynamic_capacity =
|
|
Caqti_connect_config.(get dynamic_prepare_capacity) config
|
|
end)
|
|
in
|
|
let module Connection = struct
|
|
let driver_info = driver_info
|
|
let driver_connection = None
|
|
include B
|
|
include Connection_utils.Make_convenience (System) (B)
|
|
end in
|
|
Connection.exec Q.set_timezone_to_utc () >|=
|
|
(function
|
|
| Ok () -> Ok (module Connection : CONNECTION)
|
|
| Error err -> Error (`Post_connect err)))))
|
|
end
|
|
|
|
let () =
|
|
let open Caqti_platform_unix.Driver_loader in
|
|
register "postgres" (module Connect_functor);
|
|
register "postgresql" (module Connect_functor)
|