Source file service.ml

open Lwt.Syntax
open Model

let log_src = Logs.Src.create ~doc:"database" "sihl.service.database"

module Logs = (val Logs.src_log log_src : Logs.LOG)

let ctx_key_pool : pool Core.Ctx.key = Core.Ctx.create_key ()
let find_pool ctx = Core.Ctx.find ctx_key_pool ctx
let add_pool pool ctx = Core.Ctx.add ctx_key_pool pool ctx
let remove_pool ctx = Core.Ctx.remove ctx_key_pool ctx

let print_pool_usage pool =
  let n_connections = Caqti_lwt.Pool.size pool in
  let max_connections =
    Option.value (Core.Configuration.read_int "DATABASE_POOL_SIZE") ~default:10
  in
  Logs.debug (fun m -> m "DB: Pool usage: %i/%i" n_connections max_connections)
;;

let fetch_pool () =
  match !pool_ref with
  | Some pool ->
    Logs.debug (fun m -> m "DB: Skipping pool creation, re-using existing pool");
    pool
  | None ->
    let pool_size =
      Option.value (Core.Configuration.read_int "DATABASE_POOL_SIZE") ~default:10
    in
    Logs.debug (fun m -> m "DB: Create pool with size %i" pool_size);
    Option.get ("DATABASE_URL" |> Core.Configuration.read_string)
    |> Uri.of_string
    |> Caqti_lwt.connect_pool ~max_size:pool_size
    |> (function
    | Ok pool ->
      pool_ref := Some pool;
      pool
    | Error err ->
      let msg = "DB: Failed to connect to DB pool" in
      Logs.err (fun m -> m "%s %s" msg (Caqti_error.show err));
      raise (Exception ("DB: Failed to create pool " ^ msg)))
;;

let add_pool ctx =
  let pool = fetch_pool () in
  add_pool pool ctx
;;

let query ctx f =
  match find_transaction ctx, find_connection ctx, find_pool ctx with
  | Some connection, None, None ->
    let* result = f connection in
    (match result with
    | Ok result -> Lwt.return result
    | Error error ->
      let msg = Caqti_error.show error in
      Logs.err (fun m -> m "DB: %s" msg);
      Lwt.fail (Exception msg))
  | None, Some connection, _ ->
    let* result = f connection in
    (match result with
    | Ok result -> Lwt.return result
    | Error error ->
      let msg = Caqti_error.show error in
      Logs.err (fun m -> m "DB: %s" msg);
      Lwt.fail (Exception msg))
  | None, None, pool ->
    let pool = Option.value ~default:(fetch_pool ()) pool in
    print_pool_usage pool;
    let* result = Caqti_lwt.Pool.use f pool in
    (match result with
    | Ok result -> Lwt.return result
    | Error error ->
      let msg = Caqti_error.show error in
      Logs.err (fun m -> m "DB: %s" msg);
      Lwt.fail (Exception msg))
  | _ ->
    Logs.err (fun m ->
        m
          "DB: Connection, transaction or pool found in context, this should never \
           happen and might indicate connection leaks. Please report this issue.");
    Lwt.fail (Exception "Connection and pool found")
;;

