tag2upload_service_manager/

db_migration.rs

//! Schema migration support
//!
//! Our basic approach:
//!
//! The 12-step process from s7 in <https://sqlite.org/lang_altertable.html>
//!  "Making Other Kinds Of Table Schema Changes"
//! (retrieved 2025-10-12).
//!
//! But we do it to the whole db at once, and we store the temporary copy
//! in a whole temporary database which we ATTACH.  So we don't *rename*
//! the tables.
//!
//! So the procedure is (roughly):
//!
//!  * Attach temporary database
//!  * Create tables in temporary database by running (new) schema on it
//!  * Copy data from real to temporary database,
//!    with data alteration/supplementation as needed
//!    (including sqlite_sequence table but not other sqlite_* tables).
//!  * Drop every table in the real database.
//!  * Create tables in real database by running schema on it.
//!  * Copy data from temporary database back to main database.

use crate::prelude::*;
use db_support::{get_user_version, set_user_version_stmt, table_column_names};

// TODO use more widely and move into prelude?
use bsql::{ErasedError1 as EE1, Context1 as _};

use Exception as E;
type Err = DbError<EE1>;

#[derive(Debug, Clone)]
pub struct MigrationData {
    pub schemas: Vec<Schema>,
    pub migrations: Vec<Migration>,
}

#[derive(Debug, Clone)]
pub struct Schema {
    pub version: SchemaVersion,
    pub schema: String,
}

#[derive(Debug, Clone)]
pub struct Migration {
    pub old_version: SchemaVersion,
    pub new_version: SchemaVersion,
    pub exceptions: Vec<Exception>,
}

#[derive(Debug, Clone, Copy)]
pub enum OnComplete {
    Rollback,
    Commit,
}

#[derive(Debug, Clone, Copy, Eq, PartialEq, derive_more::Display)]
pub enum Outcome {
    #[display("already version {_0}")]
    Already(SchemaVersion),
    #[display("migrated to version {_0}")]
    Migrated(SchemaVersion),
    #[display("migration tested, from {current} to {tested}")]
    Tested {
        current: SchemaVersion,
        tested: SchemaVersion,
    },
}

#[derive(Debug, Clone)]
pub enum Exception {
    ReplacementColumnValue {
        table: &'static str,
        col: &'static str,
        /// An expression for the new column value
        ///
        /// Here:
        ///   * `old` is the old *table*.
        ///   * `aside` is the new *schema* (in the temporary database)
        ///   * `main` is the old *schema* (in the main database)
        val_sql: String,
    },
    /// Also provide `ExpectedRowCountExpression` or `UncheckedRowCount`
    NewTable {
        table: &'static str,
    },
    /// After copy out hook, any SQL
    ///
    /// Runs after all the data has been copied out.
    /// There are two copies, therefore, in `main` (old schema)
    /// and `aside` (new schema).
    GeneralAfterCopyOut {
        sql: String,
    },
    /// Final hook, any SQL
    ///
    /// Runs after all the data has been copied back.
    /// The main database is the default schema.
    /// The old values are no longer accessible.
    GeneralFinal {
        sql: String,
    },
    ExpectedRowCount {
        table: &'static str,
        /// SELECT statement giving the expected number of rows
        select_sql: String,
    },
    UncheckedRowCount {
        table: &'static str,
    },
    IncreasedRowCount {
        table: &'static str,
    },
}

/// The name of the sqlite_sequence table, which needs special handling.
const SQLITE_SEQUENCE: &str = "sqlite_sequence";

#[derive(Debug, Clone, Eq, PartialEq)]
enum RowCountException {
    Unchecked,
    /// SELECT statement
    Expect(String),
    Increased,
}

#[derive(Debug, Eq, PartialEq, Default)]
struct PreprocessedExceptions {
    /// `(table, col) -> val_sql`
    replacement_col_vals: HashMap<(String, String), String>,
    /// `table -> Some(val_sql)`: `None` means unchecked
    expected_row_count: HashMap<String, RowCountException>,
    /// `table -> Some(val_sql)`: `None` means unchecked
    no_copy_tables: HashSet<String>,
    /// SQL
    general_final: String,
    /// SQL
    general_after_copy_out: String,
}

