Skip to main content

slave_framework/database/
migrations.rs

1//! Schema migrations: SurrealQL files applied in order.
2//!
3//! Each migration is a `.surql` file in [`DIR`], named
4//! `<YYYYMMDD_HHMMSS>_<name>.surql` so that names sort in the order they
5//! were written. A migration runs in one transaction together with its
6//! record in the [`TABLE`] table, so it is applied completely or not at
7//! all, and never twice. Migrations must not contain `BEGIN`/`COMMIT`.
8//!
9//! Migrations only move forward: to undo one, write another.
10
11use std::path::Path;
12
13use surrealdb::types::SurrealValue;
14
15use super::Db;
16use crate::{Error, Result};
17
18/// Where migrations live, relative to the project root.
19pub const DIR: &str = "database/migrations";
20
21/// The table recording applied migrations.
22pub const TABLE: &str = "_migration";
23
24/// A migration file.
25#[derive(Debug, Clone)]
26pub struct Migration {
27    /// The file stem, e.g. `20260101_120000_create_posts`.
28    pub name: String,
29    pub source: String,
30}
31
32/// Every `.surql` file in `dir`, in name order.
33pub fn load(dir: &Path) -> Result<Vec<Migration>> {
34    let entries = std::fs::read_dir(dir)
35        .map_err(|err| Error::Migration(format!("cannot read {}: {err}", dir.display())))?;
36    let mut migrations = Vec::new();
37    for path in entries.flatten().map(|entry| entry.path()) {
38        if path.extension().is_none_or(|ext| ext != "surql") {
39            continue;
40        }
41        let name = path
42            .file_stem()
43            .unwrap_or_default()
44            .to_string_lossy()
45            .into_owned();
46        let source = std::fs::read_to_string(&path)
47            .map_err(|err| Error::Migration(format!("cannot read {}: {err}", path.display())))?;
48        if mentions_transaction(&source) {
49            return Err(Error::Migration(format!(
50                "{name}: remove BEGIN/COMMIT; every migration already runs in a transaction"
51            )));
52        }
53        migrations.push(Migration { name, source });
54    }
55    migrations.sort_by(|a, b| a.name.cmp(&b.name));
56    Ok(migrations)
57}
58
59/// The names of the applied migrations, in order.
60pub async fn applied(db: &Db) -> Result<Vec<String>> {
61    ensure_table(db).await?;
62    let names: Vec<String> = db
63        .query(format!("SELECT VALUE name FROM {TABLE} ORDER BY name"))
64        .await?
65        .check()?
66        .take(0)?;
67    Ok(names)
68}
69
70/// Each migration in `dir` and whether it has been applied.
71pub async fn status(db: &Db, dir: &Path) -> Result<Vec<(Migration, bool)>> {
72    let applied = applied(db).await?;
73    Ok(load(dir)?
74        .into_iter()
75        .map(|migration| {
76            let done = applied.contains(&migration.name);
77            (migration, done)
78        })
79        .collect())
80}
81
82/// Applies the pending migrations in `dir`, in order, and returns their names.
83pub async fn apply(db: &Db, dir: &Path) -> Result<Vec<String>> {
84    let mut done = Vec::new();
85    for (migration, applied) in status(db, dir).await? {
86        if applied {
87            continue;
88        }
89        let statement = format!(
90            "BEGIN TRANSACTION;\n{}\n;\nCREATE {TABLE} SET name = $name, applied_at = time::now();\nCOMMIT TRANSACTION;",
91            migration.source
92        );
93        db.query(statement)
94            .bind(("name", migration.name.clone()))
95            .await?
96            .check()
97            .map_err(|err| Error::Migration(format!("{}: {err}", migration.name)))?;
98        done.push(migration.name);
99    }
100    Ok(done)
101}
102
103#[derive(Debug, Default, SurrealValue)]
104struct DatabaseInfo {
105    #[surreal(default)]
106    tables: std::collections::BTreeMap<String, String>,
107    #[surreal(default)]
108    sequences: std::collections::BTreeMap<String, String>,
109    #[surreal(default)]
110    functions: std::collections::BTreeMap<String, String>,
111    #[surreal(default)]
112    params: std::collections::BTreeMap<String, String>,
113    #[surreal(default)]
114    analyzers: std::collections::BTreeMap<String, String>,
115}
116
117/// Removes every table (with its data), sequence, function, param and
118/// analyzer from the database, including the record of applied migrations.
119pub async fn reset(db: &Db) -> Result<()> {
120    let info: Option<DatabaseInfo> = db.query("INFO FOR DB").await?.check()?.take(0)?;
121    let info = info.unwrap_or_default();
122    let mut statements = Vec::new();
123    for (prefix, names) in [
124        ("TABLE ", &info.tables),
125        ("SEQUENCE ", &info.sequences),
126        ("FUNCTION fn::", &info.functions),
127        ("PARAM $", &info.params),
128        ("ANALYZER ", &info.analyzers),
129    ] {
130        for name in names.keys() {
131            statements.push(format!("REMOVE {prefix}{name};"));
132        }
133    }
134    if !statements.is_empty() {
135        db.query(statements.join("\n")).await?.check()?;
136    }
137    Ok(())
138}
139
140/// The command line of an app's `migrate` binary:
141///
142/// - `migrate` or `migrate apply`: apply pending migrations;
143/// - `migrate status`: list migrations and whether each is applied;
144/// - `migrate fresh [--force]`: drop everything, then apply all migrations
145///   (asks to confirm unless `--force`).
146///
147/// Run from the project root, where [`DIR`] is.
148pub async fn run_cli(db: &Db, args: &[String]) -> Result<()> {
149    let dir = Path::new(DIR);
150    let command = args.first().map(String::as_str).unwrap_or("apply");
151    match command {
152        "apply" => report_applied(apply(db, dir).await?),
153        "status" => {
154            let migrations = status(db, dir).await?;
155            if migrations.is_empty() {
156                println!("No migrations in {DIR}.");
157            }
158            for (migration, applied) in migrations {
159                let state = if applied { "applied" } else { "pending" };
160                println!("{state:<8} {}", migration.name);
161            }
162        }
163        "fresh" => {
164            let force = args.iter().any(|arg| arg == "--force");
165            if !force && !confirm("Drop every table in this database and migrate again?")? {
166                println!("Cancelled.");
167                return Ok(());
168            }
169            reset(db).await?;
170            println!("Dropped every table.");
171            report_applied(apply(db, dir).await?);
172        }
173        other => {
174            return Err(Error::Migration(format!(
175                "unknown command `{other}`; use apply, status or fresh [--force]"
176            )));
177        }
178    }
179    Ok(())
180}
181
182fn report_applied(names: Vec<String>) {
183    if names.is_empty() {
184        println!("Nothing to migrate.");
185    }
186    for name in names {
187        println!("Applied {name}");
188    }
189}
190
191fn confirm(question: &str) -> Result<bool> {
192    use std::io::Write;
193    print!("{question} [y/N] ");
194    std::io::stdout().flush().map_err(Error::Server)?;
195    let mut answer = String::new();
196    std::io::stdin()
197        .read_line(&mut answer)
198        .map_err(Error::Server)?;
199    Ok(matches!(answer.trim(), "y" | "Y" | "yes"))
200}
201
202async fn ensure_table(db: &Db) -> Result<()> {
203    db.query(format!(
204        "DEFINE TABLE IF NOT EXISTS {TABLE} SCHEMAFULL;
205         DEFINE FIELD IF NOT EXISTS name ON {TABLE} TYPE string;
206         DEFINE FIELD IF NOT EXISTS applied_at ON {TABLE} TYPE datetime;
207         DEFINE INDEX IF NOT EXISTS {TABLE}_name ON {TABLE} FIELDS name UNIQUE;"
208    ))
209    .await?
210    .check()?;
211    Ok(())
212}
213
214/// Whether `source` has a `BEGIN` or `COMMIT` statement (outside comments).
215fn mentions_transaction(source: &str) -> bool {
216    source
217        .lines()
218        .map(|line| line.trim_start().to_ascii_uppercase())
219        .filter(|line| !line.starts_with("--") && !line.starts_with("//") && !line.starts_with('#'))
220        .flat_map(|line| {
221            line.split(';')
222                .map(|statement| statement.trim().to_owned())
223                .collect::<Vec<_>>()
224        })
225        .any(|statement| statement.starts_with("BEGIN") || statement.starts_with("COMMIT"))
226}
227
228#[cfg(test)]
229mod tests {
230    use super::*;
231
232    #[test]
233    fn spots_transactions_outside_comments() {
234        assert!(mentions_transaction("BEGIN TRANSACTION;\nDEFINE TABLE x;"));
235        assert!(mentions_transaction("DEFINE TABLE x; COMMIT;"));
236        assert!(!mentions_transaction("-- BEGIN here\nDEFINE TABLE x;"));
237        assert!(!mentions_transaction("DEFINE TABLE begin_at;"));
238    }
239}