From 6abe183d8efeed8f4ffda0238f6a2b4f604186df Mon Sep 17 00:00:00 2001 From: Lee Hannigan Date: Sat, 24 Oct 2026 21:56:32 +0000 Subject: [PATCH] perf(storage-cassandra): restore copy writes plain INSERTs, not the LWT put path RestoreTableFromBackup copied each backup item through put_item_impl: a read of the (empty) target row, a lightweight transaction, and a logged batch for the index rows, eight items at a time. The target is CREATING for the whole copy and the copy is its only writer, so none of that guards anything. data/restore_copy.rs holds the copy's reads and writes. read_backup_bucket returns one payload bucket's items. restore_writer resolves the target's indexes once per copy (LOCAL_QUORUM), and write_restored_item writes an item as plain INSERTs with no read and no IF condition: the base row (pk, sort key column, item_data, version) and one row per GSI and LSI the item carries keys for, issued concurrently at LOCAL_QUORUM. Index rows come from index_row_insert, so their shape and projection are the put path's. The index key type and emptiness validation is kept. version is the clock reading in microseconds, the value a later IF version = ? fence compares against. Neither the base row nor the index rows carry a client timestamp: every later index mutation on this tree (put, update, delete, transaction commit, the GSI queue worker) is stamped by the coordinator, and a restored row stamped from the host clock that ran the restore could outlive a later delete whose coordinator clock trailed that host. Nothing else is written: a restored table starts with streams disabled, TTL settings are not restored, and every GSI row is written in the copy, so there is no gsi_pending entry. copy_backup_into_target reads payload buckets two ahead of the writes and writes 64 items at a time across bucket boundaries, checking the target's ownership per item as well as per bucket, so a swept target stops taking writes within one item. The size accumulation and the item count check are unchanged. The IndexWrites enum and put_restored_item that the old copy needed are removed from put_get_item.rs. Measured on one host against a single cassandra:4.1 node with a 4 GB heap, from the CREATING response to ACTIVE: 60,000 items with no index 27.8 s before, 2.2 s after; 20,000 items with one GSI and one LSI 33.9 s before, 0.7 s after. Item counts, every item and every index row compared equal between source and target in both shapes. Direct tests: restored rows take every later write (update, conditional put and delete, transactions, across GSI and LSI rows); the index rows are coordinator-stamped; a backup item with a mistyped index key fails the restore; a recorded item count that disagrees with the payload fails the restore; the restored row shape equals the put path's. Four existing tests that probe the CREATING window use 10,000 items so the copy outlasts their probes, and the simultaneous-restore test holds a registration on the backup so exclusivity is asserted against a held registration rather than a race. Closes #404 --- crates/storage-cassandra/src/backup_engine.rs | 109 ++-- crates/storage-cassandra/src/data/index.rs | 33 +- crates/storage-cassandra/src/data/mod.rs | 1 + .../src/data/put_get_item.rs | 79 +-- .../src/data/restore_copy.rs | 442 +++++++++++++ .../tests/direct/backup_engine.rs | 138 ++-- .../tests/direct/restore_timing.rs | 270 ++++++++ .../tests/direct/restore_write.rs | 595 ++++++++++++++++++ .../tests/direct/restore_write_controls.rs | 485 ++++++++++++++ .../tests/direct_integration.rs | 6 + 10 files changed, 1989 insertions(+), 169 deletions(-) create mode 100644 crates/storage-cassandra/src/data/restore_copy.rs create mode 100644 crates/storage-cassandra/tests/direct/restore_timing.rs create mode 100644 crates/storage-cassandra/tests/direct/restore_write.rs create mode 100644 crates/storage-cassandra/tests/direct/restore_write_controls.rs diff --git a/crates/storage-cassandra/src/backup_engine.rs b/crates/storage-cassandra/src/backup_engine.rs index 7c070934..f6895bba 100755 --- a/crates/storage-cassandra/src/backup_engine.rs +++ b/crates/storage-cassandra/src/backup_engine.rs @@ -73,8 +73,8 @@ use extenddb_storage::backup_definition::{ }; use extenddb_storage::error::StorageError; use extenddb_storage::{BackupEngine, BackupRestoreSource, MetadataEngine, RestoreRequest}; -use futures::StreamExt; use futures::future::BoxFuture; +use futures::{StreamExt, TryStreamExt}; use crate::CassandraEngine; use crate::create_table::RestoreTargetColumns; @@ -87,8 +87,13 @@ use crate::delete_table::DeleteTableCaller; const BACKUP_ITEMS_PER_BUCKET: i64 = 64; const BACKUP_SCAN_PAGE_SIZE: i64 = 1_000; -/// Items of one bucket a restore copy writes concurrently. -const RESTORE_WRITE_CONCURRENCY: usize = 8; +/// Items a restore copy writes concurrently. Each write is a plain INSERT +/// per row (`data::restore_copy`), so this bounds in-flight statements per +/// index-less item; an item with indexes adds one statement per index row. +const RESTORE_WRITE_CONCURRENCY: usize = 64; + +/// Backup payload buckets a restore copy reads ahead of its writes. +const RESTORE_READ_AHEAD_BUCKETS: usize = 2; /// How often a running restore copy refreshes its target's heartbeat. pub const RESTORE_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(15); @@ -972,8 +977,8 @@ impl CassandraEngine { key_schema: job.key_schema.clone(), base_key_schema: job.key_schema.clone(), attribute_definitions: job.attribute_definitions.clone(), - // The put path reads the target's indexes from the catalog by - // table_id; these routing fields are not consulted by it. + // The restore write path reads the target's indexes from the + // catalog by table_id; these routing fields are not consulted by it. has_lsi: false, global_secondary_indexes: Vec::new(), local_secondary_indexes: Vec::new(), @@ -981,55 +986,55 @@ impl CassandraEngine { vector_indexes: Vec::new(), }; let account_keyspace = self.account_keyspace(&job.account_id); - let query = format!( - "SELECT item_data FROM {account_keyspace}.backup_items WHERE backup_arn = ? AND bucket = ?" - ); + let writer = self + .restore_writer(&key_info) + .await + .map_err(CopyError::Failed)?; + let read_bucket = |bucket: i64| { + let account_keyspace = &account_keyspace; + async move { + if heartbeat.is_lost() { + return Err(CopyError::OwnershipLost); + } + self.read_backup_bucket(account_keyspace, &job.backup_arn, bucket) + .await + .map_err(CopyError::Failed) + } + }; + // Buckets are read ahead of the writes so the copy is never waiting on + // a payload read with nothing in flight; the write stage runs + // RESTORE_WRITE_CONCURRENCY items at a time across bucket boundaries. let mut restored_count = 0_i64; let mut restored_size = 0_i64; - for bucket in 0..bucket_count(job.item_count) { - if heartbeat.is_lost() { - return Err(CopyError::OwnershipLost); - } - #[allow(clippy::cast_possible_truncation)] - let bucket_i32 = bucket as i32; - let rows = self - .session - .query_with_values( - &query, - cdrs_tokio::query_values!(job.backup_arn.as_str(), bucket_i32), - ) - .await - .map_err(|e| { - CopyError::Failed(StorageError::Internal(format!("Read backup payload: {e}"))) - })? - .response_body() - .map_err(|e| { - CopyError::Failed(StorageError::Internal(format!("Parse backup payload: {e}"))) - })? - .into_rows() - .unwrap_or_default(); - let mut items = Vec::with_capacity(rows.len()); - for row in rows { - let item_data: String = row.get_r_by_name("item_data").map_err(|e| { - CopyError::Failed(StorageError::Internal(format!("Parse backup item: {e}"))) - })?; - #[allow(clippy::cast_possible_wrap)] - let item_len = item_data.len() as i64; - restored_size = restored_size.saturating_add(item_len); - let item = serde_json::from_str(&item_data).map_err(|e| { - CopyError::Failed(StorageError::Internal(format!( - "Deserialize backup item: {e}" - ))) - })?; - items.push(item); - } - let mut writes = futures::stream::iter(items) - .map(|item| self.put_restored_item(&key_info, item)) - .buffer_unordered(RESTORE_WRITE_CONCURRENCY); - while let Some(result) = writes.next().await { - result.map_err(CopyError::Failed)?; - restored_count += 1; - } + let mut writes = futures::stream::iter(0..bucket_count(job.item_count)) + .map(read_bucket) + .buffered(RESTORE_READ_AHEAD_BUCKETS) + .map_ok(|items| futures::stream::iter(items.into_iter().map(Ok))) + .try_flatten() + .map_ok(|(item, item_len)| { + let writer = &writer; + let heartbeat = &heartbeat; + async move { + // The read stage checks ownership per bucket; with two + // buckets read ahead and 64 writes in flight, up to three + // buckets of items would otherwise still be written into + // a target the sweep has claimed. + if heartbeat.is_lost() { + return Err(CopyError::OwnershipLost); + } + self.write_restored_item(writer, &item) + .await + .map_err(CopyError::Failed)?; + Ok::(item_len) + } + }) + .try_buffer_unordered(RESTORE_WRITE_CONCURRENCY); + while let Some(result) = writes.next().await { + let item_len = result?; + #[allow(clippy::cast_possible_wrap)] + let item_len = item_len as i64; + restored_size = restored_size.saturating_add(item_len); + restored_count += 1; } if restored_count != job.item_count { return Err(CopyError::Failed(StorageError::Internal(format!( diff --git a/crates/storage-cassandra/src/data/index.rs b/crates/storage-cassandra/src/data/index.rs index 4ebd72da..8739ab93 100644 --- a/crates/storage-cassandra/src/data/index.rs +++ b/crates/storage-cassandra/src/data/index.rs @@ -506,6 +506,34 @@ pub(crate) fn insert_index_row_multi( idx_sks: &[(&str, ScalarAttributeType)], base_sks: &[(&str, ScalarAttributeType)], ) -> Result<(), StorageError> { + let (cql, qv) = index_row_insert( + account_keyspace, + idx_table, + item, + projected, + index_ks, + base_ks, + idx_sks, + base_sks, + )?; + let old = std::mem::replace(batch, BatchQueryBuilder::new()); + *batch = old.add_query(cql, qv); + Ok(()) +} + +/// The parameterized INSERT of `item`'s row in the index table `idx_table`, +/// as a statement and its bound values. +#[allow(clippy::too_many_arguments)] +pub(crate) fn index_row_insert( + account_keyspace: &str, + idx_table: &str, + item: &Item, + projected: &Item, + index_ks: &[KeySchemaElement], + base_ks: &[KeySchemaElement], + idx_sks: &[(&str, ScalarAttributeType)], + base_sks: &[(&str, ScalarAttributeType)], +) -> Result<(String, cdrs_tokio::query::QueryValues), StorageError> { let idx_pk_text = composite_pk_to_text(item, index_ks)?; let base_pk_text = composite_pk_to_text(item, base_ks)?; @@ -555,10 +583,7 @@ pub(crate) fn insert_index_row_multi( cols.join(", "), placeholders ); - let qv = cdrs_tokio::query::QueryValues::SimpleValues(values); - let old = std::mem::replace(batch, BatchQueryBuilder::new()); - *batch = old.add_query(cql, qv); - Ok(()) + Ok((cql, cdrs_tokio::query::QueryValues::SimpleValues(values))) } /// Convert a `SortKeyValue` to a `cdrs_tokio` bound `Value`. diff --git a/crates/storage-cassandra/src/data/mod.rs b/crates/storage-cassandra/src/data/mod.rs index 946dbd04..2d6e028d 100644 --- a/crates/storage-cassandra/src/data/mod.rs +++ b/crates/storage-cassandra/src/data/mod.rs @@ -21,6 +21,7 @@ pub mod index; mod put_get_item; mod query; mod query_helpers; +mod restore_copy; mod scan; pub mod transaction_ledger; mod transactions; diff --git a/crates/storage-cassandra/src/data/put_get_item.rs b/crates/storage-cassandra/src/data/put_get_item.rs index aae7567a..46e4e066 100644 --- a/crates/storage-cassandra/src/data/put_get_item.rs +++ b/crates/storage-cassandra/src/data/put_get_item.rs @@ -17,16 +17,6 @@ use super::{json_to_item, query_with_pk_sk}; use crate::CassandraEngine; use crate::stream_util::stream_record_statement; -/// How a put maintains GSIs. -#[derive(Clone, Copy, PartialEq, Eq)] -enum IndexWrites { - /// Each GSI by its propagation delay: synchronous at 0, otherwise through - /// the `gsi_pending` queue. - AsConfigured, - /// Every GSI in the item's own batch, as with a 0 delay. - AllSynchronous, -} - impl CassandraEngine { /// Implementation of `DataEngine::put_item`. /// @@ -42,68 +32,12 @@ impl CassandraEngine { condition: Option<&Expr>, maps: &ExpressionMaps, stream: Option<&StreamCapture>, - ) -> Result, StorageError> { - self.put_item_with_retry( - key_info, - item, - return_old, - condition, - maps, - stream, - IndexWrites::AsConfigured, - ) - .await - } - - /// The write a restore copy makes for one backup item: the ordinary put - /// path (same base-row write, same index maintenance helpers, same - /// projection), with every GSI written in the item's own logged batch as - /// LSIs always are, instead of through the asynchronous propagation queue. - /// A restore target becomes ACTIVE only after its last item is written, - /// and DynamoDB reports a restored table's indexes as complete at ACTIVE; - /// queueing the rows would let ACTIVE precede them. - pub(crate) async fn put_restored_item( - &self, - key_info: &TableKeyInfo, - item: Item, - ) -> Result<(), StorageError> { - self.put_item_with_retry( - key_info, - item, - false, - None, - &ExpressionMaps::default(), - None, - IndexWrites::AllSynchronous, - ) - .await - .map(|_| ()) - } - - #[allow(clippy::too_many_arguments)] - async fn put_item_with_retry( - &self, - key_info: &TableKeyInfo, - item: Item, - return_old: bool, - condition: Option<&Expr>, - maps: &ExpressionMaps, - stream: Option<&StreamCapture>, - index_writes: IndexWrites, ) -> Result, StorageError> { use super::delete_item::{TTL_CLAIM_MAX_RETRIES, ttl_claim_backoff}; for attempt in 0..=TTL_CLAIM_MAX_RETRIES { match self - .put_item_impl_inner( - key_info, - item.clone(), - return_old, - condition, - maps, - stream, - index_writes, - ) + .put_item_impl_inner(key_info, item.clone(), return_old, condition, maps, stream) .await { Err(StorageError::TransactionConflict(message)) @@ -126,7 +60,6 @@ impl CassandraEngine { unreachable!("loop returns on the final attempt") } - #[allow(clippy::too_many_arguments)] async fn put_item_impl_inner( &self, key_info: &TableKeyInfo, @@ -135,7 +68,6 @@ impl CassandraEngine { condition: Option<&Expr>, maps: &ExpressionMaps, stream: Option<&StreamCapture>, - index_writes: IndexWrites, ) -> Result, StorageError> { let data_keyspace = self.account_keyspace(&key_info.account_id); let catalog_keyspace = self.catalog_keyspace(); @@ -149,7 +81,7 @@ impl CassandraEngine { // Fetch indexes for GSI/LSI updates // Both are catalog reads with no data dependency; overlap them. - let (mut indexes, ttl_config) = futures::try_join!( + let (indexes, ttl_config) = futures::try_join!( super::index::fetch_indexes_for_table( &key_info.table_id, &self.session, @@ -174,12 +106,7 @@ impl CassandraEngine { ) .map_err(|e| StorageError::Validation(e.to_string()))?; } - if index_writes == IndexWrites::AllSynchronous { - for index in &mut indexes { - index.propagation_delay_ms = Some(0); - } - } - let sys_delay = if indexes.is_empty() || index_writes == IndexWrites::AllSynchronous { + let sys_delay = if indexes.is_empty() { 0 } else { self.gsi_default_delay_ms diff --git a/crates/storage-cassandra/src/data/restore_copy.rs b/crates/storage-cassandra/src/data/restore_copy.rs new file mode 100644 index 00000000..7787082a --- /dev/null +++ b/crates/storage-cassandra/src/data/restore_copy.rs @@ -0,0 +1,442 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! The restore copy: reading a backup's payload buckets and writing their +//! items into the target. +//! +//! A restore target is `CREATING` for the whole copy: no client request +//! reaches it, and the copy is the only writer. The ordinary put path reads +//! the row first, fences the write with a lightweight transaction and carries +//! the index rows in a logged batch; none of that protects anything here, and +//! each piece costs a round trip. This path writes an item as plain `INSERT`s +//! with no read and no `IF` condition: one for the base row and one per index +//! row, issued concurrently. +//! +//! What a write leaves behind is what an ordinary put of a new item leaves: +//! +//! * the base row's `item_data` and `version`. The version is the clock +//! reading in microseconds, the value the ordinary path gives a new row, so +//! a later `IF version = ?` fence on the restored item compares equal to +//! what that path reads back. The base row carries no explicit write +//! timestamp: ordinary writes commit it through Paxos at the coordinator's +//! clock, and a client-stamped cell ahead of that clock would shadow them; +//! * one row per GSI and LSI the item has keys for, projected as the index +//! declares, with no explicit write timestamp either. On this tree every +//! later index mutation (put, update, delete, transaction commit, the GSI +//! queue worker) is stamped by the coordinator, so a coordinator-stamped +//! restored row loses to any of them whatever the host clocks say. Once +//! index mutations carry a version-derived timestamp, the restored rows +//! take `version` as theirs so the same ordering holds for them. +//! +//! There is no stream record (a restored table starts with streams disabled), +//! no TTL work (TTL settings are not restored) and no `gsi_pending` entry +//! (every index row is written here, so the indexes are complete when the +//! table turns `ACTIVE`). + +use cdrs_tokio::consistency::Consistency; +use cdrs_tokio::query::QueryValues; +use cdrs_tokio::statement::StatementParamsBuilder; +use cdrs_tokio::types::value::Value; +use extenddb_core::types::{Item, TableKeyInfo}; +use extenddb_core::validation::IndexKeyRef; +use extenddb_storage::error::StorageError; +use extenddb_storage::util::{composite_pk_to_text, parse_sk, sk_column, sk_info}; + +use super::ddl::{all_sort_key_info, data_table_name, index_table_name}; +use super::index::{ + IndexMeta, index_row_insert, item_has_index_keys, project_item_for_index, sk_to_value, +}; +use crate::CassandraEngine; + +/// Everything a restore copy needs to write items into one target, resolved +/// once per copy rather than once per item. +pub(crate) struct RestoreWriter { + key_info: TableKeyInfo, + data_keyspace: String, + indexes: Vec, + /// The base-row `INSERT`, with a `?` for `pk`, one per sort key column, + /// `item_data` and `version`. + base_insert: String, +} + +impl CassandraEngine { + /// The items of one payload bucket of `backup_arn`, each with the byte + /// length of its stored JSON. + pub(crate) async fn read_backup_bucket( + &self, + account_keyspace: &str, + backup_arn: &str, + bucket: i64, + ) -> Result, StorageError> { + use cdrs_tokio::types::IntoRustByName; + + let query = format!( + "SELECT item_data FROM {account_keyspace}.backup_items WHERE backup_arn = ? AND bucket = ?" + ); + #[allow(clippy::cast_possible_truncation)] + let bucket = bucket as i32; + let rows = self + .session + .query_with_values(&query, cdrs_tokio::query_values!(backup_arn, bucket)) + .await + .map_err(|e| StorageError::Internal(format!("Read backup payload: {e}")))? + .response_body() + .map_err(|e| StorageError::Internal(format!("Parse backup payload: {e}")))? + .into_rows() + .unwrap_or_default(); + let mut items = Vec::with_capacity(rows.len()); + for row in rows { + let item_data: String = row + .get_r_by_name("item_data") + .map_err(|e| StorageError::Internal(format!("Parse backup item: {e}")))?; + let item: Item = serde_json::from_str(&item_data) + .map_err(|e| StorageError::Internal(format!("Deserialize backup item: {e}")))?; + items.push((item, item_data.len())); + } + Ok(items) + } + + /// Resolve the restore write path for `key_info`'s target. + /// + /// The target's indexes are read from the catalog at `LOCAL_QUORUM`: the + /// restore wrote them moments ago, and an index a lagging replica omits + /// would be left empty for every restored item. + pub(crate) async fn restore_writer( + &self, + key_info: &TableKeyInfo, + ) -> Result { + let indexes = super::index::fetch_indexes_for_table_quorum( + &key_info.table_id, + &self.session, + &self.catalog_keyspace(), + ) + .await?; + let data_keyspace = self.account_keyspace(&key_info.account_id); + let base_insert = base_insert_cql(&data_keyspace, key_info); + Ok(RestoreWriter { + key_info: key_info.clone(), + data_keyspace, + indexes, + base_insert, + }) + } + + /// Write one backup item into the restore target. + /// + /// # Errors + /// + /// [`StorageError::Validation`] if the item's index key attributes do not + /// fit the target's index key schemas, and the storage error of the first + /// write that fails. + pub(crate) async fn write_restored_item( + &self, + writer: &RestoreWriter, + item: &Item, + ) -> Result<(), StorageError> { + let key_info = &writer.key_info; + if !writer.indexes.is_empty() { + let index_refs: Vec> = writer + .indexes + .iter() + .map(|idx| IndexKeyRef { + index_name: &idx.index_name, + key_schema: &idx.key_schema, + }) + .collect(); + extenddb_core::validation::validate_index_keys( + item, + &index_refs, + &key_info.attribute_definitions, + ) + .map_err(|e| StorageError::Validation(e.to_string()))?; + } + + let version = chrono::Utc::now().timestamp_micros(); + let pk_text = composite_pk_to_text(item, &key_info.key_schema)?; + let item_text = + serde_json::to_value(item).map_err(|e| StorageError::Internal(e.to_string()))?; + let item_text = item_text.to_string(); + + let mut base_values: Vec = Vec::with_capacity(4); + base_values.push(Value::from(pk_text.as_str())); + if let Some((sk_name, sk_type)) = + sk_info(&key_info.key_schema, &key_info.attribute_definitions) + { + let sk_value = item + .get(sk_name) + .ok_or_else(|| StorageError::Internal("missing sort key".to_owned()))?; + base_values.push(sk_to_value(&parse_sk(sk_value, sk_type)?)); + } + base_values.push(item_text.as_str().into()); + base_values.push(version.into()); + + let index_rows = index_row_statements(writer, item)?; + let mut writes = Vec::with_capacity(1 + index_rows.len()); + writes.push(self.restore_write( + &writer.base_insert, + QueryValues::SimpleValues(base_values), + None, + )); + for (cql, values) in &index_rows { + // No client timestamp: every later index mutation on this tree + // (put, update, delete, transaction commit, the GSI queue worker) + // is stamped by the coordinator, so the restored row must be too, + // or a later delete whose coordinator clock trails the host that + // ran the restore would leave it in place. + writes.push(self.restore_write(cql, values.clone(), None)); + } + futures::future::try_join_all(writes).await.map(|_| ()) + } + + /// One plain write at `LOCAL_QUORUM`, at `timestamp` when given. + async fn restore_write( + &self, + cql: &str, + values: QueryValues, + timestamp: Option, + ) -> Result<(), StorageError> { + let mut params = StatementParamsBuilder::new() + .with_consistency(Consistency::LocalQuorum) + .with_values(values); + if let Some(timestamp) = timestamp { + params = params.with_timestamp(timestamp); + } + self.session + .query_with_params(cql, params.build()) + .await + .map_err(|e| StorageError::Internal(format!("Restore write: {e}")))?; + Ok(()) + } +} + +/// The base-row `INSERT` of a restored item: `pk`, the sort key column when +/// the table has one, `item_data` and `version`, each bound by `?`. +fn base_insert_cql(data_keyspace: &str, key_info: &TableKeyInfo) -> String { + let data_table = data_table_name(&key_info.table_id); + let (sort_column, sort_marker) = + match sk_info(&key_info.key_schema, &key_info.attribute_definitions) { + Some((_, sk_type)) => (format!(", {}", sk_column(sk_type)), ", ?"), + None => (String::new(), ""), + }; + format!( + "INSERT INTO {data_keyspace}.{data_table} (pk{sort_column}, item_data, version) \ + VALUES (?{sort_marker}, ?, ?)" + ) +} + +/// The index-row `INSERT`s for `item`: one per GSI and LSI whose key +/// attributes the item carries, built by the same function the ordinary put +/// path uses, so the row shape is the put path's. +fn index_row_statements( + writer: &RestoreWriter, + item: &Item, +) -> Result, StorageError> { + let key_info = &writer.key_info; + let base_sks = all_sort_key_info(&key_info.key_schema, &key_info.attribute_definitions); + let mut statements = Vec::with_capacity(writer.indexes.len()); + for idx in &writer.indexes { + if !item_has_index_keys(item, &idx.key_schema) { + continue; + } + let idx_table = index_table_name(&idx.index_id); + let idx_sks = all_sort_key_info(&idx.key_schema, &key_info.attribute_definitions); + let projected = + project_item_for_index(item, &idx.key_schema, &key_info.key_schema, &idx.projection); + statements.push(index_row_insert( + &writer.data_keyspace, + &idx_table, + item, + &projected, + &idx.key_schema, + &key_info.key_schema, + &idx_sks, + &base_sks, + )?); + } + Ok(statements) +} + +#[cfg(test)] +mod tests { + use super::*; + use extenddb_core::types::{ + AttributeDefinition, AttributeValue, KeySchemaElement, KeyType, Projection, ProjectionType, + ScalarAttributeType, + }; + + fn ks(name: &str, key_type: KeyType) -> KeySchemaElement { + KeySchemaElement { + attribute_name: name.to_owned(), + key_type, + } + } + + fn attr(name: &str, attribute_type: ScalarAttributeType) -> AttributeDefinition { + AttributeDefinition { + attribute_name: name.to_owned(), + attribute_type, + } + } + + fn key_info(key_schema: Vec) -> TableKeyInfo { + TableKeyInfo { + table_name: "T".to_owned(), + account_id: "123456789012".to_owned(), + table_id: "0a1b-2c3d".to_owned(), + key_schema: key_schema.clone(), + base_key_schema: key_schema, + attribute_definitions: vec![ + attr("id", ScalarAttributeType::S), + attr("sort", ScalarAttributeType::N), + attr("gpk", ScalarAttributeType::S), + attr("lsk", ScalarAttributeType::N), + ], + has_lsi: false, + global_secondary_indexes: Vec::new(), + local_secondary_indexes: Vec::new(), + stream_specification: None, + vector_indexes: Vec::new(), + } + } + + fn writer(key_schema: Vec, indexes: Vec) -> RestoreWriter { + let key_info = key_info(key_schema); + RestoreWriter { + base_insert: base_insert_cql("ks", &key_info), + key_info, + data_keyspace: "ks".to_owned(), + indexes, + } + } + + fn index(name: &str, index_type: &str, key_schema: Vec) -> IndexMeta { + IndexMeta { + index_name: name.to_owned(), + index_id: format!("{name}-id"), + index_type: index_type.to_owned(), + key_schema, + projection: Projection { + projection_type: ProjectionType::Include, + non_key_attributes: Some(vec!["value".to_owned()]), + }, + propagation_delay_ms: None, + } + } + + fn item(pairs: &[(&str, AttributeValue)]) -> Item { + pairs + .iter() + .map(|(k, v)| ((*k).to_owned(), v.clone())) + .collect() + } + + #[test] + fn base_insert_names_the_sort_key_column_only_when_the_table_has_one() { + let plain = writer(vec![ks("id", KeyType::Hash)], Vec::new()); + assert_eq!( + plain.base_insert, + "INSERT INTO ks.items_0a1b_2c3d (pk, item_data, version) VALUES (?, ?, ?)" + ); + let sorted = writer( + vec![ks("id", KeyType::Hash), ks("sort", KeyType::Range)], + Vec::new(), + ); + assert_eq!( + sorted.base_insert, + "INSERT INTO ks.items_0a1b_2c3d (pk, sk_n, item_data, version) VALUES (?, ?, ?, ?)" + ); + } + + #[test] + fn one_index_row_per_index_the_item_has_keys_for() { + let writer = writer( + vec![ks("id", KeyType::Hash), ks("sort", KeyType::Range)], + vec![ + index("ByGroup", "GSI", vec![ks("gpk", KeyType::Hash)]), + index( + "ByLocal", + "LSI", + vec![ks("id", KeyType::Hash), ks("lsk", KeyType::Range)], + ), + ], + ); + let full = item(&[ + ("id", AttributeValue::S("a".to_owned())), + ("sort", AttributeValue::N("1".to_owned())), + ("gpk", AttributeValue::S("g".to_owned())), + ("lsk", AttributeValue::N("7".to_owned())), + ("value", AttributeValue::S("v".to_owned())), + ("extra", AttributeValue::S("not projected".to_owned())), + ]); + let statements = index_row_statements(&writer, &full).expect("statements"); + assert_eq!(statements.len(), 2); + assert!( + statements[0] + .0 + .starts_with("INSERT INTO ks.index_ByGroup_id (") + ); + assert!( + statements[1] + .0 + .starts_with("INSERT INTO ks.index_ByLocal_id (") + ); + for (cql, values) in &statements { + assert!( + !cql.contains(" IF "), + "no condition on a restore write: {cql}" + ); + let QueryValues::SimpleValues(values) = values else { + panic!("simple values expected"); + }; + let bound = cql.matches('?').count(); + assert_eq!(bound, values.len(), "{cql}"); + } + + // No GSI key: only the LSI row. + let sparse = item(&[ + ("id", AttributeValue::S("a".to_owned())), + ("sort", AttributeValue::N("2".to_owned())), + ("lsk", AttributeValue::N("8".to_owned())), + ]); + let statements = index_row_statements(&writer, &sparse).expect("statements"); + assert_eq!(statements.len(), 1); + assert!( + statements[0] + .0 + .starts_with("INSERT INTO ks.index_ByLocal_id (") + ); + } + + #[test] + fn the_projection_is_applied_to_the_index_row() { + let writer = writer( + vec![ks("id", KeyType::Hash)], + vec![index("ByGroup", "GSI", vec![ks("gpk", KeyType::Hash)])], + ); + let full = item(&[ + ("id", AttributeValue::S("a".to_owned())), + ("gpk", AttributeValue::S("g".to_owned())), + ("value", AttributeValue::S("v".to_owned())), + ("extra", AttributeValue::S("not projected".to_owned())), + ]); + let projected = project_item_for_index( + &full, + &writer.indexes[0].key_schema, + &writer.key_info.key_schema, + &writer.indexes[0].projection, + ); + assert!(projected.contains_key("value")); + assert!(!projected.contains_key("extra")); + let statements = index_row_statements(&writer, &full).expect("statements"); + let QueryValues::SimpleValues(values) = &statements[0].1 else { + panic!("simple values expected"); + }; + // item_data is the last bound value and carries the projection. + let Some(Value::Some(bytes)) = values.last() else { + panic!("item_data must be bound"); + }; + let item_data = String::from_utf8(bytes.clone()).expect("utf-8 item_data"); + assert!(item_data.contains("\"value\""), "{item_data}"); + assert!(!item_data.contains("not projected"), "{item_data}"); + } +} diff --git a/crates/storage-cassandra/tests/direct/backup_engine.rs b/crates/storage-cassandra/tests/direct/backup_engine.rs index 30fe919a..7fbf43ad 100644 --- a/crates/storage-cassandra/tests/direct/backup_engine.rs +++ b/crates/storage-cassandra/tests/direct/backup_engine.rs @@ -29,7 +29,7 @@ use crate::helpers::{TestAccount, TestTable, setup_engine, unique_test_account}; type Item = BTreeMap; -fn item(id: &str, sort: Option<&str>, value: &str) -> Item { +pub(crate) fn item(id: &str, sort: Option<&str>, value: &str) -> Item { let mut item = BTreeMap::new(); item.insert("id".to_owned(), AttributeValue::S(id.to_owned())); if let Some(sort) = sort { @@ -39,7 +39,7 @@ fn item(id: &str, sort: Option<&str>, value: &str) -> Item { item } -fn key(id: &str, sort: Option<&str>) -> Item { +pub(crate) fn key(id: &str, sort: Option<&str>) -> Item { let mut key = BTreeMap::new(); key.insert("id".to_owned(), AttributeValue::S(id.to_owned())); if let Some(sort) = sort { @@ -48,21 +48,21 @@ fn key(id: &str, sort: Option<&str>) -> Item { key } -fn ks(name: &str, key_type: KeyType) -> KeySchemaElement { +pub(crate) fn ks(name: &str, key_type: KeyType) -> KeySchemaElement { KeySchemaElement { attribute_name: name.to_owned(), key_type, } } -fn attr(name: &str, attribute_type: ScalarAttributeType) -> AttributeDefinition { +pub(crate) fn attr(name: &str, attribute_type: ScalarAttributeType) -> AttributeDefinition { AttributeDefinition { attribute_name: name.to_owned(), attribute_type, } } -async fn activate_tables(engine: &CassandraEngine) { +pub(crate) async fn activate_tables(engine: &CassandraEngine) { tokio::time::sleep(Duration::from_millis(350)).await; engine .process_control_plane_transitions() @@ -71,7 +71,7 @@ async fn activate_tables(engine: &CassandraEngine) { } /// Create `input` in `account_id`, activate it, and return its key info. -async fn create_active_table( +pub(crate) async fn create_active_table( engine: &CassandraEngine, account_id: &str, input: CreateTableInput, @@ -90,7 +90,11 @@ async fn create_active_table( /// Make every GSI of a table synchronous, so the source's indexes are /// complete without the GSI propagation workers (not running in these tests). -async fn make_gsis_synchronous(engine: &CassandraEngine, table_id: &str, names: &[&str]) { +pub(crate) async fn make_gsis_synchronous( + engine: &CassandraEngine, + table_id: &str, + names: &[&str], +) { let query = format!( "UPDATE {}.indexes SET propagation_delay_ms = 0 WHERE table_id = ? AND index_name = ?", engine.catalog_keyspace() @@ -104,7 +108,7 @@ async fn make_gsis_synchronous(engine: &CassandraEngine, table_id: &str, names: } } -async fn put_all(engine: &CassandraEngine, key_info: &TableKeyInfo, items: Vec) { +pub(crate) async fn put_all(engine: &CassandraEngine, key_info: &TableKeyInfo, items: Vec) { let mut writes = futures::stream::iter(items) .map(|item| engine.put_item(key_info, item, false, None, &Default::default(), None)) .buffer_unordered(32); @@ -113,7 +117,11 @@ async fn put_all(engine: &CassandraEngine, key_info: &TableKeyInfo, items: Vec TableDescription { +pub(crate) async fn describe( + engine: &CassandraEngine, + account_id: &str, + name: &str, +) -> TableDescription { engine .describe_table( account_id, @@ -133,7 +141,7 @@ fn flip_last_hex(arn: &str) -> String { } /// Restore `backup_arn` into `target` the way the engine does. -async fn restore( +pub(crate) async fn restore( engine: &CassandraEngine, account_id: &str, backup_arn: &str, @@ -157,7 +165,7 @@ async fn restore( } /// Poll until `name` is ACTIVE; panic after `timeout`. -async fn wait_active( +pub(crate) async fn wait_active( engine: &CassandraEngine, account_id: &str, name: &str, @@ -178,7 +186,7 @@ async fn wait_active( } } -async fn scan_count(engine: &CassandraEngine, key_info: &TableKeyInfo) -> usize { +pub(crate) async fn scan_count(engine: &CassandraEngine, key_info: &TableKeyInfo) -> usize { let mut count = 0; let mut start: Option = None; loop { @@ -195,7 +203,7 @@ async fn scan_count(engine: &CassandraEngine, key_info: &TableKeyInfo) -> usize } /// Every item a Query on `index` for `pk_attr = value` returns, all pages. -async fn query_index( +pub(crate) async fn query_index( engine: &CassandraEngine, key_info: &TableKeyInfo, index: &str, @@ -934,7 +942,7 @@ async fn test_restore_returns_creating_and_survives_the_request() { if crate::helpers::skip_without_cassandra() { return; } - const ITEMS: usize = 1_500; + const ITEMS: usize = 10_000; let observer = setup_engine().await; let (account, _source, backup_arn) = backed_up_plain_table(&observer, "DetachedSource", ITEMS).await; @@ -990,7 +998,7 @@ async fn test_dropped_restore_request_still_reaches_active() { if crate::helpers::skip_without_cassandra() { return; } - const ITEMS: usize = 1_500; + const ITEMS: usize = 10_000; let engine = setup_engine().await; let (account, _source, backup_arn) = backed_up_plain_table(&engine, "DroppedSource", ITEMS).await; @@ -1064,7 +1072,7 @@ async fn test_backup_and_target_are_held_while_restoring() { if crate::helpers::skip_without_cassandra() { return; } - const ITEMS: usize = 1_500; + const ITEMS: usize = 10_000; let engine = setup_engine().await; let (account, source, backup_arn) = backed_up_plain_table(&engine, "HeldSource", ITEMS).await; let account_id = account.account_id.clone(); @@ -1291,7 +1299,52 @@ async fn test_a_restore_is_refused_while_a_target_of_its_backup_is_creating() { ) .await; - // Simultaneous: four restores of b2 started together, exactly one admitted. + // Simultaneous, against a held registration: a CREATING restore target + // of b2 is put in the catalog by hand and registered on b2 the way + // register_restore does, so the registration is held for as long as the + // test wants it (a real copy of b2 is over before a second target's + // schema is created). Four restores of b2 fired at once are all turned + // away as BackupInUse naming that target, and none creates a target. + // Once the registration is released, a restore of b2 is admitted. + let held_id = uuid::Uuid::new_v4().to_string(); + let catalog = engine.catalog_keyspace(); + engine + .session() + .query_with_values( + &format!( + "INSERT INTO {catalog}.tables \ + (account_id, table_name, table_id, table_status, restore_source_backup_arn) \ + VALUES (?, ?, ?, 'CREATING', ?)" + ), + cdrs_tokio::query_values!( + account_id.as_str(), + "HeldTarget", + held_id.as_str(), + b2.as_str() + ), + ) + .await + .expect("held target row"); + let registered = extenddb_storage_cassandra::cassandra_util::query_lwt( + engine.session(), + &format!( + "UPDATE {catalog}.backups_by_arn SET active_restores[?] = ? \ + WHERE account_id = ? AND backup_arn = ? \ + IF backup_status = 'AVAILABLE' AND active_restores = null" + ), + cdrs_tokio::query_values!( + held_id.as_str(), + "HeldTarget", + account_id.as_str(), + b2.as_str() + ), + ) + .await + .expect("register the held restore"); + assert!( + extenddb_storage_cassandra::cassandra_util::lwt_applied(®istered).unwrap(), + "b2 had no registration before the held one" + ); let results = futures::future::join_all((0..4).map(|i| { let engine = engine.clone(); let account_id = account_id.clone(); @@ -1308,27 +1361,38 @@ async fn test_a_restore_is_refused_while_a_target_of_its_backup_is_creating() { } })) .await; - let accepted: Vec = results - .iter() - .enumerate() - .filter_map(|(i, r)| r.is_ok().then_some(i)) - .collect(); - assert_eq!(accepted.len(), 1, "exactly one is accepted: {results:?}"); for (i, result) in results.iter().enumerate() { - if i != accepted[0] { - assert!( - matches!(result, Err(StorageError::BackupInUse(_))), - "restore {i}: {result:?}" - ); - } + assert!( + matches!(result, Err(StorageError::BackupInUse(m)) if m.contains("HeldTarget")), + "restore {i} against the held registration: {result:?}" + ); + let gone = engine + .describe_table( + &account_id, + DescribeTableInput { + table_name: format!("SimTarget{i}"), + }, + ) + .await; + assert!( + matches!(gone, Err(StorageError::TableNotFound(_))), + "refused restore {i} left a target: {gone:?}" + ); } - wait_active( - &engine, - &account_id, - &format!("SimTarget{}", accepted[0]), - Duration::from_secs(180), - ) - .await; + engine + .session() + .query_with_values( + &format!("DELETE FROM {catalog}.tables WHERE account_id = ? AND table_name = ?"), + cdrs_tokio::query_values!(account_id.as_str(), "HeldTarget"), + ) + .await + .expect("remove the held target row"); + // The registration is now stale (its target is gone); the next restore's + // backup check clears it and the restore is admitted. + restore(&engine, &account_id, &b2, "SimTarget0", &overrides) + .await + .expect("a restore once the held registration is released"); + wait_active(&engine, &account_id, "SimTarget0", Duration::from_secs(180)).await; } /// (e) A backup row with no definition (written before V005) restores as a @@ -1442,7 +1506,7 @@ async fn test_sweep_removes_abandoned_restore_target() { if crate::helpers::skip_without_cassandra() { return; } - const ITEMS: usize = 1_500; + const ITEMS: usize = 10_000; let engine = setup_engine().await; let (account, _source, backup_arn) = backed_up_plain_table(&engine, "AbandonedSource", ITEMS).await; diff --git a/crates/storage-cassandra/tests/direct/restore_timing.rs b/crates/storage-cassandra/tests/direct/restore_timing.rs new file mode 100644 index 00000000..f610d03c --- /dev/null +++ b/crates/storage-cassandra/tests/direct/restore_timing.rs @@ -0,0 +1,270 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! Restore copy timing against a live Cassandra. +//! +//! Ignored by default: the numbers from a single-node container are only +//! meaningful relative to each other on the same machine, so these run +//! manually around a change to the restore copy, never in CI: +//! +//! ```text +//! cargo test -p extenddb-storage-cassandra --test direct_integration --release -- \ +//! --ignored --test-threads=1 --nocapture restore_timing +//! ``` +//! +//! Each case backs up a table, restores it, reports the time from the +//! `CREATING` response to `ACTIVE`, and checks the restored table against the +//! source: item count, every item, and every index row. + +use std::collections::BTreeMap; +use std::time::{Duration, Instant}; + +use extenddb_core::types::{ + AttributeValue, BillingMode, CreateTableInput, GsiInput, KeyType, LsiInput, Projection, + ProjectionType, ScalarAttributeType, TableKeyInfo, TableStatus, +}; +use extenddb_storage::{BackupEngine, DataEngine, RestoreTableOverrides, TableEngine}; +use extenddb_storage_cassandra::CassandraEngine; + +use crate::backup_engine::{ + attr, create_active_table, describe, item, ks, make_gsis_synchronous, put_all, query_index, + restore, scan_count, +}; +use crate::helpers::{TestAccount, setup_engine}; + +type Item = BTreeMap; + +/// Poll `name` every 50 ms until it is ACTIVE; the time taken. +async fn time_to_active(engine: &CassandraEngine, account_id: &str, name: &str) -> Duration { + let started = Instant::now(); + loop { + let description = describe(engine, account_id, name).await; + if description.table_status == TableStatus::Active { + return started.elapsed(); + } + assert!( + started.elapsed() < Duration::from_secs(600), + "restore target {name} still {:?} after 600 s", + description.table_status + ); + tokio::time::sleep(Duration::from_millis(50)).await; + } +} + +/// Every item of `key_info`, by key, all pages. +async fn scan_all(engine: &CassandraEngine, key_info: &TableKeyInfo) -> BTreeMap { + let mut out = BTreeMap::new(); + let mut start: Option = None; + loop { + let (items, next) = engine + .scan(key_info, None, start.as_ref(), None, None, None) + .await + .expect("scan"); + for it in items { + let mut key = String::new(); + for ks in &key_info.key_schema { + key.push_str(&format!("{:?}|", it.get(&ks.attribute_name))); + } + out.insert(key, it); + } + match next { + Some(next) => start = Some(next), + None => return out, + } + } +} + +async fn back_up(engine: &CassandraEngine, account_id: &str, name: &str) -> String { + let started = Instant::now(); + let backup = engine + .create_backup(account_id, name, "timing") + .await + .expect("create backup"); + println!(" CreateBackup of {name}: {:?}", started.elapsed()); + backup.backup_arn +} + +async fn restore_and_time( + engine: &CassandraEngine, + account_id: &str, + backup_arn: &str, + target: &str, +) -> TableKeyInfo { + let initial = restore( + engine, + account_id, + backup_arn, + target, + &RestoreTableOverrides::default(), + ) + .await + .expect("restore"); + assert_eq!(initial.table_status, TableStatus::Creating); + let elapsed = time_to_active(engine, account_id, target).await; + println!(" {target} ACTIVE after {elapsed:?}"); + engine + .table_key_info(account_id, target) + .await + .expect("target key info") +} + +#[tokio::test] +#[ignore] +async fn restore_timing_60k_plain() { + if crate::helpers::skip_without_cassandra() { + return; + } + const ITEMS: usize = 60_000; + let engine = setup_engine().await; + let account = TestAccount::new(&engine, "extenddb_ttl_test").await; + let account_id = account.account_id.clone(); + let source = create_active_table( + &engine, + &account_id, + CreateTableInput { + table_name: "TimingPlain".to_owned(), + key_schema: vec![ks("id", KeyType::Hash)], + attribute_definitions: vec![attr("id", ScalarAttributeType::S)], + billing_mode: Some(BillingMode::PayPerRequest), + ..Default::default() + }, + ) + .await; + let started = Instant::now(); + put_all( + &engine, + &source, + (0..ITEMS) + .map(|i| item(&format!("item-{i:06}"), None, &"x".repeat(200))) + .collect(), + ) + .await; + println!(" seeded {ITEMS} items in {:?}", started.elapsed()); + let backup_arn = back_up(&engine, &account_id, "TimingPlain").await; + + let restored = restore_and_time(&engine, &account_id, &backup_arn, "TimingPlainTarget").await; + + let description = describe(&engine, &account_id, "TimingPlainTarget").await; + assert_eq!(description.item_count, ITEMS as i64); + assert_eq!(scan_count(&engine, &restored).await, ITEMS); + let on_source = scan_all(&engine, &source).await; + let on_target = scan_all(&engine, &restored).await; + assert_eq!(on_source.len(), ITEMS); + assert_eq!( + on_target, on_source, + "restored items differ from the source" + ); + println!(" {ITEMS} items compared equal"); +} + +#[tokio::test] +#[ignore] +async fn restore_timing_20k_indexed() { + if crate::helpers::skip_without_cassandra() { + return; + } + const ITEMS: usize = 20_000; + const GROUPS: usize = 49; + const PARTITIONS: usize = 200; + let engine = setup_engine().await; + let account = TestAccount::new(&engine, "extenddb_ttl_test").await; + let account_id = account.account_id.clone(); + let source = create_active_table( + &engine, + &account_id, + CreateTableInput { + table_name: "TimingIndexed".to_owned(), + key_schema: vec![ks("id", KeyType::Hash), ks("sort", KeyType::Range)], + attribute_definitions: vec![ + attr("id", ScalarAttributeType::S), + attr("sort", ScalarAttributeType::S), + attr("gpk", ScalarAttributeType::S), + attr("gsk", ScalarAttributeType::N), + attr("lsk", ScalarAttributeType::N), + ], + billing_mode: Some(BillingMode::PayPerRequest), + global_secondary_indexes: Some(vec![GsiInput { + index_name: "ByGroup".to_owned(), + key_schema: vec![ks("gpk", KeyType::Hash), ks("gsk", KeyType::Range)], + projection: Projection { + projection_type: ProjectionType::Include, + non_key_attributes: Some(vec!["value".to_owned()]), + }, + provisioned_throughput: None, + }]), + local_secondary_indexes: Some(vec![LsiInput { + index_name: "ByLocal".to_owned(), + key_schema: vec![ks("id", KeyType::Hash), ks("lsk", KeyType::Range)], + projection: Projection { + projection_type: ProjectionType::All, + non_key_attributes: None, + }, + }]), + ..Default::default() + }, + ) + .await; + make_gsis_synchronous(&engine, &source.table_id, &["ByGroup"]).await; + let mut items = Vec::with_capacity(ITEMS); + for i in 0..ITEMS { + let mut it = item( + &format!("p{:04}", i % PARTITIONS), + Some(&format!("s{i:06}")), + &"v".repeat(100), + ); + it.insert("lsk".to_owned(), AttributeValue::N((ITEMS - i).to_string())); + it.insert("extra".to_owned(), AttributeValue::S(format!("e{i}"))); + if i % 10 != 0 { + it.insert( + "gpk".to_owned(), + AttributeValue::S(format!("g{:03}", i % GROUPS)), + ); + it.insert("gsk".to_owned(), AttributeValue::N(i.to_string())); + } + items.push(it); + } + let started = Instant::now(); + put_all(&engine, &source, items).await; + println!(" seeded {ITEMS} indexed items in {:?}", started.elapsed()); + let backup_arn = back_up(&engine, &account_id, "TimingIndexed").await; + + let restored = restore_and_time(&engine, &account_id, &backup_arn, "TimingIndexedTarget").await; + + let description = describe(&engine, &account_id, "TimingIndexedTarget").await; + assert_eq!(description.item_count, ITEMS as i64); + let on_source = scan_all(&engine, &source).await; + let on_target = scan_all(&engine, &restored).await; + assert_eq!(on_source.len(), ITEMS); + assert_eq!( + on_target, on_source, + "restored items differ from the source" + ); + + let mut gsi_rows = 0; + for g in 0..GROUPS { + let value = AttributeValue::S(format!("g{g:03}")); + let on_source = query_index(&engine, &source, "ByGroup", "gpk", value.clone()).await; + let on_target = query_index(&engine, &restored, "ByGroup", "gpk", value).await; + assert!(!on_source.is_empty(), "source GSI g{g} is empty"); + assert_eq!( + on_target, on_source, + "GSI ByGroup g{g} differs after restore" + ); + gsi_rows += on_source.len(); + } + assert_eq!(gsi_rows, ITEMS - ITEMS / 10, "sparse GSI row count"); + let mut lsi_rows = 0; + for p in 0..PARTITIONS { + let value = AttributeValue::S(format!("p{p:04}")); + let on_source = query_index(&engine, &source, "ByLocal", "id", value.clone()).await; + let on_target = query_index(&engine, &restored, "ByLocal", "id", value).await; + assert!(!on_source.is_empty(), "source LSI p{p} is empty"); + assert_eq!( + on_target, on_source, + "LSI ByLocal p{p} differs after restore" + ); + lsi_rows += on_source.len(); + } + assert_eq!(lsi_rows, ITEMS); + println!(" {ITEMS} items, {gsi_rows} GSI rows and {lsi_rows} LSI rows compared equal"); +} diff --git a/crates/storage-cassandra/tests/direct/restore_write.rs b/crates/storage-cassandra/tests/direct/restore_write.rs new file mode 100644 index 00000000..2caec54e --- /dev/null +++ b/crates/storage-cassandra/tests/direct/restore_write.rs @@ -0,0 +1,595 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! Direct integration tests for the restore copy's write path: what a +//! restored row leaves behind, and that the ordinary write paths (`PutItem` +//! with and without a condition, `UpdateItem`, `DeleteItem`, +//! `TransactWriteItems`) read, fence and move restored items and their index +//! rows the way they do items the put path wrote. + +use std::collections::{BTreeMap, HashMap}; +use std::time::Duration; + +use cdrs_tokio::types::IntoRustByName; +use extenddb_core::expression::{CompareOp, Expr, ExpressionMaps, PathElement, UpdateAction}; +use extenddb_core::types::{ + AttributeValue, BillingMode, CreateTableInput, GsiInput, KeyType, LsiInput, Projection, + ProjectionType, ReturnValuesOnConditionCheckFailure, ScalarAttributeType, TableKeyInfo, + TableStatus, +}; +use extenddb_storage::error::StorageError; +use extenddb_storage::{ + BackupEngine, DataEngine, RestoreTableOverrides, TableEngine, TransactWriteOp, +}; +use extenddb_storage_cassandra::CassandraEngine; + +use crate::backup_engine::{ + attr, create_active_table, item, key, ks, make_gsis_synchronous, put_all, query_index, restore, + wait_active, +}; +use crate::helpers::{TestAccount, setup_engine}; + +type Item = BTreeMap; + +pub(crate) fn s(value: &str) -> AttributeValue { + AttributeValue::S(value.to_owned()) +} + +pub(crate) fn n(value: i64) -> AttributeValue { + AttributeValue::N(value.to_string()) +} + +fn path(attr: &str) -> Expr { + Expr::Path(vec![PathElement::Attribute(attr.to_owned())]) +} + +/// `attr = :v` with `:v` bound to `value`. +fn equals(attr: &str, value: AttributeValue) -> (Expr, ExpressionMaps) { + let condition = Expr::Compare { + left: Box::new(path(attr)), + op: CompareOp::Eq, + right: Box::new(Expr::Placeholder(":v".to_owned())), + }; + let mut values = HashMap::new(); + values.insert(":v".to_owned(), value); + (condition, ExpressionMaps::new(HashMap::new(), values)) +} + +/// The raw `version` column of a restored row, read with the row's text keys. +async fn row_version( + engine: &CassandraEngine, + key_info: &TableKeyInfo, + id: &str, + sort: Option<&str>, +) -> Option { + let keyspace = format!("extenddb_ttl_test_account_{}", key_info.account_id); + let table = format!("items_{}", key_info.table_id.replace('-', "_")); + let rows = match sort { + Some(sort) => { + let cql = format!("SELECT version FROM {keyspace}.{table} WHERE pk = ? AND sk_s = ?"); + engine + .session() + .query_with_values(&cql, cdrs_tokio::query_values!(id, sort)) + .await + } + None => { + let cql = format!("SELECT version FROM {keyspace}.{table} WHERE pk = ?"); + engine + .session() + .query_with_values(&cql, cdrs_tokio::query_values!(id)) + .await + } + } + .expect("read version") + .response_body() + .expect("version body") + .into_rows() + .unwrap_or_default(); + rows.into_iter() + .next() + .and_then(|row| row.get_by_name("version").ok().flatten()) +} + +/// An indexed source table (base `id` + `sort`, GSI `ByGroup` on `gpk` + +/// `gsk` projecting `value`, LSI `ByLocal` on `id` + `lsk` projecting all), +/// seeded with `items`, backed up and restored into `target`. Returns the +/// account and the restored table's key info. +pub(crate) async fn restored_indexed_table( + engine: &CassandraEngine, + items: Vec, + target: &str, +) -> (TestAccount, TableKeyInfo) { + let account = TestAccount::new(engine, "extenddb_ttl_test").await; + let account_id = account.account_id.clone(); + let source = create_active_table( + engine, + &account_id, + CreateTableInput { + table_name: format!("{target}Source"), + key_schema: vec![ks("id", KeyType::Hash), ks("sort", KeyType::Range)], + attribute_definitions: vec![ + attr("id", ScalarAttributeType::S), + attr("sort", ScalarAttributeType::S), + attr("gpk", ScalarAttributeType::S), + attr("gsk", ScalarAttributeType::N), + attr("lsk", ScalarAttributeType::N), + ], + billing_mode: Some(BillingMode::PayPerRequest), + global_secondary_indexes: Some(vec![GsiInput { + index_name: "ByGroup".to_owned(), + key_schema: vec![ks("gpk", KeyType::Hash), ks("gsk", KeyType::Range)], + projection: Projection { + projection_type: ProjectionType::Include, + non_key_attributes: Some(vec!["value".to_owned()]), + }, + provisioned_throughput: None, + }]), + local_secondary_indexes: Some(vec![LsiInput { + index_name: "ByLocal".to_owned(), + key_schema: vec![ks("id", KeyType::Hash), ks("lsk", KeyType::Range)], + projection: Projection { + projection_type: ProjectionType::All, + non_key_attributes: None, + }, + }]), + ..Default::default() + }, + ) + .await; + make_gsis_synchronous(engine, &source.table_id, &["ByGroup"]).await; + put_all(engine, &source, items).await; + let backup = engine + .create_backup(&account_id, &source.table_name, "restore-write") + .await + .expect("create backup"); + let initial = restore( + engine, + &account_id, + &backup.backup_arn, + target, + &RestoreTableOverrides::default(), + ) + .await + .expect("restore"); + assert_eq!(initial.table_status, TableStatus::Creating); + wait_active(engine, &account_id, target, Duration::from_secs(120)).await; + let restored = engine + .table_key_info(&account_id, target) + .await + .expect("target key info"); + // Restored GSIs are synchronous in these tests as well, so a later write's + // index row lands before it is queried. + make_gsis_synchronous(engine, &restored.table_id, &["ByGroup"]).await; + (account, restored) +} + +fn indexed_item(id: &str, sort: &str, group: &str, gsk: i64, lsk: i64, value: &str) -> Item { + let mut it = item(id, Some(sort), value); + it.insert("gpk".to_owned(), s(group)); + it.insert("gsk".to_owned(), n(gsk)); + it.insert("lsk".to_owned(), n(lsk)); + it +} + +/// The ids (`id|sort`) a Query on `index` for `pk_attr = value` returns. +async fn index_keys( + engine: &CassandraEngine, + key_info: &TableKeyInfo, + index: &str, + pk_attr: &str, + value: &str, +) -> Vec { + let mut keys: Vec = query_index(engine, key_info, index, pk_attr, s(value)) + .await + .into_iter() + .map(|it| { + format!( + "{}|{}", + match it.get("id") { + Some(AttributeValue::S(v)) => v.clone(), + other => format!("{other:?}"), + }, + match it.get("sort") { + Some(AttributeValue::S(v)) => v.clone(), + other => format!("{other:?}"), + } + ) + }) + .collect(); + keys.sort(); + keys +} + +#[tokio::test] +async fn test_restored_rows_carry_a_version_and_take_every_later_write() { + if crate::helpers::skip_without_cassandra() { + return; + } + let engine = setup_engine().await; + let items = vec![ + indexed_item("a", "1", "g1", 1, 10, "v-a1"), + indexed_item("a", "2", "g1", 2, 20, "v-a2"), + indexed_item("b", "1", "g2", 3, 30, "v-b1"), + indexed_item("b", "2", "g2", 4, 40, "v-b2"), + indexed_item("c", "1", "g3", 5, 50, "v-c1"), + // No GSI keys: absent from the sparse GSI, present in the LSI. + { + let mut it = item("c", Some("2"), "v-c2"); + it.insert("lsk".to_owned(), n(60)); + it + }, + ]; + let (_account, restored) = restored_indexed_table(&engine, items, "RestoreWrite").await; + let maps = ExpressionMaps::default(); + + // Every restored row carries a version in the microsecond clock domain. + for (id, sort) in [("a", "1"), ("b", "2"), ("c", "2")] { + let version = row_version(&engine, &restored, id, Some(sort)).await; + assert!( + version.is_some_and(|v| v > 1_000_000_000_000_000), + "restored row {id}/{sort} version {version:?} is not a clock reading" + ); + } + assert_eq!( + index_keys(&engine, &restored, "ByGroup", "gpk", "g1").await, + ["a|1", "a|2"] + ); + assert_eq!( + index_keys(&engine, &restored, "ByLocal", "id", "c").await, + ["c|1", "c|2"] + ); + + // UpdateItem moves a/1 from g1 to g9 and changes its LSI key; the old + // index rows go, the new ones arrive, the projection carries the new value. + let before = row_version(&engine, &restored, "a", Some("1")).await; + let mut set_values = HashMap::new(); + set_values.insert(":g".to_owned(), s("g9")); + set_values.insert(":l".to_owned(), n(11)); + set_values.insert(":val".to_owned(), s("v-a1-updated")); + let set_maps = ExpressionMaps::new(HashMap::new(), set_values); + let actions = [ + UpdateAction::Set { + path: vec![PathElement::Attribute("gpk".to_owned())], + value: Expr::Placeholder(":g".to_owned()), + }, + UpdateAction::Set { + path: vec![PathElement::Attribute("lsk".to_owned())], + value: Expr::Placeholder(":l".to_owned()), + }, + UpdateAction::Set { + path: vec![PathElement::Attribute("value".to_owned())], + value: Expr::Placeholder(":val".to_owned()), + }, + ]; + let (old, new) = engine + .update_item( + &restored, + &key("a", Some("1")), + &actions, + true, + true, + None, + &set_maps, + None, + ) + .await + .expect("update restored item"); + assert_eq!( + old.as_ref().and_then(|it| it.get("value")), + Some(&s("v-a1")) + ); + assert_eq!( + new.as_ref().and_then(|it| it.get("value")), + Some(&s("v-a1-updated")) + ); + let after = row_version(&engine, &restored, "a", Some("1")).await; + assert!( + after > before, + "version must advance: {before:?} -> {after:?}" + ); + assert_eq!( + index_keys(&engine, &restored, "ByGroup", "gpk", "g1").await, + ["a|2"] + ); + let moved = query_index(&engine, &restored, "ByGroup", "gpk", s("g9")).await; + assert_eq!(moved.len(), 1); + assert_eq!(moved[0].get("value"), Some(&s("v-a1-updated"))); + assert_eq!(moved[0].get("gsk"), Some(&n(1))); + let local = query_index(&engine, &restored, "ByLocal", "id", s("a")).await; + let lsks: Vec<&AttributeValue> = local.iter().filter_map(|it| it.get("lsk")).collect(); + assert_eq!(lsks, [&n(11), &n(20)]); + + // A conditional PutItem reads the restored image: the condition on the + // restored value holds, the condition on a wrong value fails and returns + // the restored item. + let (holds, holds_maps) = equals("value", s("v-b1")); + let mut replacement = indexed_item("b", "1", "g2", 3, 30, "v-b1-put"); + replacement.remove("gsk"); + replacement.insert("gsk".to_owned(), n(33)); + engine + .put_item( + &restored, + replacement.clone(), + false, + Some(&holds), + &holds_maps, + None, + ) + .await + .expect("conditional put over a restored item"); + assert_eq!( + engine + .get_item(&restored, &key("b", Some("1"))) + .await + .expect("get b/1"), + Some(replacement) + ); + let (fails, fails_maps) = equals("value", s("v-b2-wrong")); + let failed = engine + .put_item( + &restored, + indexed_item("b", "2", "g2", 4, 40, "never"), + false, + Some(&fails), + &fails_maps, + None, + ) + .await; + match failed { + Err(StorageError::ConditionFailed(Some(existing))) => { + assert_eq!(existing.get("value"), Some(&s("v-b2"))); + } + other => panic!("expected ConditionFailed with the restored item, got {other:?}"), + } + let in_g2 = query_index(&engine, &restored, "ByGroup", "gpk", s("g2")).await; + let gsks: Vec<&AttributeValue> = in_g2.iter().filter_map(|it| it.get("gsk")).collect(); + assert_eq!( + gsks, + [&n(4), &n(33)], + "old b/1 GSI row at gsk=3 must be gone" + ); + + // attribute_not_exists on a restored key fails: the row exists. + let not_exists = Expr::Function { + name: "attribute_not_exists".to_owned(), + args: vec![path("id")], + }; + let refused = engine + .put_item( + &restored, + indexed_item("c", "1", "g3", 5, 50, "never"), + false, + Some(¬_exists), + &maps, + None, + ) + .await; + assert!( + matches!(refused, Err(StorageError::ConditionFailed(_))), + "attribute_not_exists must see the restored row, got {refused:?}" + ); + + // DeleteItem with a condition on the restored value removes the row and + // its index rows. + let (holds, holds_maps) = equals("value", s("v-a2")); + let deleted = engine + .delete_item( + &restored, + &key("a", Some("2")), + true, + Some(&holds), + &holds_maps, + None, + ) + .await + .expect("conditional delete of a restored item"); + assert_eq!( + deleted.and_then(|it| it.get("value").cloned()), + Some(s("v-a2")) + ); + assert_eq!( + engine + .get_item(&restored, &key("a", Some("2"))) + .await + .expect("get a/2"), + None + ); + assert!( + index_keys(&engine, &restored, "ByGroup", "gpk", "g1") + .await + .is_empty() + ); + let local = query_index(&engine, &restored, "ByLocal", "id", s("a")).await; + assert_eq!(local.len(), 1, "a/2's LSI row must be gone"); + + // TransactWriteItems: a condition check on one restored item and a put + // over another commit together; the put gives c/2 its first GSI row and + // rewrites its LSI row in place. + let (check, check_maps) = equals("value", s("v-c1")); + let c2 = indexed_item("c", "2", "g7", 7, 60, "v-c2-txn"); + let c1_key = key("c", Some("1")); + let ops = vec![ + TransactWriteOp::ConditionCheck { + key_info: &restored, + key: &c1_key, + condition: &check, + maps: &check_maps, + return_values_on_ccf: ReturnValuesOnConditionCheckFailure::None, + }, + TransactWriteOp::Put { + key_info: &restored, + item: &c2, + condition: None, + maps: &maps, + return_values_on_ccf: ReturnValuesOnConditionCheckFailure::None, + stream: None, + }, + ]; + engine + .transact_write_items(&ops, None) + .await + .expect("transaction over restored items"); + assert_eq!( + engine + .get_item(&restored, &key("c", Some("2"))) + .await + .expect("get c/2"), + Some(c2) + ); + assert_eq!( + index_keys(&engine, &restored, "ByGroup", "gpk", "g7").await, + ["c|2"] + ); + let local = query_index(&engine, &restored, "ByLocal", "id", s("c")).await; + let values: Vec<&AttributeValue> = local.iter().filter_map(|it| it.get("value")).collect(); + assert_eq!(values, [&s("v-c1"), &s("v-c2-txn")]); +} + +#[tokio::test] +async fn test_restored_pk_only_rows_take_later_writes() { + if crate::helpers::skip_without_cassandra() { + return; + } + let engine = setup_engine().await; + let account = TestAccount::new(&engine, "extenddb_ttl_test").await; + let account_id = account.account_id.clone(); + let source = create_active_table( + &engine, + &account_id, + CreateTableInput { + table_name: "RestoreWritePlainSource".to_owned(), + key_schema: vec![ks("id", KeyType::Hash)], + attribute_definitions: vec![attr("id", ScalarAttributeType::S)], + billing_mode: Some(BillingMode::PayPerRequest), + ..Default::default() + }, + ) + .await; + put_all( + &engine, + &source, + (0..70) + .map(|i| item(&format!("k{i:02}"), None, &format!("v{i}"))) + .collect(), + ) + .await; + let backup = engine + .create_backup(&account_id, "RestoreWritePlainSource", "plain") + .await + .expect("create backup"); + let target = "RestoreWritePlain"; + restore( + &engine, + &account_id, + &backup.backup_arn, + target, + &RestoreTableOverrides::default(), + ) + .await + .expect("restore"); + let active = wait_active(&engine, &account_id, target, Duration::from_secs(120)).await; + assert_eq!(active.item_count, 70); + let restored = engine + .table_key_info(&account_id, target) + .await + .expect("target key info"); + let maps = ExpressionMaps::default(); + + // Items from the first and the second payload bucket (64 per bucket). + for id in ["k00", "k63", "k64", "k69"] { + assert!( + row_version(&engine, &restored, id, None) + .await + .is_some_and(|v| v > 1_000_000_000_000_000), + "{id} has no clock version" + ); + } + + // Unconditional put over a restored row returns the restored image. + let old = engine + .put_item(&restored, item("k00", None, "p0"), true, None, &maps, None) + .await + .expect("put over restored"); + assert_eq!(old.and_then(|it| it.get("value").cloned()), Some(s("v0"))); + + // UpdateItem with a condition on the restored value. + let (holds, holds_maps) = equals("value", s("v63")); + let actions = [UpdateAction::Set { + path: vec![PathElement::Attribute("value".to_owned())], + value: Expr::Placeholder(":v".to_owned()), + }]; + let before = row_version(&engine, &restored, "k63", None).await; + let (_, new) = engine + .update_item( + &restored, + &key("k63", None), + &actions, + false, + true, + Some(&holds), + &holds_maps, + None, + ) + .await + .expect("conditional update of a restored item"); + assert_eq!(new.and_then(|it| it.get("value").cloned()), Some(s("v63"))); + let after = row_version(&engine, &restored, "k63", None).await; + assert!( + after > before, + "version must advance: {before:?} -> {after:?}" + ); + + // attribute_not_exists(id) is refused on a restored key. + let not_exists = Expr::Function { + name: "attribute_not_exists".to_owned(), + args: vec![path("id")], + }; + let refused = engine + .put_item( + &restored, + item("k64", None, "never"), + false, + Some(¬_exists), + &maps, + None, + ) + .await; + assert!( + matches!(refused, Err(StorageError::ConditionFailed(_))), + "attribute_not_exists must see the restored row, got {refused:?}" + ); + + // A failing condition returns the restored item; a delete then removes it. + let (fails, fails_maps) = equals("value", s("other")); + let failed = engine + .delete_item( + &restored, + &key("k69", None), + true, + Some(&fails), + &fails_maps, + None, + ) + .await; + match failed { + Err(StorageError::ConditionFailed(Some(existing))) => { + assert_eq!(existing.get("value"), Some(&s("v69"))); + } + other => panic!("expected ConditionFailed with the restored item, got {other:?}"), + } + let deleted = engine + .delete_item(&restored, &key("k69", None), true, None, &maps, None) + .await + .expect("delete restored item"); + assert_eq!( + deleted.and_then(|it| it.get("value").cloned()), + Some(s("v69")) + ); + assert_eq!( + engine + .get_item(&restored, &key("k69", None)) + .await + .expect("get k69"), + None + ); +} diff --git a/crates/storage-cassandra/tests/direct/restore_write_controls.rs b/crates/storage-cassandra/tests/direct/restore_write_controls.rs new file mode 100644 index 00000000..d7ddb975 --- /dev/null +++ b/crates/storage-cassandra/tests/direct/restore_write_controls.rs @@ -0,0 +1,485 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! Controls on the restore copy's write path: what a restored row's write +//! timestamp is, that a backup whose payload does not fit its target or its +//! recorded count fails the restore, and that the rows the restore writes +//! have the shape the ordinary put path writes. + +use std::collections::BTreeMap; +use std::time::Duration; + +use cdrs_tokio::types::IntoRustByName; +use extenddb_core::expression::ExpressionMaps; +use extenddb_core::types::{ + AttributeValue, BillingMode, CreateTableInput, DescribeTableInput, GsiInput, KeyType, + Projection, ProjectionType, ScalarAttributeType, TableKeyInfo, TableStatus, +}; +use extenddb_storage::error::StorageError; +use extenddb_storage::{BackupEngine, DataEngine, RestoreTableOverrides, TableEngine}; +use extenddb_storage_cassandra::CassandraEngine; + +use crate::backup_engine::{ + attr, create_active_table, item, ks, make_gsis_synchronous, put_all, restore, +}; +use crate::helpers::{TestAccount, setup_engine}; +use crate::restore_write::{n, restored_indexed_table, s}; + +type Item = BTreeMap; + +/// An item carrying every key of `restored_indexed_table`'s schema. +fn indexed_item(id: &str, sort: &str, group: &str, gsk: i64, lsk: i64, value: &str) -> Item { + let mut it = item(id, Some(sort), value); + it.insert("gpk".to_owned(), s(group)); + it.insert("gsk".to_owned(), n(gsk)); + it.insert("lsk".to_owned(), n(lsk)); + it +} + +fn account_keyspace(account_id: &str) -> String { + format!("extenddb_ttl_test_account_{account_id}") +} + +fn items_table(table_id: &str) -> String { + format!("items_{}", table_id.replace('-', "_")) +} + +/// `(index_name, index table name)` for every index of `table_id`. +async fn index_tables(engine: &CassandraEngine, table_id: &str) -> Vec<(String, String)> { + let cql = format!( + "SELECT index_name, index_id FROM {}.indexes WHERE table_id = ?", + engine.catalog_keyspace() + ); + let mut out: Vec<(String, String)> = engine + .session() + .query_with_values(&cql, cdrs_tokio::query_values!(table_id)) + .await + .expect("indexes") + .response_body() + .expect("indexes body") + .into_rows() + .unwrap_or_default() + .into_iter() + .map(|row| { + let name: String = row.get_r_by_name("index_name").expect("index_name"); + let id: String = row.get_r_by_name("index_id").expect("index_id"); + (name, format!("index_{}", id.replace('-', "_"))) + }) + .collect(); + out.sort(); + out +} + +async fn rows(engine: &CassandraEngine, cql: &str) -> Vec { + engine + .session() + .query(cql) + .await + .expect("query") + .response_body() + .expect("body") + .into_rows() + .unwrap_or_default() +} + +/// Poll `target` until it is gone (the copy failed and the target was +/// removed) or `timeout` passes; panic if it turns ACTIVE. +async fn wait_removed(engine: &CassandraEngine, account_id: &str, target: &str, timeout: Duration) { + let deadline = tokio::time::Instant::now() + timeout; + loop { + let described = engine + .describe_table( + account_id, + DescribeTableInput { + table_name: target.to_owned(), + }, + ) + .await; + match described { + Err(StorageError::TableNotFound(_)) => return, + Ok(d) if d.table_status == TableStatus::Active => panic!( + "{target} turned ACTIVE (item_count {}) where the copy had to fail", + d.item_count + ), + Ok(_) => {} + Err(e) => panic!("describe {target}: {e:?}"), + } + assert!( + tokio::time::Instant::now() < deadline, + "{target} still present after {timeout:?}" + ); + tokio::time::sleep(Duration::from_millis(200)).await; + } +} + +/// Append `item` to `backup_arn`'s first payload bucket as its fourth row and +/// record `item_count` on the backup. +async fn add_backup_item( + engine: &CassandraEngine, + account_id: &str, + backup_arn: &str, + item: &Item, + item_count: i64, +) { + let payload = serde_json::to_string(item).expect("json"); + let insert = format!( + "INSERT INTO {}.backup_items (backup_arn, bucket, item_index, item_data) VALUES (?, 0, 3, ?)", + account_keyspace(account_id) + ); + engine + .session() + .query_with_values( + &insert, + cdrs_tokio::query_values!(backup_arn, payload.as_str()), + ) + .await + .expect("insert payload row"); + set_item_count(engine, account_id, backup_arn, item_count).await; +} + +async fn set_item_count(engine: &CassandraEngine, account_id: &str, backup_arn: &str, count: i64) { + let bump = format!( + "UPDATE {}.backups_by_arn SET item_count = ? WHERE account_id = ? AND backup_arn = ? IF EXISTS", + engine.catalog_keyspace() + ); + engine + .session() + .query_with_values( + &bump, + cdrs_tokio::query_values!(count, account_id, backup_arn), + ) + .await + .expect("set item_count"); +} + +/// A pk-only source table with one GSI on `gpk` (S), three items with `gpk` +/// set, backed up. Returns the account, the source key info and the ARN. +async fn backed_up_gsi_table( + engine: &CassandraEngine, + name: &str, +) -> (TestAccount, TableKeyInfo, String) { + let account = TestAccount::new(engine, "extenddb_ttl_test").await; + let account_id = account.account_id.clone(); + let source = create_active_table( + engine, + &account_id, + CreateTableInput { + table_name: name.to_owned(), + key_schema: vec![ks("id", KeyType::Hash)], + attribute_definitions: vec![ + attr("id", ScalarAttributeType::S), + attr("gpk", ScalarAttributeType::S), + ], + billing_mode: Some(BillingMode::PayPerRequest), + global_secondary_indexes: Some(vec![GsiInput { + index_name: "ByGroup".to_owned(), + key_schema: vec![ks("gpk", KeyType::Hash)], + projection: Projection { + projection_type: ProjectionType::All, + non_key_attributes: None, + }, + provisioned_throughput: None, + }]), + ..Default::default() + }, + ) + .await; + make_gsis_synchronous(engine, &source.table_id, &["ByGroup"]).await; + let items = (0..3) + .map(|i| { + let mut it = item(&format!("k{i}"), None, "v"); + it.insert("gpk".to_owned(), s("g")); + it + }) + .collect(); + put_all(engine, &source, items).await; + let backup = engine + .create_backup(&account_id, name, "control") + .await + .expect("create backup"); + (account, source, backup.backup_arn) +} + +/// One restored base row's `version`, its write time, and the write times of +/// its index rows by index name. +struct RowStamps { + version: i64, + base_wt: i64, + index_wts: Vec<(String, i64)>, +} + +/// Every restored index row's write time is the coordinator's, as the base +/// row's is. The coordinator hands each statement its own timestamp, so an +/// item's GSI row and LSI row never share one; a restore that stamped them +/// from the host clock would give both the base row's `version`. Every index +/// row's write time is also within 10 ms of its base row's. +#[tokio::test] +async fn test_restored_index_rows_take_the_coordinator_timestamp() { + if crate::helpers::skip_without_cassandra() { + return; + } + let engine = setup_engine().await; + let items = vec![ + indexed_item("a", "1", "g1", 1, 10, "v-a1"), + indexed_item("b", "1", "g2", 3, 30, "v-b1"), + { + let mut it = item("c", Some("2"), "v-c2"); + it.insert("lsk".to_owned(), n(60)); + it + }, + ]; + let (_account, restored) = restored_indexed_table(&engine, items, "RestoreStamp").await; + let keyspace = account_keyspace(&restored.account_id); + let base_table = items_table(&restored.table_id); + let indexes = index_tables(&engine, &restored.table_id).await; + assert_eq!(indexes.len(), 2, "{indexes:?}"); + let mut per_item: BTreeMap = BTreeMap::new(); + for (index_name, index_table) in indexes { + let index_rows = rows( + &engine, + &format!( + "SELECT base_pk, base_sk_s, WRITETIME(item_data) AS wt FROM {keyspace}.{index_table}" + ), + ) + .await; + assert!(!index_rows.is_empty(), "{index_name} has rows"); + for row in index_rows { + let base_pk: String = row.get_r_by_name("base_pk").expect("base_pk"); + let base_sk: String = row.get_r_by_name("base_sk_s").expect("base_sk_s"); + let index_wt: i64 = row.get_r_by_name("wt").expect("wt"); + let base = engine + .session() + .query_with_values( + &format!( + "SELECT version, WRITETIME(item_data) AS wt FROM {keyspace}.{base_table} \ + WHERE pk = ? AND sk_s = ?" + ), + cdrs_tokio::query_values!(base_pk.as_str(), base_sk.as_str()), + ) + .await + .expect("base row") + .response_body() + .expect("base body") + .into_rows() + .unwrap_or_default(); + let version: i64 = base[0].get_r_by_name("version").expect("version"); + let base_wt: i64 = base[0].get_r_by_name("wt").expect("base wt"); + per_item + .entry(format!("{base_pk}/{base_sk}")) + .or_insert(RowStamps { + version, + base_wt, + index_wts: Vec::new(), + }) + .index_wts + .push((index_name.clone(), index_wt)); + } + } + assert_eq!(per_item.len(), 3); + let mut stamped_with_version = 0; + let mut checked = 0; + for ( + key, + RowStamps { + version, + base_wt, + index_wts, + }, + ) in &per_item + { + for (index_name, index_wt) in index_wts { + assert!( + (index_wt - base_wt).abs() <= 10_000, + "{index_name} row {key}: index write time {index_wt} is not within 10 ms of the \ + base row's {base_wt}" + ); + if index_wt == version { + stamped_with_version += 1; + } + checked += 1; + } + if let [(first_name, first_wt), (second_name, second_wt)] = index_wts.as_slice() { + assert_ne!( + first_wt, second_wt, + "{key}: the {first_name} and {second_name} rows share a write time; the \ + coordinator gives each statement its own" + ); + } + } + assert_eq!(checked, 5, "2 GSI rows and 3 LSI rows"); + assert_eq!( + per_item.values().filter(|r| r.index_wts.len() == 2).count(), + 2, + "a/1 and b/1 carry a GSI row and an LSI row each" + ); + assert!( + stamped_with_version < checked, + "every index row's write time equals its base row's version: the rows were stamped from \ + the host clock" + ); +} + +/// A backup item whose GSI key attribute has the wrong type (N where the +/// index declares S) does not fit the target's index key schema: the copy +/// refuses it and the target is removed, as the ordinary put path refuses +/// the item. +#[tokio::test] +async fn test_a_mistyped_index_key_in_a_backup_item_fails_the_restore() { + if crate::helpers::skip_without_cassandra() { + return; + } + let engine = setup_engine().await; + let (account, source, backup_arn) = backed_up_gsi_table(&engine, "TypedKeySource").await; + let account_id = account.account_id.clone(); + let mut bad = item("k3", None, "v"); + bad.insert("gpk".to_owned(), n(5)); + let refused = engine + .put_item( + &source, + bad.clone(), + false, + None, + &ExpressionMaps::default(), + None, + ) + .await; + assert!( + matches!(refused, Err(StorageError::Validation(_))), + "the ordinary put refuses a mistyped index key: {refused:?}" + ); + add_backup_item(&engine, &account_id, &backup_arn, &bad, 4).await; + let initial = restore( + &engine, + &account_id, + &backup_arn, + "TypedKeyTarget", + &RestoreTableOverrides::default(), + ) + .await + .expect("restore accepted"); + assert_eq!(initial.table_status, TableStatus::Creating); + wait_removed( + &engine, + &account_id, + "TypedKeyTarget", + Duration::from_secs(90), + ) + .await; +} + +/// A backup whose recorded item count is one more than its payload does not +/// restore: the copy's count check fails it and the target is removed. +#[tokio::test] +async fn test_a_backup_item_count_above_its_payload_fails_the_restore() { + if crate::helpers::skip_without_cassandra() { + return; + } + let engine = setup_engine().await; + let account = TestAccount::new(&engine, "extenddb_ttl_test").await; + let account_id = account.account_id.clone(); + let source = create_active_table( + &engine, + &account_id, + CreateTableInput { + table_name: "CountSource".to_owned(), + key_schema: vec![ks("id", KeyType::Hash)], + attribute_definitions: vec![attr("id", ScalarAttributeType::S)], + billing_mode: Some(BillingMode::PayPerRequest), + ..Default::default() + }, + ) + .await; + put_all( + &engine, + &source, + (0..70) + .map(|i| item(&format!("k{i:02}"), None, "v")) + .collect(), + ) + .await; + let backup = engine + .create_backup(&account_id, "CountSource", "count") + .await + .expect("create backup"); + set_item_count(&engine, &account_id, &backup.backup_arn, 71).await; + let initial = restore( + &engine, + &account_id, + &backup.backup_arn, + "CountTarget", + &RestoreTableOverrides::default(), + ) + .await + .expect("restore accepted"); + assert_eq!(initial.table_status, TableStatus::Creating); + wait_removed(&engine, &account_id, "CountTarget", Duration::from_secs(90)).await; +} + +/// `SELECT JSON *` of every row of `table`, without `version`, sorted. +async fn row_shapes(engine: &CassandraEngine, keyspace: &str, table: &str) -> Vec { + let mut shapes: Vec = rows(engine, &format!("SELECT JSON * FROM {keyspace}.{table}")) + .await + .into_iter() + .map(|row| { + let json: String = row.get_r_by_name("[json]").expect("[json]"); + let mut value: serde_json::Map = + serde_json::from_str(&json).expect("row json"); + value.remove("version"); + serde_json::Value::Object(value).to_string() + }) + .collect(); + shapes.sort(); + shapes +} + +/// The rows a restore writes have the shape the ordinary put path writes: +/// the base table and each index table of a restored table hold, column by +/// column except `version`, the rows the source table holds for the same +/// items. +#[tokio::test] +async fn test_restored_rows_have_the_put_path_shape() { + if crate::helpers::skip_without_cassandra() { + return; + } + let engine = setup_engine().await; + let items = vec![ + indexed_item("a", "1", "g1", 1, 10, "v-a1"), + indexed_item("a", "2", "g1", 2, 20, "v-a2"), + indexed_item("b", "1", "g2", 3, 30, "v-b1"), + { + let mut it = item("c", Some("2"), "v-c2"); + it.insert("lsk".to_owned(), n(60)); + it.insert("extra".to_owned(), s("not projected")); + it + }, + ]; + let (account, restored) = restored_indexed_table(&engine, items, "RestoreShape").await; + let account_id = account.account_id.clone(); + let source = engine + .table_key_info(&account_id, "RestoreShapeSource") + .await + .expect("source key info"); + let keyspace = account_keyspace(&account_id); + + let on_source = row_shapes(&engine, &keyspace, &items_table(&source.table_id)).await; + let on_target = row_shapes(&engine, &keyspace, &items_table(&restored.table_id)).await; + assert_eq!(on_source.len(), 4); + assert_eq!(on_target, on_source, "base rows differ from the put path's"); + + let source_indexes = index_tables(&engine, &source.table_id).await; + let target_indexes = index_tables(&engine, &restored.table_id).await; + assert_eq!(source_indexes.len(), 2); + assert_eq!(target_indexes.len(), 2); + for ((name, source_table), (target_name, target_table)) in + source_indexes.iter().zip(&target_indexes) + { + assert_eq!(name, target_name); + let on_source = row_shapes(&engine, &keyspace, source_table).await; + let on_target = row_shapes(&engine, &keyspace, target_table).await; + assert!(!on_source.is_empty(), "{name} has rows on the source"); + assert_eq!( + on_target, on_source, + "{name} rows differ from the put path's" + ); + } +} diff --git a/crates/storage-cassandra/tests/direct_integration.rs b/crates/storage-cassandra/tests/direct_integration.rs index 43eb28b8..c023666e 100644 --- a/crates/storage-cassandra/tests/direct_integration.rs +++ b/crates/storage-cassandra/tests/direct_integration.rs @@ -39,6 +39,12 @@ mod policies; mod put_get_item; #[path = "direct/query.rs"] mod query; +#[path = "direct/restore_timing.rs"] +mod restore_timing; +#[path = "direct/restore_write.rs"] +mod restore_write; +#[path = "direct/restore_write_controls.rs"] +mod restore_write_controls; #[path = "direct/roles.rs"] mod roles; #[path = "direct/scan.rs"]