impl MigrationData {
    pub fn check(
        &self,
        target_version: SchemaVersion,
        max_version: SchemaVersion,
    ) -> Result<(), EE1> {
        expect1!(
            self.schemas.iter()
                .map(|sch| sch.version)
                .all_unique()
        )?;
        let mut mig_from_to = HashSet::new();
        for mig in &self.migrations {
            expect1!(mig.new_version == mig.old_version + 1)?;
            expect1!(mig.new_version <= max_version)?;
            expect1!(mig_from_to.insert((mig.old_version, mig.new_version)))?;
            let _: PreprocessedExceptions = mig.preprocess_exceptions()?;
        };
        expect1!(
            self.migrations.iter().next().is_none() ||
            self.get_schema(target_version).is_ok()
        )?;
        Ok(())
    }

    fn get_schema(&self, version: SchemaVersion) -> Result<&str, EE1> {
        let found = &self.schemas.iter()
            .find(|sch| sch.version == version)
            .ok_or_else(|| anyerror1!(
                "missing schema for {version}"
            ))?;
        Ok(&found.schema)
    }
}

impl Migration {
    fn preprocess_exceptions(&self) -> Result<PreprocessedExceptions, EE1> {
        let mut replacement_col_vals = HashMap::new();
        let mut expected_row_count = HashMap::new();
        let mut no_copy_tables = HashSet::new();
        let mut general_final = None;
        let mut general_after_copy_out = None;
        let os = |s: &str| -> String { s.to_owned() };

        let store_general = |g: &mut Option<String>, sql: &str| {
            expect1!(g.is_none())?;
            *g = Some(sql.into());
            Ok::<_, EE1>(())
        };
        let mut row_count = |table: &str, ex| {
            expect1!(
                expected_row_count.insert((*table).to_owned(), ex)
                    .is_none()
            )?;
            Ok::<_, EE1>(())
        };

        for e in &self.exceptions {
            match e {
                E::ReplacementColumnValue { table, col, val_sql } => expect1!(
                    replacement_col_vals.insert(
                        (os(table), os(col)),
                        val_sql.clone(),
                    )
                        .is_none()
                )?,
                E::NewTable { table } => expect1!(
                    no_copy_tables.insert(os(table))
                )?,
                E::ExpectedRowCount { table, select_sql } => row_count(
                    table,
                    RowCountException::Expect(select_sql.clone()),
                )?,
                E::UncheckedRowCount { table } => row_count(
                    table,
                    RowCountException::Unchecked,
                )?,
                E::IncreasedRowCount { table } => row_count(
                    table,
                    RowCountException::Increased,
                )?,
                E::GeneralFinal { sql } => {
                    store_general(&mut general_final, sql)?;
                },
                E::GeneralAfterCopyOut { sql } => {
                    store_general(&mut general_after_copy_out, sql)?;
                },
            }
        }

        let general_final = general_final.unwrap_or_default();
        let general_after_copy_out
            = general_after_copy_out.unwrap_or_default();

        Ok(PreprocessedExceptions {
            replacement_col_vals,
            expected_row_count,
            no_copy_tables,
            general_final,
            general_after_copy_out,
        })
    }
}

fn set_foreign_keys(conn: &rusqlite::Connection, enable: bool)
                    -> Result<(), Err>
{
    let stmt = format!("PRAGMA foreign_keys={}",
                       if enable { "ON" } else { "OFF" });
    conn.execute_batch(&stmt)
        .with_db_context(|| stmt.clone())?;
    Ok(())
}

fn exec_batch_logged(
    progress: &mut dyn io::Write,
    conn: &rusqlite::Connection,
    what: &str,
    stmt: &str
)
    -> Result<(), Err>
{
    writeln!(progress, "migration, executing statement, {what}:\n{stmt}\n")
        .context1("report progress")?;
    conn.execute_batch(stmt)
        .db_context(what)
}

impl OnComplete {
    pub fn from_commit_bool(commit: bool) -> Self {
        if commit {
            OnComplete::Commit
        } else {
            OnComplete::Rollback
        }
    }
}

impl Outcome {
    fn version_after(&self) -> SchemaVersion {
        *match self {
            Outcome::Migrated(m) => m,
            Outcome::Already(m) => m,
            Outcome::Tested { tested, .. } => tested,
        }
    }
}

