(**************************************************************************)
(*  Copyright 2014, 2015, Sebastien Mondet <seb@mondet.org>               *) (*                                                                        *) (*  Licensed under the Apache License, Version 2.0 (the "License");       *) (*  you may not use this file except in compliance with the License.      *) (*  You may obtain a copy of the License at                               *) (*                                                                        *) (*      http://www.apache.org/licenses/LICENSE-2.0                        *) (*                                                                        *) (*  Unless required by applicable law or agreed to in writing, software   *) (*  distributed under the License is distributed on an "AS IS" BASIS,     *) (*  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or       *) (*  implied.  See the License for the specific language governing         *) (*  permissions and limitations under the License.                        *)
(**************************************************************************)
open Nonstd open Pvem_lwt_unix open Pvem_lwt_unix.Deferred_result module String = StringLabels let debug = ref false let dbg fmt = ksprintf (eprintf "Trakeva_sqlite: %s\n%!") fmt      type t = {   handle: Sqlite3.db;   action_mutex: Lwt_mutex.t; } let in_posix_thread ~on_exn f =   Lwt_preemptive.detach (fun () ->   (* (fun f () -> Lwt.return (f ()))  (fun () -> *)     try `Ok (f ())     with e -> on_exn e) () (* let escape_blob s =    let b = Buffer.create (String.length s + 4) in   Buffer.add_char b '\'';   String.iter s ~f:(function     | '\'' -> Buffer.add_string b "''"     | other -> Buffer.add_char b other     );   Buffer.add_char b '\'';   Buffer.contents b *)     let escape_blob s =    let b = Buffer.create (String.length s * 2 + 4) in   Buffer.add_string b "X'";   String.iter s ~f:(fun c ->       Buffer.add_string b (sprintf "%02X" (int_of_char c));     );   Buffer.add_char b '\'';   Buffer.contents b      let default_table = "trakeva_default_table" (* https://www.sqlite.org/lang_createtable.html *) let create_table t =   sprintf "CREATE TABLE IF NOT EXISTS %s (collection BLOB, key BLOB, value BLOB, PRIMARY KEY(collection, key) ON CONFLICT REPLACE)" t let option_equals =   function   | None -> "IS NULL"   | Some s -> sprintf "= %s" (escape_blob s) let get_statement table collection key =   sprintf "SELECT value FROM %s WHERE key = %s AND collection %s" table      (escape_blob key)      (option_equals collection) let get_all_statement table collection =   sprintf "SELECT key FROM %s WHERE collection %s ORDER BY key" table     (option_equals collection) let keys_statement table collection =   sprintf "SELECT DISTINCT key FROM %s WHERE collection %s"     table (option_equals collection) (* https://www.sqlite.org/lang_insert.html *) let set_statement table collection key value =   sprintf "INSERT OR REPLACE INTO %S (collection, key, value) VALUES (%s, %s, %s)"     table     (match collection with None -> "NULL" | Some s -> escape_blob s)     (escape_blob key) (escape_blob value) (* https://www.sqlite.org/lang_delete.html *) let unset_statement table collection key =   sprintf "DELETE FROM %s WHERE key = %s AND collection %s" table     (escape_blob key)     (option_equals collection) let with_executed_statement handle statement f =   let prep = Sqlite3.prepare handle statement in   if !debug then     dbg "exec: %S\n    → counts: %d, %d"       statement (Sqlite3.column_count prep) (Sqlite3.data_count prep);   (try     let x = f prep in     let _ = Sqlite3.finalize prep in     x   with e ->     let _ = Sqlite3.finalize prep in     raise e) let is_ok_or_done_exn handle (rc: Sqlite3.Rc.t) =   let open Sqlite3.Rc in   match rc with   | OK  -> ()   | DONE -> ()   | _ -> failwith (sprintf "not ok/done: %s (global error: %s)"                      (Sqlite3.Rc.to_string rc)                      (Sqlite3.errmsg handle)) let get_row_exn prep =   match Sqlite3.step prep with   | Sqlite3.Rc.ROW -> Sqlite3.column prep 0, true   | Sqlite3.Rc.DONE -> Sqlite3.column prep 0, false   | rc -> failwith (sprintf "not a row: %s" (Sqlite3.Rc.to_string rc)) let exec_unit_exn handle statement =   with_executed_statement handle statement (fun prep ->       Sqlite3.step prep |> is_ok_or_done_exn handle) let string_option_data_exn data =   let open Sqlite3.Data in   begin match data with   | NONE  -> failwith "string_option_data_exn: none"   | NULL  -> None   | INT _ -> failwith "string_option_data_exn: int"   | FLOAT _ -> failwith "string_option_data_exn: float"   | TEXT s   | BLOB s -> Some s   end    let exec_option_exn handle statement =   with_executed_statement handle statement (fun prep ->       let first, (_ : bool) = get_row_exn prep in       string_option_data_exn first) let exec_list_exn handle statement =   with_executed_statement handle statement begin fun prep ->     let ret = ref [] in     let rec loop () =       let row, more_to_come  = get_row_exn prep in       match string_option_data_exn row with       | Some one -> ret:= one :: !ret;         if more_to_come then loop () else ()       | None -> ()     in     loop ();     !ret   end let load path =   let on_exn e = `Error (`Database (`Load path, Printexc.to_string e)) in   begin     try if Sys.getenv "TRAKEVA_SQLITE_DEBUG" = "true" then debug := true     with _ -> ()   end;   let action_mutex = Lwt_mutex.create () in   in_posix_thread ~on_exn (fun () ->       if !debug then         dbg "openning: %S" path;       (* we need the mutex `FULL: https://www.sqlite.org/threadsafe.html          the private cache is up for debate:           https://www.sqlite.org/sharedcache.html *)       let handle = Sqlite3.db_open ~mutex:`FULL ~cache:`PRIVATE path in       exec_unit_exn handle (create_table default_table);       {handle; action_mutex}     ) let close {handle} =   let on_exn e = `Error (`Database (`ClosePrintexc.to_string e)) in   in_posix_thread ~on_exn begin fun () ->     let rec loop = function     | 0 -> failwith "failed to close (busy many times)"     | n ->       if Sqlite3.db_close handle then () else (         Sqlite3.sleep 100 |> ignore;         if !debug then           dbg "closing, %d attempts left" (n - 1);         loop (n - 1)       )     in     loop 5   end open Trakeva let get ?collection t ~key =   let statement = get_statement default_table collection key in   let error_loc = `Get (Key_in_collection.create key ?collection) in   let on_exn e = `Error (`Database (error_loc , Printexc.to_string e)) in   in_posix_thread ~on_exn (fun () ->       exec_option_exn t.handle statement     ) let get_all t ~collection =   let statement = get_all_statement default_table (Some collection) in   let error_loc = `Get_all collection in   let on_exn e = `Error (`Database (error_loc , Printexc.to_string e)) in   in_posix_thread ~on_exn (fun () ->       exec_list_exn t.handle statement     ) let iterator t ~collection =   let keys_statement = keys_statement default_table (Some collection) in   let error_loc = `Iter collection in   let on_exn e = `Error (`Database (error_loc , Printexc.to_string e)) in   let state = ref None in   let rec next_exn prep =     if !debug then       dbg "exec: %S\n    → counts: %d, %d"         keys_statement (Sqlite3.column_count prep) (Sqlite3.data_count prep);     try       let row, more_to_come  = get_row_exn prep in       begin match string_option_data_exn row with       | Some one_key -> Some one_key       | None ->         let _ = Sqlite3.finalize prep in         None       end     with e ->       let _ = Sqlite3.finalize prep in       raise e   in   begin fun () ->     in_posix_thread ~on_exn (fun () ->         match !state with         | None ->           let prep = Sqlite3.prepare t.handle keys_statement in           begin match Sqlite3.prepare_tail prep with           | None  -> state := Some prep;           | Some p -> state := Some p;           end;           next_exn prep         | Some prep ->           next_exn prep       )   end let act t ~(action: Action.t) =   let rec transact (action: Action.t) =     let open Key_in_collection in     let open Action in     match action with     | Set ({ key; collection }, value) ->       let statement = set_statement default_table collection key value in       exec_unit_exn t.handle statement;       true     | Unset { key; collection } ->       let statement = unset_statement default_table collection key in       exec_unit_exn t.handle statement;       true     | Sequence l -> List.for_all l ~f:transact     | Check ({ key; collection }, opt) ->       let statement = get_statement default_table collection key in       exec_option_exn t.handle statement = opt   in   let error_loc = `Act action in   let on_exn e = `Error (`Database (error_loc , Printexc.to_string e)) in   Lwt_mutex.with_lock t.action_mutex (fun () ->        in_posix_thread ~on_exn begin fun () ->         exec_unit_exn t.handle "BEGIN TRANSACTION";         begin match transact action with         | false ->           exec_unit_exn t.handle "ROLLBACK";           `Not_done         | true ->           exec_unit_exn t.handle "COMMIT";           `Done         end       end     )