Skip to main content

rust_query/migrate/
migration.rs

1use std::{collections::HashMap, convert::Infallible, marker::PhantomData, ops::Deref};
2
3use crate::{
4    Lazy, Table, TableRow, Transaction, aggregate,
5    lower::{self, JoinableTableWithId, list_writer::Alias},
6    transaction::try_insert_private,
7};
8
9pub trait Migrateable: Table<MigrateFrom: Table<Schema = Self::FromSchema>> {
10    type Migration;
11    type FromSchema;
12    type MigrateConflict;
13
14    #[doc(hidden)]
15    fn prepare(val: Self::Migration, prev: Lazy<'_, Self::MigrateFrom>) -> Self;
16    #[doc(hidden)]
17    fn map_conflict(val: TableRow<Self::MigrateFrom>) -> Self::MigrateConflict;
18}
19
20/// Transaction type for use in migrations.
21pub struct TransactionMigrate<FromSchema: 'static> {
22    pub(super) inner: Transaction<FromSchema>,
23    pub(super) scope: lower::Scope,
24    pub(super) rename_map: HashMap<&'static str, lower::TmpTable>,
25    // creating non unique indices is delayed so that they don't need to be renamed
26    pub(super) extra_index: Vec<String>,
27}
28
29impl<FromSchema> Deref for TransactionMigrate<FromSchema> {
30    type Target = Transaction<FromSchema>;
31
32    fn deref(&self) -> &Self::Target {
33        &self.inner
34    }
35}
36
37impl<FromSchema: 'static> TransactionMigrate<FromSchema> {
38    fn new_table_name<T: Table>(&mut self) -> lower::TmpTable {
39        *self.rename_map.entry(T::NAME).or_insert_with(|| {
40            let new_table_name = self.scope.tmp_table();
41            let table = crate::schema::from_macro::Table::new::<T>().into_db();
42            self.inner
43                .execute(&table.create(lower::JoinableTable::Tmp(new_table_name)));
44            self.extra_index.extend(table.delayed_indices(T::NAME));
45            new_table_name
46        })
47    }
48
49    fn unmigrated<T: Migrateable<FromSchema = FromSchema>>(
50        &self,
51        new_name: lower::TmpTable,
52    ) -> impl Iterator<Item = TableRow<T::MigrateFrom>> {
53        self.inner.query(|rows| {
54            let old = rows.join_private::<T::MigrateFrom>();
55            rows.filter(aggregate(|rows| {
56                // manually construct Joinable because we are using the new definition in the old schema.
57                // the result type is also the old type even though it represents the new table.
58                let new = rows.join(crate::private::Joinable::new(JoinableTableWithId {
59                    name: lower::JoinableTable::Tmp(new_name),
60                    main_column: <T as Table>::ID,
61                }));
62                rows.filter(old.eq(&new));
63                rows.exists().not()
64            }));
65            rows.into_iter(old)
66        })
67    }
68
69    /// Migrate some rows to the new schema.
70    ///
71    /// This will return an error when there is a conflict.
72    /// The error type depends on the number of unique constraints that the
73    /// migration can violate:
74    /// - 0 => [Infallible]
75    /// - 1.. => [TableRow] (row in the old table that could not be migrated)
76    ///
77    /// The closure returns [Option] to indicate if each row must be kept.
78    pub fn migrate_optional<'t, T: Migrateable<FromSchema = FromSchema>>(
79        &'t mut self,
80        mut f: impl FnMut(Lazy<'t, T::MigrateFrom>) -> Option<T::Migration>,
81    ) -> Result<MigratedOptional<T>, T::MigrateConflict> {
82        let new_name = self.new_table_name::<T>();
83
84        // We will do insertions here while retrieving rows from the database.
85        // This is fine because we do not care if the query uses old or new data.
86        // The only problematic case is if sqlite decides to repeat a returned row.
87        // That would be very strange though, since we are not updating the old table.
88        // See https://sqlite.org/isolation.html for more information.
89        for row in self.unmigrated::<T>(new_name) {
90            if let Some(new) = f(self.lazy(row)) {
91                // TODO: deduplicate this self.lazy call
92                let val = T::prepare(new, self.lazy(row));
93                try_insert_private::<T>(
94                    lower::JoinableTable::Tmp(new_name),
95                    Some(row.inner.idx),
96                    val,
97                )
98                .map_err(|_| T::map_conflict(row))?;
99            };
100        }
101
102        Ok(MigratedOptional { inner: PhantomData })
103    }
104
105    /// Migrate all rows to the new schema.
106    ///
107    /// Same as [Self::migrate_optional], but it does not require wrapping all migrated
108    /// rows in [Some].
109    ///
110    /// This is most likely the variant that you want to use, unless you have a table without
111    /// unique constraint, see [Self::migrate_ok].
112    pub fn migrate<'t, T: Migrateable<FromSchema = FromSchema>>(
113        &'t mut self,
114        mut f: impl FnMut(Lazy<'t, T::MigrateFrom>) -> T::Migration,
115    ) -> Result<Migrated<'static, T>, T::MigrateConflict> {
116        self.migrate_optional(|x| Some(f(x)))
117            .map(|x| x.map_fk_err(|| unreachable!("all rows are migrated")))
118    }
119
120    /// Migrate all rows to the new schema, without unique constraint conflicts.
121    ///
122    /// Same as [Self::migrate], but can only be used when the migration is known to
123    /// never cause unique constraint conflicts.
124    pub fn migrate_ok<'t, T: Migrateable<FromSchema = FromSchema, MigrateConflict = Infallible>>(
125        &'t mut self,
126        f: impl FnMut(Lazy<'t, T::MigrateFrom>) -> T::Migration,
127    ) -> Migrated<'static, T> {
128        let Ok(res) = self.migrate(f);
129        res
130    }
131}
132
133/// [Migrated] provides a proof of migration.
134///
135/// This only needs to be provided for tables that are migrated from a previous table.
136pub struct Migrated<'t, T: Migrateable> {
137    _p: PhantomData<T>,
138    f: FkErrHandler<'t>,
139    _local: PhantomData<*const ()>,
140}
141
142impl<'t, To: Migrateable> Migrated<'t, To> {
143    #[doc(hidden)]
144    pub fn apply(self, b: &mut SchemaBuilder<'t, To::FromSchema>) {
145        b.foreign_key::<To>(self.f);
146    }
147}
148
149pub struct SchemaBuilder<'t, FromSchema: 'static> {
150    pub(super) inner: TransactionMigrate<FromSchema>,
151    pub(super) drop: Vec<String>,
152    pub(super) foreign_key: HashMap<&'static str, FkErrHandler<'t>>,
153}
154
155impl<'t, FromSchema: 'static> SchemaBuilder<'t, FromSchema> {
156    pub fn foreign_key<To: Table>(&mut self, err: FkErrHandler<'t>) {
157        self.inner.new_table_name::<To>();
158
159        self.foreign_key.insert(To::NAME, err);
160    }
161
162    pub fn create_empty<To: Table>(&mut self) {
163        self.inner.new_table_name::<To>();
164    }
165
166    pub fn drop_table<T: Table>(&mut self) {
167        self.drop.push(format!("DROP TABLE {}", Alias(T::NAME)));
168    }
169}
170
171/// Proof that a table is at least partially migrated.
172///
173/// This type can be turned into [Migrated] by providing an error
174/// handler.
175pub struct MigratedOptional<T: Migrateable> {
176    inner: PhantomData<Migrated<'static, T>>,
177}
178
179impl<T: Migrateable> MigratedOptional<T> {
180    /// The closure is called when there is a foreign key error due to some row being removed.
181    pub fn map_fk_err<'t>(self, f: impl 't + FnOnce() -> Infallible) -> Migrated<'t, T> {
182        Migrated {
183            _p: PhantomData,
184            f: FkErrHandler(Box::new(f)),
185            _local: PhantomData,
186        }
187    }
188}
189
190impl<T: Migrateable<Referer = Infallible>> MigratedOptional<T> {
191    /// The table has the `#[no_reference]` attribute, so partial migration is always ok.
192    pub fn no_reference(self) -> Migrated<'static, T> {
193        self.map_fk_err(|| unreachable!("no references exist to this table"))
194    }
195}
196
197pub(crate) struct FkErrHandler<'t>(pub Box<dyn 't + FnOnce() -> Infallible>);