(* Copyright (C) 2014--2024 Petter A. Urkedal * * 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 * and , respectively. *) [@@@alert "-caqti_private"] let dynload_library = ref None let define_loader load = dynload_library := Some load let load_library lib = (match !dynload_library with | Some load -> load lib | None -> Error (Printf.sprintf "\ Neither %s nor a dynamic loader is linked into the application." lib)) let library_name_of_scheme = function | "postgres" | "postgresql" -> "caqti-driver-postgresql" | s -> "caqti-driver-" ^ s let set_tweaks_version = function | None -> Fun.id | Some x -> Caqti_connect_config.(set tweaks_version) x let compose_subst_with_env subst env dialect = let compose_subst subst1 subst2 var = (try subst1 var with Not_found -> subst2 var) in (match subst, env with | None, None -> fun _ -> raise Not_found | Some subst, None -> subst dialect | None, Some env -> env (Caqti_driver_info.of_dialect dialect) | Some subst, Some env -> let driver_info = Caqti_driver_info.of_dialect dialect in compose_subst (subst dialect) (env driver_info)) module Make (System : System_sig.S) (Pool : Pool.S with type 'a fiber := 'a System.Fiber.t and type switch := System.Switch.t and type stdenv := System.stdenv) (Loader : Driver_loader.S with type 'a fiber := 'a System.Fiber.t and type switch := System.Switch.t and type stdenv := System.stdenv and type ('a, 'e) stream := ('a, 'e) System.Stream.t) = struct open System open System.Fiber.Infix module type CONNECTION = Caqti_connection_sig.S with type 'a fiber := 'a Fiber.t and type ('a, 'err) stream := ('a, 'err) Stream.t type connection = (module CONNECTION) let (>>=?) m f = m >>= function Ok x -> f x | Error _ as r -> Fiber.return r let (>|=?) m f = m >|= function Ok x -> (Ok (f x)) | Error _ as r -> r let (let+?) = (>|=?) module type DRIVER = Driver_loader.DRIVER with type 'a fiber := 'a Fiber.t and type ('a, 'err) stream := ('a, 'err) Stream.t and type switch := System.Switch.t and type stdenv := System.stdenv let drivers : (string, (module DRIVER)) Hashtbl.t = Hashtbl.create 11 let message_cont : (_, _, _, _) format4 = if Loader.provides_unix then "Your entry point provides both the networking and unix components." else "Your entry point provides the networking but not the unix component, \ which is required by drivers based on C bindings." let message_static = Printf.sprintf ("A suitable driver for the URI-scheme %s was not found. " ^^ message_cont) let message_dynamic = Printf.sprintf ("A suitable driver for the URI-scheme %s was not found \ after linking in %s. " ^^ message_cont) let load_driver' ~uri scheme = (match Loader.find_and_apply scheme with | Some driver -> Ok driver | None -> (match !dynload_library with | None -> let msg = message_static scheme in Error (Caqti_error.load_failed ~uri (Caqti_error.Msg msg)) | Some load -> let driver_lib = library_name_of_scheme scheme in (match load driver_lib with | Ok () -> (match Loader.find_and_apply scheme with | Some driver -> Ok driver | None -> let msg = message_dynamic scheme driver_lib in Error (Caqti_error.load_failed ~uri (Caqti_error.Msg msg))) | Error msg -> Error (Caqti_error.load_failed ~uri (Caqti_error.Msg msg))))) let load_driver uri = (match Uri.scheme uri with | None -> let msg = "Missing URI scheme." in Error (Caqti_error.load_rejected ~uri (Caqti_error.Msg msg)) | Some scheme -> (try Ok (Hashtbl.find drivers scheme) with | Not_found -> (match load_driver' ~uri scheme with | Ok driver -> Hashtbl.add drivers scheme driver; Ok driver | Error _ as r -> r))) let connect ?subst ?env ?(config = Caqti_connect_config.default) ?tweaks_version ~sw ~stdenv uri : ((module CONNECTION), _) result Fiber.t = let subst = compose_subst_with_env subst env in let config = set_tweaks_version tweaks_version config in Switch.check sw; (match load_driver uri with | Ok driver -> let module Driver = (val driver) in let+? conn = Driver.connect ~sw ~stdenv ~subst ~config uri in let module Conn = (val conn : CONNECTION) in let module Conn' = struct include Conn let disconnect = let hook = Switch.on_release_cancellable sw disconnect in fun () -> Switch.remove_hook hook; disconnect () end in (module Conn' : CONNECTION) | Error err -> Fiber.return (Error err)) let with_connection ?subst ?env ?config ?tweaks_version ~stdenv uri f = Switch.run begin fun sw -> connect ~sw ~stdenv ?subst ?env ?config ?tweaks_version uri >>=? f end let connect_pool ?pool_config ?post_connect ?subst ?env ?(config = Caqti_connect_config.default) ?tweaks_version ~sw ~stdenv uri = let subst = compose_subst_with_env subst env in let pool_config = (match pool_config with | None -> Caqti_pool_config.default_from_env () | Some pool_config -> pool_config) in let config = set_tweaks_version tweaks_version config in Switch.check sw; let check_arg cond = if not cond then invalid_arg "Caqti_connect.Make.connect_pool" in (match Caqti_pool_config.(get max_size) pool_config, Caqti_pool_config.(get max_idle_size) pool_config with | None, None -> () | Some max_size, None -> check_arg (max_size >= 0) | None, Some _ -> check_arg false | Some max_size, Some max_idle_size -> check_arg (max_size >= 0); check_arg (0 <= max_idle_size && max_idle_size <= max_size)); (match load_driver uri with | Ok driver -> let module Driver = (val driver) in let connect = (match post_connect with | None -> fun () -> (Driver.connect ~sw ~stdenv ~subst ~config uri :> (connection, _) result Fiber.t) | Some post_connect -> fun () -> (Driver.connect ~sw ~stdenv ~subst ~config uri :> (connection, _) result Fiber.t) >>=? fun conn -> post_connect conn >|=? fun () -> conn) in let disconnect (module Db : CONNECTION) = Db.disconnect () in let validate (module Db : CONNECTION) = Db.validate () in let check (module Db : CONNECTION) = Db.check in let di = Driver.driver_info in let pool_config = (match Caqti_driver_info.can_concur di, Caqti_driver_info.can_pool di, Caqti_pool_config.(get max_idle_size) pool_config with | true, true, _ -> pool_config | true, false, _ -> pool_config |> Caqti_pool_config.(set max_idle_size) 0 | false, true, Some 0 -> pool_config |> Caqti_pool_config.(set max_size) 1 |> Caqti_pool_config.(set max_idle_size) 0 | false, true, _ -> pool_config |> Caqti_pool_config.(set max_size) 1 |> Caqti_pool_config.(set max_idle_size) 1 | false, false, _ -> pool_config |> Caqti_pool_config.(set max_size) 1 |> Caqti_pool_config.(set max_idle_size) 0) in let pool = Pool.create ~config:pool_config ~validate ~check ~sw ~stdenv connect disconnect in let hook = Switch.on_release_cancellable sw (fun () -> Pool.drain pool) in Gc.finalise (fun _ -> Switch.remove_hook hook) pool; Ok pool | Error err -> Error err) end