fix(storage-cassandra): qualify migration statements with their keyspace instead of USE

crates/storage-cassandra/src/migrations.rs

@@ -30,6 +30,9 @@
 //! 2. Use the next sequential version number (e.g., V003 after V002)
 //! 3. Add the migration to the `CATALOG_MIGRATIONS` array below
 //! 4. Migrations are embedded at compile time via `include_str!()`
+//! 5. Name tables without a keyspace. The runner writes the keyspace it
+//!    migrates into each `CREATE TABLE`, `CREATE INDEX`, `ALTER TABLE` and
+//!    `INSERT INTO` statement, and rejects any other statement form.
 
 use cdrs_tokio::cluster::TcpConnectionManager;
 use cdrs_tokio::cluster::session::Session;
@@ -87,19 +90,13 @@ pub async fn run_catalog_migrations(
 ) -> OpResult<()> {
     println!("--- Running catalog migrations...");
 
-    // Set keyspace context for all subsequent queries
-    session
-        .query(format!("USE {keyspace}"))
-        .await
-        .map_err(|e| OpError::Internal(format!("Failed to USE keyspace: {e}")))?;
-
     for (filename, sql) in CATALOG_MIGRATIONS {
         if is_migration_applied(session, keyspace, filename).await? {
             println!("    {filename} — already applied, skipping.");
             continue;
         }
         println!("    Applying {filename}...");
-        execute_migration(session, sql).await?;
+        execute_migration(session, keyspace, filename, sql).await?;
         record_migration(session, keyspace, filename).await?;
     }
     println!("    Migrations applied.");
@@ -110,19 +107,13 @@ pub async fn run_catalog_migrations(
 pub async fn run_data_migrations(session: &Arc<CassandraSession>, keyspace: &str) -> OpResult<()> {
     println!("--- Running data migrations...");
 
-    // Set keyspace context for all subsequent queries
-    session
-        .query(format!("USE {keyspace}"))
-        .await
-        .map_err(|e| OpError::Internal(format!("Failed to USE keyspace: {e}")))?;
-
     for (filename, sql) in DATA_MIGRATIONS {
         if is_migration_applied(session, keyspace, filename).await? {
             println!("    {filename} — already applied, skipping.");
             continue;
         }
         println!("    Applying {filename}...");
-        execute_migration(session, sql).await?;
+        execute_migration(session, keyspace, filename, sql).await?;
         record_migration(session, keyspace, filename).await?;
     }
     println!("    Data migrations applied.");
@@ -200,22 +191,137 @@ async fn is_migration_applied(
     Ok(applied)
 }
 
-/// Execute a migration by splitting on semicolons and running each statement.
-async fn execute_migration(session: &Arc<CassandraSession>, sql: &str) -> OpResult<()> {
-    for statement in sql.split(';').filter(|s| !s.trim().is_empty()) {
-        let stmt = statement.trim();
-        if stmt.is_empty() {
-            continue;
-        }
-
+/// Execute a migration by splitting on semicolons and running each statement
+/// with the table it names written as `keyspace.table`.
+///
+/// The runners never select a keyspace with `USE` (#632). cdrs-tokio runs a
+/// `USE` on one pooled connection and copies it to the others from a
+/// background task, so a statement relying on it can run in the keyspace the
+/// connection it lands on had selected before.
+async fn execute_migration(
+    session: &Arc<CassandraSession>,
+    keyspace: &str,
+    filename: &str,
+    sql: &str,
+) -> OpResult<()> {
+    for statement in migration_statements(keyspace, filename, sql)? {
         session
-            .query(stmt)
+            .query(statement)
             .await
             .map_err(|e| OpError::Internal(format!("Migration failed: {e}")))?;
     }
     Ok(())
 }
 
+/// The statements of migration file `filename`, split on semicolons, each
+/// qualified with `keyspace` by [`qualify_statement`]. If one cannot be
+/// qualified, the file fails before any of its statements runs.
+#[doc(hidden)]
+pub fn migration_statements(keyspace: &str, filename: &str, sql: &str) -> OpResult<Vec<String>> {
+    sql.split(';')
+        .map(str::trim)
+        .filter(|statement| !statement.is_empty())
+        .map(|statement| qualify_statement(keyspace, filename, statement))
+        .collect()
+}
+
+/// `statement` with the table it names written as `keyspace.table`.
+///
+/// Recognises `CREATE TABLE [IF NOT EXISTS]`,
+/// `CREATE INDEX [IF NOT EXISTS] [name] ON`, `ALTER TABLE [IF EXISTS]` and
+/// `INSERT INTO`, each naming an unquoted table without a keyspace. Any other
+/// statement is an error naming `filename`, so no statement runs in whatever
+/// keyspace the session has selected.
+fn qualify_statement(keyspace: &str, filename: &str, statement: &str) -> OpResult<String> {
+    let body = skip_comments(statement);
+    if body.is_empty() {
+        return Err(OpError::Internal(format!(
+            "Migration file {filename} has a fragment between semicolons that holds only a \
+             comment (a semicolon inside a comment, or a comment after the last statement): {}",
+            statement.lines().next().unwrap_or_default()
+        )));
+    }
+    let at = table_name_offset(body).ok_or_else(|| {
+        OpError::Internal(format!(
+            "Migration file {filename} has a statement that names no table the runner can \
+             qualify with a keyspace: {}",
+            body.lines().next().unwrap_or_default()
+        ))
+    })?;
+    let at = statement.len() - body.len() + at;
+    Ok(format!(
+        "{}{keyspace}.{}",
+        &statement[..at],
+        &statement[at..]
+    ))
+}
+
+/// `statement` from its first character outside whitespace and `--` or `//`
+/// line comments.
+fn skip_comments(mut statement: &str) -> &str {
+    loop {
+        statement = statement.trim_start();
+        if !statement.starts_with("--") && !statement.starts_with("//") {
+            return statement;
+        }
+        statement = statement.split_once('\n').map_or("", |(_, rest)| rest);
+    }
+}
+
+/// The byte offset in `body` of the table name, for the statement forms
+/// [`qualify_statement`] recognises.
+fn table_name_offset(body: &str) -> Option<usize> {
+    fn is(word: Option<(usize, &str)>, keyword: &str) -> bool {
+        word.is_some_and(|(_, word)| word.eq_ignore_ascii_case(keyword))
+    }
+    let mut words = words(body).peekable();
+    let (first, second) = (words.next(), words.next());
+    if is(first, "CREATE") && (is(second, "TABLE") || is(second, "INDEX")) {
+        if words
+            .next_if(|&(_, w)| w.eq_ignore_ascii_case("IF"))
+            .is_some()
+            && !(is(words.next(), "NOT") && is(words.next(), "EXISTS"))
+        {
+            return None;
+        }
+        // CREATE INDEX [name] ON table: the index name is optional.
+        if is(second, "INDEX") && !is(words.next(), "ON") && !is(words.next(), "ON") {
+            return None;
+        }
+    } else if is(first, "ALTER") && is(second, "TABLE") {
+        if words
+            .next_if(|&(_, w)| w.eq_ignore_ascii_case("IF"))
+            .is_some()
+            && !is(words.next(), "EXISTS")
+        {
+            return None;
+        }
+    } else if !(is(first, "INSERT") && is(second, "INTO")) {
+        return None;
+    }
+    let (at, table) = words.next()?;
+    let mut chars = table.chars();
+    let unquoted = chars.next().is_some_and(|c| c.is_ascii_alphabetic())
+        && chars.all(|c| c.is_ascii_alphanumeric() || c == '_');
+    unquoted.then_some(at)
+}
+
+/// The words of `body` with their byte offsets, a word being a run of
+/// characters ended by whitespace or an opening parenthesis.
+fn words(body: &str) -> impl Iterator<Item = (usize, &str)> {
+    let mut at = 0;
+    std::iter::from_fn(move || {
+        let rest = &body[at..];
+        let start = at + (rest.len() - rest.trim_start().len());
+        let word = &body[start..];
+        let end = word
+            .find(|c: char| c.is_whitespace() || c == '(')
+            .unwrap_or(word.len());
+        at = start + end;
+        (end > 0).then(|| (start, &word[..end]))
+    })
+}
+
 /// Record a migration in the `schema_history` table.
 async fn record_migration(
     session: &Arc<CassandraSession>,
@@ -306,4 +412,144 @@ mod tests {
         );
         assert!(super::migration_version("005_x.cql").is_err());
     }
+
+    #[test]
+    fn qualifies_each_statement_form_with_the_keyspace() {
+        for (statement, qualified) in [
+            (
+                "CREATE TABLE IF NOT EXISTS settings (\n    key text PRIMARY KEY\n)",
+                "CREATE TABLE IF NOT EXISTS ks_a.settings (\n    key text PRIMARY KEY\n)",
+            ),
+            (
+                "create table accounts(account_id text PRIMARY KEY)",
+                "create table ks_a.accounts(account_id text PRIMARY KEY)",
+            ),
+            (
+                "CREATE INDEX IF NOT EXISTS ON tables (table_id)",
+                "CREATE INDEX IF NOT EXISTS ON ks_a.tables (table_id)",
+            ),
+            (
+                "CREATE INDEX ON tables (table_id)",
+                "CREATE INDEX ON ks_a.tables (table_id)",
+            ),
+            (
+                "CREATE INDEX idx ON tables (table_id)",
+                "CREATE INDEX idx ON ks_a.tables (table_id)",
+            ),
+            (
+                "CREATE INDEX IF NOT EXISTS idx_txn_started_at ON transaction_ledger (started_at)",
+                "CREATE INDEX IF NOT EXISTS idx_txn_started_at ON ks_a.transaction_ledger (started_at)",
+            ),
+            (
+                "ALTER TABLE tables ADD IF NOT EXISTS ttl_backfill_cursor text",
+                "ALTER TABLE ks_a.tables ADD IF NOT EXISTS ttl_backfill_cursor text",
+            ),
+            (
+                "ALTER TABLE IF EXISTS tables ADD IF NOT EXISTS c text",
+                "ALTER TABLE IF EXISTS ks_a.tables ADD IF NOT EXISTS c text",
+            ),
+            (
+                "alter table if exists tables with comment = 'x'",
+                "alter table if exists ks_a.tables with comment = 'x'",
+            ),
+            (
+                "-- Seed settings\n// two comments\nINSERT INTO settings (key) VALUES ('a')",
+                "-- Seed settings\n// two comments\nINSERT INTO ks_a.settings (key) VALUES ('a')",
+            ),
+        ] {
+            let got =
+                super::qualify_statement("ks_a", "V999__test.cql", statement).expect(statement);
+            assert_eq!(got, qualified);
+        }
+    }
+
+    #[test]
+    fn rejects_statements_it_cannot_qualify() {
+        for statement in [
+            "UPDATE settings SET value = 'x' WHERE key = 'y'",
+            "DROP TABLE settings",
+            "CREATE TABLE other.settings (key text PRIMARY KEY)",
+            "CREATE TABLE \"Settings\" (key text PRIMARY KEY)",
+            "CREATE TABLE IF EXISTS settings (key text PRIMARY KEY)",
+            "CREATE INDEX IF EXISTS idx ON tables (table_id)",
+            "ALTER TABLE IF NOT EXISTS tables ADD c text",
+            "CREATE INDEX idx settings (table_id)",
+            "-- a comment and no statement",
+        ] {
+            assert!(
+                super::qualify_statement("ks_a", "V999__test.cql", statement).is_err(),
+                "{statement}"
+            );
+        }
+    }
+
+    /// A semicolon inside a comment, or a comment after a file's last
+    /// statement, leaves a fragment that holds only a comment.
+    #[test]
+    fn a_fragment_of_only_a_comment_fails_naming_the_file() {
+        let sql = "CREATE TABLE t (k int PRIMARY KEY);\n-- a note after the last statement\n";
+        let Err(super::OpError::Internal(message)) =
+            super::migration_statements("ks_a", "V999__note.cql", sql)
+        else {
+            panic!("a fragment that holds only a comment was accepted");
+        };
+        assert!(
+            message.contains("V999__note.cql")
+                && message.contains("holds only a comment")
+                && message.ends_with("-- a note after the last statement"),
+            "{message}"
+        );
+    }
+
+    /// Every statement of every embedded migration file, as the runner sends
+    /// it, is the file's statement with the keyspace written in once, in front
+    /// of its table: after the statement's leading comments, as a word of its
+    /// own after `ON` in a `CREATE INDEX`, after `INTO` in an `INSERT INTO`,
+    /// and after `TABLE` or `EXISTS` otherwise, and not in front of the `IF`
+    /// that opens an `IF EXISTS` or `IF NOT EXISTS` clause.
+    #[test]
+    fn every_migration_statement_names_the_keyspace() {
+        let files = super::CATALOG_MIGRATIONS
+            .iter()
+            .chain(super::DATA_MIGRATIONS);
+        for (filename, sql) in files {
+            let sent = super::migration_statements("ks_a", filename, sql).expect(filename);
+            let written: Vec<&str> = sql
+                .split(';')
+                .map(str::trim)
+                .filter(|s| !s.is_empty())
+                .collect();
+            assert_eq!(sent.len(), written.len(), "{filename}");
+            for (sent, written) in sent.iter().zip(written) {
+                assert_eq!(sent.matches("ks_a.").count(), 1, "{filename}: {sent}");
+                assert_eq!(sent.replacen("ks_a.", "", 1), written, "{filename}");
+                let (before, after) = super::skip_comments(sent)
+                    .split_once("ks_a.")
+                    .unwrap_or_else(|| panic!("{filename}: keyspace in a leading comment: {sent}"));
+                let keyword = if before.ends_with(char::is_whitespace) {
+                    before.split_whitespace().next_back().unwrap_or_default()
+                } else {
+                    ""
+                };
+                let table = after
+                    .split(|c: char| c.is_whitespace() || c == '(')
+                    .next()
+                    .unwrap_or_default();
+                let mut words = super::skip_comments(written)
+                    .split_whitespace()
+                    .map(str::to_ascii_uppercase);
+                let expected: &[&str] = match (words.next().as_deref(), words.next().as_deref()) {
+                    (Some("CREATE"), Some("INDEX")) => &["ON"],
+                    (Some("INSERT"), Some("INTO")) => &["INTO"],
+                    _ => &["TABLE", "EXISTS"],
+                };
+                assert!(
+                    expected.iter().any(|k| keyword.eq_ignore_ascii_case(k))
+                        && !table.is_empty()
+                        && !table.eq_ignore_ascii_case("IF"),
+                    "{filename}: keyspace not in front of the table: {sent}"
+                );
+            }
+        }
+    }
 }

crates/storage-cassandra/tests/direct/backup_engine.rs

@@ -1573,11 +1573,23 @@ async fn test_restore_of_deleted_backup_is_not_found() {
 
 /// Catalog migrations V001-V004: the catalog shape before V005 added
 /// `backups_by_arn.definition/active_restores` and `tables.restore_*`.
-const CATALOG_THROUGH_V004: &[&str] = &[
-    include_str!("../../migrations/catalog/V001__initial_schema.cql"),
-    include_str!("../../migrations/catalog/V002__backup_restore.cql"),
-    include_str!("../../migrations/catalog/V003__account_scoped_idempotency.cql"),
-    include_str!("../../migrations/catalog/V004__ttl_backfill_cursor.cql"),
+const CATALOG_THROUGH_V004: &[(&str, &str)] = &[
+    (
+        "V001__initial_schema.cql",
+        include_str!("../../migrations/catalog/V001__initial_schema.cql"),
+    ),
+    (
+        "V002__backup_restore.cql",
+        include_str!("../../migrations/catalog/V002__backup_restore.cql"),
+    ),
+    (
+        "V003__account_scoped_idempotency.cql",
+        include_str!("../../migrations/catalog/V003__account_scoped_idempotency.cql"),
+    ),
+    (
+        "V004__ttl_backfill_cursor.cql",
+        include_str!("../../migrations/catalog/V004__ttl_backfill_cursor.cql"),
+    ),
 ];
 
 /// An engine on its own keyspace prefix (so its sweeps and workers touch no
@@ -1610,11 +1622,17 @@ async fn isolated_engine(
             .unwrap();
     }
     if through_v004 {
-        // Statement by statement, as the migration runner does, recording
-        // versions 1-4 in schema_history as the runner did for a 0.0.3 catalog.
-        session.query(format!("USE {catalog_ks}")).await.unwrap();
-        for file in CATALOG_THROUGH_V004 {
-            for statement in file.split(';').map(str::trim).filter(|s| !s.is_empty()) {
+        // Statement by statement, each qualified with the catalog keyspace by
+        // the function the migration runner uses, recording versions 1-4 in
+        // schema_history as the runner did for a 0.0.3 catalog.
+        for (filename, sql) in CATALOG_THROUGH_V004 {
+            let statements = extenddb_storage_cassandra::migrations::migration_statements(
+                &catalog_ks,
+                filename,
+                sql,
+            )
+            .expect(filename);
+            for statement in statements {
                 session.query(statement).await.expect("V001-V004 statement");
             }
         }