rydantic 0.7.0-beta.2

//! Catalog-driven model migrations. Plans are immutable and checked again under a database lock.
mod catalog;
mod plan;

use super::{ddl, Dialect, SqlBackend, SqlError, SqlModel, SqlParameter, SqlType};
use catalog::{Catalog, DbColumn};
use serde::Serialize;
use sha2::{Digest, Sha256};
use sqlx::{Pool, QueryBuilder};

/// Constant SQL defaults are independent of Python-style model default factories.
#[derive(Debug, Clone, Copy)]
pub enum SqlDefault {
    Integer(i64),
    Real(f64),
    Boolean(bool),
    Text(&'static str),
}
impl SqlDefault {
    pub(super) fn sql(self, dialect: Dialect) -> String {
        match self {
            Self::Integer(value) => value.to_string(),
            Self::Real(value) => value.to_string(),
            Self::Boolean(value) if dialect == Dialect::Postgres => value.to_string(),
            Self::Boolean(value) => i32::from(value).to_string(),
            Self::Text(value) if dialect == Dialect::MySql => {
                let hex: String = value
                    .as_bytes()
                    .iter()
                    .map(|b| format!("{b:02x}"))
                    .collect();
                format!("(CONVERT(X'{hex}' USING utf8mb4))")
            }
            Self::Text(value) if dialect == Dialect::Postgres => {
                format!("E'{}'", value.replace('\\', "\\\\").replace('\'', "''"))
            }
            Self::Text(value) => format!("'{}'", value.replace('\'', "''")),
        }
    }
    fn canonical(self, dialect: Dialect) -> String {
        match self {
            Self::Text(value) => format!("text:{value}"),
            Self::Boolean(value) if dialect == Dialect::Postgres => format!("bool:{value}"),
            Self::Boolean(value) => format!("number:{}", i32::from(value)),
            Self::Integer(value) => format!("number:{value}"),
            Self::Real(value) => format!("number:{value}"),
        }
    }
}

/// Destructive migrations require an explicit opt-in for each invocation.
#[derive(Debug, Clone, Copy, Default)]
pub struct MigrationOptions {
    pub allow_destructive: bool,
}

/// One human-readable schema change, together with its data-loss classification.
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct MigrationChange {
    pub description: String,
    pub destructive: bool,
}

/// A preview generated from the live catalog and a model's SQL metadata.
/// Statements and snapshots cannot be mutated through the public API.
#[derive(Debug, Clone)]
pub struct MigrationPlan {
    table: String,
    dialect: Dialect,
    before: Catalog,
    changes: Vec<MigrationChange>,
    statements: Vec<String>,
    target: Catalog,
    rebuild: bool,
    checks: Vec<(String, String)>,
}
impl MigrationPlan {
    pub fn table(&self) -> &str {
        &self.table
    }
    pub fn changes(&self) -> &[MigrationChange] {
        &self.changes
    }
    pub fn statements(&self) -> &[String] {
        &self.statements
    }
    pub fn is_empty(&self) -> bool {
        self.changes.is_empty()
    }
    pub fn is_destructive(&self) -> bool {
        self.changes.iter().any(|c| c.destructive)
    }
    /// Apply this exact preview, rejecting it if the live schema has changed.
    pub async fn apply<DB: SqlBackend>(
        &self,
        pool: &Pool<DB>,
        options: MigrationOptions,
    ) -> Result<MigrationReport, SqlError> {
        run::<DB>(pool, Some(self), None, options).await
    }
}

/// A successful migration receipt. No-op runs have no new history revision.
#[derive(Debug, Clone)]
pub struct MigrationReport {
    pub plan: MigrationPlan,
    pub revision: Option<i64>,
    pub checksum: String,
}

fn error(code: &'static str, message: impl Into<String>) -> SqlError {
    SqlError::Migration {
        code,
        message: message.into(),
    }
}
fn checksum(value: &impl Serialize) -> String {
    format!(
        "{:x}",
        Sha256::digest(serde_json::to_vec(value).expect("schema is JSON serializable"))
    )
}
async fn execute<DB: SqlBackend>(
    connection: &mut DB::Connection,
    sql: &str,
) -> Result<(), SqlError> {
    DB::execute(&mut QueryBuilder::new(sql), connection).await?;
    Ok(())
}
async fn rows<DB: SqlBackend>(
    connection: &mut DB::Connection,
    sql: &str,
    parameters: &[&str],
) -> Result<Vec<DB::Row>, SqlError> {
    let mut query = QueryBuilder::new(sql);
    // Catalog queries use a final placeholder generated by QueryBuilder.
    for (index, part) in sql.split("{bind}").enumerate() {
        if index == 0 {
            query = QueryBuilder::new(part);
        } else {
            DB::bind_parameter(SqlParameter::Text(parameters[index - 1].into()), &mut query);
            query.push(part);
        }
    }
    Ok(DB::fetch_all(&mut query, connection).await?)
}
fn text<DB: SqlBackend>(row: &DB::Row, name: &str) -> Result<Option<String>, SqlError> {
    Ok(match DB::read_parameter(row, name, SqlType::Text)? {
        SqlParameter::Text(value) => Some(value),
        _ => None,
    })
}
async fn scalar<DB: SqlBackend>(
    connection: &mut DB::Connection,
    sql: &str,
) -> Result<String, SqlError> {
    let result = rows::<DB>(connection, sql, &[]).await?;
    text::<DB>(
        result
            .first()
            .ok_or_else(|| error("migration_catalog", "expected a catalog result"))?,
        "value",
    )?
    .ok_or_else(|| error("migration_catalog", "catalog result was NULL"))
}

pub(super) async fn preview<T: SqlModel, DB: SqlBackend>(
    pool: &Pool<DB>,
) -> Result<MigrationPlan, SqlError> {
    ddl::check_identifiers::<T, DB>()?;
    let mut connection = pool.acquire().await?;
    let before = catalog::read::<DB>(&mut connection, T::TABLE).await?;
    plan::build::<T, DB>(before)
}

pub(super) async fn migrate<T: SqlModel, DB: SqlBackend>(
    pool: &Pool<DB>,
    options: MigrationOptions,
) -> Result<MigrationReport, SqlError> {
    ddl::check_identifiers::<T, DB>()?;
    run::<DB>(pool, None, Some((T::TABLE, plan::build::<T, DB>)), options).await
}

type Planner = fn(Catalog) -> Result<MigrationPlan, SqlError>;
struct MigrationConnection<DB: SqlBackend>(Option<sqlx::pool::PoolConnection<DB>>);
impl<DB: SqlBackend> Drop for MigrationConnection<DB> {
    fn drop(&mut self) {
        if let Some(connection) = &mut self.0 {
            connection.close_on_drop();
        }
    }
}
async fn run<DB: SqlBackend>(
    pool: &Pool<DB>,
    preview: Option<&MigrationPlan>,
    model: Option<(&str, Planner)>,
    options: MigrationOptions,
) -> Result<MigrationReport, SqlError> {
    let (table, planner) = match (preview, model) {
        (Some(plan), _) if plan.dialect != DB::DIALECT => {
            return Err(error(
                "migration_driver",
                "plan belongs to a different SQL driver",
            ))
        }
        (Some(plan), _) => (plan.table.as_str(), None),
        (_, Some((table, planner))) => (table, Some(planner)),
        _ => unreachable!(),
    };
    let mut guard = MigrationConnection(Some(pool.acquire().await?));
    let connection = guard.0.as_mut().expect("acquired connection");
    // Closing on cancellation releases advisory locks and rolls back open work.
    // SQLite foreign-key settings never leak into a reused pool connection.
    let foreign_keys = if DB::DIALECT == Dialect::Sqlite {
        let result = rows::<DB>(
            connection,
            "SELECT CAST(foreign_keys AS TEXT) AS value FROM pragma_foreign_keys",
            &[],
        )
        .await?;
        text::<DB>(&result[0], "value")?.as_deref() == Some("1")
    } else {
        false
    };
    match DB::DIALECT {
        Dialect::Sqlite => {
            execute::<DB>(connection, "PRAGMA foreign_keys = OFF").await?;
            execute::<DB>(connection, "BEGIN IMMEDIATE").await?;
        }
        Dialect::Postgres => {
            execute::<DB>(connection, "BEGIN").await?;
            execute::<DB>(connection, "SELECT pg_advisory_xact_lock(1920544372, 1)").await?;
        }
        Dialect::MySql => {
            let acquired = scalar::<DB>(connection, "SELECT CAST(GET_LOCK(CONCAT('rydantic:', LEFT(SHA2(DATABASE(), 256), 40)), 30) AS CHAR) AS value").await?;
            if acquired != "1" {
                return Err(error(
                    "migration_lock",
                    "could not acquire the database migration lock",
                ));
            }
        }
    }
    let result = apply_locked::<DB>(connection, table, planner, preview, options).await;
    if DB::DIALECT != Dialect::MySql {
        execute::<DB>(
            connection,
            if result.is_ok() { "COMMIT" } else { "ROLLBACK" },
        )
        .await?;
    }
    if DB::DIALECT == Dialect::Sqlite {
        execute::<DB>(
            connection,
            if foreign_keys {
                "PRAGMA foreign_keys = ON"
            } else {
                "PRAGMA foreign_keys = OFF"
            },
        )
        .await?;
    } else if DB::DIALECT == Dialect::MySql {
        execute::<DB>(
            connection,
            "SELECT RELEASE_LOCK(CONCAT('rydantic:', LEFT(SHA2(DATABASE(), 256), 40)))",
        )
        .await?;
    }
    drop(guard.0.take());
    result
}

async fn apply_locked<DB: SqlBackend>(
    connection: &mut DB::Connection,
    table: &str,
    planner: Option<Planner>,
    preview: Option<&MigrationPlan>,
    options: MigrationOptions,
) -> Result<MigrationReport, SqlError> {
    let mut before = catalog::read::<DB>(connection, table).await?;
    if before.exists && DB::DIALECT == Dialect::Postgres {
        execute::<DB>(
            connection,
            &format!(
                "LOCK TABLE {} IN ACCESS EXCLUSIVE MODE",
                DB::DIALECT.quote(table)
            ),
        )
        .await?;
        before = catalog::read::<DB>(connection, table).await?;
    }
    let plan = if let Some(preview) = preview {
        if preview.before != before {
            return Err(error(
                "migration_stale",
                "schema changed after this migration was planned; generate a new preview",
            ));
        }
        preview.clone()
    } else {
        planner.expect("model supplies a planner")(before)?
    };
    if plan.is_destructive() && !options.allow_destructive {
        return Err(error("migration_destructive", "migration removes columns or changes storage types; inspect migration_plan() and opt in with MigrationOptions { allow_destructive: true }"));
    }
    let digest = checksum(&plan.target);
    if plan.is_empty() {
        return Ok(MigrationReport {
            plan,
            revision: None,
            checksum: digest,
        });
    }
    if DB::DIALECT == Dialect::MySql {
        let mode = scalar::<DB>(connection, "SELECT @@SESSION.sql_mode AS value").await?;
        if !mode.contains("STRICT_ALL_TABLES") && !mode.contains("STRICT_TRANS_TABLES") {
            return Err(error(
                "migration_sql_mode",
                "MySQL automatic migrations require a strict SQL mode to reject lossy conversions",
            ));
        }
    }
    for (sql, message) in &plan.checks {
        if !rows::<DB>(connection, sql, &[]).await?.is_empty() {
            return Err(error("migration_data", message));
        }
    }
    history_table::<DB>(connection).await?;
    let revision = record::<DB>(connection, &plan, &digest, "applying").await?;
    let application = async {
        for statement in &plan.statements {
            execute::<DB>(connection, statement).await?;
        }
        if DB::DIALECT == Dialect::Sqlite
            && !rows::<DB>(connection, "PRAGMA foreign_key_check", &[])
                .await?
                .is_empty()
        {
            return Err(error(
                "migration_foreign_key",
                "migration would violate a foreign-key constraint",
            ));
        }
        let after = catalog::read::<DB>(connection, table).await?;
        if !plan::matches(&after, &plan.target) {
            return Err(error(
                "migration_verification",
                "live schema does not match the target model after migration",
            ));
        }
        Ok(())
    }
    .await;
    let status = if application.is_ok() {
        "applied"
    } else {
        "failed"
    };
    let mut query = QueryBuilder::<DB>::new("UPDATE _rydantic_migrations SET status = ");
    DB::bind_parameter(SqlParameter::Text(status.into()), &mut query);
    query.push(" WHERE table_name = ");
    DB::bind_parameter(SqlParameter::Text(table.into()), &mut query);
    query.push(" AND revision = ");
    DB::bind_parameter(SqlParameter::Integer(revision), &mut query);
    // PostgreSQL aborts the transaction on a failed statement; its journal rolls back.
    if application.is_ok() || DB::DIALECT == Dialect::MySql {
        DB::execute(&mut query, connection).await?;
    }
    application?;
    Ok(MigrationReport {
        plan,
        revision: Some(revision),
        checksum: digest,
    })
}

async fn history_table<DB: SqlBackend>(connection: &mut DB::Connection) -> Result<(), SqlError> {
    let text = if DB::DIALECT == Dialect::MySql {
        "LONGTEXT"
    } else {
        "TEXT"
    };
    execute::<DB>(connection, &format!("CREATE TABLE IF NOT EXISTS _rydantic_migrations (table_name VARCHAR(255) NOT NULL, revision BIGINT NOT NULL, checksum VARCHAR(64) NOT NULL, status VARCHAR(16) NOT NULL, before_schema {text} NOT NULL, after_schema {text} NOT NULL, changes {text} NOT NULL, applied_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, PRIMARY KEY (table_name, revision))")).await
}
async fn record<DB: SqlBackend>(
    connection: &mut DB::Connection,
    plan: &MigrationPlan,
    digest: &str,
    status: &str,
) -> Result<i64, SqlError> {
    let cast = if DB::DIALECT == Dialect::MySql {
        "CHAR"
    } else {
        "TEXT"
    };
    let sql = format!("SELECT CAST(COALESCE(MAX(revision), 0) + 1 AS {cast}) AS value FROM _rydantic_migrations WHERE table_name = {{bind}}");
    let result = rows::<DB>(connection, &sql, &[&plan.table]).await?;
    let revision = text::<DB>(&result[0], "value")?
        .unwrap_or_default()
        .parse::<i64>()
        .map_err(|_| error("migration_history", "invalid migration revision"))?;
    let mut query = QueryBuilder::<DB>::new("INSERT INTO _rydantic_migrations (table_name, revision, checksum, status, before_schema, after_schema, changes) VALUES (");
    let values = [
        SqlParameter::Text(plan.table.clone()),
        SqlParameter::Integer(revision),
        SqlParameter::Text(digest.into()),
        SqlParameter::Text(status.into()),
        SqlParameter::Text(serde_json::to_string(&plan.before).expect("schema")),
        SqlParameter::Text(serde_json::to_string(&plan.target).expect("schema")),
        SqlParameter::Text(serde_json::to_string(&plan.changes).expect("changes")),
    ];
    for (index, value) in values.into_iter().enumerate() {
        if index > 0 {
            query.push(", ");
        }
        DB::bind_parameter(value, &mut query);
    }
    query.push(")");
    DB::execute(&mut query, connection).await?;
    Ok(revision)
}