unikernel/duniverse/ocaml-caqti/caqti-driver-mariadb/lib/caqti_driver_mariadb.ml

@@ -0,0 +1,642 @@
+(* 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
+
+let (|>?) = Result.bind
+
+let cause_of_errno = function
+ | 1022 -> `Unique_violation
+ | 1048 -> `Not_null_violation
+ | 1052 -> `Integrity_constraint_violation__don't_match
+ | 1062 -> `Unique_violation
+ | 1169 -> `Unique_violation
+ | 1216 -> `Foreign_key_violation
+ | 1217 -> `Foreign_key_violation
+ | 1451 -> `Foreign_key_violation
+ | 1452 -> `Foreign_key_violation
+ | 1557 -> `Integrity_constraint_violation__don't_match
+ | 1586 -> `Unique_violation
+ | 1761 -> `Integrity_constraint_violation__don't_match
+ | 1762 -> `Integrity_constraint_violation__don't_match
+ | 1859 -> `Integrity_constraint_violation__don't_match
+ | 4025 -> `Check_violation
+ | _ -> `Unspecified__don't_match
+
+type Caqti_error.msg += Error_msg of {errno: int; error: string}
+let () =
+  let pp ppf = function
+   | Error_msg {errno; error} ->
+      Format.fprintf ppf "Error %d, %s." errno error
+   | _ ->
+      assert false
+  in
+  let cause = function
+   | Error_msg {errno; _} -> cause_of_errno errno
+   | _ -> assert false
+  in
+  Caqti_error.define_msg ~pp ~cause [%extension_constructor Error_msg]
+
+module Q = struct
+  open Caqti_template.Create
+  let set_utc =
+    direct T.(unit -->. unit) "SET time_zone = '+00:00'"
+  let set_statement_timeout =
+    direct T.(float -->. unit) "SET max_statement_time = ?"
+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)
+  module H = Connection_utils.Make_helpers (System)
+
+  let (>>=?) m f = m >>= function Ok x -> f x | Error _ as r -> Fiber.return r
+
+  let rec fold_s_list f =
+    (function
+     | [] -> Fiber.return ()
+     | x :: xs -> f x >>= fun () -> fold_s_list f xs)
+
+  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 dialect =
+    Dialect.create_mysql ~server_version:(Version.of_string_unsafe "") ()
+  let driver_info =
+    Caqti_driver_info.of_dialect dialect
+
+  (* We need to pass stdenv into the below wait, in order to implement
+   * timout for EIO, since it uses stdenv#clock.  This means that our Mdb
+   * instance becomes dependent on stdenv and will therefore be
+   * instantiated for each connection. *)
+  module Pass_stdenv
+    (Connect_env : sig val stdenv : stdenv end) =
+  struct
+    open Connect_env
+
+    module Mdb = Mariadb.Nonblocking.Make
+      (struct
+        open System
+
+        module IO = struct
+          type 'a future = 'a Fiber.t
+          let (>>=) = Fiber.Infix.(>>=)
+          let return = Fiber.return
+        end
+
+        let wait db status =
+          Mariadb.Nonblocking.fd db |> System_unix.Unix.wrap_fd @@ fun fd ->
+          System_unix.Unix.poll
+            ~stdenv
+            ~read:(Mariadb.Nonblocking.Status.read status)
+            ~write:(Mariadb.Nonblocking.Status.write status) fd
+            >|= (fun (read, write, timeout) ->
+          Mariadb.Nonblocking.Status.create ~read ~write ~timeout ())
+      end)
+
+    module Mdb_ext = struct
+      module Field = struct
+
+        let typename field =
+          (match Mdb.Field.value field with
+           | `Null -> "null"
+           | `Int _ -> "int"
+           | `Float _ -> "float"
+           | `String _ -> "string"
+           | `Bytes _ -> "bytes"
+           | `Time _ -> "time")
+
+        let bool field =
+          (match Mdb.Field.value field with
+           | `Int i -> i <> 0
+           | _ -> failwith "Mdb_ext.Field.float")
+
+        (* Lax conversion from float and string since:
+         * - Arithmetic involving integers like `1 + ?` with at least one
+         *   parameter comes back as float.
+         * - An expression like `sum(1)` comes back as a string here, presumably
+         *   as decimal from MariaDB. *)
+        let int field =
+          (match Mdb.Field.value field with
+           | `Int i -> i
+           | `Float x when fst (modf x) = 0.0 -> int_of_float x
+           | `String s -> int_of_string s
+           | _ -> failwith "Mdb_ext.Field.int")
+
+        let int32 field =
+          (match Mdb.Field.value field with
+           | `Int i -> Int32.of_int i
+           | `Float x -> Int32.of_float x   (* cf. int *)
+           | `String s -> Int32.of_string s (* cf. int *)
+           | _ -> failwith "Mdb_ext.Field.int32")
+
+        let int64 field =
+          (match Mdb.Field.value field with
+           | `Int i -> Int64.of_int i
+           | `Float x -> Int64.of_float x   (* cf. int *)
+           | `String s -> Int64.of_string s (* cf. int *)
+           | _ -> failwith "Mdb_ext.Field.int64")
+
+        let float field =
+          (match Mdb.Field.value field with
+           | `Int i -> float_of_int i
+           | `Float x -> x
+           | _ -> failwith "Mdb_ext.Field.float")
+
+        let string field =
+          (match Mdb.Field.value field with
+           | `String s -> s
+           | `Bytes s -> Bytes.to_string s
+           | _ -> failwith "Mdb_ext.Field.string")
+      end
+    end
+
+    let encode_field : type a. a Field_type.t -> a -> Mdb.Field.value =
+      fun field_type x ->
+      (match field_type with
+       | Bool -> `Int (if x then 1 else 0)
+       | Int -> `Int x
+       | Int16 -> `Int x
+       | Int32 -> `Int (Int32.to_int x)
+       | Int64 -> `Int (Int64.to_int x)
+       | Float -> `Float x
+       | String -> `String x
+       | Enum _ -> `String x
+       | Octets -> `Bytes (Bytes.of_string x)
+       | Pdate ->
+          let year, month, day = Ptime.to_date x in
+          `Time (Mdb.Time.date ~year ~month ~day ())
+       | Ptime ->
+          let (year, month, day), ((hour, minute, second), _) =
+            Ptime.to_date_time ~tz_offset_s:0 x in
+          let ps = snd (Ptime.Span.to_d_ps (Ptime.to_span x)) in
+          let tss = Int64.rem ps 1_000_000_000_000L in
+          let microsecond = Int64.to_int (Int64.div tss 1_000_000L) in
+          `Time (Mdb.Time.datetime
+                  ~year ~month ~day ~hour ~minute ~second ~microsecond ())
+       | Ptime_span ->
+          `Float (Ptime.Span.to_float_s x))
+
+    let decode_field
+        : type b. uri: Uri.t -> b Field_type.t -> Mdb.Field.t -> b =
+      fun ~uri field_type field ->
+      try match field_type with
+       | Bool -> Mdb_ext.Field.bool field
+       | Int -> Mdb_ext.Field.int field
+       | Int16 -> Mdb_ext.Field.int field
+       | Int32 -> Mdb_ext.Field.int32 field
+       | Int64 -> Mdb_ext.Field.int64 field
+       | Float -> Mdb.Field.float field
+       | String -> Mdb_ext.Field.string field
+       | Enum _ -> Mdb_ext.Field.string field
+       | Octets -> Mdb_ext.Field.string field
+       | Pdate ->
+          let t = Mdb.Field.time field in
+          let date = Mdb.Time.(year t, month t, day t) in
+          (match Ptime.of_date date with
+           | None -> failwith "Ptime.of_date"
+           | Some t -> t)
+       | Ptime ->
+          let t = Mdb.Field.time field in
+          let date = Mdb.Time.(year t, month t, day t) in
+          let time = Mdb.Time.(hour t, minute t, second t) in
+          let us = Mdb.Time.microsecond t in
+          (match Ptime.of_date_time (date, (time, 0)),
+                 Ptime.Span.of_d_ps (0, Int64.(mul (of_int us) 1_000_000L)) with
+           | None, _ | _, None -> failwith "Ptime.of_date_time"
+           | Some t, Some tss ->
+              (match Ptime.add_span t tss with
+               | None -> failwith "Ptime.of_date_time: Overflow."
+               | Some t' -> t'))
+       | Ptime_span ->
+          let t = Mdb_ext.Field.float field in
+          (match Ptime.Span.of_float_s t with
+           | None -> failwith "Ptime.Span.of_float_s"
+           | Some t -> t)
+      with
+       | Failure _ ->
+          let typename = Mdb_ext.Field.typename field in
+          let msg = Caqti_error.Msg ("Received " ^ typename ^ ".") in
+          let typ = Row_type.field field_type in
+          Request_utils.raise_decode_rejected ~uri ~typ msg
+
+    let encode_param ~uri params t v acc =
+      let write_value ~uri:_ ft fv os =
+        assert (os <> []);
+        let v = encode_field ft fv in
+        List.iter (fun j -> params.(j) <- v) (List.hd os);
+        List.tl os
+      in
+      let write_null ~uri:_ _ os = List.tl os in
+      try
+        Ok (Request_utils.encode_param ~uri {write_value; write_null} t v acc)
+      with Caqti_error.Exn (#Caqti_error.call as err) ->
+        Error err
+
+    let decode_row ~uri row_type =
+      let read_value ~uri ft (row, j) =
+        let fv = decode_field ~uri ft row.(j) in
+        (fv, (row, j + 1))
+      in
+      let skip_null n (row, j) =
+        let j' = j + n in
+        let rec check k =
+          k = j' || Mdb.Field.null_value row.(k) && check (k + 1)
+        in
+        if check j then Some (row, j') else None
+      in
+      let decode =
+        Request_utils.decode_row ~uri {read_value; skip_null} row_type
+      in
+      fun row ->
+        try
+          let (y, (_, j)) = decode (row, 0) in
+          assert (j = Row_type.length row_type);
+          Ok (Some y)
+        with Caqti_error.Exn (#Caqti_error.retrieve as err) ->
+          Error err
+
+    module Make_connection_base
+      (Connection_arg : sig
+        val subst : Query.subst
+        val uri : Uri.t
+        val db : Mdb.t
+        val dynamic_capacity : int
+      end) =
+    struct
+      open Connection_arg
+
+      let dialect =
+        Dialect.create_mysql
+          ~server_version:(Version.of_string_unsafe (Mdb.get_server_info db)) ()
+
+      let using_db_ref = ref false
+      let using_db f =
+        H.assert_single_use ~what:"MariaDB connection" using_db_ref f
+
+      let request_failed ~query (errno, error) =
+        Error (Caqti_error.request_failed ~uri ~query (Error_msg {errno; error}))
+
+      let response_failed ~query (errno, error) =
+        Error (Caqti_error.response_failed ~uri ~query (Error_msg {errno; error}))
+
+      let response_rejected ~query msg =
+        Error (Caqti_error.response_rejected ~uri ~query (Caqti_error.Msg msg))
+
+      module Response = struct
+        type ('b, +'m) t = {
+          query: string;
+          res: Mdb.Res.t;
+          row_type: 'b Row_type.t;
+        }
+
+        let reject_f ~query fmt = ksprintf (response_rejected ~query) fmt
+
+        let affected_count {res; _} = Fiber.return (Ok (Mdb.Res.affected_rows res))
+        let returned_count {res; _} = Fiber.return (Ok (Mdb.Res.num_rows res))
+
+        let decode_next_row ~query row_type =
+          let decode = decode_row ~uri row_type in
+          fun res ->
+            Mdb.Res.fetch (module Mdb.Row.Array) res >|= function
+             | Ok None as r -> r
+             | Ok (Some row) -> decode row
+             | Error err -> response_failed ~query err
+
+        let exec {res; query; _} =
+          (match Mdb.Res.num_rows res with
+           | 0 -> Fiber.return (Ok ())
+           | n -> Fiber.return (reject_f ~query "Received %d tuples for exec." n))
+
+        let find {query; res; row_type} =
+          (match Mdb.Res.num_rows res with
+           | 1 ->
+              decode_next_row ~query row_type res >|=
+              (function
+               | Ok None -> assert false
+               | Ok (Some y) -> Ok y
+               | Error _ as r -> r)
+           | n -> Fiber.return (reject_f ~query "Received %d tuples for find." n))
+
+        let find_opt {query; res; row_type} =
+          (match Mdb.Res.num_rows res with
+           | 0 -> Fiber.return (Ok None)
+           | 1 -> decode_next_row ~query row_type res
+           | n -> Fiber.return (reject_f ~query "Received %d tuples for find_opt." n))
+
+        let fold f {query; res; row_type} =
+          let decode = decode_next_row ~query row_type in
+          let rec loop acc =
+            decode res >>= function
+             | Ok None -> Fiber.return (Ok acc)
+             | Ok (Some y) -> loop (f y acc)
+             | Error _ as r -> Fiber.return r
+          in
+          loop
+
+        let fold_s f {query; res; row_type} =
+          let decode = decode_next_row ~query row_type in
+          let rec loop acc =
+            decode res >>= function
+             | Ok None -> Fiber.return (Ok acc)
+             | Ok (Some y) -> f y acc >>=? loop
+             | Error _ as r -> Fiber.return r
+          in
+          loop
+
+        let iter_s f {query; res; row_type} =
+          let decode = decode_next_row ~query row_type in
+          let rec loop () =
+            decode res >>= function
+             | Ok None -> Fiber.return (Ok ())
+             | Ok (Some y) -> f y >>=? loop
+             | Error _ as r -> Fiber.return r
+          in
+          loop ()
+
+        let to_stream {query; res; row_type} =
+          let decode = decode_next_row ~query row_type in
+          let rec loop () =
+            decode res >>= function
+             | Ok None -> Fiber.return Stream.Nil
+             | Error err -> Fiber.return (Stream.Error err)
+             | Ok (Some y) -> Fiber.return (Stream.Cons (y, loop))
+          in
+          loop
+      end
+
+      type prepared = {
+        query: string;
+        stmt: Mdb.Stmt.t;
+        param_length: int;
+        param_order: int list list;
+        quotes: Request_utils.linear_param list;
+      }
+
+      module Pcache =
+        Request_cache.Make (struct type t = prepared let weight _ = 1 end)
+
+      let pcache : Pcache.t = Pcache.create ~dynamic_capacity dialect
+
+      let pp_request_with_param ppf =
+        Request.make_pp_with_param ~subst ~dialect () ppf
+
+      let free_prepared prepared =
+        let rewrite_error = function
+         | Ok () -> Ok ()
+         | Error err ->
+            let query = sprintf "DEALLOCATE (%s)" prepared.query in
+            request_failed ~query err
+        in
+        Mdb.Stmt.close prepared.stmt >|= rewrite_error
+
+      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 prepare request =
+        deallocate_some () >>=? fun () ->
+        (match Pcache.find_and_promote pcache request with
+         | Some prepared ->
+            Fiber.return (Ok prepared)
+         | None ->
+            let templ = Request.query request dialect in
+            let templ = Query.expand ~final:true subst templ in
+            let query = Request_utils.linear_query_string templ in
+            Mdb.prepare db query >|= function
+             | Error err -> request_failed ~query err
+             | Ok stmt ->
+                let param_length = Request_utils.linear_param_length templ in
+                let param_order, quotes =
+                  Request_utils.linear_param_order templ in
+                let prepared =
+                  {query; stmt; param_length; param_order; quotes} in
+                Pcache.add pcache request prepared;
+                Ok prepared)
+
+      let call ~f req param = using_db @@ fun () ->
+        Log.debug ~src:Logging.request_log_src (fun f ->
+          f "Sending %a" pp_request_with_param (req, param))
+          >>= fun () ->
+
+        let process {query; stmt; param_length; param_order; quotes} =
+          let param_type = Request.param_type req in
+          let row_type = Request.row_type req in
+          let params = Array.make param_length `Null in
+          List.iter
+            (fun (Request_utils.Linear_param (j, t, v)) ->
+              params.(j) <- encode_field t v)
+            quotes;
+          (match encode_param ~uri params param_type param param_order with
+           | Error _ as r -> Fiber.return r
+           | Ok [] ->
+              Mdb.Stmt.execute stmt params >>=
+              (function
+               | Error err -> Fiber.return (request_failed ~query err)
+               | Ok res -> f Response.{query; res; row_type})
+           | Ok (_ :: _) -> assert false) in
+
+        (match Request.prepare_policy req with
+         | Direct ->
+            let templ = Request.query req dialect in
+            let templ = Query.expand ~final:true subst templ in
+            let query = Request_utils.linear_query_string templ in
+            Mdb.prepare db query >>=
+            (function
+             | Error err -> Fiber.return (request_failed ~query err)
+             | Ok stmt ->
+                let param_length = Request_utils.linear_param_length templ in
+                let param_order, quotes =
+                  Request_utils.linear_param_order templ in
+                let prepared =
+                  {query; stmt; param_length; param_order; quotes} in
+                process prepared >>= fun process_result ->
+                Mdb.Stmt.close stmt >>=
+                (function
+                 | Error (code, msg) ->
+                    Log.warn (fun p ->
+                      p "Ignoring error while closing statement: %d %s" code msg)
+                 | Ok () -> Fiber.return ()) >|= fun () ->
+                process_result)
+         | Dynamic | Static ->
+            let*? prepared = prepare req in
+            let* process_result = process prepared in
+            let+ () =
+              Mdb.Stmt.reset prepared.stmt >>=
+              (function
+               | Ok () -> Fiber.return ()
+               | Error (code, msg) ->
+                  Log.warn (fun p ->
+                    p "Removing statement from cache due to failed reset: %d %s"
+                      code msg) >|= fun () ->
+                  Pcache.remove_and_discard pcache req)
+            in
+            process_result)
+
+      let disconnect () = using_db @@ fun () ->
+        let close_stmt prepared =
+          Mdb.Stmt.close prepared.stmt >>=
+          (function
+           | Ok () -> Fiber.return ()
+           | Error (code, msg) ->
+              Log.warn (fun p ->
+                p "Ignoring failure during disconnect: %d %s" code msg))
+        in
+        fold_s_list close_stmt (Pcache.elements pcache) >>= fun () ->
+        Pcache.clear_and_discard pcache;
+        Mdb.close db
+
+      let validate () = using_db @@ fun () ->
+        Mdb.ping db >|= function Ok () -> true | Error _ -> false
+
+      let check f = f true (* FIXME *)
+
+      let transaction_failed query (errno, error) =
+        let msg = Error_msg {errno; error} in
+        Fiber.return (Error (Caqti_error.request_failed ~uri ~query msg))
+
+      let start () = using_db @@ fun () ->
+        Mdb.autocommit db false >>=
+        (function
+         | Ok () -> Fiber.return (Ok ())
+         | Error err -> transaction_failed "# SET autocommit = 0" err)
+
+      let commit () = using_db @@ fun () ->
+        Mdb.commit db >>= fun commit_result ->
+        Mdb.autocommit db true >>= fun autocommit_result ->
+        (match commit_result, autocommit_result with
+         | Ok (), Ok () -> Fiber.return (Ok ())
+         | Error err, _ -> transaction_failed "# COMMIT" err
+         | Ok (), Error err -> transaction_failed "# SET autocommit = 1" err)
+
+      let rollback () = using_db @@ fun () ->
+        Mdb.rollback db >>=
+        (function
+         | Ok () -> Fiber.return (Ok ())
+         | Error err -> transaction_failed "# ROLLBACK" err)
+
+      let set_statement_timeout t =
+        call ~f:Response.exec Q.set_statement_timeout
+          (match t with
+           | None -> 0.0
+           | Some t -> max 0.000001 t)
+    end
+
+    type conninfo = {
+      host: string option;
+      user: string option;
+      pass: string option;
+      port: int option;
+      db: string option;
+      flags: Mdb.flag list option;
+      config_group: string option;
+    }
+
+    let parse_uri uri =
+      let host = Uri.host uri in
+      let user = Uri.user uri in
+      let pass = Uri.password uri in
+      let port = Uri.port uri in
+      let config_group = Uri.get_query_param uri "config-group" in
+      (match Uri.path uri with
+       | "" | "/" -> Ok None
+       | path ->
+          if Filename.dirname path <> "/" then
+            let msg = Caqti_error.Msg "Bad URI path." in
+            Error (Caqti_error.connect_rejected ~uri msg)
+          else
+            Ok (Some (Filename.basename path))) |>? fun db ->
+      Ok {host; user; pass; port; db; flags = None; config_group}
+  end
+
+  let connect ~sw:_ ~stdenv ~subst ~config uri =
+    let module With_stdenv = Pass_stdenv (struct let stdenv = stdenv end) in
+    let open With_stdenv in
+
+    Fiber.return (parse_uri uri)
+      >>=? fun {host; user; pass; port; db; flags; config_group} ->
+
+    let config_group = match config_group with Some g -> g | None -> "caqti" in
+    let socket = Uri.get_query_param uri "socket" in
+    let options = [Mdb.Read_default_group config_group] in
+    Mdb.connect ?host ?user ?pass ?db ?port ?socket ?flags ~options () >>=
+    (function
+     | Ok db ->
+        let module B = Make_connection_base
+          (struct
+            let subst = subst dialect
+            let uri = uri
+            let db = db
+            let dynamic_capacity =
+              Caqti_connect_config.(get dynamic_prepare_capacity config)
+          end)
+        in
+        let module C = struct
+          let driver_info = driver_info
+          let driver_connection = None
+          include B
+          include Connection_utils.Make_convenience (System) (B)
+          include Connection_utils.Make_populate (System) (B)
+        end in
+        Mdb.set_character_set db "utf8mb4" >>=
+          (function
+           | Ok () -> Fiber.return ()
+           | Error (err_no, err_msg) ->
+              Log.warn (fun f ->
+                f "Could not enable full Unicode coverage for this MariaDB \
+                   connection. UTF-8 strings are likely limited to the BMP. \
+                   set_character_set says %S (%d)" err_msg err_no))
+          >>= fun () ->
+        (* MariaDB returns local times but without time zone, so change it to
+         * UTC for Caqti sessions. *)
+        C.call ~f:(fun resp -> C.Response.exec resp) Q.set_utc () >|=
+          (function
+           | Ok () -> Ok (module C : CONNECTION)
+           | Error err -> Error (`Post_connect err))
+     | Error (errno, error) ->
+        Fiber.return @@
+          Error (Caqti_error.connect_failed ~uri (Error_msg {errno; error})))
+end
+
+let () =
+  Caqti_platform_unix.Driver_loader.register "mariadb" (module Connect_functor)