repodb / lib / schema.ml

type table = { name : string; schema : string option }

let table ?(schema = None) name = { name; schema }

let table_name t =
  match t.schema with Some s -> s ^ "." ^ t.name | None -> t.name

type fk_action = Cascade | Restrict | SetNull | SetDefault | NoAction

type foreign_key_ref = {
  fk_table : string;
  fk_column : string;
  fk_on_delete : fk_action option;
  fk_on_update : fk_action option;
}

type column_constraint =
  | PrimaryKey
  | NotNull
  | Unique
  | Default of string
  | Check of string
  | ForeignKey of foreign_key_ref

type 'a column = {
  col_name : string;
  col_type : 'a Types.t;
  col_constraints : column_constraint list;
}

type wrapped_column = Column : 'a column -> wrapped_column

type table_def = {
  tbl_table : table;
  tbl_columns : wrapped_column list;
  tbl_primary_key : string list option;
  tbl_unique : string list list;
  tbl_checks : (string option * string) list;
}

let column ?(primary_key = false) ?(not_null = false) ?(unique = false) ?default
    ?check ?references name ty =
  let constraints = [] in
  let constraints =
    if primary_key then PrimaryKey :: constraints else constraints
  in
  let constraints = if not_null then NotNull :: constraints else constraints in
  let constraints = if unique then Unique :: constraints else constraints in
  let constraints =
    match default with
    | Some d -> Default d :: constraints
    | None -> constraints
  in
  let constraints =
    match check with Some c -> Check c :: constraints | None -> constraints
  in
  let constraints =
    match references with
    | Some r -> ForeignKey r :: constraints
    | None -> constraints
  in
  { col_name = name; col_type = ty; col_constraints = constraints }

let references ?(on_delete = None) ?(on_update = None) ~table ~column () =
  {
    fk_table = table;
    fk_column = column;
    fk_on_delete = on_delete;
    fk_on_update = on_update;
  }

let define ?(schema = None) name columns =
  let tbl = { name; schema } in
  {
    tbl_table = tbl;
    tbl_columns = List.map (fun c -> Column c) columns;
    tbl_primary_key = None;
    tbl_unique = [];
    tbl_checks = [];
  }

let with_primary_key cols def = { def with tbl_primary_key = Some cols }
let with_unique cols def = { def with tbl_unique = cols :: def.tbl_unique }

let with_check ?name expr def =
  { def with tbl_checks = (name, expr) :: def.tbl_checks }

let id_column () = column "id" Types.int64 ~primary_key:true ~not_null:true

let timestamps () =
  [
    column "inserted_at" Types.ptime ~not_null:true ~default:"NOW()";
    column "updated_at" Types.ptime ~not_null:true ~default:"NOW()";
  ]

let has_constraint c col = List.exists (fun con -> con = c) col.col_constraints
let is_primary_key col = has_constraint PrimaryKey col
let is_not_null col = has_constraint NotNull col || is_primary_key col
let is_unique col = has_constraint Unique col

let get_default col =
  List.find_map (function Default d -> Some d | _ -> None) col.col_constraints

let get_foreign_key col =
  List.find_map
    (function ForeignKey fk -> Some fk | _ -> None)
    col.col_constraints

let fk_action_to_sql = function
  | Cascade -> "CASCADE"
  | Restrict -> "RESTRICT"
  | SetNull -> "SET NULL"
  | SetDefault -> "SET DEFAULT"
  | NoAction -> "NO ACTION"

let constraint_to_sql = function
  | PrimaryKey -> "PRIMARY KEY"
  | NotNull -> "NOT NULL"
  | Unique -> "UNIQUE"
  | Default expr -> "DEFAULT " ^ expr
  | Check expr -> "CHECK (" ^ expr ^ ")"
  | ForeignKey { fk_table; fk_column; fk_on_delete; fk_on_update } -> (
      let base = Printf.sprintf "REFERENCES %s(%s)" fk_table fk_column in
      let base =
        match fk_on_delete with
        | Some od -> base ^ " ON DELETE " ^ fk_action_to_sql od
        | None -> base
      in
      match fk_on_update with
      | Some ou -> base ^ " ON UPDATE " ^ fk_action_to_sql ou
      | None -> base)

let column_to_sql (Column col) =
  let type_name = Types.sql_type_name col.col_type in
  let constraints = List.map constraint_to_sql col.col_constraints in
  String.concat " " (col.col_name :: type_name :: constraints)

let table_def_to_sql def =
  let col_defs = List.map column_to_sql def.tbl_columns in
  let pk_constraint =
    match def.tbl_primary_key with
    | Some cols ->
        [ Printf.sprintf "PRIMARY KEY (%s)" (String.concat ", " cols) ]
    | None -> []
  in
  let unique_constraints =
    List.map
      (fun cols -> Printf.sprintf "UNIQUE (%s)" (String.concat ", " cols))
      def.tbl_unique
  in
  let check_constraints =
    List.map
      (fun (name, expr) ->
        match name with
        | Some n -> Printf.sprintf "CONSTRAINT %s CHECK (%s)" n expr
        | None -> Printf.sprintf "CHECK (%s)" expr)
      def.tbl_checks
  in
  let all_parts =
    col_defs @ pk_constraint @ unique_constraints @ check_constraints
  in
  Printf.sprintf "CREATE TABLE %s (\n  %s\n)" (table_name def.tbl_table)
    (String.concat ",\n  " all_parts)