slave_framework/migration/
migrator.rs

1//! The migration engine: decides what to run and records what ran.
2
3use std::collections::{BTreeMap, HashSet};
4use std::time::{Duration, Instant};
5
6use libsql::{Connection, TransactionBehavior, Value};
7
8use super::{Registered, Schema};
9use crate::{Error, Result};
10
11/// Name of the table that records applied migrations.
12pub const TABLE: &str = "slave_migrations";
13
14/// A migration that completed successfully.
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct Outcome {
17    pub name: String,
18    pub duration: Duration,
19}
20
21/// A migration that failed. Its tracking record was not written (or removed).
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub struct Failure {
24    pub name: String,
25    pub error: String,
26    /// Whether the migration's own changes were rolled back by its transaction.
27    pub reverted: bool,
28}
29
30/// Result of [`Migrator::run`].
31#[derive(Debug, Clone, Default, PartialEq, Eq)]
32pub struct RunReport {
33    /// Batch number assigned to the migrations applied by this run.
34    pub batch: i64,
35    pub applied: Vec<Outcome>,
36    /// Migrations that were already applied and therefore not run again.
37    pub skipped: Vec<String>,
38    /// The migration that stopped the run, if any. Later migrations were not attempted.
39    pub failed: Option<Failure>,
40}
41
42/// Result of [`Migrator::rollback`].
43#[derive(Debug, Clone, Default, PartialEq, Eq)]
44pub struct RollbackReport {
45    /// The batch that was rolled back, or `None` if nothing had been applied.
46    pub batch: Option<i64>,
47    pub rolled_back: Vec<Outcome>,
48    pub failed: Option<Failure>,
49}
50
51/// Where a migration stands relative to the database.
52#[derive(Debug, Clone, PartialEq, Eq)]
53pub enum State {
54    Applied {
55        batch: i64,
56        executed_at: String,
57    },
58    Pending,
59    /// Recorded as applied, but no longer registered in the application.
60    Missing {
61        batch: i64,
62        executed_at: String,
63    },
64}
65
66/// One row of [`Migrator::status`].
67#[derive(Debug, Clone, PartialEq, Eq)]
68pub struct Status {
69    pub name: String,
70    pub state: State,
71}
72
73struct Record {
74    name: String,
75    batch: i64,
76    executed_at: String,
77}
78
79/// Runs, inspects and rolls back migrations on a single connection.
80pub struct Migrator {
81    conn: Connection,
82    migrations: Vec<Registered>,
83}
84
85impl Migrator {
86    /// Creates a migrator for `migrations`, which are run in name order.
87    pub fn new(conn: Connection, mut migrations: Vec<Registered>) -> Result<Self> {
88        migrations.sort_by_key(|m| m.name());
89        if let Some(pair) = migrations.windows(2).find(|w| w[0].name() == w[1].name()) {
90            return Err(Error::DuplicateMigration(pair[0].name().to_owned()));
91        }
92        Ok(Self { conn, migrations })
93    }
94
95    pub fn connection(&self) -> &Connection {
96        &self.conn
97    }
98
99    /// Applies every pending migration, in order, as a new batch.
100    ///
101    /// Stops at the first failure; migrations applied before it stay applied.
102    pub async fn run(&mut self) -> Result<RunReport> {
103        let records = self.records().await?;
104        let applied: HashSet<&str> = records.iter().map(|r| r.name.as_str()).collect();
105        let batch = records.iter().map(|r| r.batch).max().unwrap_or(0) + 1;
106
107        let mut report = RunReport {
108            batch,
109            ..RunReport::default()
110        };
111        let pending: Vec<Registered> = self
112            .migrations
113            .iter()
114            .filter(|m| {
115                let done = applied.contains(m.name());
116                if done {
117                    report.skipped.push(m.name().to_owned());
118                }
119                !done
120            })
121            .copied()
122            .collect();
123
124        for registered in pending {
125            let migration = registered.migration();
126            let started = Instant::now();
127            let track = (
128                format!("INSERT INTO {TABLE} (name, batch) VALUES (?1, ?2)"),
129                vec![
130                    Value::Text(registered.name().to_owned()),
131                    Value::Integer(batch),
132                ],
133            );
134            match self
135                .apply(migration.transactional(), &migration.up(), track)
136                .await
137            {
138                Ok(()) => report.applied.push(Outcome {
139                    name: registered.name().to_owned(),
140                    duration: started.elapsed(),
141                }),
142                Err(err) => {
143                    report.failed = Some(Failure {
144                        name: registered.name().to_owned(),
145                        error: err.to_string(),
146                        reverted: migration.transactional(),
147                    });
148                    break;
149                }
150            }
151        }
152        Ok(report)
153    }
154
155    /// Reverts the most recent batch, newest migration first.
156    ///
157    /// Refuses (changing nothing) if any migration in the batch is irreversible
158    /// or no longer registered.
159    pub async fn rollback(&mut self) -> Result<RollbackReport> {
160        let records = self.records().await?;
161        let Some(batch) = records.iter().map(|r| r.batch).max() else {
162            return Ok(RollbackReport::default());
163        };
164
165        let mut targets: Vec<&Record> = records.iter().filter(|r| r.batch == batch).collect();
166        targets.sort_by(|a, b| b.name.cmp(&a.name));
167
168        let mut plan = Vec::with_capacity(targets.len());
169        let (mut missing, mut irreversible) = (Vec::new(), Vec::new());
170        for record in targets {
171            match self.find(&record.name) {
172                None => missing.push(record.name.clone()),
173                Some(registered) => match registered.migration().down() {
174                    None => irreversible.push(record.name.clone()),
175                    Some(down) => plan.push((registered, down)),
176                },
177            }
178        }
179        if !missing.is_empty() {
180            return Err(Error::RollbackRefused {
181                batch,
182                reason: format!("not registered in this application: {}", missing.join(", ")),
183            });
184        }
185        if !irreversible.is_empty() {
186            return Err(Error::RollbackRefused {
187                batch,
188                reason: format!(
189                    "irreversible (down() returns None): {}",
190                    irreversible.join(", ")
191                ),
192            });
193        }
194
195        let mut report = RollbackReport {
196            batch: Some(batch),
197            ..RollbackReport::default()
198        };
199        for (registered, down) in plan {
200            let transactional = registered.migration().transactional();
201            let started = Instant::now();
202            let track = (
203                format!("DELETE FROM {TABLE} WHERE name = ?1"),
204                vec![Value::Text(registered.name().to_owned())],
205            );
206            match self.apply(transactional, &down, track).await {
207                Ok(()) => report.rolled_back.push(Outcome {
208                    name: registered.name().to_owned(),
209                    duration: started.elapsed(),
210                }),
211                Err(err) => {
212                    report.failed = Some(Failure {
213                        name: registered.name().to_owned(),
214                        error: err.to_string(),
215                        reverted: transactional,
216                    });
217                    break;
218                }
219            }
220        }
221        Ok(report)
222    }
223
224    /// Lists every registered or recorded migration in name order.
225    pub async fn status(&mut self) -> Result<Vec<Status>> {
226        let mut states: BTreeMap<String, State> = self
227            .migrations
228            .iter()
229            .map(|m| (m.name().to_owned(), State::Pending))
230            .collect();
231        for record in self.records().await? {
232            let registered = states.contains_key(&record.name);
233            let (batch, executed_at) = (record.batch, record.executed_at);
234            let state = if registered {
235                State::Applied { batch, executed_at }
236            } else {
237                State::Missing { batch, executed_at }
238            };
239            states.insert(record.name, state);
240        }
241        Ok(states
242            .into_iter()
243            .map(|(name, state)| Status { name, state })
244            .collect())
245    }
246
247    fn find(&self, name: &str) -> Option<Registered> {
248        self.migrations.iter().find(|m| m.name() == name).copied()
249    }
250
251    /// Creates the tracking table if needed and returns its rows.
252    async fn records(&self) -> Result<Vec<Record>> {
253        self.conn
254            .execute(
255                &format!(
256                    "CREATE TABLE IF NOT EXISTS {TABLE} (
257                        id INTEGER PRIMARY KEY,
258                        name TEXT NOT NULL UNIQUE,
259                        batch INTEGER NOT NULL,
260                        executed_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
261                    )"
262                ),
263                (),
264            )
265            .await?;
266
267        let mut rows = self
268            .conn
269            .query(
270                &format!("SELECT name, batch, executed_at FROM {TABLE} ORDER BY name"),
271                (),
272            )
273            .await?;
274        let mut records = Vec::new();
275        while let Some(row) = rows.next().await? {
276            records.push(Record {
277                name: row.get(0)?,
278                batch: row.get(1)?,
279                executed_at: row.get(2)?,
280            });
281        }
282        Ok(records)
283    }
284
285    /// Executes `schema` followed by the tracking statement, atomically when `transactional`.
286    async fn apply(
287        &mut self,
288        transactional: bool,
289        schema: &Schema,
290        track: (String, Vec<Value>),
291    ) -> Result<()> {
292        if !transactional {
293            return execute(&self.conn, schema, track).await;
294        }
295
296        let tx = self
297            .conn
298            .transaction_with_behavior(TransactionBehavior::Immediate)
299            .await?;
300        match execute(&tx, schema, track).await {
301            Ok(()) => Ok(tx.commit().await?),
302            Err(err) => {
303                // Report the original failure; a rollback error would only obscure it.
304                let _ = tx.rollback().await;
305                Err(err)
306            }
307        }
308    }
309}
310
311async fn execute(
312    conn: &Connection,
313    schema: &Schema,
314    (sql, params): (String, Vec<Value>),
315) -> Result<()> {
316    for (index, statement) in schema.statements().iter().enumerate() {
317        conn.execute_batch(statement)
318            .await
319            .map_err(|source| Error::Statement {
320                index: index + 1,
321                source,
322            })?;
323    }
324    conn.execute(&sql, params).await?;
325    Ok(())
326}
slave_framework/migration/migrator.rs

slave_framework/migration/
migrator.rs