/// Ensure that the database has a suitable schema
///
/// Ensures that the db file has at least `target_version`,
/// and no more than `max_version`.
/// If there is no schema at all yet, installs `target_version`.
///
/// Schema texts, and migrations, are provided via `MigrationData`.
///
/// `temp_dir` must exist.
/// Files there like `migration*` may be erased/overwritten.
/// No cleanup is done there.
///
/// Performs all necessary transaction / db locking retries.
///
/// If `on_step_complete` is `OnComplete::Rollback`,
/// cannot perform a multi-step migration.
/// If that would be needed, fails.
pub fn prepare_idempotent(
    db_file: &Path,
    temp_dir: &Path,
    timeout: &bsql::Timeout,
    progress: &mut dyn io::Write,
    target_version: SchemaVersion,
    max_version: SchemaVersion,
    migration_data: &MigrationData,
    on_step_complete: OnComplete,
) -> Result<Outcome, AE> {
    let mk_conn = || -> Result<_, Err> {
        rusqlite::Connection::open(db_file)
            .db_context("access db during schema preparation")
    };

    /// Convenience alias.  (We can't have generic closure.)
    macro_rules! retry_loop { { $f:expr } => {
        timeout.generic_retry_loop_erasederror1($f)
    } }

    loop {
        let mut conn = retry_loop!(mk_conn)?;

        let stored_version = retry_loop!(|| -> Result<_, Err> {
            let dbt = conn
                .transaction_with_behavior(
                    rusqlite::TransactionBehavior::Immediate
                )
                .db_context("start transaction, for schema setup")?;

            let stored_version = get_user_version(&dbt)
                .map_err(DbError::Sql)?;

            if stored_version == 0 {
                writeln!(progress, "initialising schema in empty database")
                    .context1("write")?;

                let target_schema = migration_data.get_schema(target_version)
                    .map_err(DbError::Other)?;

                // "execute" silently executes only the first:
                //   https://github.com/rusqlite/rusqlite/issues/397
                //   https://github.com/rusqlite/rusqlite/issues/1147
                dbt.execute_batch(&target_schema)
                    .db_context("install fresh schema")?;
                dbt
                    .execute_batch(&set_user_version_stmt(target_version))
                    .db_context("install fresh schema: set user_version")?;
                dbt.commit()
                    .db_context("install fresh schema: commit")?;
                return Ok(target_version);
            }

            Ok::<_, Err>(stored_version)
        })?;

        if stored_version >= target_version {
            if stored_version > max_version {
                Err(anyhow!(
                    "db already contains schema version {stored_version} and we only support {target_version}"
                ))?;
            }
            return Ok(Outcome::Already(stored_version));
        }

        // db is out of date, attempt a migration.

        writeln!(
            progress,
            "attempting migration from schema version {stored_version}"
        )
            .context("write")?;

        let mut conn = Some(conn);
        let outcome = retry_loop!(|| {
            let conn = conn.take().map(Ok)
                .unwrap_or_else(|| mk_conn())?;

            migration_core(
                conn,
                temp_dir,
                progress,
                stored_version,
                stored_version + 1,
                migration_data,
                on_step_complete,
            )
        })?;

        // This will have either increased the version to stored_version + 1
        // by a migration, or discovered it was already that high.
        // If we go round again, maybe we're done?

        if outcome.version_after() == target_version {
            return Ok(outcome);
        }
        match outcome {
            Outcome::Migrated(_) | // migrated, therefore to earlier, carry on
            Outcome::Already(_) => {}, // found, therefore higher, will return
            Outcome::Tested { current, tested } => Err(anyhow!(
                "needed multi-step migration: db has {current}, tested migration to {tested} OK, but target was {target_version}"
            ))?,
        }
    }
}

