Source file sql.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
(* This file is part of Dream, released under the MIT license. See LICENSE.md
   for details, or visit https://github.com/aantron/dream.

   Copyright 2021 Anton Bachin *)



module Log = Dream__server.Log
module Message = Dream_pure.Message



let log =
  Log.sub_log "dream.sql"

(* TODO Debug metadata for the pools. *)
let pool_field : (_, Caqti_error.t) Caqti_lwt_unix.Pool.t Message.field =
  Message.new_field ()

(* TODO This may not be necessary since Caqti 1.8.0. May require some messing
   around, "Enable foreign key constraint checks for SQLite3 starting at tweaks
   version 1.8." in CHANGES. *)
let foreign_keys_on =
  let open Caqti_request.Infix in
  (Caqti_type.unit ->. Caqti_type.unit) "PRAGMA foreign_keys = ON"
  [@ocaml.warning "-3"]

let post_connect (module Db : Caqti_lwt.CONNECTION) =
  match Caqti_driver_info.dialect_tag Db.driver_info with
  | `Sqlite -> Db.exec foreign_keys_on ()
  | _ -> Lwt.return (Ok ())

let sql_pool ?size uri =
    let pool_cell = ref None in
    fun inner_handler request ->

  begin match !pool_cell with
  | Some pool ->
    Message.set_field request pool_field pool;
    inner_handler request
  | None ->
    (* The correctness of this code is subtle. There is no race condition with
       two requests attempting to create a pool only because none of the code
       between checking pool_cell and setting it calls into Lwt. *)
    let parsed_uri = Uri.of_string uri in
    if Uri.scheme parsed_uri = Some "sqlite" then
      log.warning (fun log -> log ~request
        "Dream.sql_pool: \
        'sqlite' is not a valid scheme; did you mean 'sqlite3'?");
    let pool =
      let pool_config = Caqti_pool_config.create ?max_size:size () in
      Caqti_lwt_unix.connect_pool ~pool_config ~post_connect parsed_uri in
    match pool with
    | Ok pool ->
      pool_cell := Some pool;
      Message.set_field request pool_field pool;
      inner_handler request
    | Error error ->
      (* Deliberately raise an exception so that it can be communicated to any
         debug handler. *)
      let message =
        Printf.sprintf "Dream.sql_pool: cannot create pool for '%s': %s"
         uri (Caqti_error.show error) in
      log.error (fun log -> log ~request "%s" message);
      failwith message
  end

let sql request callback =
  match Message.field request pool_field with
  | None ->
    let message = "Dream.sql: no pool; did you apply Dream.sql_pool?" in
    log.error (fun log -> log ~request "%s" message);
    failwith message
  | Some pool ->
    let%lwt result =
      pool |> Caqti_lwt_unix.Pool.use (fun db ->
        (* The special exception handling is a workaround for
           https://github.com/paurkedal/ocaml-caqti/issues/68. *)
        match%lwt callback db with
        | result -> Lwt.return (Ok result)
        | exception exn -> raise exn)
    in
    Caqti_lwt.or_fail result