let atomic ctx f =
  match find_transaction ctx, find_connection ctx, find_pool ctx with
  | Some connection, None, None ->
    (* Make sure [f] can not use the pool or some other connection *)
    ctx |> remove_pool |> remove_connection |> add_transaction connection |> f
  | None, Some connection, None ->
    let (module Connection : Caqti_lwt.CONNECTION) = connection in
    let* start_result = Connection.start () in
    (match start_result with
    | Error msg ->
      Logs.debug (fun m ->
          m "DB TX: Failed to start transaction %s" (Caqti_error.show msg));
      Lwt.fail @@ Exception (Caqti_error.show msg)
    | Ok () ->
      Logs.debug (fun m -> m "DB TX: Started transaction");
      (* Remove the pool so that all subsequent queries are executed on the connection. A
         transaction can only be done only at one connection, it can not span multiple
         connections. *)
      let ctx_with_connection =
        ctx |> remove_pool |> remove_connection |> add_transaction (module Connection)
      in
      Lwt.catch
        (fun () ->
          let* result = f ctx_with_connection in
          let* commit_result = Connection.commit () in
          match commit_result with
          | Ok () ->
            Logs.debug (fun m -> m "DB TX: Successfully committed transaction");
            Lwt.return @@ result
          | Error error ->
            Logs.err (fun m ->
                m "DB TX: Failed to commit transaction %s" (Caqti_error.show error));
            Lwt.fail @@ Exception "Failed to commit transaction")
        (fun e ->
          let* rollback_result = Connection.rollback () in
          match rollback_result with
          | Ok () ->
            Logs.debug (fun m -> m "DB TX: Successfully rolled back transaction");
            Lwt.fail e
          | Error error ->
            Logs.err (fun m ->
                m "DB TX: Failed to rollback transaction %s" (Caqti_error.show error));
            Lwt.fail @@ Exception "Failed to rollback transaction"))
  | None, None, pool ->
    (* There is no transaction active, create a new one *)
    let pool = Option.value ~default:(fetch_pool ()) pool in
    print_pool_usage pool;
    let* pool_result =
      Caqti_lwt.Pool.use
        (fun connection ->
          Logs.debug (fun m -> m "DB TX: Fetched connection from pool");
          let (module Connection : Caqti_lwt.CONNECTION) = connection in
          let* start_result = Connection.start () in
          match start_result with
          | Error msg ->
            Logs.debug (fun m ->
                m "DB TX: Failed to start transaction %s" (Caqti_error.show msg));
            Lwt.return @@ Error msg
          | Ok () ->
            Logs.debug (fun m -> m "DB TX: Started transaction");
            (* Remove the pool so that all subsequent queries are executed on the
               connection. A transaction can only be done only at one connection, it can
               not span multiple connections. *)
            let ctx_with_connection =
              ctx
              |> remove_pool
              |> remove_connection
              |> add_transaction (module Connection)
            in
            Lwt.catch
              (fun () ->
                let* result = f ctx_with_connection in
                let* commit_result = Connection.commit () in
                match commit_result with
                | Ok () ->
                  Logs.debug (fun m -> m "DB TX: Successfully committed transaction");
                  Lwt.return @@ Ok result
                | Error error ->
                  Logs.err (fun m ->
                      m "DB TX: Failed to commit transaction %s" (Caqti_error.show error));
                  Lwt.fail @@ Exception "Failed to commit transaction")
              (fun e ->
                let* rollback_result = Connection.rollback () in
                match rollback_result with
                | Ok () ->
                  Logs.debug (fun m -> m "DB TX: Successfully rolled back transaction");
                  Lwt.fail e
                | Error error ->
                  Logs.err (fun m ->
                      m
                        "DB TX: Failed to rollback transaction %s"
                        (Caqti_error.show error));
                  Lwt.fail @@ Exception "Failed to rollback transaction"))
        pool
    in
    (match pool_result with
    | Ok result ->
      (* All good, return result of f ctx *)
      Lwt.return result
    | Error pool_err ->
      (* Failed to start, commit or rollback transaction *)
      Lwt.fail (Exception (pool_err |> Caqti_error.show)))
  | _ ->
    Logs.err (fun m ->
        m
          "DB: Connection, transaction or pool found in context together, this should \
           never happen and might indicate connection leaks. Please report this issue.");
    Lwt.fail (Exception "Connection and pool found")
;;

let set_fk_check_request =
  Caqti_request.exec Caqti_type.bool "SET FOREIGN_KEY_CHECKS = ?;"
;;

let set_fk_check ctx ~check =
  query ctx (fun (module Connection : Caqti_lwt.CONNECTION) ->
      Connection.exec set_fk_check_request check)
;;

let with_disabled_fk_check ctx f =
  atomic ctx (fun ctx ->
      let* () = set_fk_check ctx ~check:false in
      Lwt.finalize (fun () -> f ctx) (fun () -> set_fk_check ctx ~check:true))
;;

(* Service lifecycle *)

let start ctx = ctx |> add_pool |> Lwt.return
let stop _ = Lwt.return ()
let lifecycle = Core.Container.Lifecycle.create "db" ~start ~stop

let configure configuration =
  let configuration = Core.Configuration.make configuration in
  Core.Container.Service.create ~configuration lifecycle
;;