diff --git a/crates/storage-postgres/src/backup_engine.rs b/crates/storage-postgres/src/backup_engine.rs index 6486f522..ee02a17c 100755 --- a/crates/storage-postgres/src/backup_engine.rs +++ b/crates/storage-postgres/src/backup_engine.rs @@ -8,7 +8,6 @@ use extenddb_core::types::{ PointInTimeRecoveryDescription, SourceTableDetails, TableDescription, }; use extenddb_storage::BackupEngine; -use extenddb_storage::TableEngine; use extenddb_storage::error::StorageError; use futures::future::BoxFuture; @@ -451,16 +450,20 @@ impl BackupEngine for PostgresEngine { on_demand_throughput: None, }; - let desc = self.create_table(&account_id, create_input).await?; + // Create the table with the ACTIVE transition deferred: it enters + // CREATING with no scheduled flip, so the control-plane worker + // cannot mark it ACTIVE while the item copy below is still + // running. The explicit ACTIVE update after the copy is the only + // way this table leaves CREATING, so ACTIVE implies the restored + // data is fully present. + let desc = self + .create_table_impl(&account_id, create_input, true) + .await?; let new_table_id = &desc.table_id; let ddb_table = data_table_name(new_table_id); let ddb_table_unquoted = ddb_table.trim_matches('"'); - // Do NOT force ACTIVE — let the control plane handle the transition - // (steering rule D-2: tests run with control_plane_delay_seconds > 0). - // The table starts in CREATING and transitions to ACTIVE after the delay. - let has_sk: bool = sqlx::query_scalar( "SELECT EXISTS(SELECT 1 FROM information_schema.columns \ WHERE table_name = $1 AND column_name = 'sk')", @@ -513,10 +516,10 @@ impl BackupEngine for PostgresEngine { .await .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; - // Mark the restored table ACTIVE immediately — the data is fully - // populated and the table is ready to serve requests. This matches - // real DynamoDB behavior where restored tables become ACTIVE once - // the restore completes (the CREATING status is transient). + // Mark the restored table ACTIVE now that the copy has fully + // drained. The table was created with the transition deferred, so + // this is the first point it can become ACTIVE, by ordering rather + // than by timing. sqlx::query( "UPDATE tables SET table_status = 'ACTIVE', status_transition_at = NULL \ WHERE account_id = $1 AND table_name = $2", diff --git a/crates/storage-postgres/src/create_table.rs b/crates/storage-postgres/src/create_table.rs index b808a1fa..9ef2e39d 100755 --- a/crates/storage-postgres/src/create_table.rs +++ b/crates/storage-postgres/src/create_table.rs @@ -14,10 +14,17 @@ use crate::PostgresEngine; impl PostgresEngine { /// Core implementation of `create_table` (Fix #4: wrapped in a transaction). + /// Create a table. When `defer_active` is set (the restore path), the row + /// is written `CREATING` with **no** scheduled transition, so the + /// background control-plane worker cannot flip it to `ACTIVE` while the + /// caller is still populating it; the caller sets `ACTIVE` itself once the + /// data copy completes. Normal `CreateTable` passes `false` and gets the + /// usual timed transition. pub(crate) async fn create_table_impl( &self, account_id: &str, input: CreateTableInput, + defer_active: bool, ) -> Result { Self::validate_account_id(account_id)?; let table_id = uuid::Uuid::new_v4().to_string(); @@ -77,10 +84,12 @@ impl PostgresEngine { creation_date_time, table_arn, table_id, deletion_protection_enabled, status_transition_at, table_class, sse_specification, on_demand_throughput) VALUES ($1, $2, $3, $4, $5, $6, $7, - CASE WHEN (SELECT secs FROM delay) = 0 + CASE WHEN $14 THEN 'CREATING' + WHEN (SELECT secs FROM delay) = 0 THEN 'ACTIVE' ELSE 'CREATING' END, NOW(), $8, $9, $10, - CASE WHEN (SELECT secs FROM delay) = 0 + CASE WHEN $14 THEN NULL + WHEN (SELECT secs FROM delay) = 0 THEN NULL ELSE NOW() + make_interval(secs => (SELECT secs FROM delay)) END, @@ -100,6 +109,7 @@ impl PostgresEngine { .bind(&input.table_class) .bind(&sse_spec_json) .bind(&on_demand_json) + .bind(defer_active) .fetch_one(&mut *tx) .await .map_err(|e| match &e { diff --git a/crates/storage-postgres/src/table_engine.rs b/crates/storage-postgres/src/table_engine.rs index a3a4e1d4..09c871bc 100755 --- a/crates/storage-postgres/src/table_engine.rs +++ b/crates/storage-postgres/src/table_engine.rs @@ -22,7 +22,7 @@ impl TableEngine for PostgresEngine { input: CreateTableInput, ) -> BoxFuture<'_, Result> { let account_id = account_id.to_string(); - Box::pin(async move { self.create_table_impl(&account_id, input).await }) + Box::pin(async move { self.create_table_impl(&account_id, input, false).await }) } // H-5: Set status to DELETING with a scheduled transition to removal, diff --git a/tests/rust/src/main.rs b/tests/rust/src/main.rs index 1befdc6f..9287f713 100755 --- a/tests/rust/src/main.rs +++ b/tests/rust/src/main.rs @@ -7,6 +7,7 @@ //! Requires a running extenddb instance with `EXTENDDB_TEST_ENDPOINT` and credentials set. mod helpers; +mod restore_active_completeness; mod test_base; #[cfg(test)] @@ -18,10 +19,10 @@ mod backup_restore; #[cfg(test)] mod batch_get_item; #[cfg(test)] -mod batch_write_item; -#[cfg(test)] mod batch_transact_authz; #[cfg(test)] +mod batch_write_item; +#[cfg(test)] mod binary_sort_key; #[cfg(test)] mod capacity_throttling; diff --git a/tests/rust/src/restore_active_completeness.rs b/tests/rust/src/restore_active_completeness.rs new file mode 100644 index 00000000..cc4ea758 --- /dev/null +++ b/tests/rust/src/restore_active_completeness.rs @@ -0,0 +1,151 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! A restored table must not report ACTIVE before its data copy completes. +//! +//! DynamoDB's contract: when `DescribeTable` on a restore target first +//! returns ACTIVE, the restored data is fully present. A restore that flips +//! ACTIVE on a control-plane timer decoupled from the copy exposes an empty +//! or partial table to a client that waits-for-ACTIVE. + +use crate::test_base::*; +use aws_sdk_dynamodb::types::{ + AttributeDefinition, AttributeValue, BillingMode, KeySchemaElement, KeyType, PutRequest, + ScalarAttributeType, Select, WriteRequest, +}; + +const ITEMS: usize = 40000; + +#[tokio::test] +async fn restored_table_has_all_items_when_first_active() { + let c = client(); + let src = format!("RestoreRaceSrc_{}", ts()); + c.create_table() + .table_name(&src) + .key_schema( + KeySchemaElement::builder() + .attribute_name("pk") + .key_type(KeyType::Hash) + .build() + .unwrap(), + ) + .attribute_definitions( + AttributeDefinition::builder() + .attribute_name("pk") + .attribute_type(ScalarAttributeType::S) + .build() + .unwrap(), + ) + .billing_mode(BillingMode::PayPerRequest) + .send() + .await + .unwrap(); + wait_for_active(&c, &src).await; + + // Enough data that the restore copy takes longer than the control-plane + // transition delay, so a timer-driven ACTIVE flip would win the race. + let pad = "x".repeat(2000); + for chunk in (0..ITEMS).collect::>().chunks(25) { + let reqs: Vec = chunk + .iter() + .map(|i| { + WriteRequest::builder() + .put_request( + PutRequest::builder() + .item("pk", AttributeValue::S(format!("k{i:06}"))) + .item("d", AttributeValue::S(pad.clone())) + .build() + .unwrap(), + ) + .build() + }) + .collect(); + c.batch_write_item() + .request_items(&src, reqs) + .send() + .await + .unwrap(); + } + + let backup = c + .create_backup() + .table_name(&src) + .backup_name("restore-race-probe") + .send() + .await + .unwrap(); + let arn = backup.backup_details().unwrap().backup_arn().to_string(); + // Wait until the backup is AVAILABLE. + for _ in 0..240 { + let d = c.describe_backup().backup_arn(&arn).send().await.unwrap(); + if d.backup_description() + .and_then(|b| b.backup_details()) + .map(|b| b.backup_status().as_str() == "AVAILABLE") + .unwrap_or(false) + { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(250)).await; + } + + let dst = format!("RestoreRaceDst_{}", ts()); + // This backend's RestoreTableFromBackup blocks until the copy completes + // (real DynamoDB returns immediately), so the ACTIVE-before-data window + // is only observable by a CONCURRENT client. Race an observer task + // against the in-flight restore: the moment it sees ACTIVE, it counts. + let observer = { + let c2 = client().clone(); + let dst2 = dst.clone(); + tokio::spawn(async move { + // Wait for the table to exist, then for first ACTIVE. + loop { + match c2.describe_table().table_name(&dst2).send().await { + Ok(out) => { + let st = out.table().unwrap().table_status().unwrap().clone(); + if st.as_str() == "ACTIVE" { + break; + } + } + Err(_) => {} + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + // First ACTIVE observation: count items immediately. + let mut count = 0usize; + let mut start_key = None; + loop { + let mut req = c2.scan().table_name(&dst2).select(Select::Count); + if let Some(k) = start_key.take() { + req = req.set_exclusive_start_key(Some(k)); + } + let resp = req.send().await.unwrap(); + count += resp.count() as usize; + match resp.last_evaluated_key() { + Some(k) if !k.is_empty() => start_key = Some(k.clone()), + _ => break, + } + } + count + }) + }; + + c.restore_table_from_backup() + .target_table_name(&dst) + .backup_arn(&arn) + .send() + .await + .unwrap(); + + // The concurrent observer counted at first-ACTIVE while the restore call + // above was still in flight (or just after, if the copy was fast). + let count = observer.await.unwrap(); + + c.delete_table().table_name(&src).send().await.ok(); + c.delete_table().table_name(&dst).send().await.ok(); + + assert_eq!( + count, ITEMS, + "restored table reported ACTIVE with {count}/{ITEMS} items present — \ + ACTIVE must imply the restore copy is complete" + ); +}