Add UNIQUE constraint enforcement to ocaml-sqlite, fix OAuth identity race

+181 -37

lib/sqlite.ml

··· 41 41 mutable next_rowid : int64; 42 42 } 43 43 44 - type generic_table = { g_btree : Btree.Table.t; g_schema : schema } 44 + exception Unique_violation of string 45 + 46 + (* A unique index maps composite key values to the rowid that owns them *) 47 + type unique_index = { 48 + ui_columns : int list; (* column indices *) 49 + ui_name : string; (* for error messages *) 50 + ui_keys : (value list, int64) Hashtbl.t; 51 + } 52 + 53 + type generic_table = { 54 + g_btree : Btree.Table.t; 55 + g_schema : schema; 56 + mutable g_unique_indexes : unique_index list; 57 + } 45 58 46 59 type t = { 47 60 pager : Btree.Pager.t; ··· 239 252 col_is_rowid_alias = is_rowid_alias; 240 253 } 241 254 255 + (* Extract column names from a table-level UNIQUE(...) constraint *) 256 + let parse_unique_constraint s = 257 + let upper = String.uppercase_ascii (String.trim s) in 258 + let starts_with prefix = 259 + String.length upper >= String.length prefix 260 + && String.sub upper 0 (String.length prefix) = prefix 261 + in 262 + if starts_with "UNIQUE(" || starts_with "UNIQUE " then 263 + match String.index_opt s '(' with 264 + | None -> None 265 + | Some start -> ( 266 + match matching_paren s start with 267 + | None -> None 268 + | Some end_ -> 269 + let inner = String.sub s (start + 1) (end_ - start - 1) in 270 + let cols = String.split_on_char ',' inner |> List.map String.trim in 271 + Some cols) 272 + else None 273 + 274 + (* Check if a column definition has a UNIQUE constraint *) 275 + let column_has_unique s = 276 + if is_table_constraint s then false 277 + else 278 + let tokens = tokenize s in 279 + match tokens with 280 + | [] -> false 281 + | _ :: rest -> 282 + List.exists (fun t -> String.uppercase_ascii t = "UNIQUE") rest 283 + 242 284 let parse_create_table sql = 243 285 match String.index_opt sql '(' with 244 286 | None -> [] ··· 250 292 let parts = split_respecting_parens body in 251 293 List.filter_map parse_column_def parts) 252 294 295 + (* Parse all unique constraints from a CREATE TABLE statement. 296 + Returns a list of column-name lists. *) 297 + let parse_unique_constraints sql columns = 298 + match String.index_opt sql '(' with 299 + | None -> [] 300 + | Some start -> ( 301 + match matching_paren sql start with 302 + | None -> [] 303 + | Some body_end -> 304 + let body = String.sub sql (start + 1) (body_end - start - 1) in 305 + let parts = split_respecting_parens body in 306 + (* Table-level UNIQUE(...) constraints *) 307 + let table_level = List.filter_map parse_unique_constraint parts in 308 + (* Column-level UNIQUE constraints *) 309 + let column_level = 310 + List.filter_map 311 + (fun part -> 312 + if column_has_unique part then 313 + match tokenize part with 314 + | name :: _ -> Some [ name ] 315 + | [] -> None 316 + else None) 317 + parts 318 + in 319 + (* Deduplicate: don't add column-level if already in table-level *) 320 + let all = table_level @ column_level in 321 + (* Resolve column names to indices *) 322 + let col_names = List.map (fun c -> c.col_name) columns in 323 + List.filter_map 324 + (fun constraint_cols -> 325 + let indices = 326 + List.filter_map 327 + (fun name -> 328 + let rec find_idx i = function 329 + | [] -> None 330 + | n :: _ when n = name -> Some i 331 + | _ :: rest -> find_idx (i + 1) rest 332 + in 333 + find_idx 0 col_names) 334 + constraint_cols 335 + in 336 + if List.length indices = List.length constraint_cols then 337 + let name = String.concat ", " constraint_cols in 338 + Some (indices, name) 339 + else None) 340 + all) 341 + 342 + (* Find the index of the rowid alias column, if any *) 343 + let rowid_alias_index columns = 344 + let rec find i = function 345 + | [] -> None 346 + | c :: _ when c.col_is_rowid_alias -> Some i 347 + | _ :: rest -> find (i + 1) rest 348 + in 349 + find 0 columns 350 + 351 + (* Apply rowid substitution and trailing Vnull padding *) 352 + let fixup_values ~schema ~rowid values = 353 + let n_cols = List.length schema.columns in 354 + let len = List.length values in 355 + let values = 356 + if len < n_cols then 357 + values @ List.init (n_cols - len) (fun _ -> Btree.Record.Vnull) 358 + else values 359 + in 360 + match rowid_alias_index schema.columns with 361 + | None -> values 362 + | Some idx -> 363 + List.mapi 364 + (fun i v -> 365 + if i = idx then 366 + match v with 367 + | Btree.Record.Vnull -> Btree.Record.Vint rowid 368 + | v -> v 369 + else v) 370 + values 371 + 372 + (* Build empty unique indexes from parsed constraints *) 373 + let make_unique_indexes constraints = 374 + List.map 375 + (fun (indices, name) -> 376 + { ui_columns = indices; ui_name = name; ui_keys = Hashtbl.create 64 }) 377 + constraints 378 + 379 + (* Extract composite key values for a unique index from a row *) 380 + let extract_key ui values = 381 + List.map 382 + (fun idx -> if idx < List.length values then List.nth values idx else Vnull) 383 + ui.ui_columns 384 + 385 + (* Populate unique indexes by scanning existing rows *) 386 + let populate_unique_indexes btree schema indexes = 387 + if indexes <> [] then 388 + Btree.Table.iter btree (fun rowid payload -> 389 + let values = Btree.Record.decode payload in 390 + let values = fixup_values ~schema ~rowid values in 391 + List.iter 392 + (fun ui -> 393 + let key = extract_key ui values in 394 + Hashtbl.replace ui.ui_keys key rowid) 395 + indexes) 396 + 253 397 (* Standard kv table schema *) 254 398 let kv_columns = 255 399 [ ··· 388 532 let _page1 = Btree.Pager.allocate pager in 389 533 (* Create kv data table on page 2 *) 390 534 let kv = new_kv_table pager in 391 - let gt = { g_btree = kv.btree; g_schema = kv_schema "kv" } in 535 + let gt = 536 + { g_btree = kv.btree; g_schema = kv_schema "kv"; g_unique_indexes = [] } 537 + in 392 538 let t = 393 539 { 394 540 pager; ··· 399 545 } 400 546 in 401 547 rebuild_page1 t; 548 + Btree.Pager.sync pager; 402 549 t 403 550 404 551 (* Extract named kv tables (non-kv tables with kv schema) from all_tables *) ··· 422 569 let pager = Btree.Pager.mem ~page_size () in 423 570 let _page1 = Btree.Pager.allocate pager in 424 571 let kv = new_kv_table pager in 425 - let gt = { g_btree = kv.btree; g_schema = kv_schema "kv" } in 572 + let gt = 573 + { g_btree = kv.btree; g_schema = kv_schema "kv"; g_unique_indexes = [] } 574 + in 426 575 let t = 427 576 { 428 577 pager; ··· 442 591 (f :> Eio.File.rw_ty Eio.Resource.t) 443 592 in 444 593 let pager = Btree.Pager.v ~page_size file in 445 - if Btree.Pager.page_count pager = 0 then failwith "Database file is empty"; 594 + if Btree.Pager.page_count pager = 0 then 595 + failwith "Database file exists but is empty (delete it to recreate)"; 446 596 (* Read page 1 and validate *) 447 597 let page1 = Btree.Pager.read pager 1 in 448 598 if String.sub page1 0 16 <> magic then failwith "Not a SQLite database"; ··· 475 625 let btree = Btree.Table.open_ pager ~root_page:root in 476 626 let columns = parse_create_table sql in 477 627 let schema = { tbl_name = name; columns; sql } in 478 - { g_btree = btree; g_schema = schema }) 628 + let constraints = parse_unique_constraints sql columns in 629 + let indexes = make_unique_indexes constraints in 630 + populate_unique_indexes btree schema indexes; 631 + { g_btree = btree; g_schema = schema; g_unique_indexes = indexes }) 479 632 raw_tables 480 633 in 481 634 (* Try to find "kv" table for backward compat *) ··· 566 719 | Some gt -> gt 567 720 | None -> Fmt.failwith "No table %S found in database" name 568 721 569 - (* Find the index of the rowid alias column, if any *) 570 - let rowid_alias_index columns = 571 - let rec find i = function 572 - | [] -> None 573 - | c :: _ when c.col_is_rowid_alias -> Some i 574 - | _ :: rest -> find (i + 1) rest 575 - in 576 - find 0 columns 577 - 578 - (* Apply rowid substitution and trailing Vnull padding *) 579 - let fixup_values ~schema ~rowid values = 580 - let n_cols = List.length schema.columns in 581 - let len = List.length values in 582 - let values = 583 - if len < n_cols then 584 - values @ List.init (n_cols - len) (fun _ -> Btree.Record.Vnull) 585 - else values 586 - in 587 - match rowid_alias_index schema.columns with 588 - | None -> values 589 - | Some idx -> 590 - List.mapi 591 - (fun i v -> 592 - if i = idx then 593 - match v with 594 - | Btree.Record.Vnull -> Btree.Record.Vint rowid 595 - | v -> v 596 - else v) 597 - values 598 - 599 722 let iter_table t name ~f = 600 723 let gt = table t name in 601 724 let schema = gt.g_schema in ··· 661 784 in 662 785 let btree = Btree.Table.v t.pager in 663 786 let schema = { tbl_name = name; columns; sql } in 664 - let gt = { g_btree = btree; g_schema = schema } in 787 + let constraints = parse_unique_constraints sql columns in 788 + let indexes = make_unique_indexes constraints in 789 + let gt = { g_btree = btree; g_schema = schema; g_unique_indexes = indexes } in 665 790 t.all_tables <- t.all_tables @ [ gt ] 666 791 667 792 let next_rowid_for t name = ··· 707 832 in 708 833 (rowid, record_values) 709 834 in 835 + (* Check unique constraints before inserting *) 836 + let full_values = fixup_values ~schema:gt.g_schema ~rowid values in 837 + List.iter 838 + (fun ui -> 839 + let key = extract_key ui full_values in 840 + if Hashtbl.mem ui.ui_keys key then raise (Unique_violation ui.ui_name)) 841 + gt.g_unique_indexes; 710 842 let record = Btree.Record.encode record_values in 711 843 Btree.Table.insert gt.g_btree ~rowid record; 844 + (* Update unique indexes *) 845 + List.iter 846 + (fun ui -> 847 + let key = extract_key ui full_values in 848 + Hashtbl.replace ui.ui_keys key rowid) 849 + gt.g_unique_indexes; 712 850 rowid 713 851 714 852 (* Namespaced Tables *) ··· 737 875 | Some kv -> { parent; name; kv } 738 876 | None -> 739 877 let kv = new_kv_table parent.pager in 740 - let gt = { g_btree = kv.btree; g_schema = kv_schema name } in 878 + let gt = 879 + { 880 + g_btree = kv.btree; 881 + g_schema = kv_schema name; 882 + g_unique_indexes = []; 883 + } 884 + in 741 885 parent.named_tables <- (name, kv) :: parent.named_tables; 742 886 parent.all_tables <- parent.all_tables @ [ gt ]; 743 887 { parent; name; kv }

+8 -2

lib/sqlite.mli

··· 117 117 118 118 Create arbitrary tables and insert rows. *) 119 119 120 + exception Unique_violation of string 121 + (** Raised by {!insert} when a row would violate a [UNIQUE] constraint. The 122 + string names the constrained columns (e.g. ["provider, provider_uid"]). *) 123 + 120 124 val create_table : t -> sql:string -> unit 121 125 (** [create_table t ~sql] creates a new table from a CREATE TABLE statement. The 122 126 SQL is stored in sqlite_master and the column definitions are parsed for 123 - schema metadata. *) 127 + schema metadata. [UNIQUE] constraints (both column-level and table-level) 128 + are parsed and enforced on subsequent {!insert} calls. *) 124 129 125 130 val insert : t -> table:string -> value list -> int64 126 131 (** [insert t ~table values] inserts a row into [table] with the given column ··· 130 135 is [Vint n], the row is inserted with rowid [n]. If the value is [Vnull], 131 136 the rowid is auto-assigned. 132 137 133 - @raise Failure if the table doesn't exist. *) 138 + @raise Failure if the table doesn't exist. 139 + @raise Unique_violation if the row violates a [UNIQUE] constraint. *) 134 140 135 141 (** {1 Schema Parsing} *) 136 142

