repodb / lib / migration.ml

type column_def = {
  col_name : string;
  col_type : string;
  col_nullable : bool;
  col_primary_key : bool;
  col_unique : bool;
  col_default : string option;
  col_references : (string * string) option;
}

type index_def = {
  idx_name : string option;
  idx_table : string;
  idx_columns : string list;
  idx_unique : bool;
}

type operation =
  | Create_table of { name : string; columns : column_def list }
  | Drop_table of string
  | Alter_table of { name : string; changes : alter_change list }
  | Create_index of index_def
  | Drop_index of string
  | Execute of string

and alter_change =
  | Add_column of column_def
  | Drop_column of string
  | Rename_column of { from : string; to_ : string }
  | Alter_column of {
      name : string;
      new_type : string option;
      new_nullable : bool option;
    }

type t = {
  version : int64;
  name : string;
  up : operation list;
  down : operation list;
}

let column ?(nullable = true) ?(primary_key = false) ?(unique = false) ?default
    ?references name sql_type =
  {
    col_name = name;
    col_type = sql_type;
    col_nullable = nullable;
    col_primary_key = primary_key;
    col_unique = unique;
    col_default = default;
    col_references = references;
  }

let create_table name columns = Create_table { name; columns }
let drop_table name = Drop_table name
let add_column col = Add_column col
let drop_column name = Drop_column name
let rename_column ~from ~to_ = Rename_column { from; to_ }
let alter_table name changes = Alter_table { name; changes }

let create_index ?(unique = false) ?name table columns =
  Create_index
    {
      idx_name = name;
      idx_table = table;
      idx_columns = columns;
      idx_unique = unique;
    }

let drop_index name = Drop_index name
let execute sql = Execute sql

let timestamps () =
  [
    column "inserted_at" "TIMESTAMPTZ" ~nullable:false ~default:"NOW()";
    column "updated_at" "TIMESTAMPTZ" ~nullable:false ~default:"NOW()";
  ]

let migration ~version ~name ~up ~down = { version; name; up; down }

let column_to_sql c =
  let parts = [ c.col_name; c.col_type ] in
  let parts = if c.col_nullable then parts else parts @ [ "NOT NULL" ] in
  let parts = if c.col_primary_key then parts @ [ "PRIMARY KEY" ] else parts in
  let parts = if c.col_unique then parts @ [ "UNIQUE" ] else parts in
  let parts =
    match c.col_default with
    | Some d -> parts @ [ "DEFAULT " ^ d ]
    | None -> parts
  in
  let parts =
    match c.col_references with
    | Some (table, col) ->
        parts @ [ Printf.sprintf "REFERENCES %s(%s)" table col ]
    | None -> parts
  in
  String.concat " " parts

let operation_to_sql = function
  | Create_table { name; columns } ->
      let col_defs = List.map column_to_sql columns in
      Printf.sprintf "CREATE TABLE %s (\n  %s\n)" name
        (String.concat ",\n  " col_defs)
  | Drop_table name -> Printf.sprintf "DROP TABLE %s" name
  | Alter_table { name; changes } ->
      let change_strs =
        List.map
          (function
            | Add_column c -> Printf.sprintf "ADD COLUMN %s" (column_to_sql c)
            | Drop_column col -> Printf.sprintf "DROP COLUMN %s" col
            | Rename_column { from; to_ } ->
                Printf.sprintf "RENAME COLUMN %s TO %s" from to_
            | Alter_column { name = col; new_type; new_nullable = _ } -> (
                match new_type with
                | Some t -> Printf.sprintf "ALTER COLUMN %s TYPE %s" col t
                | None -> ""))
          changes
      in
      Printf.sprintf "ALTER TABLE %s %s" name (String.concat ", " change_strs)
  | Create_index { idx_name; idx_table; idx_columns; idx_unique } ->
      let name =
        match idx_name with
        | Some n -> n
        | None ->
            Printf.sprintf "%s_%s_index" idx_table
              (String.concat "_" idx_columns)
      in
      let unique_str = if idx_unique then "UNIQUE " else "" in
      Printf.sprintf "CREATE %sINDEX %s ON %s (%s)" unique_str name idx_table
        (String.concat ", " idx_columns)
  | Drop_index name -> Printf.sprintf "DROP INDEX %s" name
  | Execute sql -> sql

