Minimal SQLite key-value store for OCaml
0
fork

Configure Feed

Select the types of activity you want to include in your feed.

at main 1046 lines 33 kB view raw
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