extensions/ocsipersist.ml

   1 (* Ocsigen
   2  * http://www.ocsigen.org
   3  * Module ocsipersist.ml
   4  * Copyright (C) 2007 Vincent Balat - Gabriel Kerneis
   5  * Laboratoire PPS - CNRS Université Paris Diderot
   6  *
   7  * This program is free software; you can redistribute it and/or modify
   8  * it under the terms of the GNU Lesser General Public License as published by
   9  * the Free Software Foundation, with linking exception;
  10  * either version 2.1 of the License, or (at your option) any later version.
  11  *
  12  * This program is distributed in the hope that it will be useful,
  13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  15  * GNU Lesser General Public License for more details.
  16  *
  17  * You should have received a copy of the GNU Lesser General Public License
  18  * along with this program; if not, write to the Free Software
  19  * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  20  *)
  21
  22
  23 (** Module Ocsipersist: persistent data *)
  24
  25 open Lwt
  26 open Sqlite3
  27 open Printf
  28
  29 (** Data are divided into stores.
  30    Create one store for your project, where you will save all your data.
  31  *)
  32 type store = string Lwt.t
  33
  34 exception Ocsipersist_error
  35
  36
  37 (*****************************************************************************)
  38
  39 open Simplexmlparser
  40 (** getting the directory from config file *)
  41 let rec parse_global_config = function
  42   | [] -> None
  43   | (Element ("database", [("file", s)], []))::[] -> Some s
  44   | _ -> raise (Ocsigen_extensions.Error_in_config_file
  45                   ("Unexpected content inside Ocsipersist config"))
  46
  47 (* This reference is overwritten when the init function (at the end of the file)
  48    is run, which occurs when the extension is loaded *)
  49 let db_file = ref ((Ocsigen_config.get_datadir ())^"/ocsidb")
  50
  51
  52 (*****************************************************************************)
  53 (** Useful functions on database *)
  54
  55 let yield () =
  56   Thread.yield ()
  57
  58 let rec bind_safely stmt = function
  59   | [] -> stmt
  60   | (value, name)::q as l ->
  61       match Sqlite3.bind stmt (bind_parameter_index stmt name) value with
  62       | Rc.OK -> bind_safely stmt q
  63       | Rc.BUSY | Rc.LOCKED -> yield () ; bind_safely stmt l
  64       | rc -> ignore(finalize stmt) ; failwith (Rc.to_string rc)
  65
  66 let close_safely db =
  67  if not (db_close db) then
  68    ignore (Ocsigen_messages.errlog "Couldn't close database")
  69
  70 let m = Mutex.create ()
  71
  72 let exec_safely f =
  73   let aux () =
  74    let db =
  75      Mutex.lock m ;
  76      try db_open !db_file with e -> Mutex.unlock m; raise e
  77    in
  78     (try
  79       let r = f db in
  80       close_safely db ;
  81       Mutex.unlock m ;
  82       r
  83     with e -> (
  84       close_safely db ;
  85       Mutex.unlock m ;
  86       raise e))
  87   in
  88   Lwt_preemptive.detach aux ()
  89
  90 (* Référence indispensable pour les codes de retours et leur signification :
  91  * http://sqlite.org/capi3ref.html
  92  * Langage compris par SQLite : http://www.sqlite.org/lang.html
  93  *)
  94
  95 let db_create table =
  96   let sql = sprintf "CREATE TABLE IF NOT EXISTS %s (key TEXT, value BLOB,  PRIMARY KEY(key) ON CONFLICT REPLACE)" table in
  97   let create db =
  98     let stmt = prepare db sql in
  99     let rec aux () =
 100       match step stmt with
 101       | Rc.DONE -> ignore(finalize stmt)
 102       | Rc.BUSY | Rc.LOCKED ->  yield () ; aux ()
 103       | rc -> ignore(finalize stmt) ; failwith (Rc.to_string rc)
 104     in
 105     aux ()
 106   in
 107   exec_safely create >>= fun () ->
 108   return table
 109
 110 let db_remove (table, key) =
 111   let sql =  sprintf "DELETE FROM %s WHERE key = :key " table in
 112   let remove db =
 113     let stmt =  bind_safely (prepare db sql) [Data.TEXT key,":key"] in
 114     let rec aux () =
 115       match step stmt with
 116       | Rc.DONE -> ignore(finalize stmt)
 117       | Rc.BUSY | Rc.LOCKED ->  yield () ; aux ()
 118       | rc -> ignore(finalize stmt) ; failwith (Rc.to_string rc)
 119     in aux ()
 120   in
 121   exec_safely remove
 122
 123 let (db_get, db_replace, db_replace_if_exists) =
 124   let get (table, key) db =
 125     let sqlget = sprintf "SELECT value FROM %s WHERE key = :key " table in
 126     let stmt = bind_safely (prepare db sqlget) [Data.TEXT key,":key"] in
 127     let rec aux () =
 128       match step stmt with
 129       | Rc.ROW ->
 130           let value = match column stmt 0 with
 131           | Data.BLOB s -> s
 132           | _ -> assert false
 133           in
 134           ignore (finalize stmt);
 135           value
 136       | Rc.DONE -> ignore(finalize stmt) ;  raise Not_found
 137       | Rc.BUSY | Rc.LOCKED ->  yield () ; aux ()
 138       | rc -> ignore(finalize stmt) ; failwith (Rc.to_string rc)
 139     in aux ()
 140   in
 141   let replace (table, key) value db =
 142     let sqlreplace = sprintf "INSERT INTO %s VALUES ( :key , :value )" table in
 143     let stmt =
 144       bind_safely
 145         (prepare db sqlreplace)
 146         [Data.TEXT key,":key"; Data.BLOB value, ":value"]
 147     in
 148     let rec aux () =
 149       match step stmt with
 150       | Rc.DONE -> ignore(finalize stmt)
 151       | Rc.BUSY | Rc.LOCKED ->  yield () ; aux ()
 152       | rc -> ignore(finalize stmt) ; failwith (Rc.to_string rc)
 153     in aux ()
 154   in
 155   ((fun tablekey -> exec_safely (get tablekey)),
 156    (fun tablekey value -> exec_safely (replace tablekey value)),
 157    (fun tablekey value -> exec_safely
 158        (fun db -> ignore (get tablekey db); replace tablekey value db)))
 159
 160
 161 let db_iter_step table rowid =
 162   let sql =
 163     sprintf "SELECT key , value , ROWID FROM %s WHERE ROWID > :rowid" table in
 164   let iter db =
 165     let stmt = bind_safely (prepare db sql) [Data.INT rowid, ":rowid"] in
 166     let rec aux () =
 167       match step stmt with
 168       | Rc.ROW ->
 169           (match (column stmt 0,column stmt 1, column stmt 2) with
 170           | (Data.TEXT k, Data.BLOB v, Data.INT rowid) ->
 171               ignore(finalize stmt) ;
 172               Some (k, v, rowid)
 173           | _ -> assert false )
 174       | Rc.DONE -> ignore(finalize stmt) ; None
 175       | Rc.BUSY | Rc.LOCKED -> yield () ; aux ()
 176       | rc -> ignore(finalize stmt) ; failwith (Rc.to_string rc)
 177     in aux ()
 178   in
 179   exec_safely iter
 180
 181 let db_iter_block table f =
 182   let sql = sprintf "SELECT key , value FROM %s " table in
 183   let iter db =
 184     let stmt = prepare db sql in
 185     let rec aux () =
 186       match step stmt with
 187       | Rc.ROW ->
 188           (match (column stmt 0,column stmt 1) with
 189           | (Data.TEXT k, Data.BLOB v) -> f k (Marshal.from_string v 0); aux()
 190           | _ -> assert false )
 191       | Rc.DONE -> ignore(finalize stmt)
 192       | Rc.BUSY | Rc.LOCKED ->  yield () ; aux ()
 193       | rc -> ignore(finalize stmt) ; failwith (Rc.to_string rc)
 194     in aux ()
 195   in
 196   exec_safely iter
 197
 198 let db_length table =
 199   let sql = sprintf "SELECT count(*) FROM %s " table in
 200   let length db =
 201     let stmt = prepare db sql in
 202     let rec aux () =
 203       match step stmt with
 204       | Rc.ROW ->
 205           let  value = match column stmt 0 with
 206           | Data.INT s -> Int64.to_int s
 207           | _ -> assert false
 208           in
 209           ignore (finalize stmt);
 210           value
 211       | Rc.DONE -> ignore(finalize stmt) ;  raise Not_found
 212       | Rc.BUSY | Rc.LOCKED ->  yield () ; aux ()
 213       | rc -> ignore(finalize stmt) ; failwith (Rc.to_string rc)
 214     in aux ()
 215   in
 216   exec_safely length
 217
 218
 219
 220
 221
 222 (*****************************************************************************)
 223 (** Public functions: *)
 224
 225 (** Type of persistent data *)
 226 type 'a t = string * string
 227
 228 let open_store name : store =
 229   let s = "store___"^name in
 230   db_create s
 231
 232 let make_persistent_lazy ~store ~name ~default =
 233   store >>= fun store ->
 234   let pvname = (store, name) in
 235   (catch
 236      (fun () -> db_get pvname >>= (fun _ -> return ()))
 237      (function
 238        | Not_found ->
 239            let def = Marshal.to_string (default ()) []
 240            in db_replace pvname def
 241        | e -> fail e)) >>=
 242   (fun () -> return pvname)
 243
 244 let make_persistent ~store ~name ~default =
 245   make_persistent_lazy ~store ~name ~default:(fun () -> default)
 246
 247 let get (pvname : 'a t) : 'a =
 248   db_get pvname >>=
 249   (fun r -> return (Marshal.from_string r 0))
 250
 251 let set pvname v =
 252   let data = Marshal.to_string v [] in
 253   db_replace pvname data
 254
 255 (** Type of persistent tables *)
 256 type 'value table = string Lwt.t
 257
 258 (** name SHOULD NOT begin with "store___" *)
 259 let open_table name = db_create name
 260
 261 let table_name table = table
 262
 263 let find table key =
 264   table >>= fun table ->
 265   db_get (table, key) >>= fun v ->
 266   return (Marshal.from_string v 0)
 267
 268 let add table key value =
 269   table >>= fun table ->
 270   let data = Marshal.to_string value [] in
 271   db_replace (table, key) data
 272
 273 let replace_if_exists table key value =
 274   table >>= fun table ->
 275   let data = Marshal.to_string value [] in
 276   db_replace_if_exists (table, key) data
 277
 278 let remove table key =
 279   table >>= fun table ->
 280   db_remove (table, key)
 281
 282 let iter_step f table =
 283   table >>= fun table ->
 284   let rec aux rowid =
 285     db_iter_step table rowid >>=
 286     (function
 287       | None -> return ()
 288       | Some (k,v,rowid') ->
 289           f k (Marshal.from_string v 0) >>= (fun () -> aux rowid'))
 290   in
 291   aux Int64.zero
 292
 293 let fold_step f table beg =
 294   table >>= fun table ->
 295   let rec aux rowid beg =
 296     db_iter_step table rowid >>=
 297     (function
 298       | None -> return beg
 299       | Some (k, v, rowid') ->
 300           f k (Marshal.from_string v 0) beg >>= (fun res -> aux rowid' res))
 301   in
 302   aux Int64.zero beg
 303
 304 let iter_block f table =
 305   table >>= fun table ->
 306   db_iter_block table f
 307
 308 let iter_table = iter_step
 309
 310 let fold_table = fold_step
 311
 312 let length table =
 313   table >>= fun table ->
 314   db_length table
 315
 316
 317 (* Registration of the extension *)
 318
 319 let init config =
 320   db_file := Ocsigen_config.get_datadir () ^"/ocsidb";
 321   (match parse_global_config config with
 322     | None -> ()
 323     | Some d -> db_file := d
 324   );
 325   (* We check that we can access the database *)
 326   try Lwt_unix.run (exec_safely (fun _ -> ()))
 327   with e ->
 328     Ocsigen_messages.errlog
 329       (Printf.sprintf
 330          "Error opening database file '%s' when registering Ocsipersist. \
 331           Check that the directory exists, and that Ocsigen has enough \
 332           rights" !db_file);
 333     raise e
 334
 335
 336 let _ = Ocsigen_extensions.register_extension
 337   ~name:"ocsipersist"
 338   ~init_fun:init
 339   ()