let schema_migrations_table =
  Create_table
    {
      name = "schema_migrations";
      columns =
        [
          column "version" "BIGINT" ~nullable:false ~primary_key:true;
          column "name" "VARCHAR(255)" ~nullable:false;
          column "inserted_at" "TIMESTAMPTZ" ~nullable:false ~default:"NOW()";
        ];
    }

let create_schema_migrations_sql = operation_to_sql schema_migrations_table

let insert_migration_sql =
  "INSERT INTO schema_migrations (version, name) VALUES ($1, $2)"

let delete_migration_sql = "DELETE FROM schema_migrations WHERE version = $1"

let get_applied_versions_sql =
  "SELECT version FROM schema_migrations ORDER BY version"

let get_migration_status_sql =
  "SELECT version, name, inserted_at FROM schema_migrations ORDER BY version"

type migration_status = {
  applied_versions : int64 list;
  pending : t list;
  last_applied : int64 option;
}

let sort_migrations migrations =
  List.sort (fun m1 m2 -> Int64.compare m1.version m2.version) migrations

let pending_migrations ~applied_versions migrations =
  let is_applied v = List.mem v applied_versions in
  migrations
  |> List.filter (fun m -> not (is_applied m.version))
  |> sort_migrations

let generate_up_sql migration = List.map operation_to_sql migration.up
let generate_down_sql migration = List.map operation_to_sql migration.down

type migration_action =
  | Migrate of { sql : string list; version : int64; name : string }
  | Rollback of { sql : string list; version : int64 }
  | CreateSchemaTable of string
  | RecordMigration of { version : int64; name : string }
  | RemoveMigration of int64

let plan_migrate ~applied_versions ~target migrations =
  let pending = pending_migrations ~applied_versions migrations in
  let to_run =
    match target with
    | None -> pending
    | Some v -> List.filter (fun m -> m.version <= v) pending
  in
  to_run
  |> List.map (fun m ->
      [
        Migrate { sql = generate_up_sql m; version = m.version; name = m.name };
        RecordMigration { version = m.version; name = m.name };
      ])
  |> List.flatten

let plan_rollback ~applied_versions ~step migrations =
  let applied_sorted =
    List.sort (fun a b -> Int64.compare b a) applied_versions
  in
  let to_rollback =
    match step with
    | None -> ( match applied_sorted with v :: _ -> [ v ] | [] -> [])
    | Some n ->
        let rec take n lst =
          match (n, lst) with
          | 0, _ | _, [] -> []
          | n, x :: xs -> x :: take (n - 1) xs
        in
        take n applied_sorted
  in
  to_rollback
  |> List.filter_map (fun v ->
      match List.find_opt (fun m -> m.version = v) migrations with
      | Some m ->
          Some
            [
              Rollback { sql = generate_down_sql m; version = m.version };
              RemoveMigration m.version;
            ]
      | None -> None)
  |> List.flatten

let action_to_sql = function
  | CreateSchemaTable sql -> [ sql ]
  | Migrate { sql; _ } -> sql
  | Rollback { sql; _ } -> sql
  | RecordMigration { version; name } ->
      [
        Printf.sprintf
          "INSERT INTO schema_migrations (version, name) VALUES (%Ld, '%s')"
          version name;
      ]
  | RemoveMigration version ->
      [
        Printf.sprintf "DELETE FROM schema_migrations WHERE version = %Ld"
          version;
      ]

let actions_to_sql actions = actions |> List.map action_to_sql |> List.flatten

let format_status ~applied_versions ~migrations =
  let pending = pending_migrations ~applied_versions migrations in
  let buf = Buffer.create 256 in
  Buffer.add_string buf "Applied migrations:\n";
  applied_versions
  |> List.iter (fun v ->
      match List.find_opt (fun m -> m.version = v) migrations with
      | Some m ->
          Buffer.add_string buf (Printf.sprintf "  [✓] %Ld: %s\n" v m.name)
      | None ->
          Buffer.add_string buf (Printf.sprintf "  [✓] %Ld: (unknown)\n" v));
  Buffer.add_string buf "\nPending migrations:\n";
  pending
  |> List.iter (fun m ->
      Buffer.add_string buf (Printf.sprintf "  [ ] %Ld: %s\n" m.version m.name));
  Buffer.contents buf