rust_query/
migrate.rs

1pub mod config;
2mod fix_by_copy;
3pub mod migration;
4#[cfg(test)]
5mod test;
6
7use std::{
8    cell::RefCell, collections::HashMap, marker::PhantomData, mem::take, sync::atomic::AtomicI64,
9};
10
11use annotate_snippets::{Group, Renderer, renderer::DecorStyle};
12use self_cell::MutBorrow;
13
14use crate::{
15    Table, Transaction,
16    lower::{self, list_writer::Alias},
17    migrate::{
18        config::Config,
19        fix_by_copy::fix_by_copy,
20        migration::{SchemaBuilder, TransactionMigrate},
21    },
22    pool::Pool,
23    schema::{from_macro, read::read_schema},
24    transaction::{Database, OwnedTransaction, TXN, TransactionWithRows},
25};
26
27pub struct TableTypBuilder<S> {
28    pub(crate) ast: from_macro::Schema,
29    _p: PhantomData<S>,
30}
31
32impl<S> Default for TableTypBuilder<S> {
33    fn default() -> Self {
34        Self {
35            ast: Default::default(),
36            _p: Default::default(),
37        }
38    }
39}
40
41impl<S> TableTypBuilder<S> {
42    pub fn table<T: Table<Schema = S>>(&mut self) {
43        let table = from_macro::Table::new::<T>();
44        let old = self.ast.tables.insert(T::NAME, table);
45        debug_assert!(old.is_none());
46    }
47}
48
49pub trait Schema: Sized + 'static {
50    const VERSION: i64;
51    const SOURCE: &str;
52    const PATH: &str;
53    const SPAN: (usize, usize);
54    fn typs(b: &mut TableTypBuilder<Self>);
55}
56
57pub trait SchemaMigration<'a> {
58    type From: Schema;
59    type To: Schema;
60
61    fn tables(self, b: &mut SchemaBuilder<'a, Self::From>);
62}
63
64impl<S: Schema> Database<S> {
65    /// Create a [Migrator] to migrate a database.
66    ///
67    /// Returns [None] if the database `user_version` on disk is older than `S`.
68    pub fn migrator(config: Config) -> Option<Migrator<S>> {
69        let pool = Pool::new(config);
70
71        let conn = pool.pop();
72        conn.pragma_update(None, "foreign_keys", "OFF").unwrap();
73        let txn = OwnedTransaction::new(MutBorrow::new(conn), |conn| {
74            Some(
75                conn.borrow_mut()
76                    .transaction_with_behavior(rusqlite::TransactionBehavior::Exclusive)
77                    .unwrap(),
78            )
79        });
80
81        let mut user_version = Some(user_version(txn.get()).unwrap());
82
83        // check if this database is newly created
84        if schema_version(txn.get()) == 0 {
85            user_version = None;
86
87            let schema = crate::schema::from_macro::Schema::new::<S>();
88
89            for (table_name, table) in schema.tables {
90                let table = table.to_db();
91                let create = table.create(lower::JoinableTable::Table(table_name));
92                txn.get().execute(&create, []).unwrap();
93                for stmt in table.delayed_indices(table_name) {
94                    txn.get().execute(&stmt, []).unwrap();
95                }
96            }
97        } else if user_version.unwrap() < S::VERSION {
98            // We can not migrate databases older than `S`
99            return None;
100        }
101
102        debug_assert_eq!(
103            foreign_key_check(txn.get()),
104            None,
105            "foreign key constraint violated"
106        );
107
108        Some(Migrator {
109            user_version,
110            pool,
111            transaction: txn,
112            _p: PhantomData,
113        })
114    }
115}
116
117/// [Migrator] is used to apply database migrations.
118///
119/// Create it with [Database::migrator].
120///
121/// When all migrations are done, it can be turned into a [Database] instance with
122/// [Migrator::finish].
123pub struct Migrator<S> {
124    pool: Pool,
125    transaction: OwnedTransaction,
126    // Initialized to the user version when the transaction starts.
127    // This is set to None if the schema user_version is updated.
128    // Fixups are only applied if the user_version is None.
129    // Indices are fixed before this is set to None.
130    user_version: Option<i64>,
131    _p: PhantomData<S>,
132}
133
134impl<S: Schema> Migrator<S> {
135    fn with_transaction(mut self, f: impl Send + FnOnce(&'static mut Transaction<S>)) -> Self {
136        assert!(self.user_version.is_none_or(|x| x == S::VERSION));
137        let res = std::thread::scope(|s| {
138            s.spawn(|| {
139                TXN.set(Some(TransactionWithRows::new_empty(self.transaction)));
140                let txn = Transaction::new_ref();
141
142                // check if this is the first migration that is applied
143                if self.user_version.take().is_some() {
144                    // we check the schema before doing any migrations
145                    check_schema::<S>(txn)?;
146                    // fixing indices before migrations can help with migration performance
147                    fix_by_copy::<S>(txn, fix_by_copy::Detail::Indexes);
148                }
149
150                f(txn);
151
152                let transaction = TXN.take().unwrap();
153
154                Ok::<_, Renderable>(transaction.into_owner())
155            })
156            .join()
157        });
158        match res {
159            Ok(val) => self.transaction = val.unwrap_or_else(|e| e.to_panic()),
160            Err(payload) => std::panic::resume_unwind(payload),
161        }
162        self
163    }
164
165    /// Apply a database migration if the current schema is `S` and return a [Migrator] for the next schema `N`.
166    ///
167    /// This function will panic if the schema on disk does not match what is expected for its `user_version`.
168    ///
169    /// ```
170    /// # use rust_query::migration::{schema, Config};
171    /// # use rust_query::{Lazy, Database};
172    /// #[schema(Schema)]
173    /// #[version(0..=1)]
174    /// pub mod vN {
175    ///     pub struct User {
176    ///         pub name: String,
177    ///         #[version(1..)]
178    ///         pub score: i64,
179    ///     }
180    /// }
181    ///
182    /// fn main() {
183    ///     Database::migrator(Config::open_in_memory())
184    ///         .unwrap()
185    ///         .migrate(|txn| v0::migrate::Schema {
186    ///             user: txn.migrate_ok(|old: Lazy<v0::User>| v0::migrate::User {
187    ///                 score: old.name.len() as i64,
188    ///             }),
189    ///         })
190    ///         .finish()
191    ///         .unwrap();
192    /// }
193    /// ```
194    pub fn migrate<'x, M>(
195        mut self,
196        m: impl Send + FnOnce(&mut TransactionMigrate<S>) -> M,
197    ) -> Migrator<M::To>
198    where
199        M: SchemaMigration<'x, From = S>,
200    {
201        if self.user_version.is_none_or(|x| x == S::VERSION) {
202            self = self.with_transaction(|txn| {
203                let mut txn = TransactionMigrate {
204                    inner: txn.copy(),
205                    scope: Default::default(),
206                    rename_map: HashMap::new(),
207                    extra_index: Vec::new(),
208                };
209                let m = m(&mut txn);
210
211                let mut builder = SchemaBuilder {
212                    drop: vec![],
213                    foreign_key: HashMap::new(),
214                    inner: txn,
215                };
216                m.tables(&mut builder);
217                let txn = builder.inner.inner;
218
219                for drop in builder.drop {
220                    txn.execute(&drop);
221                }
222                for (to, tmp) in builder.inner.rename_map {
223                    txn.execute(&format!("ALTER TABLE main.{tmp} RENAME TO {}", Alias(to)));
224                }
225                for stmt in builder.inner.extra_index {
226                    txn.execute(&stmt);
227                }
228
229                // Change transaction schema because we are now on the new version already
230                fix_by_copy::<M::To>(&Transaction::new(), fix_by_copy::Detail::ForeignKeys);
231
232                let transaction = TXN.take().unwrap();
233                if let Some(fk) = foreign_key_check(transaction.get()) {
234                    (builder.foreign_key.remove(&*fk).unwrap())();
235                }
236
237                TXN.set(Some(transaction));
238            });
239        }
240
241        Migrator {
242            user_version: self.user_version,
243            pool: self.pool,
244            transaction: self.transaction,
245            _p: PhantomData,
246        }
247    }
248
249    /// Mutate the database as part of migrations.
250    ///
251    /// The closure will only be executed if the database got migrated to schema version `S`
252    /// by this [Migrator] instance.
253    /// If [Migrator::fixup] is used before all [Migrator::migrate], then the closures is only executed
254    /// when the database is created.
255    pub fn fixup(mut self, f: impl Send + FnOnce(&'static mut Transaction<S>)) -> Self {
256        if self.user_version.is_none() {
257            self = self.with_transaction(f);
258        }
259        self
260    }
261
262    /// Commit the migration transaction and return a [Database].
263    ///
264    /// Returns [None] if the database schema version is newer than `S`.
265    ///
266    /// This function will panic if the schema on disk does not match what is expected for its `user_version`.
267    pub fn finish(mut self) -> Option<Database<S>> {
268        if self.user_version.is_some_and(|x| x != S::VERSION) {
269            return None;
270        }
271
272        // This checks that the schema is correct and fixes indices etc
273        self = self.with_transaction(|txn| {
274            // sanity check, this should never fail
275            check_schema::<S>(txn).unwrap_or_else(|e| e.as_sanity())
276        });
277
278        // adds an sqlite_stat1 table
279        self.transaction
280            .get()
281            .execute_batch("PRAGMA optimize;")
282            .unwrap();
283
284        set_user_version(self.transaction.get(), S::VERSION).unwrap();
285        let schema_version = schema_version(self.transaction.get());
286        self.transaction.with(|x| x.commit().unwrap());
287
288        Some(Database {
289            pool: self.pool,
290            schema_version: AtomicI64::new(schema_version),
291            schema: PhantomData,
292            mut_lock: parking_lot::FairMutex::new(()),
293        })
294    }
295}
296
297impl<S> Transaction<S> {
298    #[track_caller]
299    pub(crate) fn execute(&self, sql: &str) {
300        TXN.with_borrow(|txn| txn.as_ref().unwrap().get().execute(sql, []))
301            .unwrap();
302    }
303}
304
305pub fn schema_version(conn: &rusqlite::Transaction) -> i64 {
306    conn.pragma_query_value(None, "schema_version", |r| r.get(0))
307        .unwrap()
308}
309
310// Read user version field from the SQLite db
311pub fn user_version(conn: &rusqlite::Transaction) -> Result<i64, rusqlite::Error> {
312    conn.query_row("PRAGMA user_version", [], |row| row.get(0))
313}
314
315// Set user version field from the SQLite db
316fn set_user_version(conn: &rusqlite::Transaction, v: i64) -> Result<(), rusqlite::Error> {
317    conn.pragma_update(None, "user_version", v)
318}
319
320pub(crate) fn check_schema<S: Schema>(txn: &Transaction<S>) -> Result<(), Renderable> {
321    let from_macro = crate::schema::from_macro::Schema::new::<S>();
322    let from_db = read_schema(txn);
323    let report = from_db.diff(from_macro, S::SOURCE, S::PATH, S::VERSION);
324    if report.is_empty() {
325        Ok(())
326    } else {
327        Err(Renderable(report))
328    }
329}
330
331pub struct Renderable(Vec<Group<'static>>);
332
333impl Renderable {
334    /// [Renderable] should be made into a panic on the thread of the caller.
335    pub fn to_panic(self) -> ! {
336        let renderer = RENDERER
337            .with_borrow(Clone::clone)
338            .decor_style(DecorStyle::Unicode);
339        panic!("{}", renderer.render(&self.0))
340    }
341
342    pub fn as_sanity(self) -> ! {
343        unreachable!(
344            "THIS IS A RUST-QUERY BUG {}",
345            Renderer::plain().render(&self.0)
346        );
347    }
348}
349
350thread_local! {
351    static RENDERER: RefCell<Renderer> = const { RefCell::new(Renderer::styled()) }
352}
353
354pub fn with_test_renderer<R>(f: impl FnOnce() -> R) -> R {
355    struct TestRenderGuard(Option<Renderer>);
356    impl Drop for TestRenderGuard {
357        fn drop(&mut self) {
358            RENDERER.set(take(&mut self.0).unwrap());
359        }
360    }
361    let _g = TestRenderGuard(Some(
362        RENDERER.replace(Renderer::plain().anonymized_line_numbers(true)),
363    ));
364    f()
365}
366
367fn foreign_key_check(conn: &rusqlite::Transaction) -> Option<String> {
368    let error = conn
369        .prepare("PRAGMA foreign_key_check")
370        .unwrap()
371        .query_map([], |row| row.get(2))
372        .unwrap()
373        .next();
374    error.transpose().unwrap()
375}
376
377impl<S> Transaction<S> {
378    #[cfg(test)]
379    pub(crate) fn schema(&self) -> Vec<String> {
380        TXN.with_borrow(|x| {
381            x.as_ref()
382                .unwrap()
383                .get()
384                .prepare("SELECT sql FROM 'main'.'sqlite_schema'")
385                .unwrap()
386                .query_map([], |row| row.get::<_, Option<String>>("sql"))
387                .unwrap()
388                .flat_map(|x| x.unwrap())
389                .collect()
390        })
391    }
392}
rust_query/migrate.rs

rust_query/
migrate.rs
