Minimal SQLite key-value store for OCaml
1(*---------------------------------------------------------------------------
2 Copyright (c) 2025 Thomas Gazagnaire. All rights reserved.
3 SPDX-License-Identifier: MIT
4 ---------------------------------------------------------------------------*)
5
6(* Pure OCaml B-tree backed key-value store with SQLite-compatible file format.
7
8 The file format is a valid SQLite database:
9 - Page 1: 100-byte database header + sqlite_master table B-tree
10 - Page 2+: user data tables using Table B-tree with Record encoding
11 - In-memory hashtable for O(1) key→rowid lookups *)
12
13let page_size = 4096
14let magic = "SQLite format 3\000"
15
16(* Re-export Btree.Record.value so users don't need to depend on btree *)
17
18type value = Btree.Record.value =
19 | Vnull
20 | Vint of int64
21 | Vfloat of float
22 | Vblob of string
23 | Vtext of string
24
25let pp_value = Btree.Record.pp_value
26
27(* Schema types *)
28
29type column = {
30 col_name : string;
31 col_affinity : string;
32 col_is_rowid_alias : bool;
33}
34
35type schema = { tbl_name : string; columns : column list; sql : string }
36
37(* Per-table state *)
38type kv_table = {
39 btree : Btree.Table.t;
40 keys : (string, int64) Hashtbl.t;
41 mutable next_rowid : int64;
42}
43
44exception Unique_violation of string
45
46(* A persistent unique index backed by Btree.Index *)
47type unique_index = {
48 ui_columns : int list; (* column indices *)
49 ui_name : string; (* for error messages, e.g. "provider, uid" *)
50 ui_index_name : string; (* sqlite_master name, e.g. "sqlite_autoindex_t_1" *)
51 ui_tbl_name : string; (* owning table name *)
52 ui_btree : Btree.Index.t; (* persistent B-tree index *)
53}
54
55type generic_table = {
56 g_btree : Btree.Table.t;
57 g_schema : schema;
58 g_unique_indexes : unique_index list;
59}
60
61(* Raw sqlite_master entry for schema objects we don't manage (views,
62 triggers, explicit indexes). Preserved across open/close. *)
63type raw_master_entry = {
64 rm_type : string;
65 rm_name : string;
66 rm_tbl_name : string;
67 rm_root_page : int;
68 rm_sql : Btree.Record.value;
69}
70
71type t = {
72 pager : Btree.Pager.t;
73 file : Eio.File.rw_ty Eio.Resource.t option;
74 mutable sw : Eio.Switch.t option;
75 db_path : Eio.Fs.dir_ty Eio.Path.t option;
76 data : kv_table option;
77 mutable named_tables : (string * kv_table) list;
78 mutable all_tables : generic_table list;
79 mutable extra_master : raw_master_entry list;
80 insert_rowids : (string, int64 ref) Hashtbl.t;
81}
82
83let pp ppf t =
84 let names = List.map (fun gt -> gt.g_schema.tbl_name) t.all_tables in
85 Fmt.pf ppf "sqlite(%a)" Fmt.(list ~sep:(any ",") string) names
86
87(* WAL record format: 4-byte big-endian page number + page data *)
88
89let encode_wal_page page_num data =
90 let buf = Bytes.create (4 + String.length data) in
91 Btree.Page.set_u32_be buf 0 page_num;
92 Bytes.blit_string data 0 buf 4 (String.length data);
93 Bytes.unsafe_to_string buf
94
95let decode_wal_page record =
96 if String.length record < 4 then None
97 else
98 let page_num = Btree.Page.u32_be record 0 in
99 let data = String.sub record 4 (String.length record - 4) in
100 Some (page_num, data)
101
102let wal_path (fs, name) = (fs, name ^ "-wal")
103
104let replay_wal ~file wal_path =
105 if Eio.Path.is_file wal_path then begin
106 Wal.fold wal_path ~init:() ~f:(fun () record ->
107 match decode_wal_page record with
108 | Some (page_num, data) ->
109 let offset = Optint.Int63.of_int ((page_num - 1) * page_size) in
110 Eio.File.pwrite_all file ~file_offset:offset
111 [ Cstruct.of_string data ]
112 | None -> ());
113 Eio.File.sync file;
114 Eio.Path.unlink wal_path
115 end
116
117(* CREATE TABLE parser — delegates to Lexer.parse (menhir + ocamllex) *)
118
119let parse_sql sql =
120 match Lexer.parse sql with Ok ct -> Some ct | Error _ -> None
121
122let parse_create_table sql =
123 match parse_sql sql with
124 | None -> []
125 | Some ct ->
126 List.map
127 (fun (c : Ast.column_def) ->
128 {
129 col_name = c.name;
130 col_affinity = c.affinity;
131 col_is_rowid_alias = c.is_rowid_alias;
132 })
133 ct.columns
134
135let parse_unique_constraints sql columns =
136 match parse_sql sql with
137 | None -> []
138 | Some ct ->
139 let col_names = List.map (fun c -> c.col_name) columns in
140 let find_idx name =
141 let rec go i = function
142 | [] -> None
143 | n :: _ when n = name -> Some i
144 | _ :: rest -> go (i + 1) rest
145 in
146 go 0 col_names
147 in
148 let resolve_cols names =
149 let indices = List.filter_map find_idx names in
150 if List.length indices = List.length names then
151 Some (indices, String.concat ", " names)
152 else None
153 in
154 (* Table-level UNIQUE and PRIMARY KEY constraints *)
155 let table_level =
156 List.filter_map
157 (function
158 | Ast.Tbl_unique cols | Ast.Tbl_primary_key cols ->
159 resolve_cols cols
160 | _ -> None)
161 ct.table_constraints
162 in
163 (* Column-level UNIQUE and non-rowid PRIMARY KEY constraints *)
164 let column_level =
165 List.filter_map
166 (fun (c : Ast.column_def) ->
167 if c.has_unique || (c.has_primary_key && not c.is_rowid_alias) then
168 resolve_cols [ c.name ]
169 else None)
170 ct.columns
171 in
172 table_level @ column_level
173
174(* Find the index of the rowid alias column, if any *)
175let rowid_alias_index columns =
176 let rec find i = function
177 | [] -> None
178 | c :: _ when c.col_is_rowid_alias -> Some i
179 | _ :: rest -> find (i + 1) rest
180 in
181 find 0 columns
182
183(* Apply rowid substitution and trailing Vnull padding *)
184let fixup_values ~schema ~rowid values =
185 let n_cols = List.length schema.columns in
186 let len = List.length values in
187 let values =
188 if len < n_cols then
189 values @ List.init (n_cols - len) (fun _ -> Btree.Record.Vnull)
190 else values
191 in
192 match rowid_alias_index schema.columns with
193 | None -> values
194 | Some idx ->
195 List.mapi
196 (fun i v ->
197 if i = idx then
198 match v with
199 | Btree.Record.Vnull -> Btree.Record.Vint rowid
200 | v -> v
201 else v)
202 values
203
204(* Encode the indexed column values as a Record string for Btree.Index.
205 Returns None if any indexed column is NULL — per SQL semantics,
206 NULL is never equal to NULL, so UNIQUE constraints allow multiple NULLs. *)
207let encode_index_key ui values =
208 let key_values =
209 List.map
210 (fun idx ->
211 if idx < List.length values then List.nth values idx
212 else Btree.Record.Vnull)
213 ui.ui_columns
214 in
215 if List.exists (fun v -> v = Btree.Record.Vnull) key_values then None
216 else Some (Btree.Record.encode key_values)
217
218(* Create persistent unique indexes for a table *)
219let unique_indexes pager ~tbl_name constraints =
220 List.mapi
221 (fun i (indices, col_name) ->
222 let idx_name = Fmt.str "sqlite_autoindex_%s_%d" tbl_name (i + 1) in
223 {
224 ui_columns = indices;
225 ui_name = col_name;
226 ui_index_name = idx_name;
227 ui_tbl_name = tbl_name;
228 ui_btree = Btree.Index.v pager;
229 })
230 constraints
231
232(* Standard kv table schema *)
233let kv_columns =
234 [
235 { col_name = "key"; col_affinity = "TEXT"; col_is_rowid_alias = false };
236 { col_name = "value"; col_affinity = "BLOB"; col_is_rowid_alias = false };
237 ]
238
239let table_sql name = Fmt.str "CREATE TABLE %s (key TEXT, value BLOB)" name
240
241let kv_schema name =
242 { tbl_name = name; columns = kv_columns; sql = table_sql name }
243
244(* Decode a Record payload into (key, value) *)
245let decode_kv payload =
246 match Btree.Record.decode payload with
247 | [ Btree.Record.Vtext k; Btree.Record.Vblob v ]
248 | [ Btree.Record.Vtext k; Btree.Record.Vtext v ] ->
249 Some (k, v)
250 | _ -> None
251
252(* Scan a table B-tree to build key→rowid map *)
253let scan_table btree =
254 let keys = Hashtbl.create 64 in
255 let next_rowid = ref 1L in
256 Btree.Table.iter btree (fun rowid payload ->
257 (match decode_kv payload with
258 | Some (k, _) -> Hashtbl.replace keys k rowid
259 | None -> ());
260 if rowid >= !next_rowid then next_rowid := Int64.add rowid 1L);
261 (keys, !next_rowid)
262
263(* Write the 100-byte SQLite database header *)
264let write_db_header buf ~page_count =
265 Bytes.blit_string magic 0 buf 0 16;
266 Btree.Page.set_u16_be buf 16 page_size;
267 Bytes.set_uint8 buf 18 1;
268 Bytes.set_uint8 buf 19 1;
269 Bytes.set_uint8 buf 20 0;
270 Bytes.set_uint8 buf 21 64;
271 Bytes.set_uint8 buf 22 32;
272 Bytes.set_uint8 buf 23 32;
273 Btree.Page.set_u32_be buf 24 1;
274 Btree.Page.set_u32_be buf 28 page_count;
275 Btree.Page.set_u32_be buf 32 0;
276 Btree.Page.set_u32_be buf 36 0;
277 Btree.Page.set_u32_be buf 40 1;
278 Btree.Page.set_u32_be buf 44 4;
279 Btree.Page.set_u32_be buf 48 0;
280 Btree.Page.set_u32_be buf 52 0;
281 Btree.Page.set_u32_be buf 56 1;
282 Btree.Page.set_u32_be buf 60 0;
283 Btree.Page.set_u32_be buf 64 0;
284 Btree.Page.set_u32_be buf 68 0;
285 (* Offsets 72-91: reserved for expansion, must be zero (spec section 1.2) *)
286 Bytes.fill buf 72 20 '\x00';
287 Btree.Page.set_u32_be buf 92 1;
288 Btree.Page.set_u32_be buf 96 3046000
289
290(* Build a sqlite_master table leaf cell *)
291let master_cell ~rowid ~type_ ~name ~tbl_name ~root_page ~sql =
292 let record =
293 Btree.Record.encode
294 [
295 Btree.Record.Vtext type_;
296 Btree.Record.Vtext name;
297 Btree.Record.Vtext tbl_name;
298 Btree.Record.Vint (Int64.of_int root_page);
299 sql;
300 ]
301 in
302 let payload_varint =
303 Btree.Varint.encode (Int64.of_int (String.length record))
304 in
305 let rowid_varint = Btree.Varint.encode rowid in
306 payload_varint ^ rowid_varint ^ record
307
308(* Collect all sqlite_master entries: tables, autoindexes, then extras *)
309let master_entries t =
310 let entries = ref [] in
311 let rowid = ref 1 in
312 List.iter
313 (fun gt ->
314 entries :=
315 ( !rowid,
316 gt.g_schema.tbl_name,
317 gt.g_schema.tbl_name,
318 Btree.Table.root_page gt.g_btree,
319 Btree.Record.Vtext gt.g_schema.sql,
320 "table" )
321 :: !entries;
322 incr rowid;
323 List.iter
324 (fun ui ->
325 entries :=
326 ( !rowid,
327 ui.ui_index_name,
328 ui.ui_tbl_name,
329 Btree.Index.root_page ui.ui_btree,
330 Btree.Record.Vnull,
331 "index" )
332 :: !entries;
333 incr rowid)
334 gt.g_unique_indexes)
335 t.all_tables;
336 (* Preserve explicit indexes, views, triggers *)
337 List.iter
338 (fun rm ->
339 entries :=
340 ( !rowid,
341 rm.rm_name,
342 rm.rm_tbl_name,
343 rm.rm_root_page,
344 rm.rm_sql,
345 rm.rm_type )
346 :: !entries;
347 incr rowid)
348 t.extra_master;
349 List.rev !entries
350
351(* Write page 1: db header + sqlite_master leaf table *)
352let rebuild_page1 t =
353 let buf = Bytes.create page_size in
354 write_db_header buf ~page_count:(Btree.Pager.page_count t.pager);
355 (* Leaf table header at offset 100 *)
356 Bytes.set_uint8 buf 100 0x0d;
357 Btree.Page.set_u16_be buf 101 0;
358 Bytes.set_uint8 buf 107 0;
359 let entries = master_entries t in
360 let n = List.length entries in
361 Btree.Page.set_u16_be buf 103 n;
362 (* Build cells from end of page *)
363 let cell_content_start = ref page_size in
364 let cell_ptrs = Array.make n 0 in
365 List.iteri
366 (fun i (rowid, name, tbl_name, root_page, sql, type_) ->
367 let cell =
368 master_cell ~rowid:(Int64.of_int rowid) ~type_ ~name ~tbl_name
369 ~root_page ~sql
370 in
371 let cell_len = String.length cell in
372 cell_content_start := !cell_content_start - cell_len;
373 Bytes.blit_string cell 0 buf !cell_content_start cell_len;
374 cell_ptrs.(i) <- !cell_content_start)
375 entries;
376 Btree.Page.set_u16_be buf 105 !cell_content_start;
377 (* Cell pointer array at offset 108 (100 + 8 byte leaf header) *)
378 Array.iteri
379 (fun i ptr -> Btree.Page.set_u16_be buf (108 + (i * 2)) ptr)
380 cell_ptrs;
381 Btree.Pager.write t.pager 1 (Bytes.unsafe_to_string buf)
382
383(* Initialize a new kv_table on a fresh page *)
384let new_kv_table pager =
385 let btree = Btree.Table.v pager in
386 { btree; keys = Hashtbl.create 64; next_rowid = 1L }
387
388let mkdirs_for path =
389 match Eio.Path.split path with
390 | None -> ()
391 | Some (fs, p) -> (
392 let dir = Filename.dirname p in
393 if dir <> "." && dir <> "/" then
394 try Eio.Path.mkdirs ~exists_ok:true ~perm:0o755 Eio.Path.(fs / dir)
395 with Eio.Io _ -> ())
396
397let init ~sw path =
398 mkdirs_for path;
399 let file =
400 Eio.Path.open_out ~sw ~create:(`If_missing 0o644) path |> fun f ->
401 (f :> Eio.File.rw_ty Eio.Resource.t)
402 in
403 let pager = Btree.Pager.v ~page_size file in
404 (* Allocate page 1 for db header + sqlite_master *)
405 let _page1 = Btree.Pager.allocate pager in
406 (* Create kv data table on page 2 *)
407 let kv = new_kv_table pager in
408 let gt =
409 { g_btree = kv.btree; g_schema = kv_schema "kv"; g_unique_indexes = [] }
410 in
411 let t =
412 {
413 pager;
414 file = Some file;
415 sw = Some sw;
416 db_path = Some path;
417 data = Some kv;
418 named_tables = [];
419 all_tables = [ gt ];
420 extra_master = [];
421 insert_rowids = Hashtbl.create 8;
422 }
423 in
424 rebuild_page1 t;
425 Btree.Pager.sync pager;
426 t
427
428(* Extract named kv tables (non-kv tables with kv schema) from all_tables *)
429let extract_named_kv_tables all_tables =
430 List.filter_map
431 (fun gt ->
432 let name = gt.g_schema.tbl_name in
433 if name = "kv" then None
434 else
435 match gt.g_schema.columns with
436 | [
437 { col_name = "key"; col_affinity = "TEXT"; _ };
438 { col_name = "value"; col_affinity = "BLOB"; _ };
439 ] ->
440 let keys, next_rowid = scan_table gt.g_btree in
441 Some (name, { btree = gt.g_btree; keys; next_rowid })
442 | _ -> None)
443 all_tables
444
445let in_memory () =
446 let pager = Btree.Pager.mem ~page_size () in
447 let _page1 = Btree.Pager.allocate pager in
448 let kv = new_kv_table pager in
449 let gt =
450 { g_btree = kv.btree; g_schema = kv_schema "kv"; g_unique_indexes = [] }
451 in
452 let t =
453 {
454 pager;
455 file = None;
456 sw = None;
457 db_path = None;
458 data = Some kv;
459 named_tables = [];
460 all_tables = [ gt ];
461 extra_master = [];
462 insert_rowids = Hashtbl.create 8;
463 }
464 in
465 rebuild_page1 t;
466 t
467
468(* sqlite_master row classifiers *)
469
470type master_row =
471 | Mtable of string * int * string (** name, root, sql *)
472 | Mautoindex of string * string * int (** idx, tbl, root *)
473 | Mextra of raw_master_entry
474 | Mskip
475
476let classify_index ~name ~tbl ~root sql =
477 if String.starts_with ~prefix:"sqlite_autoindex_" name then
478 Mautoindex (name, tbl, root)
479 else
480 Mextra
481 {
482 rm_type = "index";
483 rm_name = name;
484 rm_tbl_name = tbl;
485 rm_root_page = root;
486 rm_sql = sql;
487 }
488
489let classify_master_row payload =
490 match Btree.Record.decode payload with
491 | [
492 Btree.Record.Vtext "table";
493 Btree.Record.Vtext name;
494 _;
495 Btree.Record.Vint root;
496 Btree.Record.Vtext sql;
497 ] ->
498 Mtable (name, Int64.to_int root, sql)
499 | [
500 Btree.Record.Vtext "index";
501 Btree.Record.Vtext name;
502 Btree.Record.Vtext tbl;
503 Btree.Record.Vint root;
504 Btree.Record.Vtext sql;
505 ] ->
506 classify_index ~name ~tbl ~root:(Int64.to_int root)
507 (Btree.Record.Vtext sql)
508 | [
509 Btree.Record.Vtext "index";
510 Btree.Record.Vtext name;
511 Btree.Record.Vtext tbl;
512 Btree.Record.Vint root;
513 Btree.Record.Vnull;
514 ] ->
515 Mautoindex (name, tbl, Int64.to_int root)
516 | [
517 Btree.Record.Vtext type_;
518 Btree.Record.Vtext name;
519 Btree.Record.Vtext tbl_name;
520 Btree.Record.Vint root;
521 sql;
522 ]
523 when type_ = "view" || type_ = "trigger" ->
524 Mextra
525 {
526 rm_type = type_;
527 rm_name = name;
528 rm_tbl_name = tbl_name;
529 rm_root_page = Int64.to_int root;
530 rm_sql = sql;
531 }
532 | _ -> Mskip
533
534(* Parse sqlite_master into raw table, index, and other entries *)
535let read_master page1 =
536 let header = Btree.Page.parse_header page1 100 in
537 let ptrs = Btree.Page.cell_pointers page1 100 header in
538 let raw_tables = ref [] in
539 let raw_indexes = ref [] in
540 let raw_extra = ref [] in
541 for i = 0 to header.Btree.Page.cell_count - 1 do
542 let cell, _ =
543 Btree.Cell.parse_table_leaf page1 ptrs.(i) ~usable_size:page_size
544 in
545 match classify_master_row cell.Btree.Cell.payload with
546 | Mtable (name, root, sql) -> raw_tables := (name, root, sql) :: !raw_tables
547 | Mautoindex (idx, tbl, root) ->
548 raw_indexes := (idx, tbl, root) :: !raw_indexes
549 | Mextra r -> raw_extra := r :: !raw_extra
550 | Mskip -> ()
551 done;
552 (List.rev !raw_tables, List.rev !raw_indexes, List.rev !raw_extra)
553
554(* Reconnect or create unique indexes for a table *)
555let open_unique_indexes pager ~btree ~schema ~constraints ~raw_indexes =
556 let name = schema.tbl_name in
557 let tbl_indexes = List.filter (fun (_, tbl, _) -> tbl = name) raw_indexes in
558 List.mapi
559 (fun i (indices, col_name) ->
560 let expected_name = Fmt.str "sqlite_autoindex_%s_%d" name (i + 1) in
561 match
562 List.find_opt
563 (fun (idx_name, _, _) -> idx_name = expected_name)
564 tbl_indexes
565 with
566 | Some (idx_name, _, idx_root) ->
567 {
568 ui_columns = indices;
569 ui_name = col_name;
570 ui_index_name = idx_name;
571 ui_tbl_name = name;
572 ui_btree = Btree.Index.open_ pager ~root_page:idx_root;
573 }
574 | None ->
575 let ui =
576 {
577 ui_columns = indices;
578 ui_name = col_name;
579 ui_index_name = expected_name;
580 ui_tbl_name = name;
581 ui_btree = Btree.Index.v pager;
582 }
583 in
584 Btree.Table.iter btree (fun rowid payload ->
585 let values = Btree.Record.decode payload in
586 let values = fixup_values ~schema ~rowid values in
587 match encode_index_key ui values with
588 | None -> ()
589 | Some key -> Btree.Index.insert ui.ui_btree key);
590 ui)
591 constraints
592
593let open_ ~sw ?(create = true) path =
594 try
595 let file =
596 Eio.Path.open_out ~sw ~create:`Never path |> fun f ->
597 (f :> Eio.File.rw_ty Eio.Resource.t)
598 in
599 (* Replay any WAL entries from a previous crash *)
600 replay_wal ~file (wal_path path);
601 let pager = Btree.Pager.v ~page_size file in
602 if Btree.Pager.page_count pager = 0 then
603 failwith "Database file exists but is empty (delete it to recreate)";
604 let page1 = Btree.Pager.read pager 1 in
605 if String.sub page1 0 16 <> magic then failwith "Not a SQLite database";
606 let ps = Btree.Page.u16_be page1 16 in
607 if ps <> page_size then Fmt.failwith "Unsupported page size: %d" ps;
608 let raw_tables, raw_indexes, raw_extra = read_master page1 in
609 let all_tables =
610 List.map
611 (fun (name, root, sql) ->
612 let btree = Btree.Table.open_ pager ~root_page:root in
613 let columns = parse_create_table sql in
614 let schema = { tbl_name = name; columns; sql } in
615 let constraints = parse_unique_constraints sql columns in
616 let indexes =
617 open_unique_indexes pager ~btree ~schema ~constraints ~raw_indexes
618 in
619 { g_btree = btree; g_schema = schema; g_unique_indexes = indexes })
620 raw_tables
621 in
622 let data =
623 match
624 List.find_opt (fun gt -> gt.g_schema.tbl_name = "kv") all_tables
625 with
626 | None -> None
627 | Some gt ->
628 let keys, next_rowid = scan_table gt.g_btree in
629 Some { btree = gt.g_btree; keys; next_rowid }
630 in
631 let named = extract_named_kv_tables all_tables in
632 {
633 pager;
634 file = Some file;
635 sw = Some sw;
636 db_path = Some path;
637 data;
638 named_tables = named;
639 all_tables;
640 extra_master = raw_extra;
641 insert_rowids = Hashtbl.create 8;
642 }
643 with Eio.Io _ when create -> init ~sw path
644
645(* Get the kv_table, raising if no kv table exists *)
646let kv t =
647 match t.data with
648 | Some d -> d
649 | None -> failwith "No 'kv' table in this database"
650
651(* KV operations *)
652
653let find t key =
654 let d = kv t in
655 match Hashtbl.find_opt d.keys key with
656 | None -> None
657 | Some rowid -> (
658 match Btree.Table.find d.btree rowid with
659 | None -> None
660 | Some payload -> (
661 match decode_kv payload with Some (_, v) -> Some v | None -> None))
662
663let put t key value =
664 let kv = kv t in
665 let record =
666 Btree.Record.encode [ Btree.Record.Vtext key; Btree.Record.Vblob value ]
667 in
668 (match Hashtbl.find_opt kv.keys key with
669 | Some old_rowid -> Btree.Table.delete kv.btree old_rowid
670 | None -> ());
671 let rowid = kv.next_rowid in
672 kv.next_rowid <- Int64.add kv.next_rowid 1L;
673 Btree.Table.insert kv.btree ~rowid record;
674 Hashtbl.replace kv.keys key rowid
675
676let delete t key =
677 let d = kv t in
678 match Hashtbl.find_opt d.keys key with
679 | None -> ()
680 | Some rowid ->
681 Btree.Table.delete d.btree rowid;
682 Hashtbl.remove d.keys key
683
684let mem t key =
685 let d = kv t in
686 Hashtbl.mem d.keys key
687
688let iter t ~f =
689 let d = kv t in
690 Btree.Table.iter d.btree (fun _rowid payload ->
691 match decode_kv payload with Some (k, v) -> f k v | None -> ())
692
693let fold t ~init ~f =
694 let acc = ref init in
695 iter t ~f:(fun k v -> acc := f k v !acc);
696 !acc
697
698let sync t =
699 rebuild_page1 t;
700 match (t.sw, t.db_path) with
701 | Some sw, Some path ->
702 let wp = wal_path path in
703 let wal = Wal.v ~sw wp in
704 (* Write dirty pages to WAL first *)
705 Btree.Pager.iter_dirty t.pager ~f:(fun page_num data ->
706 Wal.append wal (encode_wal_page page_num data));
707 Wal.sync wal;
708 (* WAL is durable — now safe to write to database file *)
709 Btree.Pager.sync t.pager;
710 (match t.file with Some f -> Eio.File.sync f | None -> ());
711 Wal.close wal;
712 if Eio.Path.is_file wp then Eio.Path.unlink wp
713 | _ -> Btree.Pager.sync t.pager
714
715let close t =
716 (* Clear sw so sync takes the direct-write path instead of opening a new WAL
717 file. During close (and especially during Switch.on_release) the switch
718 may already be tearing down, making Wal.v ~sw fail with
719 "Switch finished!". Direct pager sync is safe here because we are doing a
720 clean shutdown, not recovering from a crash. *)
721 t.sw <- None;
722 sync t
723
724(* Transactions *)
725
726type kv_snapshot = { ks_keys : (string, int64) Hashtbl.t; ks_next : int64 }
727
728let snapshot_kv kv = { ks_keys = Hashtbl.copy kv.keys; ks_next = kv.next_rowid }
729
730type txn_snapshot = {
731 pager_snap : Btree.Pager.snapshot;
732 kv_snap : kv_snapshot option;
733 named_snaps : (string * kv_snapshot) list;
734 rowid_snaps : (string * int64) list;
735 table_roots : (Btree.Table.t * int) list;
736 index_roots : (Btree.Index.t * int) list;
737 all_tables_snap : generic_table list;
738 named_tables_snap : (string * kv_table) list;
739 extra_master_snap : raw_master_entry list;
740}
741
742let snapshot_txn t =
743 let table_roots =
744 List.map
745 (fun gt -> (gt.g_btree, Btree.Table.save_root gt.g_btree))
746 t.all_tables
747 in
748 let index_roots =
749 List.concat_map
750 (fun gt ->
751 List.map
752 (fun ui -> (ui.ui_btree, Btree.Index.save_root ui.ui_btree))
753 gt.g_unique_indexes)
754 t.all_tables
755 in
756 let rowid_snaps =
757 Hashtbl.fold (fun k v acc -> (k, !v) :: acc) t.insert_rowids []
758 in
759 {
760 pager_snap = Btree.Pager.snapshot t.pager;
761 kv_snap = Option.map snapshot_kv t.data;
762 named_snaps =
763 List.map (fun (name, kv) -> (name, snapshot_kv kv)) t.named_tables;
764 rowid_snaps;
765 table_roots;
766 index_roots;
767 all_tables_snap = t.all_tables;
768 named_tables_snap = t.named_tables;
769 extra_master_snap = t.extra_master;
770 }
771
772let restore_txn t snap =
773 Btree.Pager.rollback t.pager snap.pager_snap;
774 (* Restore B-tree root pages *)
775 List.iter
776 (fun (btree, root) -> Btree.Table.restore_root btree root)
777 snap.table_roots;
778 List.iter
779 (fun (idx, root) -> Btree.Index.restore_root idx root)
780 snap.index_roots;
781 (* Restore KV in-memory state *)
782 (match (t.data, snap.kv_snap) with
783 | Some kv, Some s ->
784 Hashtbl.reset kv.keys;
785 Hashtbl.iter (Hashtbl.replace kv.keys) s.ks_keys;
786 kv.next_rowid <- s.ks_next
787 | _ -> ());
788 List.iter
789 (fun (name, s) ->
790 match List.assoc_opt name t.named_tables with
791 | Some kv ->
792 Hashtbl.reset kv.keys;
793 Hashtbl.iter (Hashtbl.replace kv.keys) s.ks_keys;
794 kv.next_rowid <- s.ks_next
795 | None -> ())
796 snap.named_snaps;
797 (* Restore schema state *)
798 t.all_tables <- snap.all_tables_snap;
799 t.named_tables <- snap.named_tables_snap;
800 t.extra_master <- snap.extra_master_snap;
801 (* Restore insert_rowids *)
802 Hashtbl.reset t.insert_rowids;
803 List.iter
804 (fun (name, v) -> Hashtbl.replace t.insert_rowids name (ref v))
805 snap.rowid_snaps
806
807let with_transaction t f =
808 let snap = snapshot_txn t in
809 match f () with
810 | result -> result
811 | exception exn ->
812 restore_txn t snap;
813 raise exn
814
815(* Generic read API *)
816
817let tables t = List.map (fun gt -> gt.g_schema) t.all_tables
818
819let table t name =
820 match List.find_opt (fun gt -> gt.g_schema.tbl_name = name) t.all_tables with
821 | Some gt -> gt
822 | None -> Fmt.failwith "No table %S found in database" name
823
824let iter_table t name ~f =
825 let gt = table t name in
826 let schema = gt.g_schema in
827 Btree.Table.iter gt.g_btree (fun rowid payload ->
828 let values = Btree.Record.decode payload in
829 let values = fixup_values ~schema ~rowid values in
830 f rowid values)
831
832let fold_table t name ~init ~f =
833 let acc = ref init in
834 iter_table t name ~f:(fun rowid values -> acc := f rowid values !acc);
835 !acc
836
837let read_table t name =
838 fold_table t name ~init:[] ~f:(fun rowid values acc -> (rowid, values) :: acc)
839 |> List.rev
840
841(* Generic write API *)
842
843let create_table t ~sql =
844 let ast =
845 match Lexer.parse sql with
846 | Ok ct -> ct
847 | Error msg -> Fmt.failwith "Invalid CREATE TABLE statement: %s" msg
848 in
849 let name = ast.Ast.tbl_name in
850 let columns = parse_create_table sql in
851 if List.exists (fun gt -> gt.g_schema.tbl_name = name) t.all_tables then
852 Fmt.failwith "table %S already exists" name;
853 let btree = Btree.Table.v t.pager in
854 let schema = { tbl_name = name; columns; sql } in
855 let constraints = parse_unique_constraints sql columns in
856 let indexes = unique_indexes t.pager ~tbl_name:name constraints in
857 let gt = { g_btree = btree; g_schema = schema; g_unique_indexes = indexes } in
858 t.all_tables <- t.all_tables @ [ gt ]
859
860let next_rowid_for t name =
861 match Hashtbl.find_opt t.insert_rowids name with
862 | Some r -> r
863 | None ->
864 let gt = table t name in
865 let max_rowid = ref 0L in
866 Btree.Table.iter gt.g_btree (fun rowid _payload ->
867 if rowid > !max_rowid then max_rowid := rowid);
868 let r = ref (Int64.add !max_rowid 1L) in
869 Hashtbl.replace t.insert_rowids name r;
870 r
871
872let insert t ~table:name values =
873 let gt = table t name in
874 let next = next_rowid_for t name in
875 (* Determine rowid: for INTEGER PRIMARY KEY, use explicit value or auto *)
876 let rowid, record_values =
877 match rowid_alias_index gt.g_schema.columns with
878 | None ->
879 let rowid = !next in
880 next := Int64.add rowid 1L;
881 (rowid, values)
882 | Some idx ->
883 let pk_value =
884 if idx < List.length values then List.nth values idx else Vnull
885 in
886 let rowid =
887 match pk_value with
888 | Vint n ->
889 (* INTEGER PRIMARY KEY is implicitly UNIQUE — reject duplicates *)
890 if Btree.Table.find gt.g_btree n <> None then
891 raise
892 (Unique_violation
893 (Fmt.str "INTEGER PRIMARY KEY (rowid %Ld)" n));
894 if Int64.add n 1L > !next then next := Int64.add n 1L;
895 n
896 | _ ->
897 let rowid = !next in
898 next := Int64.add rowid 1L;
899 rowid
900 in
901 (* Store Vnull for the INTEGER PRIMARY KEY column in the record,
902 since SQLite stores the rowid separately *)
903 let record_values =
904 List.mapi (fun i v -> if i = idx then Vnull else v) values
905 in
906 (rowid, record_values)
907 in
908 (* Check unique constraints before inserting *)
909 let full_values = fixup_values ~schema:gt.g_schema ~rowid values in
910 List.iter
911 (fun ui ->
912 match encode_index_key ui full_values with
913 | None -> () (* NULL columns — UNIQUE allows multiple NULLs *)
914 | Some key ->
915 if Btree.Index.mem ui.ui_btree key then
916 raise (Unique_violation ui.ui_name))
917 gt.g_unique_indexes;
918 let record = Btree.Record.encode record_values in
919 Btree.Table.insert gt.g_btree ~rowid record;
920 (* Update persistent unique indexes *)
921 List.iter
922 (fun ui ->
923 match encode_index_key ui full_values with
924 | None -> ()
925 | Some key -> Btree.Index.insert ui.ui_btree key)
926 gt.g_unique_indexes;
927 rowid
928
929let delete_row t ~table:name rowid =
930 let gt = table t name in
931 (* Remove from unique indexes first *)
932 (match Btree.Table.find gt.g_btree rowid with
933 | None -> ()
934 | Some payload ->
935 let values = Btree.Record.decode payload in
936 let values = fixup_values ~schema:gt.g_schema ~rowid values in
937 List.iter
938 (fun ui ->
939 match encode_index_key ui values with
940 | None -> ()
941 | Some key -> Btree.Index.delete ui.ui_btree key)
942 gt.g_unique_indexes);
943 Btree.Table.delete gt.g_btree rowid
944
945(* Namespaced Tables *)
946
947module Table = struct
948 type db = t
949 type t = { kv : kv_table }
950
951 let valid_name name =
952 String.length name > 0
953 && (let first = name.[0] in
954 (first >= 'a' && first <= 'z')
955 || (first >= 'A' && first <= 'Z')
956 || first = '_')
957 && String.for_all
958 (fun c ->
959 (c >= 'a' && c <= 'z')
960 || (c >= 'A' && c <= 'Z')
961 || (c >= '0' && c <= '9')
962 || c = '_')
963 name
964
965 let create parent ~name =
966 if not (valid_name name) then Fmt.invalid_arg "Invalid table name: %S" name;
967 match List.assoc_opt name parent.named_tables with
968 | Some kv -> { kv }
969 | None ->
970 (* Check if a table with this name already exists (e.g. the default
971 "kv" table or a table created via create_table). If so, reuse it
972 as a named kv table rather than creating a duplicate. *)
973 let existing_gt =
974 List.find_opt
975 (fun gt -> gt.g_schema.tbl_name = name)
976 parent.all_tables
977 in
978 let kv =
979 match existing_gt with
980 | Some gt ->
981 let keys, next_rowid = scan_table gt.g_btree in
982 { btree = gt.g_btree; keys; next_rowid }
983 | None ->
984 let kv = new_kv_table parent.pager in
985 let gt =
986 {
987 g_btree = kv.btree;
988 g_schema = kv_schema name;
989 g_unique_indexes = [];
990 }
991 in
992 parent.all_tables <- parent.all_tables @ [ gt ];
993 kv
994 in
995 parent.named_tables <- (name, kv) :: parent.named_tables;
996 { kv }
997
998 (* Scan the B-tree for a key, returning (rowid, value) if found.
999 This is the authoritative lookup — no stale cache. *)
1000 let btree_find_key kv key =
1001 let result = ref None in
1002 Btree.Table.iter kv.btree (fun rowid payload ->
1003 if !result = None then
1004 match decode_kv payload with
1005 | Some (k, v) when k = key -> result := Some (rowid, v)
1006 | _ -> ());
1007 !result
1008
1009 (* Compute the next available rowid by scanning the B-tree. *)
1010 let btree_max_rowid kv =
1011 let max_id = ref 0L in
1012 Btree.Table.iter kv.btree (fun rowid _payload ->
1013 if rowid > !max_id then max_id := rowid);
1014 Int64.add !max_id 1L
1015
1016 let find t key =
1017 match btree_find_key t.kv key with Some (_, v) -> Some v | None -> None
1018
1019 let put t key value =
1020 let kv = t.kv in
1021 let record =
1022 Btree.Record.encode [ Btree.Record.Vtext key; Btree.Record.Vblob value ]
1023 in
1024 (* Delete existing entry if present — read from B-tree, not cache. *)
1025 (match btree_find_key kv key with
1026 | Some (old_rowid, _) -> Btree.Table.delete kv.btree old_rowid
1027 | None -> ());
1028 (* Use authoritative next rowid from B-tree. *)
1029 let rowid = btree_max_rowid kv in
1030 kv.next_rowid <- Int64.add rowid 1L;
1031 Btree.Table.insert kv.btree ~rowid record;
1032 Hashtbl.replace kv.keys key rowid
1033
1034 let delete t key =
1035 match btree_find_key t.kv key with
1036 | None -> ()
1037 | Some (rowid, _) ->
1038 Btree.Table.delete t.kv.btree rowid;
1039 Hashtbl.remove t.kv.keys key
1040
1041 let mem t key = Option.is_some (btree_find_key t.kv key)
1042
1043 let iter t ~f =
1044 Btree.Table.iter t.kv.btree (fun _rowid payload ->
1045 match decode_kv payload with Some (k, v) -> f k v | None -> ())
1046end