Repository: libsql/libsql
File: libsql-server/src/schema/mod.rs (Rust)

Code before:
use crate::query_analysis::StmtKind;

// validate program is valid for migration, and return whether foreign keys should be disabled
pub fn validate_migration(migration: &mut Program) -> Result<bool, Error> {
    let mut steps = migration.steps.iter_mut().peekable();
    let mut explicit_tx = false;
    let mut disable_foreign_key = false;
    // skip pragmas prologue
    while steps.next_if(|s| s.query.stmt.is_pragma()).is_some() {
        disable_foreign_key = true;
    }

    // first step can be a BEGIN
    if let Some(step) = steps.next() {
        if matches!(step.query.stmt.kind, StmtKind::TxnBegin) {
            // neutralize step
            step.query.stmt.stmt = r#"SELECT 'neutralized txn begin'"#.into();
            explicit_tx = true;
        }
    }

    // skip all steps that are not tx items
    while steps.next_if(|s| !s.query.stmt.kind.is_txn()).is_some() {}

    // last stmt can be a tx commit
    if let Some(step) = steps.next() {
        if matches!(step.query.stmt.kind, StmtKind::TxnEnd) {
            if !explicit_tx {
                // transaction is closed but was never opened
                return Err(Error::MigrationContainsTransactionStatements);
            }
            // neutralize step
            step.query.stmt.stmt = r#"SELECT 'neutralized txn end'"#.into();
        }
    }

    // validate pragma epilogue
    if steps.by_ref().any(|s| !s.query.stmt.is_pragma()) {
        // only accept pragmas after tx end
        return Err(Error::MigrationContainsTransactionStatements);
    }

    Ok(disable_foreign_key)
}


Review comment (performance):
This is not a regression but while we're at this we might fix the transactional batch case when both commit and rollback are added at the end of hrana batch. This code currently expects only one transaction end statement at the end but for our transactional batches we actually have two so we should probably handle both 1 and 2 transaction end statements here.

Code after:
use crate::query_analysis::StmtKind;

// validate program is valid for migration, and return whether foreign keys should be disabled
pub fn validate_migration(migration: &mut Program) -> Result<bool, Error> {
    let mut steps = migration.steps_mut().unwrap().iter_mut().peekable();
    let mut explicit_tx = false;
    let mut disable_foreign_key = false;
    // skip pragmas prologue
    while steps.next_if(|s| s.query.stmt.is_pragma()).is_some() {
        disable_foreign_key = true;
    }

    // first step can be a BEGIN
    if let Some(step) = steps.next() {
        if matches!(step.query.stmt.kind, StmtKind::TxnBegin) {
            // neutralize step
            step.query.stmt.stmt = r#"SELECT 'neutralized txn begin'"#.into();
            explicit_tx = true;
        }
    }

    // skip all steps that are not tx items
    while steps.next_if(|s| !s.query.stmt.kind.is_txn()).is_some() {}

    // last stmt can be a tx commit
    while let Some(step) = steps.next_if(|s| s.query.stmt.kind.is_txn()) {
        if matches!(step.query.stmt.kind, StmtKind::TxnEnd) {
            if !explicit_tx {
                // transaction is closed but was never opened
                return Err(Error::MigrationContainsTransactionStatements);
            }
            // neutralize step
            step.query.stmt.stmt = r#"SELECT 'neutralized txn component'"#.into();
        }
    }

    // validate pragma epilogue
    if steps.by_ref().any(|s| !s.query.stmt.is_pragma()) {
        // only accept pragmas after tx end
        return Err(Error::MigrationContainsTransactionStatements);
    }

    Ok(disable_foreign_key)
}