Source file caqti_connection.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
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
(* Copyright (C) 2019--2020  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 OCaml static compilation 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
 * along with this library.  If not, see <http://www.gnu.org/licenses/>.
 *)

open Caqti_common_priv
open Printf

module Make_helpers
  (System : Caqti_driver_sig.System_common) =
struct
  open System

  let assert_single_use ~what in_use f =
    if !in_use then
      failwith ("Invalid concurrent usage of " ^ what ^ " detected.");
    assert (not !in_use);
    in_use := true;
    f () >|= fun y ->
    assert !in_use;
    in_use := false;
    y
end

module Make_convenience
  (System : Caqti_driver_sig.System_common)
  (C : Caqti_connection_sig.Base
        with type 'a future := 'a System.future
         and type ('a, 'err) stream := ('a, 'err) System.Stream.t) =
struct
  open System
  module Response = C.Response

  let exec q p = C.call ~f:Response.exec q p
  let find q p = C.call ~f:Response.find q p
  let find_opt q p = C.call ~f:Response.find_opt q p
  let fold q f p acc = C.call ~f:(fun resp -> Response.fold f resp acc) q p
  let fold_s q f p acc = C.call ~f:(fun resp -> Response.fold_s f resp acc) q p
  let iter_s q f p = C.call ~f:(fun resp -> Response.iter_s f resp) q p
  let collect_list q p =
    let f resp = Response.fold List.cons resp [] >|= Result.map List.rev in
    C.call ~f q p
  let rev_collect_list q p =
    let f resp = Response.fold List.cons resp [] in
    C.call ~f q p

  let exec_with_affected_count q p =
    let f response =
      Response.exec response >>= fun execResult ->
      match execResult with
      | Ok () -> Response.affected_count response
      | Error x -> return (Error x) in
    C.call ~f q p
end

module Make_populate
  (System : Caqti_driver_sig.System_common)
  (C : Caqti_connection_sig.Base
        with type 'a future := 'a System.future
         and type ('a, 'e) stream := ('a, 'e) System.Stream.t) =
struct
  open System
  let (>>=?) m f = m >>= function Ok x -> f x | Error _ as r -> return r

  let populate ~table ~columns row_type data =

    let request =
      let columns_tuple = String.concat ", " columns in
      let q =
        let open Caqti_query in
        S[L(sprintf "INSERT INTO %s (%s) VALUES (" table columns_tuple);
          concat ", " (List.mapi (fun i _ -> P i) columns); L")"]
      in
      Caqti_request.create row_type Caqti_type.unit Caqti_mult.zero (fun _ -> q)
    in

    C.start () >>=? fun () ->
    Stream.iter_s ~f:(C.call ~f:C.Response.exec request) data >>= fun res ->
    C.deallocate request >>= fun _ ->
    (match res with
     | Ok () ->
        C.commit ()
     | Error (`Congested err) ->
        C.rollback () >>=? fun () ->
        return (Error (`Congested err))
     | Error err ->
        return (Error err))
end