sqlite3_utils.ml

(* This file is free software, part of bender-ocaml. See file "license" for more details. *)

(** {1 Utils for SQLite} *)

module Data = Sqlite3.Data
module Rc = Sqlite3.Rc
type db = Sqlite3.db

exception RcError of Sqlite3.Rc.t
exception Type_error of Data.t

let () = Printexc.register_printer
    (function
      | RcError rc -> Some ("sqlite error: " ^ Sqlite3.Rc.to_string rc)
      | _ -> None)

type t = db

let check_ret_exn = function
  | Sqlite3.Rc.DONE
  | Sqlite3.Rc.OK -> ()
  | rc -> raise (RcError rc)

let check_ret x = function
  | Sqlite3.Rc.DONE | Sqlite3.Rc.OK -> Ok x
  | rc -> Error rc

let err_string = function
  | Ok x -> Ok x
  | Error rc -> Error (Rc.to_string rc)

(* on "busy", wait 300ms before failing *)
let setup_timeout ?(ms=300) db : unit =
  Sqlite3.busy_timeout db ms

let bind_ stmt i d : unit = check_ret_exn (Sqlite3.bind stmt i d)

(** Parameters passed to a statement *)
module Ty = struct
  type _ arg =
    | Int : int arg
    | Int64 : int64 arg
    | Float : float arg
    | String : [`Blob|`Text|`Both] -> string arg
    | Data : Data.t arg
    | Nullable : 'a arg -> 'a option arg

  type (_,_) t =
    | [] : ('res, 'res) t
    | (::) : 'a arg * ('b, 'res) t -> ('a -> 'b, 'res) t
    (* TODO: add this (with distinct count for args and row)
       | Raw : (Data.t array, 'res) t
    *)

  let nil = []
  let int = Int
  let int64 = Int64
  let float = Float
  let text = String `Text
  let blob = String `Blob
  let any_str = String `Both
  let data = Data

  let nullable
    : type a. a arg -> a option arg
    = function
      | Nullable _ -> invalid_arg "Sqlite3_utils.Ty.nullable can't be nested"
      | x -> Nullable x

  let (@>) x y = (::) (x,y)
  let p1 x = [x]
  let p2 x y = [x;y]
  let p3 x y z = [x;y;z]
  let p4 x y z w = [x;y;z;w]
  let p5 x y z w1 w2 = [x;y;z;w1;w2]
  let p6 x y z w1 w2 w3 = [x;y;z;w1;w2;w3]

  let rec (@>>)
    : type a b res. (a, b) t -> (b, res) t -> (a, res) t =
    fun ty1 ty2 ->
    match ty1 with
      | [] -> ty2
      | (::) (x, ty1') -> (::) (x, ty1' @>> ty2)

  let id x = x
  let mkp2 x y = x,y
  let mkp3 x y z = x,y,z
  let mkp4 x y z w = x,y,z,w
  let mkp5 x y z w1 w2 = x,y,z,w1,w2
  let mkp6 x y z w1 w2 w3 = x,y,z,w1,w2,w3

  let rec count : type a r. (a, r) t -> int
    = function
    | [] -> 0
    | (::) (_, tl) -> 1 + count tl

  (* translate parameters *)
  let rec tr_args
    : type a res. Sqlite3.stmt -> int -> (a,res) t -> (unit->res) -> a
    = fun stmt i p cb ->
      match p with
      | [] -> cb()
      | Int :: k ->
        (fun x -> bind_ stmt i (Data.INT (Int64.of_int x)); tr_args stmt (i+1) k cb)
      | Int64 :: k ->
        (fun x -> bind_ stmt i (Data.INT x); tr_args stmt (i+1) k cb)
      | String `Text :: k ->
        (fun x -> bind_ stmt i (Data.TEXT x); tr_args stmt (i+1) k cb)
      | String (`Blob|`Both) :: k ->
        (fun x -> bind_ stmt i (Data.BLOB x); tr_args stmt (i+1) k cb)
      | Float :: k ->
        (fun x -> bind_ stmt i (Data.FLOAT x); tr_args stmt (i+1) k cb)
      | Data :: k ->
        (fun x -> bind_ stmt i x; tr_args stmt (i+1) k cb)
      | Nullable p1 :: k ->
        (function
          | None -> bind_ stmt i Data.NULL; tr_args stmt (i+1) k cb
          | Some x -> tr_args stmt i (p1 :: k) cb x
        )

  (* translate results *)
  let rec tr_row
    : type a res. (int->Data.t) -> int -> (a,res) t -> a -> res
    = fun get i ty f -> match ty with
      | [] -> f
      | (::) (ty, k) ->
        let data = get i in
        tr_row_data get data i ty k f
  and tr_row_data
    : type a b res. (int->Data.t) -> Data.t -> int -> a arg -> (b,res) t -> (a->b) -> res
    = fun get data i ty k f ->
      match ty with
      | Data -> tr_row get (i+1) k (f data)
      | Int64 ->
        (match data with
         | Data.INT x -> tr_row get (i+1) k (f x)
         | d -> raise (Type_error d))
      | Int ->
        (match data with
         | Data.INT x as d ->
           let x = try Int64.to_int x with _ -> raise (Type_error d) in
           tr_row get (i+1) k (f x)
         | d -> raise (Type_error d))
      | Float ->
        (match data with
         | Data.FLOAT x -> tr_row get (i+1) k (f x)
         | d -> raise (Type_error d))
      | String kind ->
        (match data, kind with
         | Data.BLOB x, `Blob -> tr_row get (i+1) k (f x)
         | Data.TEXT x, `Text -> tr_row get (i+1) k (f x)
         | (Data.BLOB x | Data.TEXT x), `Both -> tr_row get (i+1) k (f x)
         | d, _ -> raise (Type_error d))
      | Nullable ty' ->
        (match data with
         | Data.NONE | Data.NULL -> tr_row get (i+1) k (f None)
         | _ -> tr_row_data get data i ty' k (fun x -> f (Some x)))
end

module Cursor = struct
  type 'a t = {
    db: db;
    stmt: Sqlite3.stmt;
    read: Sqlite3.stmt -> 'a;
    mutable cur: 'a option;
    mutable last_row_id: int64;
  }

  let ignore _ = ()

  let next_ self : unit =
    match Sqlite3.step self.stmt with
    | Sqlite3.Rc.DONE ->
      self.cur <- None;
    | Sqlite3.Rc.ROW ->
      let x = self.read self.stmt in
      self.last_row_id <- Sqlite3.last_insert_rowid self.db;
      self.cur <- Some x
    | rc -> raise (RcError rc)

  let make_ db stmt read =
    let self = { db; stmt; cur=None; read; last_row_id=0L; } in
    next_ self;
    self

  let opt_map_ f = function None -> None | Some x -> Some (f x)

  let get_one c = match c.cur with
    | None -> Error Rc.NOTFOUND
    | Some x -> Ok x

  let get_one_exn c = match get_one c with
    | Ok x -> x
    | Error rc -> raise (RcError rc)

  let map ~f c : _ t =
    {c with
     read=(fun stmt -> f (c.read stmt));
     cur=opt_map_ f c.cur;
    }

  let make db stmt ty f =
    let read stmt = Ty.tr_row (Sqlite3.column stmt) 0 ty f in
    make_ db stmt read

  let make_raw db stmt : Data.t array t =
    make_ db stmt Sqlite3.row_data

  (* next value in the cursor *)
  let next self : _ option =
    match self.cur with
    | None -> None
    | Some _ as x ->
      next_ self;
      x

  (* next value in the cursor *)
  let next_with_last_row_id self : _ option =
    match self.cur with
    | None -> None
    | Some x ->
      let last = self.last_row_id in
      next_ self;
      Some (x, last)

  let rec iter ~f self = match self.cur with
    | None -> ()
    | Some res ->
      f res;
      next_ self;
      iter ~f self

  let rec iter_with_last_row_id ~f self = match self.cur with
    | None -> ()
    | Some res ->
      f self.last_row_id res;
      next_ self;
      iter_with_last_row_id ~f self

  let to_seq self =
    let rec get_next () =
      let n = lazy (
        match self.cur with
        | None -> Seq.Nil
        | Some x ->
          next_ self;
          let tl = get_next () in
          Seq.Cons (x, tl)
      ) in
      fun () -> Lazy.force n
    in
    get_next ()

  (* convert a cursor into a list of answers *)
  let to_list_rev (c:'a t) : 'a list =
    let rec aux acc c = match next c with
      | None -> acc
      | Some d -> aux (d::acc) c
    in
    aux [] c

  let to_list c = List.rev (to_list_rev c)
end

let finally_ ~hok ~herr x f =
  try
    let res = f x in
    hok x;
    res
  with e ->
    herr x;
    raise e

let finalize_check_ stmt = check_ret_exn @@ Sqlite3.finalize stmt
let finalize_nocheck_ stmt =
  try ignore (Sqlite3.finalize stmt : Rc.t) with Sqlite3.Error _ -> ()

let db_close_rec_ db =
  while not (Sqlite3.db_close db) do () done

let with_db ?(wal=false) ?mode ?mutex ?cache ?vfs ?timeout str f =
  let db = Sqlite3.db_open ?mode ?mutex ?cache ?vfs str in
  if wal then check_ret_exn @@ Sqlite3.exec db "pragma journal_mode=WAL;";
  (match timeout with Some ms -> setup_timeout db ~ms | None -> ());
  finally_ ~hok:db_close_rec_ ~herr:db_close_rec_ db f

let with_stmt db str ~f =
  let stmt = Sqlite3.prepare db str in
  finally_ ~hok:finalize_check_ ~herr:finalize_nocheck_ stmt f

let check_arity_params_ stmt n : unit =
  if Sqlite3.bind_parameter_count stmt <> n then (
    invalid_arg
      (Format.sprintf "wrong number of parameters: expected %d, got %d"
         (Sqlite3.bind_parameter_count stmt) n);
  )

let check_arity_res_ stmt n : unit =
  if Sqlite3.column_count stmt <> n then (
    invalid_arg
      (Format.sprintf "wrong number of columns in result: expected %d, got %d"
         (Sqlite3.column_count stmt) n);
  )

let exec_get_column_names self str : _ list =
  let stmt = Sqlite3.prepare self str in
  let n = Sqlite3.column_count stmt in
  List.init n (Sqlite3.column_name stmt)

(*$R
  let db = Sqlite3.db_open ":memory:" in
  exec0_exn db "create table foo (a text, b text, c int);";
  let l = exec_get_column_names db "select * from foo;" in
  assert_equal ~printer:(String.concat ";") ["a"; "b"; "c"] l;

  let l2 = exec_get_column_names db "select b, a, c as c2 from foo;" in
  assert_equal ~printer:(String.concat ";") ["b"; "a"; "c2"] l2;
*)

let exec0_exn db str : unit = check_ret_exn @@ Sqlite3.exec db str
let exec0 db str : _ result = check_ret () @@ Sqlite3.exec db str

(* execute statement, return cursor *)
let exec_raw_exn db str ~f =
  with_stmt db str
    ~f:(fun stmt ->
        check_arity_params_ stmt 0;
        f (Cursor.make_raw db stmt))

let exec_raw db str ~f =
  try Ok (exec_raw_exn db str ~f)
  with RcError c -> Error c

(* execute statement parametrized by the array of arguments *)
let exec_raw_args_exn db str a ~f =
  with_stmt db str
    ~f:(fun stmt ->
        check_arity_params_ stmt (Array.length a);
        Array.iteri (fun i x -> check_ret_exn (Sqlite3.bind stmt (i+1) x)) a;
        f (Cursor.make_raw db stmt))

let exec_raw_args db str a ~f =
  try Ok (exec_raw_args_exn db str a ~f)
  with RcError c -> Error c

(* execute statement parametrized by the array of arguments *)
let exec_ db str ~ty ~f =
  let params, ty_r, f_r = ty in
  let stmt = Sqlite3.prepare db str in
  (* initial checks. Returns a function that will, or not, call
     a continuation [k] after all parameters are bound. *)
  let f0 =
    try
      check_arity_params_ stmt (Ty.count params);
      fun ~k -> k (Ok ())
    with
    | RcError e ->
      fun ~k ->
        finalize_nocheck_ stmt;
        k (Error e)
    | e ->
      fun ~k:_ ->
        finalize_nocheck_ stmt;
        raise e
  in
  (* caution, bind starts at [1] *)
  Ty.tr_args stmt 1 params
    (fun () ->
       f0 ~k:(function
           | Error e -> f (Error e)
           | Ok () ->
             try
               check_arity_res_ stmt (Ty.count ty_r);
               let c = Cursor.make db stmt ty_r f_r in
               let r = f (Ok c) in
               finalize_check_ stmt;
               r
             with
             | RcError e ->
               finalize_nocheck_ stmt;
               f (Error e)
             | e ->
               finalize_nocheck_ stmt;
               raise e))

let exec db str ~ty ~f =
  exec_ db str ~ty
    ~f:(function
        | Ok x ->
          (try Ok (f x) with RcError e -> Error e)
        | Error e -> Error e)

let exec_exn db str ~ty ~f =
  exec_ db str ~ty
    ~f:(function
        | Ok x -> f x
        | Error e -> raise (RcError e))

let exec_no_params_exn db str ~ty ~f =
  with_stmt db str
    ~f:(fun stmt ->
        check_arity_params_ stmt 0;
        let ty_r, f_r = ty in
        f (Cursor.make db stmt ty_r f_r))

let exec_no_params db str ~ty ~f =
  exec_no_params_exn db str ~ty
    ~f:(fun c -> try (Ok (f c)) with RcError rc -> Error rc)

let exec_no_cursor_ db str ~ty ~f =
  let stmt = Sqlite3.prepare db str in
  let f0 =
    try
      check_arity_params_ stmt (Ty.count ty);
      fun ~k -> k (Ok ())
    with
    | RcError e ->
      fun ~k ->
        finalize_nocheck_ stmt;
        k (Error e)
    | e ->
      fun ~k:_ ->
        finalize_nocheck_ stmt;
        raise e
  in
  (* caution, bind starts at [1] *)
  Ty.tr_args stmt 1 ty
    (fun () ->
       f0 ~k:(function
           | Error e -> f (Error e)
           | Ok () ->
             try
               (* just execute one step *)
               check_ret_exn @@ Sqlite3.step stmt;
               finalize_check_ stmt;
               f (Ok ())
             with
             | RcError e ->
               finalize_nocheck_ stmt;
               f (Error e)
             | e ->
               finalize_nocheck_ stmt;
               raise e))

let exec_no_cursor_exn db str ~ty =
  exec_no_cursor_ db str ~ty
    ~f:(function
        | Ok () -> ()
        | Error e -> raise (RcError e))

let exec_no_cursor db str ~ty =
  exec_no_cursor_ db str ~ty ~f:(fun x -> x)

(* From [ocaml-sqlite3EZ](https://github.com/mlin/ocaml-sqlite3EZ),
   with some changes. Compatible license (MIT) *)
let transact db f =
  exec0_exn db "BEGIN;";
  try
    let y = f db in
    exec0_exn db "COMMIT;";
    y
  with
  | e ->
    exec0_exn db "ROLLBACK;";
    raise e

(* From [ocaml-sqlite3EZ](https://github.com/mlin/ocaml-sqlite3EZ),
   with some changes. Compatible license (MIT) *)
let atomically db f =
  exec0_exn db "SAVEPOINT a;";
  try
    let y = f db in
    exec0_exn db "RELEASE a;";
    y
  with
  | exn ->
    exec0_exn db "RELEASE a;";
    exec0_exn db "ROLLBACK TO a;";
    raise exn

(*$inject
  let with_test_db f : unit =
    with_db ":memory:" (fun db ->
        exec0_exn db
          "create table person(name text, job text, age int); ";
        exec0_exn db
          "insert into person values
            ('john', 'PHB', 42), ('alice', 'hacker', 20),
            ('bob', 'caller', 20), ('eve', 'spy', 99);";
        exec0_exn db
          "create table friendorfoe(p1 text, p2 text, level int); ";
        exec0_exn db
          "insert into friendorfoe values
            ('john', 'bob', 1), ('bob', 'alice', 1),
            ('bob', 'eve', -2), ('alice', 'eve', -5) ; ";
        f db)
*)

(* update test *)
(*$R
  with_test_db (fun db ->
        exec0_exn db
          "insert or ignore into friendorfoe (p1, p2, level)
            select p2, p1, level from friendorfoe ;";
        let l = exec_exn db ~ty:Ty.(p1 text, p2 text int, mkp2)
          "select p2, level from friendorfoe where p1 = ? order by p2;"
          "eve" ~f:Cursor.to_list
        in
        assert_equal ~printer:Q.Print.(list @@ pair string int) [
          "alice", -5;
          "bob", -2;
        ] l;)
*)


(* basic cursor test *)
(*$R
  with_test_db (fun db ->
        exec_no_params_exn db
          "select name, job from person where age=20 order by name;"
          ~ty:Ty.([text; text], mkp2)
          ~f:(fun c ->
              match Cursor.next c with
              | None -> assert_failure "cursor too short"
              | Some (n1,j1) ->
                assert_equal "alice" n1;
                assert_equal "hacker" j1;
                match Cursor.next c with
                | None -> assert_failure "cursor too short"
                | Some (n2,j2) ->
                  assert_equal "bob" n2;
                  assert_equal "caller" j2;
                  match Cursor.next c with
                  | None -> ()
                  | Some _ -> assert_failure "cursor too long"
            ))
*)

(* recursive query test *)
(*$R
  let q = "with recursive fib(a,b,c) as
    ( values (1,1,1),(2,1,2) UNION select a+1, c, b+c from fib where a<100)
    select a, c from fib where a<= ?;"
  in
  let l_expect =  [
    1,1;
    2,2;
    3,3;
    4,5;
    5,8;
    6,13;
    7,21;
    8,34;
    9,55;
    10,89;
    11,144;
    12,233;
    13,377;
    14,610;
    15,987;
  ] in
  with_db ":memory:" (fun db ->
      let l =
        exec_exn db q ~ty:Ty.([int], [int;int], (fun x y->x,y))
          15 ~f:Cursor.to_list
      in
      assert_equal ~printer:Q.Print.(list @@ pair int int) l_expect l
    )
*)

(* regression for #1 *)
(*$R
  let db = Sqlite3.db_open ":memory:" in
  exec0_exn db "CREATE TABLE ids(id INTEGER PRIMARY KEY)";
  let insert db =
    (exec_no_cursor
      db
       "INSERT INTO ids (id) VALUES(1)"
       ~ty:(Ty.nil)
     : (unit, Sqlite3.Rc.t) result)
  in
  assert_bool "isok" (match insert db with Ok _ -> true | Error _ -> false);
  assert_bool "must not raise" (try ignore (insert db:_ result); true with _ -> false);
  ()

*)