slave_framework/database/
migrations.rs1use std::path::Path;
12
13use surrealdb::types::SurrealValue;
14
15use super::Db;
16use crate::{Error, Result};
17
18pub const DIR: &str = "database/migrations";
20
21pub const TABLE: &str = "_migration";
23
24#[derive(Debug, Clone)]
26pub struct Migration {
27 pub name: String,
29 pub source: String,
30}
31
32pub 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
59pub 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
70pub 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
82pub 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
117pub 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
140pub 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
214fn 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}