pub fn migration_core<'w>(
    mut main_conn: rusqlite::Connection,
    temp_dir: &Path,
    progress: &'w mut dyn io::Write,
    old_version: SchemaVersion,
    new_version: SchemaVersion,
    migration_data: &MigrationData,
    on_complete: OnComplete,
) -> Result<Outcome, Err> {
    let temp_dir = temp_dir.to_str()
        .ok_or_else(|| anyerror1!("temp dir path not utf-8 {temp_dir:?}"))?;
    let tmp_db = format!("{temp_dir}/migration-aside.db");

    macro_rules! progressln { { $($m:tt)* } => {
        writeln!(progress, $($m)*).context1("report progress")
    } }

    progressln!("migration, preparing")?;

    //----- Find the migration in the migration data -----

    migration_data.check(new_version, SchemaVersion::MAX)?;

    let migration = migration_data
        .migrations.iter()
        .find(|exc|
              exc.old_version == old_version &&
              exc.new_version == new_version)
        .ok_or_else(|| anyerror1!(
            "unsupported migration {old_version} => {new_version}"
        ))?;

    let new_schema = migration_data.get_schema(new_version)?;

    let mut es = migration.preprocess_exceptions()
        .context1("preprocess exceptions")?;

    //----- Clear out things that look like our temporary db -----

    let preclean_glob = format!("{tmp_db}*");
    for preclean in glob::glob(&preclean_glob)
        .with_context1(|| preclean_glob.clone())
        .context1("glob tmp dirs")?
    {
        let preclean = preclean.context1("identify path to clean")?;
        match fs::remove_file(&preclean) {
            Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()),
            other => other,
        }
            .with_context1(|| preclean.to_string_lossy().to_string())
            .context1("clean tmp db file")?;
    }

    //----- Attach a temporary db, containing the new schema -----

    // We can't run the schema on the temporary database when we
    // have it ATTACHed because it doesn't use a schema name,
    // and sqlite has no way to control the default schema.
    let tmp_conn = rusqlite::Connection::open(&tmp_db)
        .with_db_context(|| format!("open tmp db {tmp_db:?}"))?;
    tmp_conn.execute_batch(&new_schema)
        .db_context("execute current schema")?;
    tmp_conn.close()
        .map_err(|(_, e)| e)
        .db_context("close on tmp db current schema")?;

    // We *consume* the connection, so we don't have to worry about
    // leaving this temporary db attached to the caller's connection.
    main_conn.execute("ATTACH DATABASE ? AS aside", [&tmp_db])
        .db_context("attach temp database")?;

    //----- turn off foreign key checking! -----

    set_foreign_keys(&main_conn, false)?;

    //----- start the transaction -----

    let dbt = main_conn
        .transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)
        .db_context("start transaction")?;

    //----- check the schema version in the db -----

    let found_old_version = get_user_version(&dbt)
        .map_err(DbError::Sql)?;
    if found_old_version >= new_version {
        return Ok(Outcome::Already(found_old_version));
    }
    if found_old_version != old_version {
        Err(anyerror1!(
 "running schema migration from {old_version} but db has {found_old_version}"
        ))?;
    }

    //----- determine the names of the tables we're migrating -----

    // table names in new schema
    // including sqlite_sequence (first) but not any other sqlite_*
    let tables = {
        let mut tables = dbt
            // get all the tables names from the schema
            .prepare(r#"
                SELECT name FROM aside.sqlite_schema WHERE type = 'table'
                ORDER BY name ASC
            "#)
            .db_context("prepare find tables")?
            .query_map([], |row| row.get::<_, String>(0))
            .db_context("find tables")?
            .collect::<Result<Vec<_>, _>>()
            .db_context("convert table name")?
            .into_iter()
            // - filter out sqlite_*,
            // - prepend an ordering override, putting sqlite_sequence first
            .filter_map(|table| Some(match &*table {
                SQLITE_SEQUENCE => (0, table),
                s if s.starts_with("sqlite_") => return None,
                _ => (1, table),
            }))
            .collect_vec();

        // sort and collect again
        tables.sort();
        tables.into_iter().map(|(_, t)| t).collect_vec()
    };

    //----- preprocess the exceptions -----

    let expected_row_count: HashMap<_, (u64, Option<_>)> = {
        let mut out = HashMap::new();
        for table in &tables {
            let get_exp = |sql: &str| -> Result<u64, Err> {
                let exp = dbt.query_one(sql, [], |row| row.get(0))
                    .db_context("query existing row count")?;
                Ok(exp)
            };
            let get_current = || get_exp(&format!(
                "SELECT count(*) FROM main.{table}"
            ));
            let exp = match es.expected_row_count.remove(&**table) {
                None => {
                    let e = get_current()?;
                    (e, Some(e))
                },
                Some(RowCountException::Expect(exp)) => {
                    let e = get_exp(&exp)?;
                    (e, Some(e))
                },
                Some(RowCountException::Increased) => {
                    let e = get_current()?;
                    (e, None)
                },
                Some(RowCountException::Unchecked) => continue,
            };
            out.insert(table, exp.into());
        }
        out
    };

    //----- copy data to temp db, transforming it to new schema -----

    progressln!("migration, copying data out")?;

    for table in &tables {
        if es.no_copy_tables.remove(&*table) { continue };

        let new_cols = table_column_names(&dbt, &format!("aside.{table}"))
            .map_err(DbError::Sql)?;
        let new_col_names = new_cols.iter().map(|s| &**s)
            .intersperse(", ").collect::<String>();

        let mut s = String::new();

        s += &format!(r#"
            INSERT INTO aside.{table}
                     ( {new_col_names} )
               SELECT
"#);

        let new_col_settings = new_cols.iter().map(|col| {
            let col_id = (table.clone(), col.clone());
            let new_val = match es.replacement_col_vals.remove(&col_id) {
                None => format!("old.{col}"),
                Some(val_sql) => format!("{val_sql}"),
            };
            format!(
r#"                      {new_val} AS {col}"#)
        }).collect_vec();

        s.extend(
            new_col_settings.iter().map(|s| &**s)
                .intersperse(",\n")
        );
        s += &format!(r#"
                 FROM main.{table} AS old
"#);

        exec_batch_logged(progress, &dbt, "copy data out", &s)?;
    }

    exec_batch_logged(progress, &dbt, "GeneralAfterCopyOut",
                      &mem::take(&mut es.general_after_copy_out))?;

    progressln!("migration, turning around")?;

    //----- delete every table we're migrating in the main db -----

    // We only delete tables we know we're migrating, so we can't
    // accidentally delete any *other* tables that we somehow missed
    // earlier and whose data we aren't copying.
    //
    // Deleting the tables deletes all their indices, and their
    // corresponding rows in sqlite_sequence.

    for table in &tables {
        if table == SQLITE_SEQUENCE {
            // We don't delete this table.  (We do save/restore its rows.)
            continue
        }

        exec_batch_logged(progress, &dbt, "clear old table", &format!(r#"
            DROP TABLE IF EXISTS main.{table}
"#
        ))?;
    }

    //----- install the new schema in the main db -----

    // The deletion above should have cleared out the schema.
    // So this creation should run without trouble.
    dbt.execute_batch(&new_schema)
        .db_context("set up new schema in main")?;

    //----- copy the data back from the temp db to the main db -----

    progressln!("migration, copying data back")?;

    // It's already in the right format, so we can just
    // INSERT INTO SELECT without any ado.
    for table in &tables {
        exec_batch_logged(progress, &dbt, "copy data in", &format!(
r#"    INSERT INTO main.{table} SELECT * FROM aside.{table}"#
        ))?;
    }

    //----- migration-specific fixups -----

    exec_batch_logged(progress, &dbt, "GeneralFinal",
                      &mem::take(&mut es.general_final))?;

    //----- consistency checks -----

    progressln!("migration, checking")?;

    // Perform and reenable foreign key checks.
    exec_batch_logged(progress, &dbt, "check FK", r#"
        PRAGMA foreign_key_check;
"#)?;
    // This oughtn't to won't make any difference, since it's per-connection
    // and we're about to throw the connection away.  Let's do it anyway.
    set_foreign_keys(&dbt, true)?;

    // Check that we have the right number of rows
    expect1!(expected_row_count.iter().next().is_some())?;
    for (table, (exp_min, exp_max)) in expected_row_count {
        let got: u64 = dbt.query_one(
            &format!("SELECT count(*) FROM main.{table}"),
            [], |row| row.get(0),
        ).db_context("get new row count")?;
        (got >= exp_min).then_some(()).ok_or_else(|| anyerror1!(
 "table {table}: expected >= {exp_min} rows after migration but had {got}"
        ))?;
        if let Some(exp_max) = exp_max {
            (got <= exp_max).then_some(()).ok_or_else(|| anyerror1!(
 "table {table}: expected <= {exp_max} rows after migration but had {got}"
            ))?;
        }
    }

    expect1!(es == PreprocessedExceptions::default(),
             "unused exception(s) {es:#?}")?;

    //----- completion -----

    // Set the user version, recording that this is the new schema.
    dbt.execute_batch(&set_user_version_stmt(new_version))
        .db_context("set user version after migration")?;

    // Commit or rollback, as instructued.
    match on_complete {
        OnComplete::Commit => {
            progressln!("migration complete, committing")?;
            dbt.commit().db_context("commit")?;
            progressln!("migrated to schema version {new_version}")?;
            Ok(Outcome::Migrated(new_version))
        },
        OnComplete::Rollback => {
            progressln!("migration check, apparently successful")?;
            Ok(Outcome::Tested {
                current: old_version,
                tested: new_version,
            })
        }
    }
}

#[cfg(test)]
mod db_migr_test;