+115

test/test_sqlite.ml

··· 1013 1013 Alcotest.fail "should have raised" 1014 1014 with Failure _ -> () 1015 1015 1016 + (* ── Unique constraint tests ─────────────────────────────────────── *) 1017 + 1018 + let test_unique_column_level () = 1019 + with_temp_db @@ fun _fs db -> 1020 + Sqlite.create_table db 1021 + ~sql:"CREATE TABLE t (id INTEGER PRIMARY KEY, email TEXT UNIQUE)"; 1022 + let _ = 1023 + Sqlite.insert db ~table:"t" [ Sqlite.Vnull; Sqlite.Vtext "a@b.com" ] 1024 + in 1025 + (try 1026 + let _ = 1027 + Sqlite.insert db ~table:"t" [ Sqlite.Vnull; Sqlite.Vtext "a@b.com" ] 1028 + in 1029 + Alcotest.fail "should have raised Unique_violation" 1030 + with Sqlite.Unique_violation cols -> 1031 + Alcotest.(check string) "column name" "email" cols); 1032 + (* Different email should succeed *) 1033 + let _ = 1034 + Sqlite.insert db ~table:"t" [ Sqlite.Vnull; Sqlite.Vtext "c@d.com" ] 1035 + in 1036 + () 1037 + 1038 + let test_unique_table_level () = 1039 + with_temp_db @@ fun _fs db -> 1040 + Sqlite.create_table db 1041 + ~sql: 1042 + "CREATE TABLE t (id INTEGER PRIMARY KEY, provider TEXT, uid TEXT, \ 1043 + UNIQUE(provider, uid))"; 1044 + let _ = 1045 + Sqlite.insert db ~table:"t" 1046 + [ Sqlite.Vnull; Sqlite.Vtext "github"; Sqlite.Vtext "123" ] 1047 + in 1048 + (try 1049 + let _ = 1050 + Sqlite.insert db ~table:"t" 1051 + [ Sqlite.Vnull; Sqlite.Vtext "github"; Sqlite.Vtext "123" ] 1052 + in 1053 + Alcotest.fail "should have raised Unique_violation" 1054 + with Sqlite.Unique_violation cols -> 1055 + Alcotest.(check string) "columns" "provider, uid" cols); 1056 + (* Same provider, different uid should succeed *) 1057 + let _ = 1058 + Sqlite.insert db ~table:"t" 1059 + [ Sqlite.Vnull; Sqlite.Vtext "github"; Sqlite.Vtext "456" ] 1060 + in 1061 + (* Different provider, same uid should succeed *) 1062 + let _ = 1063 + Sqlite.insert db ~table:"t" 1064 + [ Sqlite.Vnull; Sqlite.Vtext "google"; Sqlite.Vtext "123" ] 1065 + in 1066 + () 1067 + 1068 + let test_unique_composite () = 1069 + with_temp_db @@ fun _fs db -> 1070 + Sqlite.create_table db 1071 + ~sql:"CREATE TABLE t (a TEXT, b TEXT, c TEXT, UNIQUE(a, b))"; 1072 + let _ = 1073 + Sqlite.insert db ~table:"t" 1074 + [ Sqlite.Vtext "x"; Sqlite.Vtext "y"; Sqlite.Vtext "z1" ] 1075 + in 1076 + try 1077 + let _ = 1078 + Sqlite.insert db ~table:"t" 1079 + [ Sqlite.Vtext "x"; Sqlite.Vtext "y"; Sqlite.Vtext "z2" ] 1080 + in 1081 + Alcotest.fail "should have raised Unique_violation" 1082 + with Sqlite.Unique_violation _ -> () 1083 + 1084 + let test_unique_allows_distinct () = 1085 + with_temp_db @@ fun _fs db -> 1086 + Sqlite.create_table db ~sql:"CREATE TABLE t (name TEXT UNIQUE, age INTEGER)"; 1087 + let _ = 1088 + Sqlite.insert db ~table:"t" [ Sqlite.Vtext "alice"; Sqlite.Vint 30L ] 1089 + in 1090 + let _ = Sqlite.insert db ~table:"t" [ Sqlite.Vtext "bob"; Sqlite.Vint 30L ] in 1091 + let rows = Sqlite.read_table db "t" in 1092 + Alcotest.(check int) "two rows" 2 (List.length rows) 1093 + 1094 + let test_unique_persists () = 1095 + Eio_main.run @@ fun env -> 1096 + let cwd = Eio.Stdenv.cwd env in 1097 + let tmp_dir = Eio.Path.(cwd / "_build" / "test_sqlite") in 1098 + (try Eio.Path.mkdirs ~exists_ok:true ~perm:0o755 tmp_dir with Eio.Io _ -> ()); 1099 + let path = 1100 + Eio.Path.(tmp_dir / Fmt.str "unique_%d.db" (Random.int 1_000_000)) 1101 + in 1102 + (* Create and insert *) 1103 + Eio.Switch.run (fun sw -> 1104 + let db = Sqlite.open_ ~sw ~create:true path in 1105 + Sqlite.create_table db 1106 + ~sql:"CREATE TABLE t (id INTEGER PRIMARY KEY, email TEXT UNIQUE)"; 1107 + let _ = 1108 + Sqlite.insert db ~table:"t" [ Sqlite.Vnull; Sqlite.Vtext "a@b.com" ] 1109 + in 1110 + Sqlite.close db); 1111 + (* Reopen and verify constraint is enforced *) 1112 + Eio.Switch.run (fun sw -> 1113 + let db = Sqlite.open_ ~sw ~create:false path in 1114 + (try 1115 + let _ = 1116 + Sqlite.insert db ~table:"t" [ Sqlite.Vnull; Sqlite.Vtext "a@b.com" ] 1117 + in 1118 + Alcotest.fail "should have raised Unique_violation after reopen" 1119 + with Sqlite.Unique_violation _ -> ()); 1120 + Sqlite.close db) 1121 + 1016 1122 let suite = 1017 1123 ( "sqlite", 1018 1124 List.concat ··· 1130 1236 test_insert_coexists_with_kv; 1131 1237 Alcotest.test_case "insert nonexistent" `Quick 1132 1238 test_insert_nonexistent_table; 1239 + ]; 1240 + [ 1241 + Alcotest.test_case "unique column-level" `Quick 1242 + test_unique_column_level; 1243 + Alcotest.test_case "unique table-level" `Quick test_unique_table_level; 1244 + Alcotest.test_case "unique composite" `Quick test_unique_composite; 1245 + Alcotest.test_case "unique allows distinct" `Quick 1246 + test_unique_allows_distinct; 1247 + Alcotest.test_case "unique persists" `Quick test_unique_persists; 1133 1248 ]; 1134 1249 ] )

Configure Feed

Configure Feed