From b924f3df48a7cc87d2501db00e1f15dcd15b0179 Mon Sep 17 00:00:00 2001 From: Scott Robinson Date: Sun, 4 Oct 2026 23:24:33 +0000 Subject: [PATCH 1/4] fix(postgres): back up and restore tables with a sort key CreateBackup decided whether a table had a sort key by looking for a column named `sk`. Data tables name their sort key columns by type (`sk_s`, `sk_n`, `sk_b`), so the check was false for every table with a sort key: the backup recorded no sort key, and RestoreTableFromBackup then failed its first insert on the NOT NULL sort key column. The failure came after the target had been created CREATING with no scheduled transition, so the target stayed CREATING forever and its name could not be reused. `item_data` holds the whole item, key attributes included, so neither side needs the physical key columns: - CreateBackup records `pk` and `item_data` only. - RestoreTableFromBackup derives every key column from the item (all HASH parts into `pk`, each RANGE part into its typed column) and writes it with a plain INSERT inside one data transaction. Two backup rows for one key fail the restore instead of collapsing into one item. - The restored table reports its item count and size when it turns ACTIVE, instead of 0 bytes until the size worker runs. - The flip to ACTIVE happens only from CREATING, so a DeleteTable issued while the copy runs is not undone. - If the copy fails, the target is removed by table id, synchronously, whatever the control-plane delay, and the error is returned. A failed restore leaves no CREATING or DELETING table and frees the name. Backups written by earlier binaries already carry the full item in `item_data`, so they restore correctly with this change; no migration is needed. Tests: - tests/test_backup_restore_fidelity.py round-trips a hash-only table and S, N, and B sort-key tables and compares key schema, attribute definitions, and every item attribute for attribute, by Scan and by point read. It waits for the backup to be AVAILABLE so it also holds against the service. On main the three composite-key cases fail on PostgreSQL; all four pass on PostgreSQL and SQLite with this change. - crates/storage-postgres/tests/backup_restore.rs: a backup in the pre-fix row shape restores with its sort keys and size; a restore that fails mid-copy removes its target under a non-zero control-plane delay; a backup with two rows for one key fails; a two-HASH, two-RANGE key backup restores every key part. Added to the PostgreSQL storage-level CI step. Readiness assessment P0-2 (sort-key defect). GSI and LSI restore follow separately. Signed-off-by: Scott Robinson --- .github/workflows/integration.yml | 12 +- crates/storage-postgres/src/backup_engine.rs | 260 ++++++---- .../storage-postgres/tests/backup_restore.rs | 449 ++++++++++++++++++ tests/test_backup_restore_fidelity.py | 216 +++++++++ 4 files changed, 835 insertions(+), 102 deletions(-) create mode 100644 crates/storage-postgres/tests/backup_restore.rs create mode 100644 tests/test_backup_restore_fidelity.py diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 5d1ab88ad..ed6cb9360 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -296,13 +296,17 @@ jobs: run: devtools/run-tests --extenddb --rust-integration --release # The control plane for vector indexes is not reachable over the wire while - # this backend declares no vector search capability, so its tests drive the - # storage layer directly against this job's PostgreSQL. They build their own - # throwaway databases; the connection string is the server, not a database. + # this backend declares no vector search capability, and the backup restore + # cases here need catalog rows the current binary would not write, so these + # tests drive the storage layer directly against this job's PostgreSQL. They + # build their own throwaway databases; the connection string is the server, + # not a database. - name: Run PostgreSQL storage-level tests env: EXTENDDB_TEST_PG_CONNECTION_STRING: postgresql://postgres:devpass@127.0.0.1:5432 - run: cargo test --release -p extenddb-storage-postgres --test vector_control_plane + run: >- + cargo test --release -p extenddb-storage-postgres + --test vector_control_plane --test backup_restore # The daemonized server logs to syslog; dump it so server-side failures # are diagnosable from the job log. diff --git a/crates/storage-postgres/src/backup_engine.rs b/crates/storage-postgres/src/backup_engine.rs index a93239850..9ea057971 100755 --- a/crates/storage-postgres/src/backup_engine.rs +++ b/crates/storage-postgres/src/backup_engine.rs @@ -4,7 +4,7 @@ //! Backup and point-in-time recovery implementation for `PostgreSQL` storage. use extenddb_core::types::{ - BackupDescription, BackupDetails, BackupSummary, ContinuousBackupsDescription, + BackupDescription, BackupDetails, BackupSummary, ContinuousBackupsDescription, Item, PointInTimeRecoveryDescription, SourceTableDetails, TableDescription, }; use extenddb_storage::BackupEngine; @@ -12,7 +12,9 @@ use extenddb_storage::error::StorageError; use futures::future::BoxFuture; use crate::PostgresEngine; -use crate::data::data_table_name; +use extenddb_storage::util::{SortKeyValue, composite_pk_to_text, parse_sk, sk_column_n}; + +use crate::data::{all_sort_key_info, data_table_name}; /// Current epoch milliseconds, used as the leading component of a backup id. fn epoch_millis() -> u128 { @@ -41,6 +43,137 @@ fn pg_timestamp_to_epoch(ts: time::OffsetDateTime) -> f64 { ts.unix_timestamp() as f64 } +impl PostgresEngine { + /// Copy a backup's items into a freshly created restore target and mark + /// it ACTIVE. + /// + /// Every key column is derived from the item itself, so the copy does not + /// depend on how the source table's columns were laid out. Each row is a + /// plain INSERT: two backup rows for one key are a corrupt backup and fail + /// the restore rather than silently collapsing into one item. The copy is + /// one data transaction, so a failure part-way leaves the target empty. + async fn copy_backup_items( + &self, + desc: &TableDescription, + backup_arn: &str, + ) -> Result<(), StorageError> { + let items: Vec<(serde_json::Value,)> = + sqlx::query_as("SELECT item_data FROM backup_items WHERE backup_arn = $1") + .bind(backup_arn) + .fetch_all(&self.pool) + .await + .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; + + let ddb_table = data_table_name(&desc.table_id); + let sort_keys = all_sort_key_info(&desc.key_schema, &desc.attribute_definitions); + let mut cols = vec!["pk".to_owned()]; + cols.extend( + sort_keys + .iter() + .enumerate() + .map(|(i, &(_, t))| sk_column_n(i, t)), + ); + cols.push("item_data".to_owned()); + let placeholders: Vec = (1..=cols.len()).map(|i| format!("${i}")).collect(); + let insert_sql = format!( + "INSERT INTO {ddb_table} ({}) VALUES ({})", + cols.join(", "), + placeholders.join(", ") + ); + + let mut tx = self + .data_pool + .begin() + .await + .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; + for (item_json,) in &items { + let item: Item = serde_json::from_value(item_json.clone()) + .map_err(|e| StorageError::Internal(format!("Parse backup item: {e}")))?; + let mut query = + sqlx::query(&insert_sql).bind(composite_pk_to_text(&item, &desc.key_schema)?); + for &(name, sk_type) in &sort_keys { + let value = item.get(name).ok_or_else(|| { + StorageError::Internal(format!( + "backup {backup_arn} has an item without sort key {name}" + )) + })?; + query = match parse_sk(value, sk_type)? { + SortKeyValue::S(v) => query.bind(v), + SortKeyValue::N(v) => query.bind(v), + SortKeyValue::B(v) => query.bind(v), + }; + } + query + .bind(item_json) + .execute(&mut *tx) + .await + .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; + } + tx.commit() + .await + .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; + + #[allow(clippy::cast_possible_wrap)] + let item_count = items.len() as i64; + let (table_size,): (i64,) = sqlx::query_as(&format!( + "SELECT COALESCE(pg_total_relation_size('{ddb_table}'), 0)" + )) + .fetch_one(&self.data_pool) + .await + .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; + + // Mark the restored table ACTIVE now that the copy has committed. The + // table was created with the transition deferred, so this is the first + // point it can become ACTIVE, by ordering rather than by timing. + // + // Only from CREATING: a client may have deleted the target while the + // copy ran (DeleteTable accepts a CREATING table), and that delete + // wins. Flipping a DELETING row back to ACTIVE would cancel the + // delete and leave a table with no scheduled removal. + let activated = sqlx::query( + "UPDATE tables SET item_count = $1, table_size_bytes = $2, table_status = 'ACTIVE', \ + status_transition_at = NULL WHERE table_id = $3 AND table_status = 'CREATING'", + ) + .bind(item_count) + .bind(table_size) + .bind(&desc.table_id) + .execute(&self.pool) + .await + .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; + if activated.rows_affected() == 0 { + tracing::info!( + "restore of {backup_arn} into {} finished after the target was deleted; \ + leaving the delete in place", + desc.table_name + ); + } + Ok(()) + } + + /// Remove a restore target whose copy failed. + /// + /// Keyed by table id, not name, so it can only ever remove the table this + /// restore created, and synchronous regardless of the control-plane + /// delay: the ordinary DeleteTable path would leave the name held in + /// DELETING for that long. The target was created with no secondary or + /// vector indexes, so its data table is the only physical object. + async fn abort_restore(&self, table_id: &str) -> Result<(), StorageError> { + sqlx::query("DELETE FROM tables WHERE table_id = $1 AND table_status = 'CREATING'") + .bind(table_id) + .execute(&self.pool) + .await + .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; + sqlx::query(&format!( + "DROP TABLE IF EXISTS {}", + data_table_name(table_id) + )) + .execute(&self.data_pool) + .await + .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; + Ok(()) + } +} + impl BackupEngine for PostgresEngine { fn create_backup( &self, @@ -92,34 +225,17 @@ impl BackupEngine for PostgresEngine { id = backup_id() ); - // Snapshot items from the data table. + // Snapshot items from the data table. `item_data` is the complete + // item, key attributes included, so it is all a restore needs: the + // typed key columns (`sk_s`, `sk_n`, `sk_b`) are derived from it on + // the way back in. Reading them here instead would tie the backup + // to this table's physical column layout. let ddb_table = data_table_name(&table_id); - let ddb_table_unquoted = ddb_table.trim_matches('"'); - let has_sk: bool = sqlx::query_scalar( - "SELECT EXISTS(SELECT 1 FROM information_schema.columns \ - WHERE table_name = $1 AND column_name = 'sk')", - ) - .bind(ddb_table_unquoted) - .fetch_one(&self.data_pool) - .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; - - let items: Vec<(String, Option, serde_json::Value)> = if has_sk { - sqlx::query_as(&format!("SELECT pk, sk, item_data FROM {ddb_table}")) + let items: Vec<(String, serde_json::Value)> = + sqlx::query_as(&format!("SELECT pk, item_data FROM {ddb_table}")) .fetch_all(&self.data_pool) .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))? - } else { - sqlx::query_as::<_, (String, serde_json::Value)>(&format!( - "SELECT pk, item_data FROM {ddb_table}" - )) - .fetch_all(&self.data_pool) - .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))? - .into_iter() - .map(|(pk, data)| (pk, None, data)) - .collect() - }; + .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; #[allow(clippy::cast_possible_wrap)] let actual_count = items.len() as i64; @@ -182,14 +298,13 @@ impl BackupEngine for PostgresEngine { .await .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; - for (pk, sk, item_data) in &items { + // `backup_items.sk` stays NULL: restore reads keys from `item_data`. + for (pk, item_data) in &items { sqlx::query( - "INSERT INTO backup_items (backup_arn, pk, sk, item_data) \ - VALUES ($1, $2, $3, $4)", + "INSERT INTO backup_items (backup_arn, pk, item_data) VALUES ($1, $2, $3)", ) .bind(&backup_arn) .bind(pk) - .bind(sk.as_deref()) .bind(item_data) .execute(&mut *tx) .await @@ -539,76 +654,25 @@ impl BackupEngine for PostgresEngine { .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('"'); - - let has_sk: bool = sqlx::query_scalar( - "SELECT EXISTS(SELECT 1 FROM information_schema.columns \ - WHERE table_name = $1 AND column_name = 'sk')", - ) - .bind(ddb_table_unquoted) - .fetch_one(&self.data_pool) - .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; - - let items: Vec<(String, Option, serde_json::Value)> = - sqlx::query_as("SELECT pk, sk, item_data FROM backup_items WHERE backup_arn = $1") - .bind(&backup_arn) - .fetch_all(&self.pool) - .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; - - for (pk, sk, item_data) in &items { - if has_sk { - sqlx::query(&format!( - "INSERT INTO {ddb_table} (pk, sk, item_data) VALUES ($1, $2, $3)" - )) - .bind(pk) - .bind(sk.as_deref()) - .bind(item_data) - .execute(&self.data_pool) - .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; - } else { - sqlx::query(&format!( - "INSERT INTO {ddb_table} (pk, item_data) VALUES ($1, $2)" - )) - .bind(pk) - .bind(item_data) - .execute(&self.data_pool) - .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; + // From here on a failure must not leave the target behind: it is + // CREATING with no scheduled transition, so nothing else would ever + // move it on, and the name would stay taken. Remove it and report + // the failure. + if let Err(e) = self.copy_backup_items(&desc, &backup_arn).await { + tracing::error!( + "restore of {backup_arn} into {target_table_name} failed, \ + removing the partial table: {e}" + ); + if let Err(cleanup) = self.abort_restore(&desc.table_id).await { + tracing::error!( + "could not remove partially restored table {target_table_name} \ + ({}): {cleanup}", + desc.table_id + ); } + return Err(e); } - #[allow(clippy::cast_possible_wrap)] - let item_count = items.len() as i64; - - sqlx::query( - "UPDATE tables SET item_count = $1 WHERE account_id = $2 AND table_name = $3", - ) - .bind(item_count) - .bind(&account_id) - .bind(&target_table_name) - .execute(&self.pool) - .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; - - // 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", - ) - .bind(&account_id) - .bind(&target_table_name) - .execute(&self.pool) - .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; - // Return CREATING — the API response shows the initial status, // but the table is already ACTIVE by the time the caller polls. Ok(desc) diff --git a/crates/storage-postgres/tests/backup_restore.rs b/crates/storage-postgres/tests/backup_restore.rs new file mode 100644 index 000000000..dd0399ea6 --- /dev/null +++ b/crates/storage-postgres/tests/backup_restore.rs @@ -0,0 +1,449 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 +//! Storage-level tests for PostgreSQL backup restore. +//! +//! The wire suite (`tests/test_backup_restore_fidelity.py`) checks round trips +//! on every backend. Two cases are only reachable here, because they need a +//! `backup_items` row the current binary would not write: +//! +//! - a backup written before the sort-key fix, whose rows carry the item with +//! `sk` NULL, must restore with its sort keys (they are read from the item); +//! - a backup whose copy fails part-way must not leave the target table +//! behind in CREATING, where nothing would ever move it on and its name +//! would stay taken; +//! - a backup with two rows for one key must fail rather than collapse them; +//! - a multi-part key backup must restore every key part (multi-part base +//! keys are a preview gated by `allow_multipart_table_keys`, not reachable +//! with the default configuration). +//! +//! Each test builds a throwaway catalog and data database and drops them when +//! it passes; a failing test leaves `eddb_bkup_*` behind for inspection. +//! Requires `EXTENDDB_TEST_PG_CONNECTION_STRING` (a base URL with no database +//! component); without it every test reports a skip and passes. + +use extenddb_core::types::{ + AttributeDefinition, BillingMode, CreateTableInput, DescribeTableInput, Item, KeySchemaElement, + KeyType, ScalarAttributeType, TableKeyInfo, +}; +use extenddb_storage::error::StorageError; +use extenddb_storage::{BackupEngine, DataEngine, TableEngine}; +use extenddb_storage_postgres::{PostgresConfig, PostgresEngine}; +use sqlx::PgPool; +use sqlx::postgres::PgPoolOptions; + +const ACCOUNT: &str = "123456789012"; +const REGION: &str = "us-east-1"; + +/// A catalog database and a separate data database, the layout `extenddb +/// init` creates. They must be separate here: the catalog schema carries a +/// legacy `stream_shards` table with a foreign key to `tables`, which would +/// shadow the data schema's table if both were applied to one database. +struct Scratch { + engine: PostgresEngine, + /// The data database (`stream_shards`, `stream_records`). + db: PgPool, + catalog: PgPool, + admin: PgPool, + db_names: [String; 2], +} + +impl Scratch { + async fn cleanup(self) { + let Scratch { + engine, + db, + catalog, + admin, + db_names, + } = self; + drop(engine); + db.close().await; + catalog.close().await; + for name in db_names { + sqlx::query(&format!("DROP DATABASE IF EXISTS \"{name}\" WITH (FORCE)")) + .execute(&admin) + .await + .expect("drop a scratch database"); + } + admin.close().await; + } +} + +fn base_conn() -> Option { + let conn = std::env::var("EXTENDDB_TEST_PG_CONNECTION_STRING").ok()?; + (!conn.trim().is_empty()).then(|| conn.trim_end_matches('/').to_owned()) +} + +async fn connect(url: &str) -> PgPool { + PgPoolOptions::new() + .max_connections(2) + .connect(url) + .await + .expect("connect to a scratch database") +} + +async fn apply(pool: &PgPool, migrations: &[&str]) { + for sql in migrations { + sqlx::raw_sql(sql) + .execute(pool) + .await + .expect("apply a shipped migration"); + } +} + +async fn scratch() -> Scratch { + let base = base_conn().expect("caller checks base_conn() first"); + let stem = format!("eddb_bkup_{}", uuid::Uuid::new_v4().simple())[..24].to_owned(); + let (catalog_name, data_name) = (format!("{stem}_c"), format!("{stem}_d")); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect(&format!("{base}/postgres")) + .await + .expect("connect to the postgres maintenance database"); + for name in [&catalog_name, &data_name] { + sqlx::query(&format!("CREATE DATABASE \"{name}\"")) + .execute(&admin) + .await + .expect("create a scratch database"); + } + let catalog_url = format!("{base}/{catalog_name}"); + let data_url = format!("{base}/{data_name}"); + let catalog = connect(&catalog_url).await; + let db = connect(&data_url).await; + apply( + &catalog, + &[ + include_str!("../migrations/001_schema.sql"), + include_str!("../migrations/002_vector_indexes.sql"), + ], + ) + .await; + apply( + &db, + &[ + include_str!("../data_migrations/001_data_schema.sql"), + include_str!("../data_migrations/002_gsi_pending.sql"), + include_str!("../data_migrations/003_idempotency_account_scope.sql"), + include_str!("../data_migrations/004_vector_index_state.sql"), + ], + ) + .await; + sqlx::query("UPDATE settings SET value = '0' WHERE key = 'control_plane_delay_seconds'") + .execute(&catalog) + .await + .expect("pin the control-plane delay to zero"); + sqlx::query( + "INSERT INTO settings (key, value) VALUES ('data_database_connection_string', $1) \ + ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value", + ) + .bind(&data_url) + .execute(&catalog) + .await + .expect("point the catalog at the data database"); + sqlx::query("INSERT INTO accounts (account_id, account_name) VALUES ($1, $2)") + .bind(ACCOUNT) + .bind(format!("acct-{stem}")) + .execute(&catalog) + .await + .expect("seed the account row"); + let engine = PostgresEngine::new( + &PostgresConfig { + connection_string: catalog_url, + pool_size: 10, + max_item_size_bytes: 400_000, + }, + REGION, + ) + .await + .expect("open a PostgresEngine on the scratch databases"); + Scratch { + engine, + db, + catalog, + admin, + db_names: [catalog_name, data_name], + } +} + +fn composite_table(name: &str) -> CreateTableInput { + let key = |name: &str, key_type| KeySchemaElement { + attribute_name: name.to_owned(), + key_type, + }; + let attr = |name: &str, attribute_type| AttributeDefinition { + attribute_name: name.to_owned(), + attribute_type, + }; + CreateTableInput { + table_name: name.to_owned(), + key_schema: vec![key("pk", KeyType::Hash), key("sk", KeyType::Range)], + attribute_definitions: vec![ + attr("pk", ScalarAttributeType::S), + attr("sk", ScalarAttributeType::N), + ], + billing_mode: Some(BillingMode::PayPerRequest), + ..Default::default() + } +} + +/// A backup of an empty composite-key table, then `rows` written straight +/// into `backup_items` in the pre-fix shape: `sk` NULL, the item in +/// `item_data`. +async fn legacy_backup(s: &Scratch, source: &str, rows: &[serde_json::Value]) -> String { + legacy_backup_of(s, composite_table(source), rows).await +} + +async fn legacy_backup_of( + s: &Scratch, + source: CreateTableInput, + rows: &[serde_json::Value], +) -> String { + let source_name = source.table_name.clone(); + s.engine + .create_table(ACCOUNT, source) + .await + .expect("create the source table"); + let backup = s + .engine + .create_backup(ACCOUNT, &source_name, "legacy") + .await + .expect("back up the empty source"); + for item in rows { + sqlx::query( + "INSERT INTO backup_items (backup_arn, pk, sk, item_data) VALUES ($1, $2, NULL, $3)", + ) + .bind(&backup.backup_arn) + .bind(item["pk"]["S"].as_str().unwrap_or("x")) + .bind(item) + .execute(&s.catalog) + .await + .expect("write a legacy backup row"); + } + backup.backup_arn +} + +async fn table_status(s: &Scratch, name: &str) -> Option { + sqlx::query_scalar("SELECT table_status FROM tables WHERE account_id = $1 AND table_name = $2") + .bind(ACCOUNT) + .bind(name) + .fetch_optional(&s.catalog) + .await + .expect("read the table status") +} + +#[tokio::test] +async fn legacy_composite_backup_restores_with_sort_keys() { + if base_conn().is_none() { + eprintln!("SKIP legacy_composite_backup_restores_with_sort_keys: no PostgreSQL"); + return; + } + let s = scratch().await; + let rows: Vec = (0..12) + .map(|i| { + serde_json::json!({ + "pk": {"S": format!("p{}", i % 3)}, + "sk": {"N": format!("{}", i * 10 - 50)}, + "v": {"S": format!("value-{i}")}, + }) + }) + .collect(); + let arn = legacy_backup(&s, "legacy_src", &rows).await; + + s.engine + .restore_table_from_backup(ACCOUNT, "legacy_dst", &arn) + .await + .expect("restore a pre-fix backup"); + assert_eq!( + table_status(&s, "legacy_dst").await.as_deref(), + Some("ACTIVE") + ); + + let desc = s + .engine + .describe_table( + ACCOUNT, + DescribeTableInput { + table_name: "legacy_dst".to_owned(), + }, + ) + .await + .expect("describe the restored table"); + assert_eq!(desc.item_count, 12); + assert!( + desc.table_size_bytes > 0, + "a restored table reports its size as soon as it is ACTIVE" + ); + let key_info = TableKeyInfo { + table_name: "legacy_dst".to_owned(), + account_id: ACCOUNT.to_owned(), + table_id: desc.table_id.clone(), + key_schema: desc.key_schema.clone(), + base_key_schema: desc.key_schema.clone(), + attribute_definitions: desc.attribute_definitions.clone(), + ..Default::default() + }; + + // Every row is addressable by its full key, so the sort key landed in its + // typed column rather than being dropped. + for row in &rows { + let item: Item = serde_json::from_value(row.clone()).expect("item from json"); + let key: Item = item + .iter() + .filter(|(k, _)| *k == "pk" || *k == "sk") + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + let got = s + .engine + .get_item(&key_info, &key) + .await + .expect("point read"); + assert_eq!(got.as_ref(), Some(&item)); + } + + s.cleanup().await; +} + +#[tokio::test] +async fn failed_restore_removes_the_partial_table() { + if base_conn().is_none() { + eprintln!("SKIP failed_restore_removes_the_partial_table: no PostgreSQL"); + return; + } + let s = scratch().await; + let rows = vec![ + serde_json::json!({"pk": {"S": "a"}, "sk": {"N": "1"}}), + // No sort key: the copy cannot place this item. + serde_json::json!({"pk": {"S": "b"}}), + ]; + let arn = legacy_backup(&s, "broken_src", &rows).await; + // A non-zero control-plane delay sends an ordinary DeleteTable through + // DELETING; cleanup of a failed restore must not depend on that. + sqlx::query("UPDATE settings SET value = '5' WHERE key = 'control_plane_delay_seconds'") + .execute(&s.catalog) + .await + .expect("set a non-zero control-plane delay"); + + let err = s + .engine + .restore_table_from_backup(ACCOUNT, "broken_dst", &arn) + .await + .expect_err("a backup row without its sort key cannot restore"); + assert!(matches!(err, StorageError::Internal(_)), "{err:?}"); + + // Neither CREATING, DELETING, nor half-filled: the target and its data + // table are gone and the name is free. + assert_eq!(table_status(&s, "broken_dst").await, None); + let leftover: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pg_tables WHERE schemaname = 'public' AND tablename LIKE '\\_ddb\\_%'", + ) + .fetch_one(&s.db) + .await + .expect("count data tables"); + assert_eq!(leftover, 1, "only the source table's data table remains"); + s.engine + .create_table(ACCOUNT, composite_table("broken_dst")) + .await + .expect("the target name is free again"); + + s.cleanup().await; +} + +#[tokio::test] +async fn duplicate_backup_rows_fail_the_restore() { + if base_conn().is_none() { + eprintln!("SKIP duplicate_backup_rows_fail_the_restore: no PostgreSQL"); + return; + } + let s = scratch().await; + // 1 and 1.0 are one N key value: a backup carrying both is corrupt, and + // the restore must say so rather than keep one and report two. + let rows = vec![ + serde_json::json!({"pk": {"S": "a"}, "sk": {"N": "1"}, "v": {"S": "first"}}), + serde_json::json!({"pk": {"S": "a"}, "sk": {"N": "1.0"}, "v": {"S": "second"}}), + ]; + let arn = legacy_backup(&s, "dup_src", &rows).await; + let err = s + .engine + .restore_table_from_backup(ACCOUNT, "dup_dst", &arn) + .await + .expect_err("two rows for one key cannot restore"); + assert!(matches!(err, StorageError::Internal(_)), "{err:?}"); + assert_eq!(table_status(&s, "dup_dst").await, None); + s.cleanup().await; +} + +#[tokio::test] +async fn multipart_key_backup_restores_every_key_part() { + if base_conn().is_none() { + eprintln!("SKIP multipart_key_backup_restores_every_key_part: no PostgreSQL"); + return; + } + let s = scratch().await; + let key = |name: &str, key_type| KeySchemaElement { + attribute_name: name.to_owned(), + key_type, + }; + let attr = |name: &str, attribute_type| AttributeDefinition { + attribute_name: name.to_owned(), + attribute_type, + }; + let input = CreateTableInput { + table_name: "multi_src".to_owned(), + key_schema: vec![ + key("h1", KeyType::Hash), + key("h2", KeyType::Hash), + key("r1", KeyType::Range), + key("r2", KeyType::Range), + ], + attribute_definitions: vec![ + attr("h1", ScalarAttributeType::S), + attr("h2", ScalarAttributeType::N), + attr("r1", ScalarAttributeType::S), + attr("r2", ScalarAttributeType::B), + ], + billing_mode: Some(BillingMode::PayPerRequest), + ..Default::default() + }; + // Rows differ only in the second HASH or the second RANGE attribute, so a + // restore that keyed on the first of each would collapse them. + let rows: Vec = [("1", "AA=="), ("2", "AA=="), ("1", "AQ==")] + .iter() + .map(|(h2, r2)| { + serde_json::json!({ + "h1": {"S": "same"}, "h2": {"N": h2}, + "r1": {"S": "same"}, "r2": {"B": r2}, + }) + }) + .collect(); + let arn = legacy_backup_of(&s, input, &rows).await; + s.engine + .restore_table_from_backup(ACCOUNT, "multi_dst", &arn) + .await + .expect("restore a multi-part key backup"); + let desc = s + .engine + .describe_table( + ACCOUNT, + DescribeTableInput { + table_name: "multi_dst".to_owned(), + }, + ) + .await + .expect("describe the restored table"); + // Checked on the physical rows: on this backend the PutItem and GetItem + // paths address only the first HASH and RANGE attribute of a base table + // (multi-part base keys are off by default), so a point read cannot be + // the oracle here. Each row must carry both HASH parts in `pk` and both + // RANGE parts in their typed columns, with nothing left NULL. + let (rows_n, pks, r1s, r2s, nulls): (i64, i64, i64, i64, i64) = sqlx::query_as(&format!( + "SELECT count(*), count(DISTINCT pk), count(DISTINCT sk_s), count(DISTINCT sk2_b), \ + count(*) FILTER (WHERE sk_s IS NULL OR sk2_b IS NULL) FROM \"_ddb_{}\"", + desc.table_id + )) + .fetch_one(&s.db) + .await + .expect("read restored rows"); + assert_eq!((rows_n, pks, r1s, r2s, nulls), (3, 2, 1, 2, 0)); + assert_eq!(desc.item_count, 3); + s.cleanup().await; +} diff --git a/tests/test_backup_restore_fidelity.py b/tests/test_backup_restore_fidelity.py new file mode 100644 index 000000000..9d00652a2 --- /dev/null +++ b/tests/test_backup_restore_fidelity.py @@ -0,0 +1,216 @@ +# Copyright 2026 ExtendDB contributors +# SPDX-License-Identifier: Apache-2.0 + +"""Backup and restore reproduce the table they were taken from. + +Every backend, every key shape. A restore must reach ACTIVE with the source's +key schema, attribute definitions, and exactly the source's items, attribute +for attribute. Comparing whole items rather than counts is what catches a +backend that restores the right number of rows under the wrong keys. + +Readiness assessment P0-2: on PostgreSQL every table with a sort key backed up +without its sort keys, and restoring it left the target in CREATING forever. +""" + +from __future__ import annotations + +import time +import uuid + +import pytest +from botocore.exceptions import ClientError + +from conftest import _poll_interval, wait_for_active, wait_for_deleted + +ITEMS_PER_TABLE = 60 +# The service takes minutes to make a backup AVAILABLE and to restore a table. +BACKUP_TIMEOUT_S = 900.0 +RESTORE_TIMEOUT_S = 1800.0 + + +def _create_backup(client, table_name: str) -> str: + """CreateBackup, then wait for the backup to be AVAILABLE. + + The service creates a backup asynchronously and refuses to restore one + that is still CREATING; ExtendDB returns it AVAILABLE. A just-created + table can also refuse CreateBackup for a short while, which is retried. + """ + deadline = time.monotonic() + 60 + while True: + try: + arn = client.create_backup( + TableName=table_name, BackupName=f"{table_name}-bkp" + )["BackupDetails"]["BackupArn"] + break + except ClientError as e: + code = e.response["Error"]["Code"] + retryable = code in ("ContinuousBackupsUnavailableException", "TableInUseException") + if not retryable or time.monotonic() >= deadline: + raise + time.sleep(0.5) + deadline = time.monotonic() + BACKUP_TIMEOUT_S + while True: + status = client.describe_backup(BackupArn=arn)["BackupDescription"][ + "BackupDetails" + ]["BackupStatus"] + if status == "AVAILABLE": + return arn + assert status == "CREATING", f"backup {arn} is {status}" + assert time.monotonic() < deadline, f"backup {arn} not AVAILABLE in time" + time.sleep(_poll_interval()) + + +def _scan_all(client, table_name: str) -> list[dict]: + items: list[dict] = [] + kwargs: dict = {"TableName": table_name, "ConsistentRead": True} + while True: + resp = client.scan(**kwargs) + items += resp["Items"] + if "LastEvaluatedKey" not in resp: + return items + kwargs["ExclusiveStartKey"] = resp["LastEvaluatedKey"] + + +def _canonical(item: dict) -> str: + """Order-independent, type-preserving form of a wire item for comparison.""" + import json + + def norm(v): + if isinstance(v, dict): + return {k: norm(v[k]) for k in sorted(v)} + if isinstance(v, list): + return [norm(x) for x in v] + if isinstance(v, (bytes, bytearray)): + return {"__bytes__": bytes(v).hex()} + return v + + # Set members have no order on the wire. + def sort_sets(av): + if isinstance(av, dict): + out = {} + for t, val in av.items(): + if t in ("SS", "NS"): + out[t] = sorted(val) + elif t == "BS": + out[t] = sorted(bytes(b).hex() for b in val) + elif t == "M": + out[t] = {k: sort_sets(x) for k, x in val.items()} + elif t == "L": + out[t] = [sort_sets(x) for x in val] + else: + out[t] = val + return out + return av + + return json.dumps(norm({k: sort_sets(v) for k, v in item.items()}), sort_keys=True) + + +def _drop(client, *names: str) -> None: + for name in names: + try: + client.delete_table(TableName=name) + except client.exceptions.ResourceNotFoundException: + continue + wait_for_deleted(client, name) + + +def _sort_value(kind: str, i: int) -> dict: + if kind == "S": + return {"S": f"sort-{i:04d}"} + if kind == "N": + # Negative, fractional, and large magnitudes in one column. + return {"N": str((i - 30) * 1.5 if i % 3 else (i - 30) * 10**20)} + return {"B": i.to_bytes(2, "big") + b"\x00\xff"} + + +def _item(i: int, sort_kind: str | None) -> dict: + # With a sort key, several items share a partition; without one, every + # partition key must be distinct or the puts overwrite each other. + item: dict = { + "pk": {"S": f"part-{i % 7}" if sort_kind else f"part-{i}"}, + "str": {"S": f"value-{i}"}, + "num": {"N": str(i)}, + "nested": {"M": {"list": {"L": [{"N": "1"}, {"S": "two"}, {"BOOL": i % 2 == 0}]}}}, + "tags": {"SS": [f"t{i}", "shared"]}, + "blob": {"B": bytes([i % 256]) * 3}, + } + if i % 5 == 0: + item["maybe"] = {"NULL": True} + if sort_kind: + item["sk"] = _sort_value(sort_kind, i) + return item + + +def _round_trip(client, sort_kind: str | None) -> None: + source = f"restore-fid-{uuid.uuid4().hex[:10]}" + restored = f"{source}-r" + key_schema = [{"AttributeName": "pk", "KeyType": "HASH"}] + attr_defs = [{"AttributeName": "pk", "AttributeType": "S"}] + if sort_kind: + key_schema.append({"AttributeName": "sk", "KeyType": "RANGE"}) + attr_defs.append({"AttributeName": "sk", "AttributeType": sort_kind}) + client.create_table( + TableName=source, + KeySchema=key_schema, + AttributeDefinitions=attr_defs, + BillingMode="PAY_PER_REQUEST", + ) + backup_arn = None + try: + wait_for_active(client, source) + for i in range(ITEMS_PER_TABLE): + client.put_item(TableName=source, Item=_item(i, sort_kind)) + # Compare against what the source serves, not what was sent: numbers + # come back canonicalized (`-42.0` reads as `-42`), here as in DynamoDB. + expected = _scan_all(client, source) + assert len(expected) == ITEMS_PER_TABLE + + backup_arn = _create_backup(client, source) + resp = client.restore_table_from_backup(TargetTableName=restored, BackupArn=backup_arn) + assert resp["TableDescription"]["TableName"] == restored + wait_for_active(client, restored, timeout=RESTORE_TIMEOUT_S) + + table = client.describe_table(TableName=restored)["Table"] + assert table["KeySchema"] == key_schema + assert sorted(table["AttributeDefinitions"], key=lambda a: a["AttributeName"]) == sorted( + attr_defs, key=lambda a: a["AttributeName"] + ) + + got = sorted(_canonical(i) for i in _scan_all(client, restored)) + want = sorted(_canonical(i) for i in expected) + assert len(got) == len(want), f"restored {len(got)} of {len(want)} items" + assert got == want + + # Point reads by full key, so a restore that stored the right items + # under the wrong key columns fails here even if Scan looked right. + for item in expected[:10]: + key = {"pk": item["pk"]} + if sort_kind: + key["sk"] = item["sk"] + fetched = client.get_item(TableName=restored, Key=key, ConsistentRead=True) + assert _canonical(fetched.get("Item", {})) == _canonical(item) + + # The restored table takes writes like any other. + client.put_item(TableName=restored, Item=_item(ITEMS_PER_TABLE, sort_kind)) + finally: + # Each cleanup step runs even if an earlier one fails. + try: + _drop(client, restored) + finally: + try: + _drop(client, source) + finally: + if backup_arn: + try: + client.delete_backup(BackupArn=backup_arn) + except client.exceptions.BackupNotFoundException: + pass + + +def test_restore_hash_only_table(dynamodb_client): + _round_trip(dynamodb_client, None) + + +@pytest.mark.parametrize("sort_kind", ["S", "N", "B"]) +def test_restore_composite_key_table(dynamodb_client, sort_kind): + _round_trip(dynamodb_client, sort_kind) From 1d981b89b3d5dcd3a81012815d86672da6eab47b Mon Sep 17 00:00:00 2001 From: Scott Robinson Date: Mon, 5 Oct 2026 21:36:23 +0000 Subject: [PATCH 2/4] fix(backup): restore indexes and throughput, bound memory, recover abandoned restores Follows the sort-key fix. Restores on PostgreSQL and SQLite now reproduce what the service reproduces, and the copy is bounded and crash-safe. Restored, from a new per-backup table definition: - global and local secondary indexes, with their key schemas, projections, and GSI throughput, filled during the copy; - billing mode and provisioned throughput (was always 5/5); - table class, SSE specification, on-demand throughput. Not restored, as on the service: streams, TTL, tags, deletion protection. Catalog 0.0.4. PostgreSQL migration 003 and the SQLite schema add `backup_definitions`, one row per backup holding that definition as JSON in the wire's shape behind a version marker (`extenddb_storage::backup_definition`). Backups taken before the upgrade have no row and restore as before: keys and items, no secondary indexes, 5/5 for a provisioned table. The migration runner now writes the compiled catalog version after every walk, as #335 also does, so a replayed earlier migration cannot leave the version behind the schema. Consistency: - PostgreSQL CreateBackup holds the source table row FOR SHARE from its first catalog read until its data snapshot is taken, and locks the source data table before releasing it. UpdateTable and DeleteTable take the row FOR UPDATE, so the definition recorded is the one in force when the items were read. A GSI whose catalog row exists but whose data table is not yet built is left out, with any attribute definition only it used. - SQLite CreateBackup reads the table, its definition, and its items in one write transaction. - DeleteBackup removes a backup's rows and marks it DELETED in one transaction; a restore reads from a snapshot (PostgreSQL) or re-checks under the write lock per batch (SQLite), so it copies all of a backup or fails with BackupNotFoundException. - DeleteTable on a restore target returns ResourceInUseException, as the service does for a table it is still creating, instead of deleting it out from under the copy. Memory: items are streamed and buffered up to 500 items or 4 MiB of stored JSON, whichever comes first, on both sides of both backends. Measured on PostgreSQL: 100,000 items of 4 KB (about 400 MB), peak server RSS 87 MB during the backup and 109 MB during the restore; 150 items of about 2.9 MB JSON each (435 MB), 147 MB and 168 MB with MALLOC_ARENA_MAX=1 (glibc's per-thread arenas otherwise retain freed memory and RSS reads higher). Concurrency: SQLite restores commit per batch, so other writers wait for a batch rather than the whole restore (PutItem p50 137 ms, max 276 ms during a 60,000-item restore). SQLite backups still hold the write lock for their duration, as they must to be consistent. PostgreSQL allows four concurrent restores per process; a fifth waits up to 30 s and is then refused with LimitExceededException. Failure and crash: - A failed copy removes the target by table id, synchronously: it is claimed CREATING -> DELETING in one statement (so removal and the flip to ACTIVE are mutually exclusive), its data and index tables are dropped, and its row deleted. If a drop fails the row stays DELETING and the control plane finishes the job. - PostgreSQL: a restore holds a session advisory lock keyed by account and target name from before the target is created until the copy ends, on a dedicated connection with idle_session_timeout disabled. The control-plane pass removes targets that are CREATING with no scheduled transition, older than 60 s, and whose lock is free. Verified with kill -9 mid-restore: the target is removed about 60 s later and the name can be restored into again. - SQLite: the same targets are removed at startup, before serving; this assumes one server process per database file, which the backend's in-process write lock already requires. Verified with kill -9. Refused rather than restored wrongly: backups of tables with vector indexes (as before), backups of multi-part-key tables (a preview the item paths address only by their first parts), and definitions with an unknown version. RestoreTableFromBackup now validates TargetTableName. Tests: - tests/test_backup_restore_fidelity.py: hash-only tables with S, N, and B keys; composite keys with S, N, and B sort keys; a provisioned table with four GSIs (ALL, INCLUDE, KEYS_ONLY; S and N hash keys, N and B range keys; sparse) and an LSI, comparing definitions, throughput, IndexStatus, every index's contents, projected attribute sets against the projection definitions, and ordered ranged index reads in both directions. Waits for backups and restores as the service needs. 7/7 on PostgreSQL and SQLite; the index test fails on the previous commit. - crates/storage-postgres/tests/backup_restore.rs (14): legacy row shape, pre-0.0.4 backup after replaying migration 003, failed-copy cleanup under a non-zero control-plane delay, duplicate rows, multi-part refusal, indexes and throughput, table class/SSE/on-demand, >1,000 rows with oversized items, the abandoned-restore sweep (grace, held lock, unowned), a name whose lock is held, the backup barrier against an uncommitted UpdateTable (fails on the previous approach), a half-built GSI, and DeleteTable during a restore. Run one at a time, and scratch databases are dropped even when a test panics. - crates/storage-sqlite/src/backup.rs restore_tests (9): indexes and throughput, the startup sweep, failed-copy cleanup, multi-part refusal, table class/SSE/on-demand, 730 rows across batches with physical row and index-row equality against the write path, a database with the exact 0.0.3 schema (testdata/schema_0_0_3.sql) holding a 0.0.3-format backup, byte-bounded batch cuts, and DeleteTable during a restore. - tests/test_cli_vector_catalog_migration.py covers 0.0.2 -> 0.0.4 and the replay case. Not changed here: RestoreSummary appears only in the RestoreTableFromBackup response and not in later DescribeTable output (pre-existing; needs a persisted column). DeleteBackup during a restore fails the restore with BackupNotFoundException rather than refusing with BackupInUseException. Vector index restore is still refused. Readiness assessment P0-2. Signed-off-by: Scott Robinson --- crates/engine/src/backup.rs | 4 + .../migrations/003_backup_definitions.sql | 27 + crates/storage-postgres/src/backup_engine.rs | 942 +++++++++--- crates/storage-postgres/src/delete_table.rs | 23 + crates/storage-postgres/src/lib.rs | 2 +- crates/storage-postgres/src/migrations.rs | 28 +- crates/storage-postgres/src/worker_store.rs | 13 + .../storage-postgres/tests/backup_restore.rs | 817 +++++++++- .../storage-postgres/tests/key_collation.rs | 1 + .../tests/vector_control_plane.rs | 1 + crates/storage-sqlite/src/backup.rs | 1329 +++++++++++++++-- crates/storage-sqlite/src/data/mod.rs | 4 +- crates/storage-sqlite/src/delete_table.rs | 23 + crates/storage-sqlite/src/lib.rs | 14 + crates/storage-sqlite/src/schema.rs | 14 +- .../storage-sqlite/testdata/schema_0_0_3.sql | 400 +++++ crates/storage/src/backup_definition.rs | 316 ++++ crates/storage/src/lib.rs | 1 + docs/getting-started.md | 6 +- docs/manuals/01-architecture-guide.md | 2 +- docs/manuals/04-quickstart-setup-guide.md | 4 +- docs/manuals/07-upgrade-manual.md | 27 +- docs/manuals/08-install-linux.md | 2 +- docs/manuals/09-install-macos.md | 2 +- tests/test_backup_restore_fidelity.py | 349 ++++- tests/test_cli_vector_catalog_migration.py | 53 +- 26 files changed, 3953 insertions(+), 451 deletions(-) create mode 100644 crates/storage-postgres/migrations/003_backup_definitions.sql create mode 100644 crates/storage-sqlite/testdata/schema_0_0_3.sql create mode 100644 crates/storage/src/backup_definition.rs diff --git a/crates/engine/src/backup.rs b/crates/engine/src/backup.rs index 510190d1e..36d42d098 100755 --- a/crates/engine/src/backup.rs +++ b/crates/engine/src/backup.rs @@ -145,6 +145,7 @@ pub(crate) async fn handle_restore_table_from_backup( .to_owned(), ) })?; + extenddb_core::validation::validate_table_name(target_table_name, &ctx.limits)?; let backup_arn = backup_arn_field(&body, &ctx.account_id)?; let mut desc = ctx @@ -292,6 +293,9 @@ fn storage_err_to_dynamo(e: extenddb_storage::error::StorageError) -> DynamoDbEr extenddb_storage::error::StorageError::Unsupported(msg) => { DynamoDbError::ValidationException(msg) } + extenddb_storage::error::StorageError::LimitExceeded(msg) => { + DynamoDbError::LimitExceededException(msg) + } other => { tracing::error!(internal_error = %other, "backup storage error"); DynamoDbError::InternalServerError("Internal server error".to_owned()) diff --git a/crates/storage-postgres/migrations/003_backup_definitions.sql b/crates/storage-postgres/migrations/003_backup_definitions.sql new file mode 100644 index 000000000..654537494 --- /dev/null +++ b/crates/storage-postgres/migrations/003_backup_definitions.sql @@ -0,0 +1,27 @@ +-- Copyright 2026 ExtendDB contributors +-- SPDX-License-Identifier: Apache-2.0 +-- Migration 003: record a backup's table definition (catalog version 0.0.4). +-- +-- A backup kept the source table's key schema, attribute definitions, and +-- billing mode, and nothing else, so a restored table came back without its +-- global and local secondary indexes, with 5/5 provisioned throughput, and +-- without its table class or encryption settings. One row per backup holds +-- those, in the wire's own shape behind a version marker (see +-- `extenddb_storage::backup_definition`). A backup taken before this +-- migration has no row and restores as before. +-- +-- Written to tolerate a replay, like 002: the runner applies a migration and +-- records it in `schema_history` as two separate commits, so a crash in +-- between leaves this file applied but unrecorded, and the next migrate runs +-- it again. CREATE TABLE IF NOT EXISTS and the version UPDATE are idempotent. + +BEGIN; + +CREATE TABLE IF NOT EXISTS backup_definitions ( + backup_arn TEXT PRIMARY KEY REFERENCES backups(backup_arn) ON DELETE CASCADE, + definition JSONB NOT NULL +); + +UPDATE settings SET value = '0.0.4' WHERE key = 'catalog_version'; + +COMMIT; diff --git a/crates/storage-postgres/src/backup_engine.rs b/crates/storage-postgres/src/backup_engine.rs index 9ea057971..c20bfbc84 100755 --- a/crates/storage-postgres/src/backup_engine.rs +++ b/crates/storage-postgres/src/backup_engine.rs @@ -4,17 +4,24 @@ //! Backup and point-in-time recovery implementation for `PostgreSQL` storage. use extenddb_core::types::{ - BackupDescription, BackupDetails, BackupSummary, ContinuousBackupsDescription, Item, - PointInTimeRecoveryDescription, SourceTableDetails, TableDescription, + BackupDescription, BackupDetails, BackupSummary, ContinuousBackupsDescription, GsiInput, Item, + KeySchemaElement, LsiInput, PointInTimeRecoveryDescription, Projection, ScalarAttributeType, + SourceTableDetails, TableDescription, }; use extenddb_storage::BackupEngine; +use extenddb_storage::backup_definition::{ + BACKUP_DEFINITION_VERSION, BackupTableDefinition, COPY_BATCH_BYTES, COPY_BATCH_ITEMS, + ensure_single_part_base_key, throughput_from_catalog, +}; use extenddb_storage::error::StorageError; +use extenddb_storage::util::{SortKeyValue, composite_pk_to_text, parse_sk, sk_column_n}; use futures::future::BoxFuture; use crate::PostgresEngine; -use extenddb_storage::util::{SortKeyValue, composite_pk_to_text, parse_sk, sk_column_n}; - -use crate::data::{all_sort_key_info, data_table_name}; +use crate::data::{ + all_sort_key_info, data_table_name, index_table_name, insert_index_row_multi, + item_has_index_keys, project_item_for_index, +}; /// Current epoch milliseconds, used as the leading component of a backup id. fn epoch_millis() -> u128 { @@ -43,93 +50,323 @@ fn pg_timestamp_to_epoch(ts: time::OffsetDateTime) -> f64 { ts.unix_timestamp() as f64 } +/// Seed for the 64-bit restore lock key. Restore locks use the single-`bigint` +/// advisory lock form, which PostgreSQL keeps in a key space separate from the +/// two-`int4` form the migration and vector-build locks use, so they cannot +/// meet; the seed keeps this key space distinct from any other `bigint` user. +const RESTORE_LOCK_SEED: i64 = 0x0045_4452; // 'E', 'D', 'R' + +/// Restores in flight at once in this process. Each holds one dedicated +/// connection for its lock on top of its pooled ones, so this bounds the +/// connections restores can open outside the configured pools. A restore +/// beyond the limit waits for a slot. +static RESTORE_SLOTS: tokio::sync::Semaphore = tokio::sync::Semaphore::const_new(4); + +/// How long a restore waits for a slot before it is refused with +/// `LimitExceededException`. +const RESTORE_SLOT_WAIT_SECS: u64 = 30; + +/// How old an unowned restore target must be before it is treated as +/// abandoned. The owner takes its lock milliseconds after creating the +/// target, so this only has to cover that gap with a wide margin. +const ABANDONED_RESTORE_GRACE_SECS: i32 = 60; + +/// A secondary index of a restore target, as the copy needs it. +struct RestoreIndex { + table: String, + key_schema: Vec, + projection: Projection, +} + +/// Session-scoped ownership of one restore, held for the life of the copy. +/// +/// A `pg_try_advisory_lock` on a dedicated connection: it dies with the +/// connection, so a crashed process's claim disappears on its own and the +/// abandoned-restore sweep can tell a dead restore from a running one. +/// +/// If only this connection is lost while the copy carries on (its backend is +/// terminated, say), the sweep may claim the target after the grace period. +/// That cannot leave a wrong table: the claim and the restore's flip to +/// ACTIVE are both conditional on CREATING, so exactly one wins, and the +/// restore then fails with "deleted while it was being restored". +struct RestoreOwner { + _conn: sqlx::PgConnection, + _slot: Option>, +} + +/// The advisory-lock key for a restore into `table_name` in `account_id`. +/// +/// Keyed by name rather than table id so the restore can hold it before the +/// target exists: the sweep must never see an unowned target, including in +/// the window while `create_table_impl` is still running DDL. The key is a +/// 64-bit `hashtextextended`, so a collision is negligible, and one would +/// only make the sweep skip a table while the other name's lock is held, +/// never remove a table it should not. +fn restore_lock_key(account_id: &str, table_name: &str) -> String { + format!("{account_id}/{table_name}") +} + +async fn try_restore_lock( + pool: &sqlx::PgPool, + account_id: &str, + table_name: &str, + slot: Option>, +) -> Result, StorageError> { + let options = pool.connect_options(); + let mut conn = ::connect_with(&options) + .await + .map_err(|e| StorageError::Internal(format!("restore lock connection: {e}")))?; + // This session sits idle for the whole copy. A server-side + // idle_session_timeout would end it and release the lock while the + // restore is still running, so it is disabled for this session only. + sqlx::query("SET idle_session_timeout = 0") + .execute(&mut conn) + .await + .map_err(|e| StorageError::Internal(format!("restore lock session: {e}")))?; + let taken: bool = sqlx::query_scalar("SELECT pg_try_advisory_lock(hashtextextended($1, $2))") + .bind(restore_lock_key(account_id, table_name)) + .bind(RESTORE_LOCK_SEED) + .fetch_one(&mut conn) + .await + .map_err(|e| StorageError::Internal(format!("restore lock: {e}")))?; + Ok(taken.then_some(RestoreOwner { + _conn: conn, + _slot: slot, + })) +} + +/// Insert one buffered batch of backup rows in one statement and clear the +/// buffers. Returns the number written. +async fn insert_backup_batch( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + backup_arn: &str, + pks: &mut Vec, + datas: &mut Vec, +) -> Result { + if pks.is_empty() { + return Ok(0); + } + let n = i64::try_from(pks.len()).unwrap_or(i64::MAX); + sqlx::query( + "INSERT INTO backup_items (backup_arn, pk, item_data) \ + SELECT $1, pk, item_data::jsonb FROM UNNEST($2::text[], $3::text[]) AS t(pk, item_data)", + ) + .bind(backup_arn) + .bind(&*pks) + .bind(&*datas) + .execute(&mut **tx) + .await + .map_err(db_err)?; + pks.clear(); + datas.clear(); + Ok(n) +} + +fn db_err(e: sqlx::Error) -> StorageError { + StorageError::Internal(format!("Database error: {e}")) +} + impl PostgresEngine { - /// Copy a backup's items into a freshly created restore target and mark - /// it ACTIVE. + /// Read the parts of a table's definition a restore recreates. + async fn capture_table_definition( + conn: &mut sqlx::PgConnection, + table_id: &str, + ) -> Result<(BackupTableDefinition, Vec<(String, String)>), StorageError> { + let (billing_mode, pt, table_class, sse, on_demand): ( + String, + Option, + Option, + Option, + Option, + ) = sqlx::query_as( + "SELECT billing_mode, provisioned_throughput, table_class, sse_specification, \ + on_demand_throughput FROM tables WHERE table_id = $1", + ) + .bind(table_id) + .fetch_one(&mut *conn) + .await + .map_err(db_err)?; + + let rows: Vec<( + String, + String, + String, + serde_json::Value, + serde_json::Value, + Option, + )> = sqlx::query_as( + "SELECT index_name, index_id, index_type, key_schema, projection, provisioned_throughput \ + FROM indexes WHERE table_id = $1 ORDER BY index_name", + ) + .bind(table_id) + .fetch_all(&mut *conn) + .await + .map_err(db_err)?; + let mut gsis = Vec::new(); + let mut lsis = Vec::new(); + let mut gsi_ids = Vec::new(); + for (index_name, index_id, index_type, ks, proj, ipt) in rows { + let key_schema: Vec = serde_json::from_value(ks) + .map_err(|e| StorageError::Internal(format!("index key schema: {e}")))?; + let projection: Projection = serde_json::from_value(proj) + .map_err(|e| StorageError::Internal(format!("index projection: {e}")))?; + if index_type == "LSI" { + lsis.push(LsiInput { + index_name, + key_schema, + projection, + }); + } else { + gsi_ids.push((index_name.clone(), index_id)); + gsis.push(GsiInput { + index_name, + key_schema, + projection, + provisioned_throughput: throughput_from_catalog(ipt.as_ref()), + }); + } + } + + let vector_index_names: Vec = sqlx::query_scalar( + "SELECT index_name FROM vector_indexes WHERE table_id = $1 ORDER BY index_name", + ) + .bind(table_id) + .fetch_all(&mut *conn) + .await + .map_err(db_err)?; + + Ok(( + BackupTableDefinition { + version: BACKUP_DEFINITION_VERSION, + billing_mode, + provisioned_throughput: throughput_from_catalog(pt.as_ref()), + global_secondary_indexes: gsis, + local_secondary_indexes: lsis, + table_class, + sse_specification: sse, + on_demand_throughput: on_demand + .map(serde_json::from_value) + .transpose() + .map_err(|e| StorageError::Internal(format!("on-demand throughput: {e}")))?, + vector_index_names, + }, + gsi_ids, + )) + } + + /// Copy a backup's items into a freshly created restore target, populate + /// its secondary indexes, and mark it ACTIVE. /// /// Every key column is derived from the item itself, so the copy does not - /// depend on how the source table's columns were laid out. Each row is a - /// plain INSERT: two backup rows for one key are a corrupt backup and fail - /// the restore rather than silently collapsing into one item. The copy is - /// one data transaction, so a failure part-way leaves the target empty. + /// depend on how the source table's columns were laid out. Rows are plain + /// INSERTs: two backup rows for one key are a corrupt backup and fail the + /// restore rather than silently collapsing into one item. Items are read + /// as a stream and written in batches, so memory is bounded by the batch + /// size; the whole copy is one data transaction, so a failure part-way + /// leaves the target empty. The target is CREATING throughout, which + /// refuses every data-plane request, so nothing else writes to it. async fn copy_backup_items( &self, desc: &TableDescription, backup_arn: &str, ) -> Result<(), StorageError> { - let items: Vec<(serde_json::Value,)> = - sqlx::query_as("SELECT item_data FROM backup_items WHERE backup_arn = $1") - .bind(backup_arn) - .fetch_all(&self.pool) - .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; + use futures::TryStreamExt; let ddb_table = data_table_name(&desc.table_id); let sort_keys = all_sort_key_info(&desc.key_schema, &desc.attribute_definitions); - let mut cols = vec!["pk".to_owned()]; - cols.extend( - sort_keys - .iter() - .enumerate() - .map(|(i, &(_, t))| sk_column_n(i, t)), - ); - cols.push("item_data".to_owned()); - let placeholders: Vec = (1..=cols.len()).map(|i| format!("${i}")).collect(); - let insert_sql = format!( - "INSERT INTO {ddb_table} ({}) VALUES ({})", - cols.join(", "), - placeholders.join(", ") - ); + let index_rows: Vec<(String, serde_json::Value, serde_json::Value)> = sqlx::query_as( + "SELECT index_id, key_schema, projection FROM indexes WHERE table_id = $1", + ) + .bind(&desc.table_id) + .fetch_all(&self.pool) + .await + .map_err(db_err)?; + let mut indexes = Vec::with_capacity(index_rows.len()); + for (index_id, ks, proj) in index_rows { + indexes.push(RestoreIndex { + table: index_table_name(&index_id), + key_schema: serde_json::from_value(ks) + .map_err(|e| StorageError::Internal(format!("index key schema: {e}")))?, + projection: serde_json::from_value(proj) + .map_err(|e| StorageError::Internal(format!("index projection: {e}")))?, + }); + } - let mut tx = self - .data_pool - .begin() + // The backup's rows are read from one catalog snapshot in which the + // backup is still AVAILABLE. DeleteBackup removes the rows and marks + // the backup DELETED in one transaction, so a concurrent delete is + // either wholly visible here (the restore fails) or not at all (the + // restore copies every row). + let mut snapshot = self + .pool + .begin_with("BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY") .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; - for (item_json,) in &items { - let item: Item = serde_json::from_value(item_json.clone()) + .map_err(db_err)?; + let available: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM backups WHERE backup_arn = $1 \ + AND backup_status = 'AVAILABLE')", + ) + .bind(backup_arn) + .fetch_one(&mut *snapshot) + .await + .map_err(db_err)?; + if !available { + return Err(StorageError::Validation(format!( + "Backup not found: {backup_arn}" + ))); + } + let mut tx = self.data_pool.begin().await.map_err(db_err)?; + let mut rows = sqlx::query_scalar::<_, String>( + "SELECT item_data::text FROM backup_items WHERE backup_arn = $1", + ) + .bind(backup_arn) + .fetch(&mut *snapshot); + let mut batch: Vec<(Item, String)> = Vec::with_capacity(COPY_BATCH_ITEMS); + let mut batch_bytes: usize = 0; + let mut item_count: i64 = 0; + while let Some(item_json) = rows.try_next().await.map_err(db_err)? { + let item: Item = serde_json::from_str(&item_json) .map_err(|e| StorageError::Internal(format!("Parse backup item: {e}")))?; - let mut query = - sqlx::query(&insert_sql).bind(composite_pk_to_text(&item, &desc.key_schema)?); - for &(name, sk_type) in &sort_keys { - let value = item.get(name).ok_or_else(|| { - StorageError::Internal(format!( - "backup {backup_arn} has an item without sort key {name}" - )) - })?; - query = match parse_sk(value, sk_type)? { - SortKeyValue::S(v) => query.bind(v), - SortKeyValue::N(v) => query.bind(v), - SortKeyValue::B(v) => query.bind(v), - }; + // The parsed item is kept for its keys and index projections; the + // budget counts its stored text, which bounds both forms. + batch_bytes += item_json.len(); + batch.push((item, item_json)); + if batch.len() == COPY_BATCH_ITEMS || batch_bytes >= COPY_BATCH_BYTES { + self.write_restore_batch( + &mut tx, &ddb_table, desc, &sort_keys, &indexes, &batch, backup_arn, + ) + .await?; + item_count += i64::try_from(batch.len()).unwrap_or(i64::MAX); + batch.clear(); + batch_bytes = 0; } - query - .bind(item_json) - .execute(&mut *tx) - .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; } - tx.commit() - .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; + drop(rows); + snapshot.commit().await.map_err(db_err)?; + if !batch.is_empty() { + self.write_restore_batch( + &mut tx, &ddb_table, desc, &sort_keys, &indexes, &batch, backup_arn, + ) + .await?; + item_count += i64::try_from(batch.len()).unwrap_or(i64::MAX); + } + tx.commit().await.map_err(db_err)?; - #[allow(clippy::cast_possible_wrap)] - let item_count = items.len() as i64; let (table_size,): (i64,) = sqlx::query_as(&format!( "SELECT COALESCE(pg_total_relation_size('{ddb_table}'), 0)" )) .fetch_one(&self.data_pool) .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; + .map_err(db_err)?; // Mark the restored table ACTIVE now that the copy has committed. The // table was created with the transition deferred, so this is the first // point it can become ACTIVE, by ordering rather than by timing. // - // Only from CREATING: a client may have deleted the target while the - // copy ran (DeleteTable accepts a CREATING table), and that delete - // wins. Flipping a DELETING row back to ACTIVE would cancel the - // delete and leave a table with no scheduled removal. + // Only from CREATING. DeleteTable refuses a restore target, but the + // abandoned-restore sweep can claim one whose lock session was lost + // (moving it to DELETING); that claim wins, and flipping the row back + // to ACTIVE would leave a table whose data tables are being dropped. let activated = sqlx::query( "UPDATE tables SET item_count = $1, table_size_bytes = $2, table_status = 'ACTIVE', \ status_transition_at = NULL WHERE table_id = $3 AND table_status = 'CREATING'", @@ -139,38 +376,199 @@ impl PostgresEngine { .bind(&desc.table_id) .execute(&self.pool) .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; + .map_err(db_err)?; if activated.rows_affected() == 0 { - tracing::info!( - "restore of {backup_arn} into {} finished after the target was deleted; \ - leaving the delete in place", + // The target was claimed for removal while the copy ran; the + // copy's rows go with it. + return Err(StorageError::TableNotFound(format!( + "restore of {backup_arn} into {} did not complete: the table was deleted \ + while it was being restored", desc.table_name - ); + ))); } Ok(()) } - /// Remove a restore target whose copy failed. + /// Write one batch of restored items: one multi-row INSERT into the base + /// table, then each item's row in every secondary index it belongs to. + #[allow(clippy::too_many_arguments)] + async fn write_restore_batch( + &self, + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + ddb_table: &str, + desc: &TableDescription, + sort_keys: &[(&str, ScalarAttributeType)], + indexes: &[RestoreIndex], + batch: &[(Item, String)], + backup_arn: &str, + ) -> Result<(), StorageError> { + let mut cols = vec!["pk".to_owned()]; + cols.extend( + sort_keys + .iter() + .enumerate() + .map(|(i, &(_, t))| sk_column_n(i, t)), + ); + cols.push("item_data".to_owned()); + let width = cols.len(); + // The last column is the item, bound as text and cast in SQL. + let tuples: Vec = (0..batch.len()) + .map(|r| { + let ps: Vec = (1..=width) + .map(|c| { + if c == width { + format!("${}::jsonb", r * width + c) + } else { + format!("${}", r * width + c) + } + }) + .collect(); + format!("({})", ps.join(", ")) + }) + .collect(); + let sql = format!( + "INSERT INTO {ddb_table} ({}) VALUES {}", + cols.join(", "), + tuples.join(", ") + ); + let mut query = sqlx::query(&sql); + for (item, item_json) in batch { + query = query.bind(composite_pk_to_text(item, &desc.key_schema)?); + for &(name, sk_type) in sort_keys { + let value = item.get(name).ok_or_else(|| { + StorageError::Internal(format!( + "backup {backup_arn} has an item without sort key {name}" + )) + })?; + query = match parse_sk(value, sk_type)? { + SortKeyValue::S(v) => query.bind(v), + SortKeyValue::N(v) => query.bind(v), + SortKeyValue::B(v) => query.bind(v), + }; + } + query = query.bind(item_json.as_str()); + } + query.execute(&mut **tx).await.map_err(db_err)?; + + if indexes.is_empty() { + return Ok(()); + } + let base_sks = all_sort_key_info(&desc.key_schema, &desc.attribute_definitions); + for idx in indexes { + let idx_sks = all_sort_key_info(&idx.key_schema, &desc.attribute_definitions); + for (item, _) in batch { + if !item_has_index_keys(item, &idx.key_schema) { + continue; + } + let projected = project_item_for_index( + item, + &idx.key_schema, + &desc.key_schema, + &idx.projection, + ); + insert_index_row_multi( + tx, + &idx.table, + item, + &projected, + &idx.key_schema, + &desc.key_schema, + &idx_sks, + &base_sks, + ) + .await?; + } + } + Ok(()) + } + + /// Remove a restore target whose copy failed or was abandoned. /// - /// Keyed by table id, not name, so it can only ever remove the table this + /// Keyed by table id, not name, so it can only ever remove the table that /// restore created, and synchronous regardless of the control-plane /// delay: the ordinary DeleteTable path would leave the name held in - /// DELETING for that long. The target was created with no secondary or - /// vector indexes, so its data table is the only physical object. - async fn abort_restore(&self, table_id: &str) -> Result<(), StorageError> { - sqlx::query("DELETE FROM tables WHERE table_id = $1 AND table_status = 'CREATING'") + /// DELETING for that long. + /// + /// The target is first claimed by moving it from CREATING to DELETING in + /// one statement. That is what makes the removal and the restore's own + /// flip to ACTIVE mutually exclusive: both are conditional on CREATING, so + /// exactly one wins. The claim also schedules the row for the ordinary + /// DELETING removal, so if a drop below fails, the row stays DELETING with + /// its index rows and the control plane finishes the job later. + async fn abort_restore(&self, table_id: &str) -> Result { + let mut tx = self.pool.begin().await.map_err(db_err)?; + let claimed = sqlx::query( + "UPDATE tables SET table_status = 'DELETING', status_transition_at = NOW() \ + WHERE table_id = $1 AND table_status = 'CREATING' AND status_transition_at IS NULL", + ) + .bind(table_id) + .execute(&mut *tx) + .await + .map_err(db_err)? + .rows_affected() + > 0; + let index_ids: Vec = + sqlx::query_scalar("SELECT index_id FROM indexes WHERE table_id = $1") + .bind(table_id) + .fetch_all(&mut *tx) + .await + .map_err(db_err)?; + tx.commit().await.map_err(db_err)?; + if !claimed { + return Ok(false); + } + let mut drops = vec![data_table_name(table_id)]; + drops.extend(index_ids.iter().map(|id| index_table_name(id))); + for t in drops { + sqlx::query(&format!("DROP TABLE IF EXISTS {t}")) + .execute(&self.data_pool) + .await + .map_err(db_err)?; + } + sqlx::query("DELETE FROM tables WHERE table_id = $1 AND table_status = 'DELETING'") .bind(table_id) .execute(&self.pool) .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; - sqlx::query(&format!( - "DROP TABLE IF EXISTS {}", - data_table_name(table_id) - )) - .execute(&self.data_pool) + .map_err(db_err)?; + Ok(true) + } + + /// Remove restore targets whose restore died with its process. + /// + /// A restore target is the only table that sits CREATING with no + /// scheduled transition. One older than the grace period whose restore + /// lock is free has no live owner: the process that created it crashed or + /// was killed mid-copy, and nothing else would ever move it on. Removing + /// it frees the name; the client sees the table disappear, as it would + /// after a failed restore. A target whose lock is held is being copied by + /// a live process, here or on another instance, and is left alone. + /// + /// Returns the names of the tables removed. + pub(crate) async fn sweep_abandoned_restores(&self) -> Result, StorageError> { + let candidates: Vec<(String, String, String)> = sqlx::query_as( + "SELECT table_id, account_id, table_name FROM tables \ + WHERE table_status = 'CREATING' AND status_transition_at IS NULL \ + AND creation_date_time < NOW() - make_interval(secs => $1)", + ) + .bind(f64::from(ABANDONED_RESTORE_GRACE_SECS)) + .fetch_all(&self.pool) .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; - Ok(()) + .map_err(db_err)?; + let mut removed = Vec::new(); + for (table_id, account_id, table_name) in candidates { + let Some(_owner) = try_restore_lock(&self.pool, &account_id, &table_name, None).await? + else { + continue; + }; + if self.abort_restore(&table_id).await? { + tracing::warn!( + "removed table {table_name} ({table_id}): its restore did not finish and \ + no process owns it" + ); + removed.push(table_name); + } + } + Ok(removed) } } @@ -185,6 +583,14 @@ impl BackupEngine for PostgresEngine { let table_name = table_name.to_string(); let backup_name = backup_name.to_string(); Box::pin(async move { + // The table's definition and its items must describe one instant. + // They live in different databases, so no single snapshot covers + // both. Instead the table row is held FOR SHARE from the first + // catalog read until the data snapshot has been taken: UpdateTable + // and DeleteTable both take it FOR UPDATE, so no definition change + // can commit in between, and the definition read here is the one + // in force when the data snapshot starts. + let mut meta = self.pool.begin().await.map_err(db_err)?; // Verify table exists and get metadata. let row: ( String, @@ -199,11 +605,12 @@ impl BackupEngine for PostgresEngine { "SELECT table_id, table_arn, key_schema, attribute_definitions, \ billing_mode, table_size_bytes, item_count, \ COALESCE(provisioned_throughput::text, '{}') \ - FROM tables WHERE account_id = $1 AND table_name = $2 AND table_status = 'ACTIVE'", + FROM tables WHERE account_id = $1 AND table_name = $2 AND table_status = 'ACTIVE' \ + FOR SHARE", ) .bind(&account_id) .bind(&table_name) - .fetch_optional(&self.pool) + .fetch_optional(&mut *meta) .await .map_err(|e| StorageError::Internal(format!("Database error: {e}")))? .ok_or_else(|| StorageError::TableNotFound(format!("Table not found: {table_name}")))?; @@ -212,9 +619,9 @@ impl BackupEngine for PostgresEngine { table_id, _table_arn, key_schema, - attr_defs, + mut attr_defs, billing_mode, - size_bytes, + _size_bytes, _item_count, _prov, ) = row; @@ -225,21 +632,6 @@ impl BackupEngine for PostgresEngine { id = backup_id() ); - // Snapshot items from the data table. `item_data` is the complete - // item, key attributes included, so it is all a restore needs: the - // typed key columns (`sk_s`, `sk_n`, `sk_b`) are derived from it on - // the way back in. Reading them here instead would tie the backup - // to this table's physical column layout. - let ddb_table = data_table_name(&table_id); - let items: Vec<(String, serde_json::Value)> = - sqlx::query_as(&format!("SELECT pk, item_data FROM {ddb_table}")) - .fetch_all(&self.data_pool) - .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; - - #[allow(clippy::cast_possible_wrap)] - let actual_count = items.len() as i64; - // Snapshot the source table's vector index configuration alongside // its key schema. Restore refuses a backup whose snapshot is // non-empty rather than silently dropping a declared index, and it @@ -265,51 +657,170 @@ impl BackupEngine for PostgresEngine { FROM vector_indexes WHERE table_id = $1", ) .bind(&table_id) - .fetch_optional(&self.pool) + .fetch_optional(&mut *meta) .await .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; - // Wrap all catalog-side writes in a single transaction so a crash - // cannot leave a backup marked AVAILABLE with partial items. - let mut tx = self - .pool - .begin() + let (mut definition, gsi_ids) = + Self::capture_table_definition(&mut meta, &table_id).await?; + + // Take the data snapshot while the barrier holds. A REPEATABLE + // READ transaction's snapshot is fixed by its first statement, + // which here also locks the source table ACCESS SHARE, so a + // DeleteTable released by the barrier cannot drop it under the + // read. + let mut snapshot = self + .data_pool + .begin_with("BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY") .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; + .map_err(db_err)?; + sqlx::query(&format!( + "LOCK TABLE {} IN ACCESS SHARE MODE", + data_table_name(&table_id) + )) + .execute(&mut *snapshot) + .await + .map_err(db_err)?; + meta.commit().await.map_err(db_err)?; + + // UpdateTable commits a new index's catalog row before it creates + // and fills the index's data table, and removes the row again if + // that fails. An index whose data table is not there yet is + // therefore not part of the table as of this snapshot; leave it + // out rather than record an index that may never exist. + let gsi_ids_len = gsi_ids.len(); + // The ids were read under the barrier, so this checks exactly the + // indexes the definition describes, against the data snapshot. + let mut present = Vec::with_capacity(definition.global_secondary_indexes.len()); + for (gsi, (name, index_id)) in std::mem::take(&mut definition.global_secondary_indexes) + .into_iter() + .zip(gsi_ids) + { + debug_assert_eq!(gsi.index_name, name); + let exists: bool = sqlx::query_scalar("SELECT to_regclass($1) IS NOT NULL") + .bind(index_table_name(&index_id)) + .fetch_one(&mut *snapshot) + .await + .map_err(db_err)?; + if exists { + present.push(gsi); + } + } + let omitted = present.len() < gsi_ids_len; + definition.global_secondary_indexes = present; + // Leaving an index out can leave an attribute definition no key + // uses, which CreateTable would refuse on the service; drop those. + if omitted { + let mut used: std::collections::HashSet = + serde_json::from_value::>(key_schema.clone()) + .map_err(|e| StorageError::Internal(format!("key schema: {e}")))? + .into_iter() + .map(|k| k.attribute_name) + .collect(); + for k in definition + .global_secondary_indexes + .iter() + .flat_map(|g| &g.key_schema) + .chain( + definition + .local_secondary_indexes + .iter() + .flat_map(|l| &l.key_schema), + ) + { + used.insert(k.attribute_name.clone()); + } + if let Some(defs) = attr_defs.as_array_mut() { + defs.retain(|d| { + d.get("AttributeName") + .and_then(serde_json::Value::as_str) + .is_some_and(|n| used.contains(n)) + }); + } + } + let definition = definition.to_json()?; + + // All catalog-side writes are one transaction, so a crash cannot + // leave a backup AVAILABLE with partial items. The items are read + // from one REPEATABLE READ snapshot of the data table, streamed, + // and written in batches, so the backup is consistent as of one + // instant and memory is bounded by the batch size. + let mut tx = self.pool.begin().await.map_err(db_err)?; sqlx::query( "INSERT INTO backups (backup_arn, backup_name, table_id, table_name, account_id, \ backup_status, backup_size_bytes, item_count, key_schema, attribute_definitions, \ billing_mode, vector_indexes) \ - VALUES ($1, $2, $3, $4, $5, 'AVAILABLE', $6, $7, $8, $9, $10, $11)", + VALUES ($1, $2, $3, $4, $5, 'AVAILABLE', 0, 0, $6, $7, $8, $9)", ) .bind(&backup_arn) .bind(&backup_name) .bind(&table_id) .bind(&table_name) .bind(&account_id) - .bind(size_bytes) - .bind(actual_count) .bind(&key_schema) .bind(&attr_defs) .bind(&billing_mode) .bind(&vector_indexes) .execute(&mut *tx) .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; - - // `backup_items.sk` stays NULL: restore reads keys from `item_data`. - for (pk, item_data) in &items { - sqlx::query( - "INSERT INTO backup_items (backup_arn, pk, item_data) VALUES ($1, $2, $3)", - ) + .map_err(db_err)?; + sqlx::query("INSERT INTO backup_definitions (backup_arn, definition) VALUES ($1, $2)") .bind(&backup_arn) - .bind(pk) - .bind(item_data) + .bind(&definition) .execute(&mut *tx) .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; - } + .map_err(db_err)?; + + // `backup_items.sk` stays NULL: `item_data` is the whole item, key + // attributes included, and restore derives every key column from + // it. Reading the typed columns here instead would tie the backup + // to this table's physical layout. + let (actual_count, size_bytes) = { + use futures::TryStreamExt; + let ddb_table = data_table_name(&table_id); + let mut count: i64 = 0; + let mut size: i64 = 0; + let mut pks: Vec = Vec::with_capacity(COPY_BATCH_ITEMS); + let mut datas: Vec = Vec::with_capacity(COPY_BATCH_ITEMS); + let mut batch_bytes: usize = 0; + // Rows come back as JSON text and are buffered as text, so the + // batch budget counts what is actually held. Each item is + // parsed only long enough to size it. + let select_sql = format!("SELECT pk, item_data::text FROM {ddb_table}"); + { + let mut rows = + sqlx::query_as::<_, (String, String)>(&select_sql).fetch(&mut *snapshot); + while let Some((pk, data)) = rows.try_next().await.map_err(db_err)? { + let item: Item = serde_json::from_str(&data) + .map_err(|e| StorageError::Internal(format!("Parse item: {e}")))?; + let item_bytes = extenddb_core::types::item_size_bytes(&item); + drop(item); + size += i64::try_from(item_bytes).unwrap_or(i64::MAX); + batch_bytes += data.len(); + pks.push(pk); + datas.push(data); + if pks.len() == COPY_BATCH_ITEMS || batch_bytes >= COPY_BATCH_BYTES { + batch_bytes = 0; + count += + insert_backup_batch(&mut tx, &backup_arn, &mut pks, &mut datas) + .await?; + } + } + } + count += insert_backup_batch(&mut tx, &backup_arn, &mut pks, &mut datas).await?; + snapshot.commit().await.map_err(db_err)?; + (count, size) + }; + sqlx::query( + "UPDATE backups SET item_count = $1, backup_size_bytes = $2 WHERE backup_arn = $3", + ) + .bind(actual_count) + .bind(size_bytes) + .bind(&backup_arn) + .execute(&mut *tx) + .await + .map_err(db_err)?; // Read back the creation timestamp assigned by the database. let created_at: time::OffsetDateTime = @@ -499,16 +1010,29 @@ impl BackupEngine for PostgresEngine { // reported missing here and the writes below never run. let desc = self.describe_backup(&account_id, &backup_arn).await?; - // The account predicate is repeated on both writes rather than - // relying on the lookup above, so the statements are correct on - // their own terms. + // One transaction, so a restore reading its snapshot sees either + // the whole backup or a DELETED one, never AVAILABLE with its rows + // gone. The account predicate is repeated on every write rather + // than relying on the lookup above, so the statements are correct + // on their own terms. + let mut tx = self.pool.begin().await.map_err(db_err)?; sqlx::query( "DELETE FROM backup_items WHERE backup_arn = $1 AND EXISTS (\ SELECT 1 FROM backups b WHERE b.backup_arn = $1 AND b.account_id = $2)", ) .bind(&backup_arn) .bind(&account_id) - .execute(&self.pool) + .execute(&mut *tx) + .await + .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; + + sqlx::query( + "DELETE FROM backup_definitions WHERE backup_arn = $1 AND EXISTS (\ + SELECT 1 FROM backups b WHERE b.backup_arn = $1 AND b.account_id = $2)", + ) + .bind(&backup_arn) + .bind(&account_id) + .execute(&mut *tx) .await .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; @@ -518,10 +1042,12 @@ impl BackupEngine for PostgresEngine { ) .bind(&backup_arn) .bind(&account_id) - .execute(&self.pool) + .execute(&mut *tx) .await .map_err(|e| StorageError::Internal(format!("Database error: {e}")))?; + tx.commit().await.map_err(db_err)?; + Ok(BackupDescription { backup_details: BackupDetails { backup_status: "DELETED".to_owned(), @@ -542,37 +1068,36 @@ impl BackupEngine for PostgresEngine { let target_table_name = target_table_name.to_string(); let backup_arn = backup_arn.to_string(); Box::pin(async move { + #[allow(clippy::type_complexity)] let backup_row: ( - String, serde_json::Value, serde_json::Value, String, Option, Option, ) = sqlx::query_as( - "SELECT table_name, key_schema, attribute_definitions, billing_mode, \ - provisioned_throughput, vector_indexes \ - FROM backups \ - WHERE backup_arn = $1 AND account_id = $2 AND backup_status = 'AVAILABLE'", + "SELECT b.key_schema, b.attribute_definitions, b.billing_mode, \ + b.vector_indexes, d.definition \ + FROM backups b LEFT JOIN backup_definitions d ON d.backup_arn = b.backup_arn \ + WHERE b.backup_arn = $1 AND b.account_id = $2 AND b.backup_status = 'AVAILABLE'", ) .bind(&backup_arn) .bind(&account_id) .fetch_optional(&self.pool) .await - .map_err(|e| StorageError::Internal(format!("Database error: {e}")))? + .map_err(db_err)? .ok_or_else(|| StorageError::Validation(format!("Backup not found: {backup_arn}")))?; - let (_orig_table, ks_json, ad_json, billing, _prov, vector_indexes) = backup_row; + let (ks_json, ad_json, billing, vector_indexes, definition) = backup_row; - // Restore rebuilds the target from the snapshot's key schema, - // attribute definitions and billing mode, and does not carry indexes - // across. For a source table that had vector indexes that would be - // silent loss of a declared index: the restored table would answer - // every request except a search, and the client would only find out - // on the first one. The service preserves vector indexes through + // Vector indexes are not restored on this backend. For a source + // table that had them, a restore without them would be silent loss + // of a declared index: the restored table would answer every + // request except a search, and the client would only find out on + // the first one. The service preserves vector indexes through // backup and restore (measured 2026-08-19), so the conformant end - // state is to restore them; until then this refuses, because a typed - // refusal is recoverable and silent loss is not. + // state is to restore them; until then this refuses, because a + // typed refusal is recoverable and silent loss is not. // // A NULL snapshot is a backup taken before the column existed, which // cannot have carried vector indexes: this backend could not create @@ -584,14 +1109,6 @@ impl BackupEngine for PostgresEngine { Some(snapshot) => { let version = snapshot.get("Version").and_then(serde_json::Value::as_u64); if version != Some(1) { - // Version skew, not a fault: a newer build wrote a shape - // this one does not know. Reported the same way as the - // refusal below, so a client gets a reason rather than a - // 500 and the operator gets no error-level noise. Nearly - // unreachable, because a shape change would normally come - // with a catalog version bump that the startup gate - // refuses first; the gap is a future build adding a member - // without any schema change. return Err(StorageError::Unsupported(format!( "backup {backup_arn} carries a vector index snapshot this build \ cannot read (version {version:?})" @@ -610,65 +1127,112 @@ impl BackupEngine for PostgresEngine { storage backend" ))); } + let definition = definition + .map(|d| BackupTableDefinition::from_json(d, &backup_arn)) + .transpose()?; + if let Some(d) = &definition { + d.ensure_restorable(&backup_arn)?; + } - let key_schema: Vec = - serde_json::from_value(ks_json) - .map_err(|e| StorageError::Internal(format!("Parse key schema: {e}")))?; + let key_schema: Vec = serde_json::from_value(ks_json) + .map_err(|e| StorageError::Internal(format!("Parse key schema: {e}")))?; + ensure_single_part_base_key(&key_schema, &backup_arn)?; let attr_defs: Vec = serde_json::from_value(ad_json) .map_err(|e| StorageError::Internal(format!("Parse attr defs: {e}")))?; - let billing_mode = if billing == "PAY_PER_REQUEST" { - Some(extenddb_core::types::BillingMode::PayPerRequest) - } else { - Some(extenddb_core::types::BillingMode::Provisioned) - }; - - let create_input = extenddb_core::types::CreateTableInput { + let mut create_input = extenddb_core::types::CreateTableInput { table_name: target_table_name.clone(), key_schema, attribute_definitions: attr_defs, - billing_mode, - provisioned_throughput: Some(extenddb_core::types::ProvisionedThroughput { - read_capacity_units: 5, - write_capacity_units: 5, - }), - global_secondary_indexes: None, - local_secondary_indexes: None, - stream_specification: None, - tags: None, - deletion_protection_enabled: None, - sse_specification: None, - table_class: None, - on_demand_throughput: None, ..Default::default() }; + match definition { + Some(d) => d.apply_to(&mut create_input), + None => { + // A backup taken before table definitions were recorded + // carries only its keys and billing mode, so it restores as + // it always has: no secondary indexes, and 5/5 throughput for + // a provisioned table since the source's is unknown. + let on_demand = billing == "PAY_PER_REQUEST"; + create_input.billing_mode = Some(if on_demand { + extenddb_core::types::BillingMode::PayPerRequest + } else { + extenddb_core::types::BillingMode::Provisioned + }); + create_input.provisioned_throughput = + (!on_demand).then_some(extenddb_core::types::ProvisionedThroughput { + read_capacity_units: 5, + write_capacity_units: 5, + }); + } + } + + // Ownership first, before the target exists, so the + // abandoned-restore sweep never sees this target without an owner, + // including while the DDL below runs. It is held until the copy + // ends; a crash anywhere releases it with the connection, and the + // sweep then removes the target. A lock already held means another + // restore into this name is in flight right now; let it report the + // name as taken, as CreateTable would. + let slot = match tokio::time::timeout( + std::time::Duration::from_secs(RESTORE_SLOT_WAIT_SECS), + RESTORE_SLOTS.acquire(), + ) + .await + { + Ok(slot) => { + slot.map_err(|e| StorageError::Internal(format!("restore slot: {e}")))? + } + Err(_) => { + return Err(StorageError::LimitExceeded( + "Too many restores are in progress on this server; retry later".to_owned(), + )); + } + }; + let Some(_owner) = + try_restore_lock(&self.pool, &account_id, &target_table_name, Some(slot)).await? + else { + return Err(StorageError::TableAlreadyExists(target_table_name.clone())); + }; // 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. + // data is fully present. Its secondary indexes are created with it, + // empty, and filled by the copy. let desc = self .create_table_impl(&account_id, create_input, true) .await?; // From here on a failure must not leave the target behind: it is - // CREATING with no scheduled transition, so nothing else would ever - // move it on, and the name would stay taken. Remove it and report - // the failure. - if let Err(e) = self.copy_backup_items(&desc, &backup_arn).await { - tracing::error!( - "restore of {backup_arn} into {target_table_name} failed, \ - removing the partial table: {e}" - ); - if let Err(cleanup) = self.abort_restore(&desc.table_id).await { - tracing::error!( - "could not remove partially restored table {target_table_name} \ - ({}): {cleanup}", + // CREATING with no scheduled transition, so nothing else would + // move it on, and the name would stay taken. + let copied = self.copy_backup_items(&desc, &backup_arn).await; + if let Err(e) = copied { + match self.abort_restore(&desc.table_id).await { + // Not ours to remove: it was already claimed for removal + // while it was being copied into, and that is why the copy + // failed. Report it as such rather than as a server fault a + // client would retry. + Ok(false) => { + return Err(StorageError::TableNotFound(format!( + "restore of {backup_arn} into {target_table_name} did not \ + complete: the table was deleted while it was being restored" + ))); + } + Ok(true) => tracing::error!( + "restore of {backup_arn} into {target_table_name} failed and the \ + partial table was removed: {e}" + ), + Err(cleanup) => tracing::error!( + "restore of {backup_arn} into {target_table_name} failed ({e}), and \ + the partial table ({}) could not be removed yet; the control plane \ + retries: {cleanup}", desc.table_id - ); + ), } return Err(e); } diff --git a/crates/storage-postgres/src/delete_table.rs b/crates/storage-postgres/src/delete_table.rs index 5a55231fe..2f56c318a 100755 --- a/crates/storage-postgres/src/delete_table.rs +++ b/crates/storage-postgres/src/delete_table.rs @@ -47,6 +47,29 @@ impl PostgresEngine { return Err(StorageError::DeletionProtected(row.table_arn.clone())); } + // A restore target (CREATING with no scheduled transition) is being + // filled by RestoreTableFromBackup. The service refuses to delete a + // table it is still creating, and deleting this one would make the + // restore fail part-way; refuse until the restore finishes. An + // abandoned target is removed by the restore sweep, not here. + if row.table_status == "CREATING" { + let restoring: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM tables WHERE table_id = $1 \ + AND table_status = 'CREATING' AND status_transition_at IS NULL)", + ) + .bind(&row.table_id) + .fetch_one(&mut *tx) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + if restoring { + return Err(StorageError::IndexesInUse(format!( + "Attempt to change a resource which is still in use: Table is being \ + restored: {}", + row.table_name + ))); + } + } + // Fetch indexes for the response description. let index_rows: Vec = sqlx::query_as( r"SELECT index_name, index_id, index_type, key_schema, projection, diff --git a/crates/storage-postgres/src/lib.rs b/crates/storage-postgres/src/lib.rs index 0d7a51c0d..63646cc1f 100755 --- a/crates/storage-postgres/src/lib.rs +++ b/crates/storage-postgres/src/lib.rs @@ -135,7 +135,7 @@ use sqlx::postgres::PgPoolOptions; /// /// The tuple is the single source of truth. Use `CATALOG_VERSION.to_string()` /// wherever a string representation is needed. -pub const CATALOG_VERSION: CatalogVersion = CatalogVersion::new(0, 0, 3); +pub const CATALOG_VERSION: CatalogVersion = CatalogVersion::new(0, 0, 4); /// Minimum number of connections allowed per pool. /// diff --git a/crates/storage-postgres/src/migrations.rs b/crates/storage-postgres/src/migrations.rs index 45af16167..17f8c9c92 100755 --- a/crates/storage-postgres/src/migrations.rs +++ b/crates/storage-postgres/src/migrations.rs @@ -16,8 +16,18 @@ pub(crate) const CATALOG_MIGRATIONS: &[(&str, &str)] = &[ "002_vector_indexes.sql", include_str!("../../storage-postgres/migrations/002_vector_indexes.sql"), ), + ( + "003_backup_definitions.sql", + include_str!("../../storage-postgres/migrations/003_backup_definitions.sql"), + ), ]; +/// The runner's convergence write, executed after the walk over +/// [`CATALOG_MIGRATIONS`]. Same statement shape as the files' own in-file +/// writes, parameterized on the compiled version. +pub(crate) const SET_CATALOG_VERSION_SQL: &str = + "UPDATE settings SET value = $1 WHERE key = 'catalog_version'"; + /// Run catalog migrations, skipping already-applied ones. pub(crate) async fn run_catalog_migrations(pool: &PgPool) -> OpResult<()> { println!("--- Running catalog migrations..."); @@ -41,6 +51,20 @@ pub(crate) async fn run_catalog_migrations(pool: &PgPool) -> OpResult<()> { // another migration lands. record_migration(pool, filename).await?; } + // The runner owns the final version write. Each file still writes the + // version it introduces, but a replay can re-apply an EARLIER file while a + // later one stays recorded and skipped: the re-applied file's in-file write + // then leaves the stored version behind the schema actually present, the + // startup gate refuses the catalog, and a second migrate finds nothing + // unrecorded and writes nothing, stranding the deployment. Converging on + // the compiled version after every completed walk closes that gap. The + // in-file writes stay until #221 moves version ownership into the runner + // entirely. + sqlx::query(SET_CATALOG_VERSION_SQL) + .bind(crate::CATALOG_VERSION.to_string()) + .execute(pool) + .await + .map_err(|e| OpError::Internal(format!("Write catalog_version: {e}")))?; println!(" Migrations applied."); Ok(()) } @@ -311,10 +335,10 @@ mod tests { fn the_migration_count_and_the_catalog_version_agree() { assert_eq!( CATALOG_MIGRATIONS.len(), - 2, + 3, "a catalog migration was added or removed; update CATALOG_VERSION and this count" ); - assert_eq!(CATALOG_VERSION.to_string(), "0.0.3"); + assert_eq!(CATALOG_VERSION.to_string(), "0.0.4"); } /// The version the binary expects must be the version the schema writes. diff --git a/crates/storage-postgres/src/worker_store.rs b/crates/storage-postgres/src/worker_store.rs index 372fbf49d..235e77ce2 100755 --- a/crates/storage-postgres/src/worker_store.rs +++ b/crates/storage-postgres/src/worker_store.rs @@ -165,6 +165,19 @@ impl PostgresEngine { .map_err(|e| StorageError::Internal(e.to_string()))?; } + // CREATING with no owner → removed. A restore whose process died + // mid-copy leaves its target CREATING with no scheduled transition, + // which nothing above would ever move on. Logged and skipped on + // failure, so it cannot hold up the transitions above. + match self.sweep_abandoned_restores().await { + Ok(names) => { + for name in names { + transitions.push((name, "abandoned restore → removed")); + } + } + Err(e) => tracing::warn!("abandoned restore sweep failed: {e}"), + } + Ok(transitions) } } diff --git a/crates/storage-postgres/tests/backup_restore.rs b/crates/storage-postgres/tests/backup_restore.rs index dd0399ea6..0ae45af1f 100644 --- a/crates/storage-postgres/tests/backup_restore.rs +++ b/crates/storage-postgres/tests/backup_restore.rs @@ -12,13 +12,14 @@ //! behind in CREATING, where nothing would ever move it on and its name //! would stay taken; //! - a backup with two rows for one key must fail rather than collapse them; -//! - a multi-part key backup must restore every key part (multi-part base -//! keys are a preview gated by `allow_multipart_table_keys`, not reachable -//! with the default configuration). +//! - a backup of a multi-part key table (a preview gated by +//! `enable_multipart_keys`) must be refused, not restored wrongly; +//! - the abandoned-restore sweep, which needs a crash to reach. //! //! Each test builds a throwaway catalog and data database and drops them when -//! it passes; a failing test leaves `eddb_bkup_*` behind for inspection. -//! Requires `EXTENDDB_TEST_PG_CONNECTION_STRING` (a base URL with no database +//! it ends, passing or failing. The tests run one at a time: each opens its +//! own pools, and running them all at once can exhaust the server's +//! connection limit. Requires `EXTENDDB_TEST_PG_CONNECTION_STRING` (a base URL with no database //! component); without it every test reports a skip and passes. use extenddb_core::types::{ @@ -40,32 +41,72 @@ const REGION: &str = "us-east-1"; /// shadow the data schema's table if both were applied to one database. struct Scratch { engine: PostgresEngine, - /// The data database (`stream_shards`, `stream_records`). + /// The data database. db: PgPool, catalog: PgPool, admin: PgPool, db_names: [String; 2], + _guard: DbGuard, + _serial: tokio::sync::MutexGuard<'static, ()>, } +/// Drops the scratch databases if the test panics, from the moment they are +/// created, including during setup. +struct DbGuard { + names: Vec, +} + +/// One test at a time; see the module comment. +static SERIAL: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + impl Scratch { + /// Close the pools and drop the scratch databases. Also run, on a fresh + /// runtime, if a test panics before calling it. async fn cleanup(self) { - let Scratch { - engine, - db, - catalog, - admin, - db_names, - } = self; - drop(engine); - db.close().await; - catalog.close().await; - for name in db_names { - sqlx::query(&format!("DROP DATABASE IF EXISTS \"{name}\" WITH (FORCE)")) - .execute(&admin) - .await - .expect("drop a scratch database"); + self.db.close().await; + self.catalog.close().await; + drop_databases(&self.admin, &self.db_names).await; + self.admin.close().await; + } +} + +async fn drop_databases(admin: &PgPool, names: &[String]) { + for name in names { + let _ = sqlx::query(&format!("DROP DATABASE IF EXISTS \"{name}\" WITH (FORCE)")) + .execute(admin) + .await; + } +} + +impl Drop for DbGuard { + fn drop(&mut self) { + if !std::thread::panicking() { + return; } - admin.close().await; + // A failed test: drop its databases on a separate thread and runtime, + // since the test's own runtime is unwinding. + let (Some(base), names) = (base_conn(), self.names.clone()) else { + return; + }; + let _ = std::thread::spawn(move || { + let Ok(rt) = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + else { + return; + }; + rt.block_on(async { + if let Ok(admin) = PgPoolOptions::new() + .max_connections(1) + .connect(&format!("{base}/postgres")) + .await + { + drop_databases(&admin, &names).await; + admin.close().await; + } + }); + }) + .join(); } } @@ -92,6 +133,7 @@ async fn apply(pool: &PgPool, migrations: &[&str]) { } async fn scratch() -> Scratch { + let serial = SERIAL.lock().await; let base = base_conn().expect("caller checks base_conn() first"); let stem = format!("eddb_bkup_{}", uuid::Uuid::new_v4().simple())[..24].to_owned(); let (catalog_name, data_name) = (format!("{stem}_c"), format!("{stem}_d")); @@ -100,6 +142,9 @@ async fn scratch() -> Scratch { .connect(&format!("{base}/postgres")) .await .expect("connect to the postgres maintenance database"); + let guard = DbGuard { + names: vec![catalog_name.clone(), data_name.clone()], + }; for name in [&catalog_name, &data_name] { sqlx::query(&format!("CREATE DATABASE \"{name}\"")) .execute(&admin) @@ -115,6 +160,7 @@ async fn scratch() -> Scratch { &[ include_str!("../migrations/001_schema.sql"), include_str!("../migrations/002_vector_indexes.sql"), + include_str!("../migrations/003_backup_definitions.sql"), ], ) .await; @@ -132,6 +178,10 @@ async fn scratch() -> Scratch { .execute(&catalog) .await .expect("pin the control-plane delay to zero"); + sqlx::query("UPDATE settings SET value = '0' WHERE key = 'index_propagation_delay_ms'") + .execute(&catalog) + .await + .expect("make index propagation synchronous, as no queue worker runs here"); sqlx::query( "INSERT INTO settings (key, value) VALUES ('data_database_connection_string', $1) \ ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value", @@ -149,7 +199,7 @@ async fn scratch() -> Scratch { let engine = PostgresEngine::new( &PostgresConfig { connection_string: catalog_url, - pool_size: 10, + pool_size: 4, max_item_size_bytes: 400_000, }, REGION, @@ -162,6 +212,8 @@ async fn scratch() -> Scratch { catalog, admin, db_names: [catalog_name, data_name], + _guard: guard, + _serial: serial, } } @@ -373,9 +425,9 @@ async fn duplicate_backup_rows_fail_the_restore() { } #[tokio::test] -async fn multipart_key_backup_restores_every_key_part() { +async fn multipart_key_backup_is_refused() { if base_conn().is_none() { - eprintln!("SKIP multipart_key_backup_restores_every_key_part: no PostgreSQL"); + eprintln!("SKIP multipart_key_backup_is_refused: no PostgreSQL"); return; } let s = scratch().await; @@ -393,57 +445,710 @@ async fn multipart_key_backup_restores_every_key_part() { key("h1", KeyType::Hash), key("h2", KeyType::Hash), key("r1", KeyType::Range), - key("r2", KeyType::Range), ], attribute_definitions: vec![ attr("h1", ScalarAttributeType::S), attr("h2", ScalarAttributeType::N), attr("r1", ScalarAttributeType::S), - attr("r2", ScalarAttributeType::B), ], billing_mode: Some(BillingMode::PayPerRequest), ..Default::default() }; - // Rows differ only in the second HASH or the second RANGE attribute, so a - // restore that keyed on the first of each would collapse them. - let rows: Vec = [("1", "AA=="), ("2", "AA=="), ("1", "AQ==")] - .iter() - .map(|(h2, r2)| { - serde_json::json!({ - "h1": {"S": "same"}, "h2": {"N": h2}, - "r1": {"S": "same"}, "r2": {"B": r2}, - }) - }) - .collect(); + let rows = vec![serde_json::json!({"h1": {"S": "a"}, "h2": {"N": "1"}, "r1": {"S": "x"}})]; let arn = legacy_backup_of(&s, input, &rows).await; - s.engine + // Multi-part base keys are a preview the item paths address only by their + // first parts, so no restore of one can be right; it must refuse, before + // creating anything. + let err = s + .engine .restore_table_from_backup(ACCOUNT, "multi_dst", &arn) .await - .expect("restore a multi-part key backup"); + .expect_err("a multi-part key backup is refused"); + assert!(matches!(err, StorageError::Unsupported(_)), "{err:?}"); + assert_eq!(table_status(&s, "multi_dst").await, None); + s.cleanup().await; +} + +/// A provisioned (pk S, sk N) table with one GSI and one LSI, holding 30 +/// items, some of them outside each index. +async fn indexed_source(s: &Scratch, name: &str) -> String { + let input: CreateTableInput = serde_json::from_value(serde_json::json!({ + "TableName": name, + "KeySchema": [ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"} + ], + "AttributeDefinitions": [ + {"AttributeName": "pk", "AttributeType": "S"}, + {"AttributeName": "sk", "AttributeType": "N"}, + {"AttributeName": "g", "AttributeType": "S"}, + {"AttributeName": "l", "AttributeType": "S"} + ], + "BillingMode": "PROVISIONED", + "ProvisionedThroughput": {"ReadCapacityUnits": 7, "WriteCapacityUnits": 9}, + "GlobalSecondaryIndexes": [{ + "IndexName": "gi", + "KeySchema": [{"AttributeName": "g", "KeyType": "HASH"}], + "Projection": {"ProjectionType": "KEYS_ONLY"}, + "ProvisionedThroughput": {"ReadCapacityUnits": 3, "WriteCapacityUnits": 4} + }], + "LocalSecondaryIndexes": [{ + "IndexName": "li", + "KeySchema": [ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "l", "KeyType": "RANGE"} + ], + "Projection": {"ProjectionType": "ALL"} + }] + })) + .expect("input"); + let desc = s + .engine + .create_table(ACCOUNT, input) + .await + .expect("create the source table"); + let key_info = TableKeyInfo { + table_name: name.to_owned(), + account_id: ACCOUNT.to_owned(), + table_id: desc.table_id.clone(), + key_schema: desc.key_schema.clone(), + base_key_schema: desc.key_schema.clone(), + attribute_definitions: desc.attribute_definitions.clone(), + ..Default::default() + }; + for i in 0..30 { + let mut item = + serde_json::json!({"pk": {"S": format!("p{}", i % 3)}, "sk": {"N": i.to_string()}}); + if i % 2 == 0 { + item["g"] = serde_json::json!({"S": format!("g{}", i % 4)}); + } + if i % 3 == 0 { + item["l"] = serde_json::json!({"S": format!("l{i}")}); + } + let item: Item = serde_json::from_value(item).expect("item"); + s.engine + .put_item( + &key_info, + item, + false, + None, + &extenddb_core::expression::ExpressionMaps::default(), + None, + ) + .await + .expect("put"); + } + desc.table_id +} + +/// `(index_name, index_id)` of a table's secondary indexes, by name. +async fn index_ids(s: &Scratch, table_id: &str) -> Vec<(String, String)> { + sqlx::query_as( + "SELECT index_name, index_id FROM indexes WHERE table_id = $1 ORDER BY index_name", + ) + .bind(table_id) + .fetch_all(&s.catalog) + .await + .expect("read indexes") +} + +async fn data_rows(s: &Scratch, physical_id: &str) -> i64 { + sqlx::query_scalar(&format!("SELECT count(*) FROM \"_ddb_{physical_id}\"")) + .fetch_one(&s.db) + .await + .expect("count rows") +} + +async fn data_table_exists(s: &Scratch, physical_id: &str) -> bool { + sqlx::query_scalar("SELECT to_regclass($1) IS NOT NULL") + .bind(format!("public.\"_ddb_{physical_id}\"")) + .fetch_one(&s.db) + .await + .expect("regclass") +} + +async fn table_id_of(s: &Scratch, name: &str) -> Option { + sqlx::query_scalar("SELECT table_id FROM tables WHERE account_id = $1 AND table_name = $2") + .bind(ACCOUNT) + .bind(name) + .fetch_optional(&s.catalog) + .await + .expect("table id") +} + +#[tokio::test] +async fn restore_recreates_indexes_and_throughput() { + if base_conn().is_none() { + eprintln!("SKIP restore_recreates_indexes_and_throughput: no PostgreSQL"); + return; + } + let s = scratch().await; + let src = indexed_source(&s, "idx_src").await; + let backup = s + .engine + .create_backup(ACCOUNT, "idx_src", "b") + .await + .expect("backup"); let desc = s + .engine + .restore_table_from_backup(ACCOUNT, "idx_dst", &backup.backup_arn) + .await + .expect("restore"); + assert_eq!(table_status(&s, "idx_dst").await.as_deref(), Some("ACTIVE")); + + let restored = s .engine .describe_table( ACCOUNT, DescribeTableInput { - table_name: "multi_dst".to_owned(), + table_name: "idx_dst".to_owned(), }, ) .await - .expect("describe the restored table"); - // Checked on the physical rows: on this backend the PutItem and GetItem - // paths address only the first HASH and RANGE attribute of a base table - // (multi-part base keys are off by default), so a point read cannot be - // the oracle here. Each row must carry both HASH parts in `pk` and both - // RANGE parts in their typed columns, with nothing left NULL. - let (rows_n, pks, r1s, r2s, nulls): (i64, i64, i64, i64, i64) = sqlx::query_as(&format!( - "SELECT count(*), count(DISTINCT pk), count(DISTINCT sk_s), count(DISTINCT sk2_b), \ - count(*) FILTER (WHERE sk_s IS NULL OR sk2_b IS NULL) FROM \"_ddb_{}\"", - desc.table_id - )) - .fetch_one(&s.db) + .expect("describe"); + let pt = restored.provisioned_throughput; + assert_eq!((pt.read_capacity_units, pt.write_capacity_units), (7, 9)); + let gsis = restored.global_secondary_indexes.expect("gsis"); + assert_eq!(gsis.len(), 1); + let gpt = gsis[0] + .provisioned_throughput + .as_ref() + .expect("gsi throughput"); + assert_eq!((gpt.read_capacity_units, gpt.write_capacity_units), (3, 4)); + assert_eq!(restored.local_secondary_indexes.map(|l| l.len()), Some(1)); + + // Each index holds exactly the rows the source's does. + let src_idx = index_ids(&s, &src).await; + let dst_idx = index_ids(&s, &desc.table_id).await; + assert_eq!( + dst_idx.iter().map(|(n, _)| n.as_str()).collect::>(), + ["gi", "li"] + ); + for ((_, a), (_, b)) in src_idx.iter().zip(&dst_idx) { + let want = data_rows(&s, a).await; + assert!(want > 0); + assert_eq!(data_rows(&s, b).await, want); + } + assert_eq!(data_rows(&s, &desc.table_id).await, 30); + s.cleanup().await; +} + +#[tokio::test] +async fn abandoned_restore_is_removed_only_when_unowned_and_old() { + if base_conn().is_none() { + eprintln!("SKIP abandoned_restore_is_removed_only_when_unowned_and_old: no PostgreSQL"); + return; + } + let s = scratch().await; + indexed_source(&s, "ab_src").await; + let backup = s + .engine + .create_backup(ACCOUNT, "ab_src", "b") + .await + .expect("backup"); + let desc = s + .engine + .restore_table_from_backup(ACCOUNT, "ab_dst", &backup.backup_arn) + .await + .expect("restore"); + let indexes = index_ids(&s, &desc.table_id).await; + + // The state a process killed mid-copy leaves: the target CREATING with no + // scheduled transition, its copy transaction rolled back. + let crash = |age: &'static str| { + let catalog = s.catalog.clone(); + let id = desc.table_id.clone(); + async move { + sqlx::query(&format!( + "UPDATE tables SET table_status = 'CREATING', status_transition_at = NULL, \ + creation_date_time = NOW() - INTERVAL '{age}' WHERE table_id = $1" + )) + .bind(&id) + .execute(&catalog) + .await + .expect("simulate a crash"); + } + }; + + // Inside the grace period: a restore that has not taken its lock yet. + crash("1 second").await; + s.engine + .process_control_plane_transitions() + .await + .expect("sweep"); + assert_eq!( + table_status(&s, "ab_dst").await.as_deref(), + Some("CREATING") + ); + + // Old, but its lock is held: a live restore on another instance. The key + // is the account and name, so the lock covers the target before it exists. + crash("10 minutes").await; + let mut owner = s.catalog.acquire().await.expect("connection"); + sqlx::query("SELECT pg_advisory_lock(hashtextextended($2, $1))") + .bind(0x0045_4452_i64) + .bind(format!("{ACCOUNT}/ab_dst")) + .execute(&mut *owner) + .await + .expect("hold the restore lock"); + s.engine + .process_control_plane_transitions() + .await + .expect("sweep"); + assert_eq!( + table_status(&s, "ab_dst").await.as_deref(), + Some("CREATING") + ); + sqlx::query("SELECT pg_advisory_unlock(hashtextextended($2, $1))") + .bind(0x0045_4452_i64) + .bind(format!("{ACCOUNT}/ab_dst")) + .execute(&mut *owner) + .await + .expect("release the restore lock"); + drop(owner); + + // Old and unowned: removed, with its data and index tables. + let transitions = s + .engine + .process_control_plane_transitions() + .await + .expect("sweep"); + assert!( + transitions.iter().any(|(n, _)| n == "ab_dst"), + "{transitions:?}" + ); + assert_eq!(table_status(&s, "ab_dst").await, None); + assert!(!data_table_exists(&s, &desc.table_id).await); + for (_, id) in &indexes { + assert!( + !data_table_exists(&s, id).await, + "index table {id} left behind" + ); + } + + // An ACTIVE table is never a candidate, and the name is free again. + s.engine + .restore_table_from_backup(ACCOUNT, "ab_dst", &backup.backup_arn) + .await + .expect("restore again"); + sqlx::query("UPDATE tables SET creation_date_time = NOW() - INTERVAL '10 minutes'") + .execute(&s.catalog) + .await + .expect("age every table"); + s.engine + .process_control_plane_transitions() + .await + .expect("sweep"); + assert_eq!(table_status(&s, "ab_dst").await.as_deref(), Some("ACTIVE")); + s.cleanup().await; +} + +#[tokio::test] +async fn failed_restore_with_indexes_leaves_no_tables() { + if base_conn().is_none() { + eprintln!("SKIP failed_restore_with_indexes_leaves_no_tables: no PostgreSQL"); + return; + } + let s = scratch().await; + indexed_source(&s, "fi_src").await; + let backup = s + .engine + .create_backup(ACCOUNT, "fi_src", "b") + .await + .expect("backup"); + sqlx::query("INSERT INTO backup_items (backup_arn, pk, item_data) VALUES ($1, 'x', $2)") + .bind(&backup.backup_arn) + .bind(serde_json::json!({"pk": {"S": "x"}})) + .execute(&s.catalog) + .await + .expect("add a row without its sort key"); + let before: i64 = + sqlx::query_scalar("SELECT count(*) FROM pg_tables WHERE schemaname = 'public'") + .fetch_one(&s.db) + .await + .expect("tables"); + s.engine + .restore_table_from_backup(ACCOUNT, "fi_dst", &backup.backup_arn) + .await + .expect_err("a row without its sort key cannot restore"); + assert_eq!(table_id_of(&s, "fi_dst").await, None); + let after: i64 = + sqlx::query_scalar("SELECT count(*) FROM pg_tables WHERE schemaname = 'public'") + .fetch_one(&s.db) + .await + .expect("tables"); + assert_eq!( + after, before, + "the target's data and index tables are dropped" + ); + s.cleanup().await; +} + +#[tokio::test] +async fn restore_keeps_table_class_sse_and_on_demand_limits() { + if base_conn().is_none() { + eprintln!("SKIP restore_keeps_table_class_sse_and_on_demand_limits: no PostgreSQL"); + return; + } + let s = scratch().await; + let input: CreateTableInput = serde_json::from_value(serde_json::json!({ + "TableName": "cls_src", + "KeySchema": [{"AttributeName": "pk", "KeyType": "HASH"}], + "AttributeDefinitions": [{"AttributeName": "pk", "AttributeType": "S"}], + "BillingMode": "PAY_PER_REQUEST", + "TableClass": "STANDARD_INFREQUENT_ACCESS", + "SSESpecification": {"Enabled": true, "SSEType": "KMS"}, + "OnDemandThroughput": {"MaxReadRequestUnits": 100, "MaxWriteRequestUnits": 50} + })) + .expect("input"); + let src = s.engine.create_table(ACCOUNT, input).await.expect("create"); + let backup = s + .engine + .create_backup(ACCOUNT, "cls_src", "b") + .await + .expect("backup"); + s.engine + .restore_table_from_backup(ACCOUNT, "cls_dst", &backup.backup_arn) + .await + .expect("restore"); + let read = |name: &'static str| { + let catalog = s.catalog.clone(); + async move { + sqlx::query_as::< + _, + ( + String, + Option, + Option, + Option, + ), + >( + "SELECT billing_mode, table_class, sse_specification, on_demand_throughput \ + FROM tables WHERE account_id = $1 AND table_name = $2", + ) + .bind(ACCOUNT) + .bind(name) + .fetch_one(&catalog) + .await + .expect("row") + } + }; + let want = read("cls_src").await; + assert_eq!(want.1.as_deref(), Some("STANDARD_INFREQUENT_ACCESS")); + assert!(want.2.is_some() && want.3.is_some(), "{want:?}"); + assert_eq!(read("cls_dst").await, want); + drop(src); + s.cleanup().await; +} + +#[tokio::test] +async fn restore_crosses_batch_boundaries() { + if base_conn().is_none() { + eprintln!("SKIP restore_crosses_batch_boundaries: no PostgreSQL"); + return; + } + let s = scratch().await; + let src = indexed_source(&s, "many_src").await; + // More rows than one item batch, and a few large enough that the byte + // budget, not the row count, ends some batches. + let desc_rows: Vec = (0..1_234) + .map(|i| { + let mut v = serde_json::json!({ + "pk": {"S": format!("bulk{}", i % 11)}, + "sk": {"N": format!("{}", 1000 + i)}, + "g": {"S": format!("g{}", i % 4)}, + }); + if i % 97 == 0 { + v["pad"] = serde_json::json!({"S": "x".repeat(300_000)}); + } + v + }) + .collect(); + let backup = s + .engine + .create_backup(ACCOUNT, "many_src", "b") + .await + .expect("backup"); + sqlx::query( + "INSERT INTO backup_items (backup_arn, pk, item_data) \ + SELECT $1, '', d FROM UNNEST($2::jsonb[]) AS t(d)", + ) + .bind(&backup.backup_arn) + .bind(&desc_rows) + .execute(&s.catalog) + .await + .expect("add rows to the backup"); + let desc = s + .engine + .restore_table_from_backup(ACCOUNT, "many_dst", &backup.backup_arn) + .await + .expect("restore"); + assert_eq!(data_rows(&s, &desc.table_id).await, 30 + 1_234); + let src_gsi = &index_ids(&s, &src).await[0].1; + let dst_gsi = &index_ids(&s, &desc.table_id).await[0].1; + assert_eq!( + data_rows(&s, dst_gsi).await, + data_rows(&s, src_gsi).await + 1_234 + ); + s.cleanup().await; +} + +#[tokio::test] +async fn restore_into_a_name_whose_restore_lock_is_held_is_refused() { + if base_conn().is_none() { + eprintln!("SKIP restore_into_a_name_whose_restore_lock_is_held_is_refused: no PostgreSQL"); + return; + } + let s = scratch().await; + indexed_source(&s, "lk_src").await; + let backup = s + .engine + .create_backup(ACCOUNT, "lk_src", "b") + .await + .expect("backup"); + let mut other = s.catalog.acquire().await.expect("connection"); + sqlx::query("SELECT pg_advisory_lock(hashtextextended($2, $1))") + .bind(0x0045_4452_i64) + .bind(format!("{ACCOUNT}/lk_dst")) + .execute(&mut *other) + .await + .expect("another restore owns the name"); + let err = s + .engine + .restore_table_from_backup(ACCOUNT, "lk_dst", &backup.backup_arn) + .await + .expect_err("the name is being restored into already"); + assert!( + matches!(err, StorageError::TableAlreadyExists(_)), + "{err:?}" + ); + assert_eq!(table_status(&s, "lk_dst").await, None); + sqlx::query("SELECT pg_advisory_unlock_all()") + .execute(&mut *other) + .await + .expect("release"); + drop(other); + s.cleanup().await; +} + +/// A backup taken before catalog 0.0.4 has no definition row. Rolling the +/// catalog back to that shape, replaying migration 003 twice (as an upgrade +/// interrupted between applying and recording it would), and restoring must +/// give the table as restores did before: keys and items, no secondary +/// indexes, and 5/5 throughput for a provisioned table. +#[tokio::test] +async fn pre_0_0_4_backup_restores_after_migrating() { + if base_conn().is_none() { + eprintln!("SKIP pre_0_0_4_backup_restores_after_migrating: no PostgreSQL"); + return; + } + let s = scratch().await; + indexed_source(&s, "old_src").await; + let backup = s + .engine + .create_backup(ACCOUNT, "old_src", "b") + .await + .expect("backup"); + sqlx::query("DROP TABLE backup_definitions") + .execute(&s.catalog) + .await + .expect("roll the catalog back to 0.0.3"); + for _ in 0..2 { + sqlx::raw_sql(include_str!("../migrations/003_backup_definitions.sql")) + .execute(&s.catalog) + .await + .expect("apply migration 003"); + } + let version: String = + sqlx::query_scalar("SELECT value FROM settings WHERE key = 'catalog_version'") + .fetch_one(&s.catalog) + .await + .expect("version"); + assert_eq!(version, "0.0.4"); + + let desc = s + .engine + .restore_table_from_backup(ACCOUNT, "old_dst", &backup.backup_arn) + .await + .expect("restore a pre-0.0.4 backup"); + assert_eq!(table_status(&s, "old_dst").await.as_deref(), Some("ACTIVE")); + assert!(index_ids(&s, &desc.table_id).await.is_empty()); + assert_eq!(data_rows(&s, &desc.table_id).await, 30); + let restored = s + .engine + .describe_table( + ACCOUNT, + DescribeTableInput { + table_name: "old_dst".to_owned(), + }, + ) + .await + .expect("describe"); + assert_eq!( + ( + restored.provisioned_throughput.read_capacity_units, + restored.provisioned_throughput.write_capacity_units + ), + (5, 5) + ); + s.cleanup().await; +} + +/// The table row is held FOR SHARE until the data snapshot is taken, so a +/// definition change cannot commit between the two: an UpdateTable blocked +/// on the row waits for the backup's barrier, and the backup records the +/// definition in force when its items were read. +#[tokio::test] +async fn backup_definition_and_items_share_one_instant() { + if base_conn().is_none() { + eprintln!("SKIP backup_definition_and_items_share_one_instant: no PostgreSQL"); + return; + } + let s = scratch().await; + indexed_source(&s, "bar_src").await; + let table_id = table_id_of(&s, "bar_src").await.expect("table"); + + // Hold the row the way UpdateTable does, and change the billing mode + // without committing yet. + let mut writer = s.catalog.begin().await.expect("writer"); + sqlx::query("SELECT 1 FROM tables WHERE table_id = $1 FOR UPDATE") + .bind(&table_id) + .execute(&mut *writer) + .await + .expect("lock the row"); + sqlx::query("UPDATE tables SET billing_mode = 'PAY_PER_REQUEST' WHERE table_id = $1") + .bind(&table_id) + .execute(&mut *writer) + .await + .expect("change the definition"); + + // The backup must wait for the writer rather than read around it. + let arn = { + let backup = s.engine.create_backup(ACCOUNT, "bar_src", "b"); + tokio::pin!(backup); + let early = tokio::time::timeout(std::time::Duration::from_millis(500), &mut backup).await; + assert!( + early.is_err(), + "CreateBackup read the definition past an uncommitted change" + ); + writer.commit().await.expect("commit the change"); + backup.await.expect("backup").backup_arn + }; + + let def: serde_json::Value = + sqlx::query_scalar("SELECT definition FROM backup_definitions WHERE backup_arn = $1") + .bind(&arn) + .fetch_one(&s.catalog) + .await + .expect("definition"); + assert_eq!(def["BillingMode"], "PAY_PER_REQUEST"); + s.cleanup().await; +} + +/// UpdateTable commits a new GSI's catalog row before it builds the index's +/// data table. A backup taken in between leaves that index out (it is not +/// part of the table yet), drops the attribute definition only it used, and +/// restores. +#[tokio::test] +async fn backup_leaves_out_a_gsi_still_being_built() { + if base_conn().is_none() { + eprintln!("SKIP backup_leaves_out_a_gsi_still_being_built: no PostgreSQL"); + return; + } + let s = scratch().await; + let src = indexed_source(&s, "unb_src").await; + let gi = index_ids(&s, &src) + .await + .into_iter() + .find(|(n, _)| n == "gi") + .expect("gi") + .1; + sqlx::query(&format!("DROP TABLE \"_ddb_{gi}\"")) + .execute(&s.db) + .await + .expect("make the GSI look half-built"); + let backup = s + .engine + .create_backup(ACCOUNT, "unb_src", "b") + .await + .expect("backup"); + let def: serde_json::Value = + sqlx::query_scalar("SELECT definition FROM backup_definitions WHERE backup_arn = $1") + .bind(&backup.backup_arn) + .fetch_one(&s.catalog) + .await + .expect("definition"); + assert_eq!(def["GlobalSecondaryIndexes"], serde_json::json!([])); + let attrs: serde_json::Value = + sqlx::query_scalar("SELECT attribute_definitions FROM backups WHERE backup_arn = $1") + .bind(&backup.backup_arn) + .fetch_one(&s.catalog) + .await + .expect("attrs"); + let names: Vec<&str> = attrs + .as_array() + .expect("array") + .iter() + .filter_map(|a| a["AttributeName"].as_str()) + .collect(); + assert!(!names.contains(&"g"), "{names:?}"); + assert!(names.contains(&"l"), "the LSI still uses l: {names:?}"); + + let desc = s + .engine + .restore_table_from_backup(ACCOUNT, "unb_dst", &backup.backup_arn) + .await + .expect("restore"); + assert_eq!(table_status(&s, "unb_dst").await.as_deref(), Some("ACTIVE")); + let names: Vec = index_ids(&s, &desc.table_id) + .await + .into_iter() + .map(|(n, _)| n) + .collect(); + assert_eq!(names, ["li"]); + s.cleanup().await; +} + +#[tokio::test] +async fn delete_table_refuses_a_restore_in_progress() { + if base_conn().is_none() { + eprintln!("SKIP delete_table_refuses_a_restore_in_progress: no PostgreSQL"); + return; + } + let s = scratch().await; + indexed_source(&s, "dt_src").await; + let backup = s + .engine + .create_backup(ACCOUNT, "dt_src", "b") + .await + .expect("backup"); + let desc = s + .engine + .restore_table_from_backup(ACCOUNT, "dt_dst", &backup.backup_arn) + .await + .expect("restore"); + sqlx::query( + "UPDATE tables SET table_status = 'CREATING', status_transition_at = NULL \ + WHERE table_id = $1", + ) + .bind(&desc.table_id) + .execute(&s.catalog) .await - .expect("read restored rows"); - assert_eq!((rows_n, pks, r1s, r2s, nulls), (3, 2, 1, 2, 0)); - assert_eq!(desc.item_count, 3); + .expect("back to in progress"); + let err = s + .engine + .delete_table( + ACCOUNT, + extenddb_core::types::DeleteTableInput { + table_name: "dt_dst".to_owned(), + }, + ) + .await + .expect_err("refused"); + assert!(matches!(err, StorageError::IndexesInUse(_)), "{err:?}"); + assert_eq!( + table_status(&s, "dt_dst").await.as_deref(), + Some("CREATING") + ); s.cleanup().await; } diff --git a/crates/storage-postgres/tests/key_collation.rs b/crates/storage-postgres/tests/key_collation.rs index 263972dc9..9d1ba4a58 100644 --- a/crates/storage-postgres/tests/key_collation.rs +++ b/crates/storage-postgres/tests/key_collation.rs @@ -93,6 +93,7 @@ async fn scratch() -> Scratch { for sql in [ include_str!("../migrations/001_schema.sql"), include_str!("../migrations/002_vector_indexes.sql"), + include_str!("../migrations/003_backup_definitions.sql"), include_str!("../data_migrations/001_data_schema.sql"), include_str!("../data_migrations/002_gsi_pending.sql"), include_str!("../data_migrations/003_idempotency_account_scope.sql"), diff --git a/crates/storage-postgres/tests/vector_control_plane.rs b/crates/storage-postgres/tests/vector_control_plane.rs index 336011702..cafd23732 100644 --- a/crates/storage-postgres/tests/vector_control_plane.rs +++ b/crates/storage-postgres/tests/vector_control_plane.rs @@ -160,6 +160,7 @@ async fn scratch(pgvector: Pgvector) -> Scratch { for sql in [ include_str!("../migrations/001_schema.sql"), include_str!("../migrations/002_vector_indexes.sql"), + include_str!("../migrations/003_backup_definitions.sql"), include_str!("../data_migrations/001_data_schema.sql"), include_str!("../data_migrations/002_gsi_pending.sql"), include_str!("../data_migrations/003_idempotency_account_scope.sql"), diff --git a/crates/storage-sqlite/src/backup.rs b/crates/storage-sqlite/src/backup.rs index b79422a21..a4e3fc52a 100644 --- a/crates/storage-sqlite/src/backup.rs +++ b/crates/storage-sqlite/src/backup.rs @@ -3,23 +3,33 @@ //! `BackupEngine` for the SQLite backend. //! -//! A backup snapshots every item's `item_data` into `backup_items`. Restore -//! recreates the table via `create_table` and upserts the snapshot under the -//! engine write lock. `RestoreTableToPointInTime` is implemented as a +//! A backup snapshots every item's `item_data` into `backup_items` and the +//! table's definition (secondary indexes, billing mode and throughput, table +//! class, encryption) into `backup_definitions`. Restore recreates the table +//! with that definition and copies the snapshot, filling the secondary +//! indexes as it goes, under the engine write lock. `RestoreTableToPointInTime` is implemented as a //! snapshot-then-restore (then discard the temporary backup), matching the //! PostgreSQL backend's behaviour. use extenddb_core::types::{ AttributeDefinition, BackupDescription, BackupDetails, BackupSummary, BillingMode, - ContinuousBackupsDescription, CreateTableInput, KeySchemaElement, - PointInTimeRecoveryDescription, ProvisionedThroughput, SourceTableDetails, TableDescription, - TableKeyInfo, + ContinuousBackupsDescription, CreateTableInput, GsiInput, Item, KeySchemaElement, LsiInput, + PointInTimeRecoveryDescription, Projection, ProvisionedThroughput, SourceTableDetails, + TableDescription, }; use extenddb_storage::BackupEngine; +use extenddb_storage::backup_definition::{ + BACKUP_DEFINITION_VERSION, BackupTableDefinition, COPY_BATCH_BYTES, + ensure_single_part_base_key, throughput_from_catalog, +}; use extenddb_storage::error::StorageError; +use extenddb_storage::util::{composite_pk_to_text, parse_sk, sk_column_n}; use futures::future::BoxFuture; -use crate::data::{data_table_name, upsert_item_in_tx}; +use crate::data::{ + BoundValue, all_sort_key_info, data_table_name, index_table_name, insert_index_row_multi, + item_has_index_keys, project_item_for_index, sk_bound, +}; use crate::sqlite_util::parse_timestamp; use crate::store::SqliteEngine; @@ -46,6 +56,415 @@ fn ts_to_epoch(s: &str) -> f64 { .unwrap_or(0.0) } +/// Rows considered per read. Each batch is cut to [`COPY_BATCH_BYTES`] of +/// stored JSON (at least one row) before its rows are fetched, and written out +/// before the next is read, so a backup or a restore holds at most one batch. +const COPY_BATCH: i64 = 500; + +/// Read the next batch of `(rowid, item_data)` from `table` after +/// `last_rowid`, restricted by `filter` (a fixed SQL condition with one `?` +/// bound to `filter_arg`, or `None`). +/// +/// Two statements: the first reads only row lengths, so the batch can be cut +/// to the byte budget before any item text is loaded. Both run in the +/// caller's write transaction, so no row appears or moves between them. +async fn next_batch( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + table: &str, + filter: Option<(&str, &str)>, + last_rowid: i64, +) -> Result, StorageError> { + let (cond, arg) = filter.map_or(("1 = 1", None), |(c, a)| (c, Some(a))); + let sizes_sql = format!( + "SELECT rowid, length(CAST(item_data AS BLOB)) FROM {table} \ + WHERE {cond} AND rowid > ? ORDER BY rowid LIMIT ?" + ); + let mut q = sqlx::query_as::<_, (i64, i64)>(&sizes_sql); + if let Some(a) = arg { + q = q.bind(a); + } + let sizes = q + .bind(last_rowid) + .bind(COPY_BATCH) + .fetch_all(&mut **tx) + .await + .map_err(db_err)?; + let Some(&(first, _)) = sizes.first() else { + return Ok(Vec::new()); + }; + let mut upto = first; + let mut bytes: usize = 0; + for &(rowid, len) in &sizes { + let len = usize::try_from(len).unwrap_or(usize::MAX); + if rowid != first && bytes.saturating_add(len) > COPY_BATCH_BYTES { + break; + } + bytes = bytes.saturating_add(len); + upto = rowid; + } + let rows_sql = format!( + "SELECT rowid, item_data FROM {table} \ + WHERE {cond} AND rowid > ? AND rowid <= ? ORDER BY rowid" + ); + let mut q = sqlx::query_as::<_, (i64, String)>(&rows_sql); + if let Some(a) = arg { + q = q.bind(a); + } + q.bind(last_rowid) + .bind(upto) + .fetch_all(&mut **tx) + .await + .map_err(db_err) +} + +fn db_err(e: sqlx::Error) -> StorageError { + StorageError::Internal(e.to_string()) +} + +fn parse_json(s: &str, what: &str) -> Result { + serde_json::from_str(s).map_err(|e| StorageError::Internal(format!("{what}: {e}"))) +} + +/// A secondary index of a restore target, as the copy needs it. +struct RestoreIndex { + table: String, + key_schema: Vec, + projection: Projection, +} + +impl SqliteEngine { + /// Read the parts of a table's definition a restore recreates. + async fn capture_table_definition( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + table_id: &str, + ) -> Result { + let (billing_mode, pt, table_class, sse, on_demand): ( + String, + Option, + Option, + Option, + Option, + ) = sqlx::query_as( + "SELECT billing_mode, provisioned_throughput, table_class, sse_specification, \ + on_demand_throughput FROM tables WHERE table_id = ?", + ) + .bind(table_id) + .fetch_one(&mut **tx) + .await + .map_err(db_err)?; + + let rows: Vec<(String, String, String, String, Option)> = sqlx::query_as( + "SELECT index_name, index_type, key_schema, projection, provisioned_throughput \ + FROM indexes WHERE table_id = ? ORDER BY index_name", + ) + .bind(table_id) + .fetch_all(&mut **tx) + .await + .map_err(db_err)?; + let mut gsis = Vec::new(); + let mut lsis = Vec::new(); + for (index_name, index_type, ks, proj, ipt) in rows { + let key_schema: Vec = parse_json(&ks, "index key schema")?; + let projection: Projection = parse_json(&proj, "index projection")?; + if index_type == "LSI" { + lsis.push(LsiInput { + index_name, + key_schema, + projection, + }); + } else { + let ipt: Option = ipt + .as_deref() + .map(|s| parse_json(s, "index throughput")) + .transpose()?; + gsis.push(GsiInput { + index_name, + key_schema, + projection, + provisioned_throughput: throughput_from_catalog(ipt.as_ref()), + }); + } + } + + let vector_index_names: Vec = sqlx::query_scalar( + "SELECT index_name FROM vector_indexes WHERE table_id = ? ORDER BY index_name", + ) + .bind(table_id) + .fetch_all(&mut **tx) + .await + .map_err(db_err)?; + + let pt: Option = pt + .as_deref() + .map(|s| parse_json(s, "throughput")) + .transpose()?; + Ok(BackupTableDefinition { + version: BACKUP_DEFINITION_VERSION, + billing_mode, + provisioned_throughput: throughput_from_catalog(pt.as_ref()), + global_secondary_indexes: gsis, + local_secondary_indexes: lsis, + table_class, + sse_specification: sse.as_deref().map(|s| parse_json(s, "sse")).transpose()?, + on_demand_throughput: on_demand + .as_deref() + .map(|s| parse_json(s, "on-demand throughput")) + .transpose()?, + vector_index_names, + }) + } + + /// Copy a backup's items into a freshly created restore target and fill + /// its secondary indexes, inside the caller's write transaction. + /// + /// Every key column is derived from the item (all HASH parts into `pk`, + /// each RANGE part into its typed column), and each row is a plain + /// INSERT, so two backup rows for one key fail the restore instead of + /// collapsing into one item. Items are read and written in byte-bounded + /// batches, each in its own write transaction, and the target is flipped + /// to ACTIVE in the last one. + async fn copy_backup_items( + &self, + desc: &TableDescription, + backup_arn: &str, + ) -> Result<(), StorageError> { + let ddb_table = data_table_name(&desc.table_id); + let sort_keys = all_sort_key_info(&desc.key_schema, &desc.attribute_definitions); + let mut cols = vec!["pk".to_owned()]; + cols.extend( + sort_keys + .iter() + .enumerate() + .map(|(i, &(_, t))| sk_column_n(i, t)), + ); + cols.push("item_data".to_owned()); + let insert_sql = format!( + "INSERT INTO {ddb_table} ({}) VALUES ({})", + cols.join(", "), + vec!["?"; cols.len()].join(", ") + ); + + let index_rows: Vec<(String, String, String)> = sqlx::query_as( + "SELECT index_id, key_schema, projection FROM indexes WHERE table_id = ?", + ) + .bind(&desc.table_id) + .fetch_all(&self.pool) + .await + .map_err(db_err)?; + let mut indexes = Vec::with_capacity(index_rows.len()); + for (index_id, ks, proj) in index_rows { + indexes.push(RestoreIndex { + table: index_table_name(&index_id), + key_schema: parse_json(&ks, "index key schema")?, + projection: parse_json(&proj, "index projection")?, + }); + } + let base_sks = all_sort_key_info(&desc.key_schema, &desc.attribute_definitions); + + let mut last_rowid: i64 = 0; + let mut count: i64 = 0; + loop { + // One write transaction per batch, so other writers on the server + // wait for a batch rather than the whole restore. The target is + // CREATING throughout, which refuses every data-plane request, so + // no one sees it part-filled. Each batch first re-checks, under + // the write lock, that the backup is still AVAILABLE (DeleteBackup + // removes its rows and marks it DELETED in one transaction) and + // that the target is still this restore's. + let _writer = self.write_lock.lock().await; + let mut tx = self + .pool + .begin_with("BEGIN IMMEDIATE") + .await + .map_err(db_err)?; + Self::ensure_restore_can_continue(&mut tx, desc, backup_arn).await?; + let batch = next_batch( + &mut tx, + "backup_items", + Some(("backup_arn = ?", backup_arn)), + last_rowid, + ) + .await?; + let Some(&(tail, _)) = batch.last() else { + let (table_size,): (i64,) = sqlx::query_as(&format!( + "SELECT COALESCE(SUM(length(item_data)), 0) FROM {ddb_table}" + )) + .fetch_one(&mut *tx) + .await + .map_err(db_err)?; + // Only from CREATING, which the check above confirmed under + // the same lock. + sqlx::query( + "UPDATE tables SET item_count = ?, table_size_bytes = ?, \ + table_status = 'ACTIVE', status_transition_at = NULL \ + WHERE table_id = ? AND table_status = 'CREATING'", + ) + .bind(count) + .bind(table_size) + .bind(&desc.table_id) + .execute(&mut *tx) + .await + .map_err(db_err)?; + tx.commit().await.map_err(db_err)?; + return Ok(()); + }; + last_rowid = tail; + for (_, item_json) in &batch { + let item: Item = parse_json(item_json, "backup item")?; + let mut values = vec![BoundValue::Text(composite_pk_to_text( + &item, + &desc.key_schema, + )?)]; + for &(name, sk_type) in &sort_keys { + let value = item.get(name).ok_or_else(|| { + StorageError::Internal(format!( + "backup {backup_arn} has an item without sort key {name}" + )) + })?; + values.push(sk_bound(&parse_sk(value, sk_type)?)); + } + values.push(BoundValue::Text(item_json.clone())); + let mut q = sqlx::query(&insert_sql); + for v in values { + q = crate::data::bind_bound!(q, v); + } + q.execute(&mut *tx).await.map_err(db_err)?; + + for idx in &indexes { + if !item_has_index_keys(&item, &idx.key_schema) { + continue; + } + let idx_sks = all_sort_key_info(&idx.key_schema, &desc.attribute_definitions); + let projected = project_item_for_index( + &item, + &idx.key_schema, + &desc.key_schema, + &idx.projection, + ); + insert_index_row_multi( + &mut tx, + &idx.table, + &item, + &projected, + &idx.key_schema, + &desc.key_schema, + &idx_sks, + &base_sks, + ) + .await?; + } + count += 1; + } + tx.commit().await.map_err(db_err)?; + } + } + + /// Fail a restore whose backup was deleted or whose target is no longer + /// the one it created. + async fn ensure_restore_can_continue( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + desc: &TableDescription, + backup_arn: &str, + ) -> Result<(), StorageError> { + let available: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM backups WHERE backup_arn = ? \ + AND backup_status = 'AVAILABLE')", + ) + .bind(backup_arn) + .fetch_one(&mut **tx) + .await + .map_err(db_err)?; + if !available { + return Err(StorageError::Validation(format!( + "Backup not found: {backup_arn}" + ))); + } + let still_target: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM tables WHERE table_id = ? \ + AND table_status = 'CREATING' AND status_transition_at IS NULL)", + ) + .bind(&desc.table_id) + .fetch_one(&mut **tx) + .await + .map_err(db_err)?; + if !still_target { + return Err(StorageError::TableNotFound(format!( + "restore of {backup_arn} into {} did not complete: the table was deleted \ + while it was being restored", + desc.table_name + ))); + } + Ok(()) + } + + /// Remove a restore target whose copy failed or was abandoned: its + /// catalog row (cascading its index rows) and every data table. Only a + /// target still CREATING with no scheduled transition is removed, and by + /// table id, so this can never touch a table the restore did not create + /// or one a client already deleted. + async fn abort_restore(&self, table_id: &str) -> Result { + let _writer = self.write_lock.lock().await; + let mut tx = self + .pool + .begin_with("BEGIN IMMEDIATE") + .await + .map_err(db_err)?; + let index_ids: Vec = + sqlx::query_scalar("SELECT index_id FROM indexes WHERE table_id = ?") + .bind(table_id) + .fetch_all(&mut *tx) + .await + .map_err(db_err)?; + let removed = sqlx::query( + "DELETE FROM tables WHERE table_id = ? AND table_status = 'CREATING' \ + AND status_transition_at IS NULL", + ) + .bind(table_id) + .execute(&mut *tx) + .await + .map_err(db_err)? + .rows_affected() + > 0; + if removed { + for id in &index_ids { + Self::drop_index_data_table(&mut tx, id).await?; + } + Self::drop_data_table(&mut tx, table_id).await?; + } + tx.commit().await.map_err(db_err)?; + Ok(removed) + } + + /// Remove restore targets left by a process that died mid-restore. + /// + /// A restore target is the only table that sits CREATING with no + /// scheduled transition. Called once at startup, before any request is + /// served, so no restore of this process can be in flight and every such + /// table is abandoned. + /// + /// This assumes one server process per database file, which the backend + /// already depends on: writers are serialized by an in-process lock, so a + /// second process writing the same file is unsupported regardless. A + /// second process started against the file while the first is mid-restore + /// would remove that restore's target. + pub(crate) async fn sweep_abandoned_restores(&self) -> Result, StorageError> { + let candidates: Vec<(String, String)> = sqlx::query_as( + "SELECT table_id, table_name FROM tables \ + WHERE table_status = 'CREATING' AND status_transition_at IS NULL", + ) + .fetch_all(&self.pool) + .await + .map_err(db_err)?; + let mut removed = Vec::new(); + for (table_id, table_name) in candidates { + if self.abort_restore(&table_id).await? { + removed.push(table_name); + } + } + Ok(removed) + } +} + impl BackupEngine for SqliteEngine { fn create_backup( &self, @@ -57,17 +476,29 @@ impl BackupEngine for SqliteEngine { let table_name = table_name.to_owned(); let backup_name = backup_name.to_owned(); Box::pin(async move { - let row: Option<(String, String, String, String, i64)> = sqlx::query_as( - "SELECT table_id, key_schema, attribute_definitions, billing_mode, table_size_bytes \ + // One write transaction for everything: the table lookup, its + // definition, the items, and the backup row. Every writer is held + // off by the write lock, so the backup is consistent as of one + // instant, including against a concurrent UpdateTable. Items are + // read in rowid batches on the same connection, so memory is + // bounded by the batch size. + let _writer = self.write_lock.lock().await; + let mut tx = self + .pool + .begin_with("BEGIN IMMEDIATE") + .await + .map_err(db_err)?; + let row: Option<(String, String, String, String)> = sqlx::query_as( + "SELECT table_id, key_schema, attribute_definitions, billing_mode \ FROM tables WHERE account_id = ? AND table_name = ? AND table_status = 'ACTIVE'", ) .bind(&account_id) .bind(&table_name) - .fetch_optional(&self.pool) + .fetch_optional(&mut *tx) .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + .map_err(db_err)?; - let (table_id, key_schema, attr_defs, billing_mode, size_bytes) = + let (table_id, key_schema, attr_defs, billing_mode) = row.ok_or_else(|| StorageError::TableNotFound(table_name.clone()))?; let backup_arn = format!( @@ -76,50 +507,68 @@ impl BackupEngine for SqliteEngine { id = backup_id() ); - let ddb_table = data_table_name(&table_id); - let items: Vec<(String,)> = - sqlx::query_as(&format!("SELECT item_data FROM {ddb_table}")) - .fetch_all(&self.pool) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - let item_count = i64::try_from(items.len()).unwrap_or(i64::MAX); - - let _writer = self.write_lock.lock().await; - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await + let definition = Self::capture_table_definition(&mut tx, &table_id).await?; + let definition = serde_json::to_string(&definition.to_json()?) .map_err(|e| StorageError::Internal(e.to_string()))?; sqlx::query( "INSERT INTO backups (backup_arn, backup_name, table_id, table_name, account_id, \ backup_status, backup_size_bytes, item_count, key_schema, attribute_definitions, \ - billing_mode) VALUES (?, ?, ?, ?, ?, 'AVAILABLE', ?, ?, ?, ?, ?)", + billing_mode) VALUES (?, ?, ?, ?, ?, 'AVAILABLE', 0, 0, ?, ?, ?)", ) .bind(&backup_arn) .bind(&backup_name) .bind(&table_id) .bind(&table_name) .bind(&account_id) - .bind(size_bytes) - .bind(item_count) .bind(&key_schema) .bind(&attr_defs) .bind(&billing_mode) .execute(&mut *tx) .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - - for (item_data,) in &items { - sqlx::query( - "INSERT INTO backup_items (backup_arn, pk, sk, item_data) VALUES (?, '', NULL, ?)", - ) + .map_err(db_err)?; + sqlx::query("INSERT INTO backup_definitions (backup_arn, definition) VALUES (?, ?)") .bind(&backup_arn) - .bind(item_data) + .bind(&definition) .execute(&mut *tx) .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + .map_err(db_err)?; + + let ddb_table = data_table_name(&table_id); + let mut last_rowid: i64 = 0; + let mut item_count: i64 = 0; + let mut size_bytes: i64 = 0; + loop { + let batch = next_batch(&mut tx, &ddb_table, None, last_rowid).await?; + let Some(&(tail, _)) = batch.last() else { + break; + }; + last_rowid = tail; + for (_, item_data) in &batch { + let item: Item = parse_json(item_data, "item")?; + size_bytes += i64::try_from(extenddb_core::types::item_size_bytes(&item)) + .unwrap_or(i64::MAX); + sqlx::query( + "INSERT INTO backup_items (backup_arn, pk, sk, item_data) \ + VALUES (?, '', NULL, ?)", + ) + .bind(&backup_arn) + .bind(item_data) + .execute(&mut *tx) + .await + .map_err(db_err)?; + item_count += 1; + } } + sqlx::query( + "UPDATE backups SET item_count = ?, backup_size_bytes = ? WHERE backup_arn = ?", + ) + .bind(item_count) + .bind(size_bytes) + .bind(&backup_arn) + .execute(&mut *tx) + .await + .map_err(db_err)?; let created_at: (String,) = sqlx::query_as("SELECT created_at FROM backups WHERE backup_arn = ?") @@ -276,16 +725,31 @@ impl BackupEngine for SqliteEngine { // past a concurrent locked writer's busy_timeout. let _writer = self.write_lock.lock().await; - // The account predicate is repeated on both writes rather than - // relying on the lookup above, so the statements are correct on - // their own terms. + // One transaction, so a restore never sees the backup AVAILABLE + // with its rows gone. The account predicate is repeated on every + // write rather than relying on the lookup above, so the + // statements are correct on their own terms. + let mut tx = self + .pool + .begin_with("BEGIN IMMEDIATE") + .await + .map_err(db_err)?; sqlx::query( "DELETE FROM backup_items WHERE backup_arn = ?1 AND EXISTS (\ SELECT 1 FROM backups b WHERE b.backup_arn = ?1 AND b.account_id = ?2)", ) .bind(&backup_arn) .bind(&account_id) - .execute(&self.pool) + .execute(&mut *tx) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + sqlx::query( + "DELETE FROM backup_definitions WHERE backup_arn = ?1 AND EXISTS (\ + SELECT 1 FROM backups b WHERE b.backup_arn = ?1 AND b.account_id = ?2)", + ) + .bind(&backup_arn) + .bind(&account_id) + .execute(&mut *tx) .await .map_err(|e| StorageError::Internal(e.to_string()))?; sqlx::query( @@ -294,9 +758,10 @@ impl BackupEngine for SqliteEngine { ) .bind(&backup_arn) .bind(&account_id) - .execute(&self.pool) + .execute(&mut *tx) .await .map_err(|e| StorageError::Internal(e.to_string()))?; + tx.commit().await.map_err(db_err)?; Ok(BackupDescription { backup_details: BackupDetails { backup_status: "DELETED".to_owned(), @@ -317,116 +782,90 @@ impl BackupEngine for SqliteEngine { let target_table_name = target_table_name.to_owned(); let backup_arn = backup_arn.to_owned(); Box::pin(async move { - let row: Option<(String, String, String)> = sqlx::query_as( - "SELECT key_schema, attribute_definitions, billing_mode \ - FROM backups WHERE backup_arn = ? AND account_id = ? \ - AND backup_status = 'AVAILABLE'", + let row: Option<(String, String, String, Option)> = sqlx::query_as( + "SELECT b.key_schema, b.attribute_definitions, b.billing_mode, d.definition \ + FROM backups b LEFT JOIN backup_definitions d ON d.backup_arn = b.backup_arn \ + WHERE b.backup_arn = ? AND b.account_id = ? AND b.backup_status = 'AVAILABLE'", ) .bind(&backup_arn) .bind(&account_id) .fetch_optional(&self.pool) .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + .map_err(db_err)?; - let (ks, ad, billing) = row.ok_or_else(|| { + let (ks, ad, billing, definition) = row.ok_or_else(|| { StorageError::Validation(format!("Backup not found: {backup_arn}")) })?; - let key_schema: Vec = - serde_json::from_str(&ks).map_err(|e| StorageError::Internal(e.to_string()))?; - let attr_defs: Vec = - serde_json::from_str(&ad).map_err(|e| StorageError::Internal(e.to_string()))?; - let billing_mode = Some(if billing == "PAY_PER_REQUEST" { - BillingMode::PayPerRequest - } else { - BillingMode::Provisioned - }); + let definition = definition + .map(|d| { + BackupTableDefinition::from_json( + parse_json(&d, "table definition")?, + &backup_arn, + ) + }) + .transpose()?; + if let Some(d) = &definition { + d.ensure_restorable(&backup_arn)?; + } + let key_schema: Vec = parse_json(&ks, "key schema")?; + ensure_single_part_base_key(&key_schema, &backup_arn)?; + let attr_defs: Vec = parse_json(&ad, "attribute definitions")?; - let create_input = CreateTableInput { + let mut create_input = CreateTableInput { table_name: target_table_name.clone(), - key_schema: key_schema.clone(), - attribute_definitions: attr_defs.clone(), - billing_mode, - provisioned_throughput: Some(ProvisionedThroughput { - read_capacity_units: 5, - write_capacity_units: 5, - }), - global_secondary_indexes: None, - local_secondary_indexes: None, - stream_specification: None, - tags: None, - deletion_protection_enabled: None, - sse_specification: None, - table_class: None, - on_demand_throughput: None, + key_schema, + attribute_definitions: attr_defs, ..Default::default() }; + match definition { + Some(d) => d.apply_to(&mut create_input), + None => { + // A backup taken before table definitions were recorded + // carries only its keys and billing mode, so it restores as + // it always has: no secondary indexes, and 5/5 throughput for + // a provisioned table since the source's is unknown. + let on_demand = billing == "PAY_PER_REQUEST"; + create_input.billing_mode = Some(if on_demand { + BillingMode::PayPerRequest + } else { + BillingMode::Provisioned + }); + create_input.provisioned_throughput = + (!on_demand).then_some(ProvisionedThroughput { + read_capacity_units: 5, + write_capacity_units: 5, + }); + } + } // `defer_active`: the target is written CREATING with no scheduled // transition, so the control-plane worker cannot report it ACTIVE - // while the copy below is still running. The explicit ACTIVE update - // after the copy commits is the only flip. + // while the copy below is still running, and every data-plane + // request against it is refused. Its secondary indexes are created + // with it, empty, and filled by the copy. let desc = self .create_table_impl(&account_id, create_input, true) .await?; - let key_info = TableKeyInfo { - table_name: target_table_name.clone(), - account_id: account_id.clone(), - table_id: desc.table_id.clone(), - base_key_schema: key_schema.clone(), - key_schema, - attribute_definitions: attr_defs, - has_lsi: false, - // The restored table is created above without secondary - // indexes, so there is no index metadata to carry. - global_secondary_indexes: Vec::new(), - local_secondary_indexes: Vec::new(), - stream_specification: None, - ..Default::default() - }; - let items: Vec<(String,)> = - sqlx::query_as("SELECT item_data FROM backup_items WHERE backup_arn = ?") - .bind(&backup_arn) - .fetch_all(&self.pool) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - let item_count = i64::try_from(items.len()).unwrap_or(i64::MAX); - - let _writer = self.write_lock.lock().await; - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - for (item_json,) in &items { - let item: extenddb_core::types::Item = serde_json::from_str(item_json) - .map_err(|e| StorageError::Internal(e.to_string()))?; - upsert_item_in_tx(&mut tx, &key_info, &item).await?; + // The copy commits in batches and flips the target ACTIVE in its + // last one; until then the target is CREATING and refuses every + // data-plane request. A failure part-way is cleaned up here, a + // crash part-way by the startup sweep. + let copied = self.copy_backup_items(&desc, &backup_arn).await; + if let Err(e) = copied { + tracing::error!( + "restore of {backup_arn} into {target_table_name} failed, \ + removing the partial table: {e}" + ); + if let Err(cleanup) = self.abort_restore(&desc.table_id).await { + tracing::error!( + "could not remove partially restored table {target_table_name} \ + ({}); the next startup will: {cleanup}", + desc.table_id + ); + } + return Err(e); } - sqlx::query("UPDATE tables SET item_count = ? WHERE account_id = ? AND table_name = ?") - .bind(item_count) - .bind(&account_id) - .bind(&target_table_name) - .execute(&mut *tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - // Mark the restored table ACTIVE immediately: the data is fully - // populated and ready to serve. This matches the Postgres backend - // and real DynamoDB, where a restored table becomes ACTIVE once the - // restore completes (the CREATING status is transient) rather than - // waiting for the control-plane transition poller. - sqlx::query( - "UPDATE tables SET table_status = 'ACTIVE', status_transition_at = NULL \ - WHERE account_id = ? AND table_name = ?", - ) - .bind(&account_id) - .bind(&target_table_name) - .execute(&mut *tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - tx.commit() - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; Ok(desc) }) @@ -542,3 +981,631 @@ impl BackupEngine for SqliteEngine { }) } } + +#[cfg(test)] +mod restore_tests { + use extenddb_core::types::{Item, TableKeyInfo}; + use extenddb_storage::{BackupEngine, DataEngine}; + use serde_json::json; + + use crate::SqliteEngine; + + const ACCOUNT: &str = "000000000000"; + + async fn engine() -> SqliteEngine { + let engine = SqliteEngine::new(":memory:", 1, "us-east-1", 409_600) + .await + .expect("engine"); + crate::schema::apply(&engine.pool).await.expect("schema"); + sqlx::query("UPDATE settings SET value = '0' WHERE key = 'control_plane_delay_seconds'") + .execute(&engine.pool) + .await + .expect("delay"); + sqlx::query("UPDATE settings SET value = '0' WHERE key = 'index_propagation_delay_ms'") + .execute(&engine.pool) + .await + .expect("propagation"); + sqlx::query("INSERT INTO accounts (account_id, account_name) VALUES (?, 'default')") + .bind(ACCOUNT) + .execute(&engine.pool) + .await + .expect("account"); + engine + } + + /// A provisioned two-part-key table with one GSI and one LSI. + async fn indexed_table(engine: &SqliteEngine, name: &str) -> String { + let input: extenddb_core::types::CreateTableInput = serde_json::from_value(json!({ + "TableName": name, + "KeySchema": [ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"} + ], + "AttributeDefinitions": [ + {"AttributeName": "pk", "AttributeType": "S"}, + {"AttributeName": "sk", "AttributeType": "N"}, + {"AttributeName": "g", "AttributeType": "S"}, + {"AttributeName": "l", "AttributeType": "S"} + ], + "BillingMode": "PROVISIONED", + "ProvisionedThroughput": {"ReadCapacityUnits": 7, "WriteCapacityUnits": 9}, + "GlobalSecondaryIndexes": [{ + "IndexName": "gi", + "KeySchema": [{"AttributeName": "g", "KeyType": "HASH"}], + "Projection": {"ProjectionType": "KEYS_ONLY"}, + "ProvisionedThroughput": {"ReadCapacityUnits": 3, "WriteCapacityUnits": 4} + }], + "LocalSecondaryIndexes": [{ + "IndexName": "li", + "KeySchema": [ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "l", "KeyType": "RANGE"} + ], + "Projection": {"ProjectionType": "ALL"} + }] + })) + .expect("input"); + let desc = engine + .create_table_impl(ACCOUNT, input, false) + .await + .expect("create"); + let key_info = TableKeyInfo { + table_name: name.to_owned(), + account_id: ACCOUNT.to_owned(), + table_id: desc.table_id.clone(), + key_schema: desc.key_schema.clone(), + base_key_schema: desc.key_schema.clone(), + attribute_definitions: desc.attribute_definitions.clone(), + ..Default::default() + }; + for i in 0..30 { + let mut item = json!({"pk": {"S": format!("p{}", i % 3)}, "sk": {"N": i.to_string()}}); + if i % 2 == 0 { + item["g"] = json!({"S": format!("g{}", i % 4)}); + } + if i % 3 == 0 { + item["l"] = json!({"S": format!("l{i}")}); + } + let item: Item = serde_json::from_value(item).expect("item"); + engine + .put_item( + &key_info, + item, + false, + None, + &extenddb_core::expression::ExpressionMaps::default(), + None, + ) + .await + .expect("put"); + } + desc.table_id + } + + async fn count(engine: &SqliteEngine, table: &str) -> i64 { + sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {table}")) + .fetch_one(&engine.pool) + .await + .expect("count") + } + + async fn index_tables(engine: &SqliteEngine, table_id: &str) -> Vec<(String, String)> { + sqlx::query_as( + "SELECT index_name, index_id FROM indexes WHERE table_id = ? ORDER BY index_name", + ) + .bind(table_id) + .fetch_all(&engine.pool) + .await + .expect("indexes") + } + + #[tokio::test] + async fn restore_recreates_indexes_and_throughput() { + let engine = engine().await; + let src = indexed_table(&engine, "src").await; + let backup = engine + .create_backup(ACCOUNT, "src", "bkp") + .await + .expect("backup"); + let desc = engine + .restore_table_from_backup(ACCOUNT, "dst", &backup.backup_arn) + .await + .expect("restore"); + + let (status, pt): (String, Option) = sqlx::query_as( + "SELECT table_status, provisioned_throughput FROM tables WHERE table_id = ?", + ) + .bind(&desc.table_id) + .fetch_one(&engine.pool) + .await + .expect("row"); + assert_eq!(status, "ACTIVE"); + let pt: serde_json::Value = serde_json::from_str(&pt.expect("throughput")).expect("json"); + assert_eq!(pt["ReadCapacityUnits"], 7); + assert_eq!(pt["WriteCapacityUnits"], 9); + + // Same indexes, each holding the same number of rows as the source's. + let src_idx = index_tables(&engine, &src).await; + let dst_idx = index_tables(&engine, &desc.table_id).await; + assert_eq!( + src_idx.iter().map(|(n, _)| n.as_str()).collect::>(), + vec!["gi", "li"] + ); + assert_eq!( + dst_idx.iter().map(|(n, _)| n.clone()).collect::>(), + src_idx.iter().map(|(n, _)| n.clone()).collect::>() + ); + for ((_, s), (_, d)) in src_idx.iter().zip(&dst_idx) { + let (s, d) = ( + crate::data::index_table_name(s), + crate::data::index_table_name(d), + ); + assert_eq!(count(&engine, &d).await, count(&engine, &s).await); + assert!(count(&engine, &d).await > 0); + } + assert_eq!( + count(&engine, &crate::data::data_table_name(&desc.table_id)).await, + 30 + ); + } + + #[tokio::test] + async fn startup_sweep_removes_an_abandoned_restore() { + let engine = engine().await; + indexed_table(&engine, "src").await; + let backup = engine + .create_backup(ACCOUNT, "src", "bkp") + .await + .expect("backup"); + let desc = engine + .restore_table_from_backup(ACCOUNT, "dst", &backup.backup_arn) + .await + .expect("restore"); + let index_ids: Vec = index_tables(&engine, &desc.table_id) + .await + .into_iter() + .map(|(_, id)| id) + .collect(); + + // The state a process killed mid-copy leaves: the target CREATING with + // no scheduled transition (its copy transaction rolled back). + sqlx::query( + "UPDATE tables SET table_status = 'CREATING', status_transition_at = NULL \ + WHERE table_id = ?", + ) + .bind(&desc.table_id) + .execute(&engine.pool) + .await + .expect("simulate crash"); + + let removed = engine.sweep_abandoned_restores().await.expect("sweep"); + assert_eq!(removed, vec!["dst".to_owned()]); + let rows: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM tables WHERE table_name = 'dst'") + .fetch_one(&engine.pool) + .await + .expect("rows"); + assert_eq!(rows, 0); + let mut leftover = vec![desc.table_id.clone()]; + leftover.extend(index_ids); + for id in leftover { + let exists: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?)", + ) + .bind(format!("_ddb_{id}")) + .fetch_one(&engine.pool) + .await + .expect("exists"); + assert!(!exists, "data table _ddb_{id} left behind"); + } + + // A live table is never a candidate, and the name is free again. + assert!( + engine + .sweep_abandoned_restores() + .await + .expect("sweep") + .is_empty() + ); + engine + .restore_table_from_backup(ACCOUNT, "dst", &backup.backup_arn) + .await + .expect("restore again"); + } + + #[tokio::test] + async fn failed_copy_removes_the_target() { + let engine = engine().await; + indexed_table(&engine, "src").await; + let backup = engine + .create_backup(ACCOUNT, "src", "bkp") + .await + .expect("backup"); + // Corrupt one backup row so the copy fails part-way. + sqlx::query( + "UPDATE backup_items SET item_data = '{\"pk\": {\"S\": \"x\"}}' \ + WHERE rowid = (SELECT MAX(rowid) FROM backup_items WHERE backup_arn = ?)", + ) + .bind(&backup.backup_arn) + .execute(&engine.pool) + .await + .expect("corrupt"); + let before: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM sqlite_master WHERE type = 'table'") + .fetch_one(&engine.pool) + .await + .expect("tables"); + engine + .restore_table_from_backup(ACCOUNT, "dst", &backup.backup_arn) + .await + .expect_err("an item without its sort key cannot restore"); + let rows: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM tables WHERE table_name = 'dst'") + .fetch_one(&engine.pool) + .await + .expect("rows"); + assert_eq!(rows, 0); + let after: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM sqlite_master WHERE type = 'table'") + .fetch_one(&engine.pool) + .await + .expect("tables"); + assert_eq!( + after, before, + "the target's data and index tables are dropped" + ); + } + + async fn create(engine: &SqliteEngine, input: serde_json::Value) -> String { + let input: extenddb_core::types::CreateTableInput = + serde_json::from_value(input).expect("input"); + engine + .create_table_impl(ACCOUNT, input, false) + .await + .expect("create") + .table_id + } + + #[tokio::test] + async fn multipart_key_backup_is_refused() { + let engine = engine().await; + create( + &engine, + json!({ + "TableName": "m", + "KeySchema": [ + {"AttributeName": "a", "KeyType": "HASH"}, + {"AttributeName": "b", "KeyType": "HASH"} + ], + "AttributeDefinitions": [ + {"AttributeName": "a", "AttributeType": "S"}, + {"AttributeName": "b", "AttributeType": "S"} + ], + "BillingMode": "PAY_PER_REQUEST" + }), + ) + .await; + let backup = engine + .create_backup(ACCOUNT, "m", "bkp") + .await + .expect("backup"); + let err = engine + .restore_table_from_backup(ACCOUNT, "m2", &backup.backup_arn) + .await + .expect_err("refused"); + assert!( + matches!(err, extenddb_storage::error::StorageError::Unsupported(_)), + "{err:?}" + ); + let rows: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM tables WHERE table_name = 'm2'") + .fetch_one(&engine.pool) + .await + .expect("rows"); + assert_eq!(rows, 0); + } + + #[tokio::test] + async fn restore_keeps_table_class_sse_and_on_demand_limits() { + let engine = engine().await; + create( + &engine, + json!({ + "TableName": "c", + "KeySchema": [{"AttributeName": "pk", "KeyType": "HASH"}], + "AttributeDefinitions": [{"AttributeName": "pk", "AttributeType": "S"}], + "BillingMode": "PAY_PER_REQUEST", + "TableClass": "STANDARD_INFREQUENT_ACCESS", + "SSESpecification": {"Enabled": true, "SSEType": "KMS"}, + "OnDemandThroughput": {"MaxReadRequestUnits": 100, "MaxWriteRequestUnits": 50} + }), + ) + .await; + let backup = engine + .create_backup(ACCOUNT, "c", "bkp") + .await + .expect("backup"); + engine + .restore_table_from_backup(ACCOUNT, "c2", &backup.backup_arn) + .await + .expect("restore"); + let read = |name: &'static str| { + let pool = engine.pool.clone(); + async move { + let row: (String, Option, Option, Option) = sqlx::query_as( + "SELECT billing_mode, table_class, sse_specification, on_demand_throughput \ + FROM tables WHERE table_name = ?", + ) + .bind(name) + .fetch_one(&pool) + .await + .expect("row"); + let json = |v: Option| { + v.map(|s| serde_json::from_str::(&s).expect("json")) + }; + (row.0, row.1, json(row.2), json(row.3)) + } + }; + let want = read("c").await; + assert_eq!(want.1.as_deref(), Some("STANDARD_INFREQUENT_ACCESS")); + assert!(want.2.is_some() && want.3.is_some(), "{want:?}"); + assert_eq!(read("c2").await, want); + } + + #[tokio::test] + async fn restore_crosses_batch_boundaries_and_keeps_n_key_order() { + let engine = engine().await; + let src = indexed_table(&engine, "src").await; + let (src_gsi_before,) = (index_tables(&engine, &src).await[0].1.clone(),); + let base = crate::data::data_table_name(&src); + // Add rows straight into the source, well past one read batch, with N + // sort keys whose encoded order differs from their text order. + let key_info = TableKeyInfo { + table_name: "src".to_owned(), + account_id: ACCOUNT.to_owned(), + table_id: src.clone(), + key_schema: serde_json::from_value(json!([ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"} + ])) + .expect("ks"), + ..Default::default() + }; + let mut key_info = key_info; + key_info.base_key_schema = key_info.key_schema.clone(); + key_info.attribute_definitions = serde_json::from_value(json!([ + {"AttributeName": "pk", "AttributeType": "S"}, + {"AttributeName": "sk", "AttributeType": "N"}, + {"AttributeName": "g", "AttributeType": "S"}, + {"AttributeName": "l", "AttributeType": "S"} + ])) + .expect("ad"); + // More rows than one read batch (500), so the restore commits + // several batches. + for i in 0..700 { + let n = if i % 2 == 0 { + format!("-{i}.5") + } else { + format!("{}", i * 1000) + }; + let item: Item = serde_json::from_value(json!({ + "pk": {"S": "bulk"}, "sk": {"N": n}, "g": {"S": "gbulk"} + })) + .expect("item"); + engine + .put_item( + &key_info, + item, + false, + None, + &extenddb_core::expression::ExpressionMaps::default(), + None, + ) + .await + .expect("put"); + } + let backup = engine + .create_backup(ACCOUNT, "src", "bkp") + .await + .expect("backup"); + let desc = engine + .restore_table_from_backup(ACCOUNT, "dst", &backup.backup_arn) + .await + .expect("restore"); + let dst = crate::data::data_table_name(&desc.table_id); + assert_eq!(count(&engine, &dst).await, count(&engine, &base).await); + assert_eq!(count(&engine, &dst).await, 730); + // The physical rows are identical to the ones the write path made, + // sort-key encoding included, so key order is preserved. + let rows = |t: String| { + let pool = engine.pool.clone(); + async move { + sqlx::query_scalar::<_, String>(&format!( + "SELECT pk || '|' || sk_n || '|' || item_data FROM {t} ORDER BY pk, sk_n" + )) + .fetch_all(&pool) + .await + .expect("rows") + } + }; + assert_eq!(rows(dst).await, rows(base).await); + let dst_gsi = index_tables(&engine, &desc.table_id).await[0].1.clone(); + let gsi_rows = |id: String| { + let pool = engine.pool.clone(); + async move { + sqlx::query_scalar::<_, String>(&format!( + "SELECT pk || '|' || base_pk || '|' || base_sk_n || '|' || item_data \ + FROM {} ORDER BY 1", + crate::data::index_table_name(&id) + )) + .fetch_all(&pool) + .await + .expect("gsi rows") + } + }; + assert_eq!(gsi_rows(dst_gsi).await, gsi_rows(src_gsi_before).await); + } + + /// A database with the schema catalog 0.0.3 shipped, holding a backup in + /// the row shape 0.0.3 wrote (`pk` empty, `sk` NULL, no definition), is + /// brought to 0.0.4 by re-applying the schema, which is what + /// `extenddb migrate` does on this backend; the old backup then restores + /// as it did before: keys and items, no secondary indexes. + #[tokio::test] + async fn a_0_0_3_database_migrates_and_its_backups_restore() { + let engine = SqliteEngine::new(":memory:", 1, "us-east-1", 409_600) + .await + .expect("engine"); + sqlx::raw_sql(include_str!("../testdata/schema_0_0_3.sql")) + .execute(&engine.pool) + .await + .expect("0.0.3 schema"); + sqlx::query("UPDATE settings SET value = '0' WHERE key = 'control_plane_delay_seconds'") + .execute(&engine.pool) + .await + .expect("delay"); + sqlx::query("INSERT INTO accounts (account_id, account_name) VALUES (?, 'default')") + .bind(ACCOUNT) + .execute(&engine.pool) + .await + .expect("account"); + assert!(engine.check_catalog_version().await.is_err()); + + // A backup exactly as the 0.0.3 binary wrote one. + let arn = format!("arn:aws:dynamodb:us-east-1:{ACCOUNT}:table/old/backup/1-00000000"); + sqlx::query( + "INSERT INTO backups (backup_arn, backup_name, table_id, table_name, account_id, \ + backup_status, backup_size_bytes, item_count, key_schema, attribute_definitions, \ + billing_mode) VALUES (?, 'b', 'gone', 'old', ?, 'AVAILABLE', 0, 3, ?, ?, \ + 'PROVISIONED')", + ) + .bind(&arn) + .bind(ACCOUNT) + .bind(r#"[{"AttributeName":"pk","KeyType":"HASH"},{"AttributeName":"sk","KeyType":"RANGE"}]"#) + .bind(r#"[{"AttributeName":"pk","AttributeType":"S"},{"AttributeName":"sk","AttributeType":"N"}]"#) + .execute(&engine.pool) + .await + .expect("old backup row"); + for i in 0..3 { + sqlx::query( + "INSERT INTO backup_items (backup_arn, pk, sk, item_data) VALUES (?, '', NULL, ?)", + ) + .bind(&arn) + .bind(format!( + r#"{{"pk":{{"S":"a"}},"sk":{{"N":"{i}"}},"v":{{"S":"x{i}"}}}}"# + )) + .execute(&engine.pool) + .await + .expect("old backup item"); + } + + crate::schema::apply(&engine.pool).await.expect("migrate"); + engine + .check_catalog_version() + .await + .expect("0.0.4 after migrate"); + let desc = engine + .restore_table_from_backup(ACCOUNT, "new", &arn) + .await + .expect("restore a 0.0.3 backup"); + assert!(index_tables(&engine, &desc.table_id).await.is_empty()); + assert_eq!( + count(&engine, &crate::data::data_table_name(&desc.table_id)).await, + 3 + ); + let pt: String = + sqlx::query_scalar("SELECT provisioned_throughput FROM tables WHERE table_id = ?") + .bind(&desc.table_id) + .fetch_one(&engine.pool) + .await + .expect("throughput"); + let pt: serde_json::Value = serde_json::from_str(&pt).expect("json"); + assert_eq!( + ( + pt["ReadCapacityUnits"].as_i64(), + pt["WriteCapacityUnits"].as_i64() + ), + (Some(5), Some(5)) + ); + } + + #[tokio::test] + async fn batches_are_cut_by_stored_bytes() { + let engine = engine().await; + sqlx::query("CREATE TABLE t (item_data TEXT NOT NULL)") + .execute(&engine.pool) + .await + .expect("table"); + // One row over the budget on its own, then small rows, then two rows + // that together exceed it. + let big = "x".repeat(super::COPY_BATCH_BYTES + 10); + let half = "y".repeat(super::COPY_BATCH_BYTES / 2 + 10); + let mut rows = vec![big.clone()]; + rows.extend((0..10).map(|i| format!("small-{i}"))); + rows.push(half.clone()); + rows.push(half.clone()); + for r in &rows { + sqlx::query("INSERT INTO t (item_data) VALUES (?)") + .bind(r) + .execute(&engine.pool) + .await + .expect("insert"); + } + let mut tx = engine.pool.begin().await.expect("tx"); + let mut last = 0; + let mut seen: Vec = Vec::new(); + let mut batches: Vec = Vec::new(); + loop { + let batch = super::next_batch(&mut tx, "t", None, last) + .await + .expect("batch"); + let Some(&(tail, _)) = batch.last() else { + break; + }; + let bytes: usize = batch.iter().map(|(_, d)| d.len()).sum(); + assert!( + batch.len() == 1 || bytes <= super::COPY_BATCH_BYTES, + "{bytes}" + ); + batches.push(batch.len()); + seen.extend(batch.into_iter().map(|(_, d)| d)); + last = tail; + } + assert_eq!(seen, rows, "every row once, in order"); + // The oversized row alone; then small rows and one half; then the other half. + assert_eq!(batches, vec![1, 11, 1]); + } + + #[tokio::test] + async fn delete_table_refuses_a_restore_in_progress() { + let engine = engine().await; + indexed_table(&engine, "src").await; + let backup = engine + .create_backup(ACCOUNT, "src", "bkp") + .await + .expect("backup"); + let desc = engine + .restore_table_from_backup(ACCOUNT, "dst", &backup.backup_arn) + .await + .expect("restore"); + // Back to the state the target is in while its copy runs. + sqlx::query( + "UPDATE tables SET table_status = 'CREATING', status_transition_at = NULL \ + WHERE table_id = ?", + ) + .bind(&desc.table_id) + .execute(&engine.pool) + .await + .expect("in progress"); + let err = extenddb_storage::TableEngine::delete_table( + &engine, + ACCOUNT, + extenddb_core::types::DeleteTableInput { + table_name: "dst".to_owned(), + }, + ) + .await + .expect_err("refused"); + assert!( + matches!(err, extenddb_storage::error::StorageError::IndexesInUse(_)), + "{err:?}" + ); + } +} diff --git a/crates/storage-sqlite/src/data/mod.rs b/crates/storage-sqlite/src/data/mod.rs index 20d5db839..b1be9b6e2 100644 --- a/crates/storage-sqlite/src/data/mod.rs +++ b/crates/storage-sqlite/src/data/mod.rs @@ -42,9 +42,9 @@ mod update_item; pub(crate) mod vector_index; pub(crate) use index::{ - PendingApplyContext, apply_pending_context, insert_index_row_multi, project_item_for_index, + PendingApplyContext, apply_pending_context, insert_index_row_multi, item_has_index_keys, + project_item_for_index, }; -pub(crate) use tx_helpers::upsert_item_in_tx; /// Quoted SQL identifier for a virtual DynamoDB table's data table. pub(crate) fn data_table_name(table_id: &str) -> String { diff --git a/crates/storage-sqlite/src/delete_table.rs b/crates/storage-sqlite/src/delete_table.rs index 4621b387d..72533abea 100644 --- a/crates/storage-sqlite/src/delete_table.rs +++ b/crates/storage-sqlite/src/delete_table.rs @@ -37,6 +37,29 @@ impl SqliteEngine { return Err(StorageError::DeletionProtected(row.table_arn.clone())); } + // A restore target (CREATING with no scheduled transition) is being + // filled by RestoreTableFromBackup. The service refuses to delete a + // table it is still creating, and deleting this one would make the + // restore fail part-way; refuse until the restore finishes. An + // abandoned target is removed by the restore sweep, not here. + if row.table_status == "CREATING" { + let restoring: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM tables WHERE table_id = ? \ + AND table_status = 'CREATING' AND status_transition_at IS NULL)", + ) + .bind(&row.table_id) + .fetch_one(&self.pool) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + if restoring { + return Err(StorageError::IndexesInUse(format!( + "Attempt to change a resource which is still in use: Table is being \ + restored: {}", + row.table_name + ))); + } + } + let index_rows: Vec = sqlx::query_as(&format!( "SELECT {INDEX_COLUMNS} FROM indexes WHERE table_id = ?" )) diff --git a/crates/storage-sqlite/src/lib.rs b/crates/storage-sqlite/src/lib.rs index 54ffc0aa6..2bc7e0e4d 100644 --- a/crates/storage-sqlite/src/lib.rs +++ b/crates/storage-sqlite/src/lib.rs @@ -241,6 +241,20 @@ fn sqlite_server_components_factory( Err(e) => tracing::error!("Failed to recover control plane transitions: {e}"), } + // Remove any restore target a prior crash left mid-copy. Before the + // server takes requests, so no restore can be in flight. + match engine.sweep_abandoned_restores().await { + Ok(names) => { + for name in &names { + tracing::warn!( + "removed table '{name}': its restore did not finish before the last \ + shutdown" + ); + } + } + Err(e) => tracing::error!("Failed to sweep abandoned restores: {e}"), + } + // Rebuild any GSI left mid-backfill (status CREATING) by a prior crash. match engine.reconcile_incomplete_gsis().await { Ok(n) if n > 0 => tracing::info!("Reconciled {n} incomplete GSI(s) at startup"), diff --git a/crates/storage-sqlite/src/schema.rs b/crates/storage-sqlite/src/schema.rs index f28532e8e..a558b777f 100644 --- a/crates/storage-sqlite/src/schema.rs +++ b/crates/storage-sqlite/src/schema.rs @@ -27,7 +27,7 @@ use sqlx::SqlitePool; /// Compiled-in catalog version. Single source of truth for the SQLite backend; /// mirrors the PostgreSQL backend's `CATALOG_VERSION`. pub const CATALOG_VERSION: extenddb_core::version::CatalogVersion = - extenddb_core::version::CatalogVersion::new(0, 0, 3); + extenddb_core::version::CatalogVersion::new(0, 0, 4); /// Complete catalog schema, applied once on a fresh database. /// @@ -375,6 +375,16 @@ CREATE TABLE IF NOT EXISTS backup_items ( CREATE INDEX IF NOT EXISTS idx_backup_items_arn ON backup_items (backup_arn); +-- A backup's table definition beyond its keys: secondary indexes, billing +-- mode and throughput, table class, encryption (catalog 0.0.4). JSON in the +-- wire's shape behind a version marker; see +-- `extenddb_storage::backup_definition`. A backup taken before 0.0.4 has no +-- row and restores as before. +CREATE TABLE IF NOT EXISTS backup_definitions ( + backup_arn TEXT PRIMARY KEY REFERENCES backups(backup_arn) ON DELETE CASCADE, + definition TEXT NOT NULL +); + -- Continuous backups / PITR status. CREATE TABLE IF NOT EXISTS continuous_backups ( account_id TEXT NOT NULL, @@ -426,7 +436,7 @@ INSERT OR IGNORE INTO seq_counters (name, value) -- never advance, so a migration could add objects and still leave the server -- refusing to start on a version mismatch. Must stay in step with -- `CATALOG_VERSION` above; they are checked against each other in a test. -INSERT INTO settings (key, value) VALUES ('catalog_version', '0.0.3') +INSERT INTO settings (key, value) VALUES ('catalog_version', '0.0.4') ON CONFLICT(key) DO UPDATE SET value = excluded.value; INSERT OR IGNORE INTO settings (key, value) VALUES ('control_plane_delay_seconds', '0.25'); INSERT OR IGNORE INTO settings (key, value) VALUES ('index_propagation_delay_ms', '10'); diff --git a/crates/storage-sqlite/testdata/schema_0_0_3.sql b/crates/storage-sqlite/testdata/schema_0_0_3.sql new file mode 100644 index 000000000..172ffbaf8 --- /dev/null +++ b/crates/storage-sqlite/testdata/schema_0_0_3.sql @@ -0,0 +1,400 @@ +-- Copyright 2026 ExtendDB contributors +-- SPDX-License-Identifier: Apache-2.0 +-- The SQLite catalog schema exactly as catalog 0.0.3 shipped it (main at +-- c178814), for the upgrade test in src/backup.rs. Do not edit. + +-- Accounts — multi-account support (REQ-AUTH-005). +CREATE TABLE IF NOT EXISTS accounts ( + account_id TEXT PRIMARY KEY, + account_name TEXT NOT NULL UNIQUE, + created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')) +); + +-- Table metadata. +CREATE TABLE IF NOT EXISTS tables ( + account_id TEXT NOT NULL REFERENCES accounts(account_id) ON DELETE CASCADE, + table_name TEXT NOT NULL, + key_schema TEXT NOT NULL, + attribute_definitions TEXT NOT NULL, + billing_mode TEXT NOT NULL DEFAULT 'PAY_PER_REQUEST', + provisioned_throughput TEXT, + stream_specification TEXT, + table_status TEXT NOT NULL DEFAULT 'CREATING', + creation_date_time TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')), + table_size_bytes INTEGER NOT NULL DEFAULT 0, + item_count INTEGER NOT NULL DEFAULT 0, + table_arn TEXT NOT NULL, + table_id TEXT NOT NULL, + ttl_attribute TEXT, + deletion_protection_enabled INTEGER NOT NULL DEFAULT 0, + status_transition_at TEXT, + stream_label TEXT, + ttl_index_ready INTEGER NOT NULL DEFAULT 0, + table_class TEXT, + sse_specification TEXT, + on_demand_throughput TEXT, + PRIMARY KEY (account_id, table_name), + CONSTRAINT tables_table_id_unique UNIQUE (table_id) +); + +CREATE INDEX IF NOT EXISTS idx_tables_pending_transition + ON tables (status_transition_at) + WHERE status_transition_at IS NOT NULL; + +-- Index metadata. index_id is always supplied by the engine (no DB-side UUID). +CREATE TABLE IF NOT EXISTS indexes ( + table_id TEXT NOT NULL, + index_id TEXT NOT NULL, + index_name TEXT NOT NULL, + index_type TEXT NOT NULL, + key_schema TEXT NOT NULL, + projection TEXT NOT NULL, + index_status TEXT NOT NULL DEFAULT 'ACTIVE', + provisioned_throughput TEXT, + propagation_delay_ms INTEGER, + PRIMARY KEY (table_id, index_name), + CONSTRAINT indexes_table_id_fkey + FOREIGN KEY (table_id) REFERENCES tables(table_id) ON DELETE CASCADE, + CONSTRAINT chk_propagation_delay_ms_non_negative + CHECK (propagation_delay_ms IS NULL OR propagation_delay_ms >= 0) +); + +-- Vector index metadata. Kept out of `indexes` deliberately: a vector index is +-- not described by a key schema, so reusing that table's `key_schema` column +-- would mean storing something meaningless in a NOT NULL column. The engine +-- supplies index_id, as it does for GSIs. +-- +-- `search_schema` is nullable because the HASH element is optional (measured +-- against the live service): with one the search is partition-scoped and +-- SearchConditionExpression is required, without one it spans the table. +-- +-- `backfilling` mirrors the measured lifecycle: false while CREATING before the +-- scan starts, true while it runs, and the member is absent once ACTIVE. Stored +-- as an integer so the ACTIVE state is representable as NULL rather than as a +-- third boolean value. +CREATE TABLE IF NOT EXISTS vector_indexes ( + table_id TEXT NOT NULL, + index_id TEXT NOT NULL, + index_name TEXT NOT NULL, + dimensions INTEGER NOT NULL, + distance_function TEXT NOT NULL, + vector_attribute TEXT NOT NULL, + search_schema TEXT, + projection TEXT NOT NULL, + index_status TEXT NOT NULL DEFAULT 'CREATING', + backfilling INTEGER, + -- Items the backfill skipped because their stored bytes cannot enter the + -- index (unparseable row, malformed or wrong-dimension vector). NULL until + -- a backfill has completed; 0 afterwards when nothing was skipped. Kept so + -- an operator can see that an ACTIVE index deliberately omits rows, rather + -- than the build looping forever on them or dying part-way. + skipped_item_count INTEGER, + PRIMARY KEY (table_id, index_name), + CONSTRAINT vector_indexes_table_id_fkey + FOREIGN KEY (table_id) REFERENCES tables(table_id) ON DELETE CASCADE, + CONSTRAINT chk_vector_dimensions_positive CHECK (dimensions > 0), + CONSTRAINT chk_vector_backfilling_bool + CHECK (backfilling IS NULL OR backfilling IN (0, 1)), + -- An ACTIVE index must not carry the member at all, which is what the + -- service does. Enforced here as well as in core, so a bug in the backend + -- cannot persist a state the wire contract forbids. + CONSTRAINT chk_vector_active_has_no_backfilling + CHECK (index_status <> 'ACTIVE' OR backfilling IS NULL) +); + +CREATE UNIQUE INDEX IF NOT EXISTS idx_vector_indexes_index_id + ON vector_indexes (index_id); + +-- Resource tags. +CREATE TABLE IF NOT EXISTS tags ( + resource_arn TEXT NOT NULL, + tag_key TEXT NOT NULL, + tag_value TEXT NOT NULL, + PRIMARY KEY (resource_arn, tag_key) +); + +-- Migration tracking. +CREATE TABLE IF NOT EXISTS schema_history ( + filename TEXT PRIMARY KEY, + applied_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')) +); + +-- Settings (catalog version, data database name, runtime config). +CREATE TABLE IF NOT EXISTS settings ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL +); + +-- Stream shards. +CREATE TABLE IF NOT EXISTS stream_shards ( + shard_id TEXT PRIMARY KEY, + table_id TEXT NOT NULL REFERENCES tables(table_id) ON DELETE CASCADE, + parent_shard_id TEXT, + starting_sequence_number TEXT NOT NULL, + ending_sequence_number TEXT, + created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')) +); + +CREATE INDEX IF NOT EXISTS idx_stream_shards_table ON stream_shards (table_id); + +-- Stream records. +CREATE TABLE IF NOT EXISTS stream_records ( + shard_id TEXT NOT NULL REFERENCES stream_shards(shard_id) ON DELETE CASCADE, + sequence_number TEXT NOT NULL, + table_id TEXT NOT NULL, + event_name TEXT NOT NULL, + record_data TEXT NOT NULL, + created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')), + PRIMARY KEY (shard_id, sequence_number) +); + +CREATE INDEX IF NOT EXISTS idx_stream_records_created ON stream_records (created_at); + +-- Admin users. +CREATE TABLE IF NOT EXISTS admin_users ( + admin_name TEXT PRIMARY KEY, + password_hash TEXT NOT NULL, + created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')) +); + +-- IAM users. +CREATE TABLE IF NOT EXISTS iam_users ( + account_id TEXT NOT NULL REFERENCES accounts(account_id) ON DELETE CASCADE, + user_name TEXT NOT NULL, + user_arn TEXT NOT NULL UNIQUE, + password_hash TEXT, + created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')), + PRIMARY KEY (account_id, user_name) +); + +-- IAM user tags. +CREATE TABLE IF NOT EXISTS iam_user_tags ( + account_id TEXT NOT NULL, + user_name TEXT NOT NULL, + tag_key TEXT NOT NULL, + tag_value TEXT NOT NULL, + PRIMARY KEY (account_id, user_name, tag_key), + FOREIGN KEY (account_id, user_name) + REFERENCES iam_users(account_id, user_name) ON DELETE CASCADE +); + +-- Access keys. +CREATE TABLE IF NOT EXISTS access_keys ( + access_key_id TEXT PRIMARY KEY, + secret_key_encrypted BLOB NOT NULL, + account_id TEXT NOT NULL, + user_name TEXT NOT NULL, + is_active INTEGER NOT NULL DEFAULT 1, + created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')), + FOREIGN KEY (account_id, user_name) + REFERENCES iam_users(account_id, user_name) ON DELETE CASCADE +); + +-- IAM groups. +CREATE TABLE IF NOT EXISTS iam_groups ( + account_id TEXT NOT NULL REFERENCES accounts(account_id) ON DELETE CASCADE, + group_name TEXT NOT NULL, + group_arn TEXT NOT NULL UNIQUE, + created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')), + PRIMARY KEY (account_id, group_name) +); + +-- IAM group membership. +CREATE TABLE IF NOT EXISTS iam_group_members ( + account_id TEXT NOT NULL, + group_name TEXT NOT NULL, + user_name TEXT NOT NULL, + PRIMARY KEY (account_id, group_name, user_name), + FOREIGN KEY (account_id, group_name) + REFERENCES iam_groups(account_id, group_name) ON DELETE CASCADE, + FOREIGN KEY (account_id, user_name) + REFERENCES iam_users(account_id, user_name) ON DELETE CASCADE +); + +-- IAM roles. +CREATE TABLE IF NOT EXISTS iam_roles ( + account_id TEXT NOT NULL REFERENCES accounts(account_id) ON DELETE CASCADE, + role_name TEXT NOT NULL, + role_arn TEXT NOT NULL UNIQUE, + trust_policy TEXT NOT NULL, + permissions_boundary_arn TEXT, + created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')), + PRIMARY KEY (account_id, role_name) +); + +-- IAM role tags. +CREATE TABLE IF NOT EXISTS iam_role_tags ( + account_id TEXT NOT NULL, + role_name TEXT NOT NULL, + tag_key TEXT NOT NULL, + tag_value TEXT NOT NULL, + PRIMARY KEY (account_id, role_name, tag_key), + FOREIGN KEY (account_id, role_name) + REFERENCES iam_roles(account_id, role_name) ON DELETE CASCADE +); + +-- IAM sessions. +CREATE TABLE IF NOT EXISTS iam_sessions ( + session_token TEXT PRIMARY KEY, + access_key_id TEXT NOT NULL UNIQUE, + secret_key_encrypted BLOB NOT NULL, + account_id TEXT NOT NULL, + role_name TEXT NOT NULL, + session_name TEXT NOT NULL, + session_tags TEXT, + session_policy TEXT, + expires_at TEXT NOT NULL, + created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')), + FOREIGN KEY (account_id, role_name) + REFERENCES iam_roles(account_id, role_name) ON DELETE CASCADE +); + +-- IAM policies. +CREATE TABLE IF NOT EXISTS iam_policies ( + account_id TEXT NOT NULL REFERENCES accounts(account_id) ON DELETE CASCADE, + principal_type TEXT NOT NULL CHECK (principal_type IN ('user', 'group', 'role')), + principal_name TEXT NOT NULL, + policy_name TEXT NOT NULL, + policy_document TEXT NOT NULL, + created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')), + PRIMARY KEY (account_id, principal_type, principal_name, policy_name) +); + +-- IAM permissions boundaries. +CREATE TABLE IF NOT EXISTS iam_permissions_boundaries ( + account_id TEXT NOT NULL REFERENCES accounts(account_id) ON DELETE CASCADE, + principal_type TEXT NOT NULL CHECK (principal_type IN ('user', 'role')), + principal_name TEXT NOT NULL, + policy_document TEXT NOT NULL, + PRIMARY KEY (account_id, principal_type, principal_name) +); + +-- Idempotency tokens for TransactWriteItems. +CREATE TABLE IF NOT EXISTS idempotency_tokens ( + account_id TEXT NOT NULL, + token TEXT NOT NULL, + fingerprint TEXT NOT NULL, + created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')), + PRIMARY KEY (account_id, token) +); + +CREATE INDEX IF NOT EXISTS idx_idempotency_tokens_created ON idempotency_tokens (created_at); + +-- Metrics (1-minute aggregation). ±9e999 store as ±Infinity in SQLite REAL, +-- matching PostgreSQL's Infinity / -Infinity seed for min / max. +CREATE TABLE IF NOT EXISTS metrics ( + bucket TEXT NOT NULL, + metric TEXT NOT NULL, + table_name TEXT NOT NULL DEFAULT '', + index_name TEXT NOT NULL DEFAULT '', + operation TEXT NOT NULL DEFAULT '', + sum REAL NOT NULL DEFAULT 0, + count INTEGER NOT NULL DEFAULT 0, + min REAL NOT NULL DEFAULT 9e999, + max REAL NOT NULL DEFAULT -9e999, + PRIMARY KEY (bucket, metric, table_name, index_name, operation) +); + +CREATE INDEX IF NOT EXISTS idx_metrics_bucket ON metrics (bucket); + +-- Login attempt tracking. +CREATE TABLE IF NOT EXISTS login_attempts ( + principal TEXT NOT NULL, + attempted_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')), + success INTEGER NOT NULL, + source_ip TEXT +); + +CREATE INDEX IF NOT EXISTS idx_login_attempts_principal_time + ON login_attempts (principal, attempted_at DESC); + +CREATE INDEX IF NOT EXISTS idx_login_attempts_source_ip_time + ON login_attempts (source_ip, attempted_at DESC) + WHERE source_ip IS NOT NULL; + +-- Backup metadata. +CREATE TABLE IF NOT EXISTS backups ( + backup_arn TEXT PRIMARY KEY, + backup_name TEXT NOT NULL, + table_id TEXT NOT NULL, + table_name TEXT NOT NULL, + account_id TEXT NOT NULL, + backup_status TEXT NOT NULL DEFAULT 'AVAILABLE', + backup_type TEXT NOT NULL DEFAULT 'USER', + backup_size_bytes INTEGER NOT NULL DEFAULT 0, + item_count INTEGER NOT NULL DEFAULT 0, + key_schema TEXT NOT NULL, + attribute_definitions TEXT NOT NULL, + billing_mode TEXT NOT NULL DEFAULT 'PAY_PER_REQUEST', + provisioned_throughput TEXT, + stream_specification TEXT, + created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')) +); + +CREATE INDEX IF NOT EXISTS idx_backups_table ON backups (account_id, table_name); + +-- Backup items. +CREATE TABLE IF NOT EXISTS backup_items ( + backup_arn TEXT NOT NULL REFERENCES backups(backup_arn) ON DELETE CASCADE, + pk TEXT NOT NULL, + sk TEXT, + item_data TEXT NOT NULL +); + +CREATE INDEX IF NOT EXISTS idx_backup_items_arn ON backup_items (backup_arn); + +-- Continuous backups / PITR status. +CREATE TABLE IF NOT EXISTS continuous_backups ( + account_id TEXT NOT NULL, + table_name TEXT NOT NULL, + pitr_enabled INTEGER NOT NULL DEFAULT 0, + earliest_restorable TEXT, + latest_restorable TEXT, + PRIMARY KEY (account_id, table_name) +); + +-- Persistent queue for async GSI propagation. A row is inserted inside the +-- base write transaction (zero crash window) and consumed by a background +-- worker once `ready_at` has passed; survives process crash/restart. +-- +-- One row per async index: each row is self-describing — `index_context` +-- carries the base key schema, attribute definitions, and the single target +-- index definition captured at enqueue, so the worker applies with zero +-- catalog reads. `worker_partition` is a stable hash of the base table key; +-- all updates to a given base item share a partition and `ready_at` is kept +-- monotonically non-decreasing within it, so the worker (which drains in `id` +-- order) preserves per-key FIFO even with randomized propagation jitter. +CREATE TABLE IF NOT EXISTS gsi_pending ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + table_id TEXT NOT NULL, + worker_partition INTEGER NOT NULL, + old_item TEXT, + new_item TEXT, + index_context TEXT NOT NULL, + ready_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')) +); + +CREATE INDEX IF NOT EXISTS idx_gsi_pending_claim + ON gsi_pending (worker_partition, ready_at, id); + +-- Monotonic counters. Replaces PostgreSQL sequences. The stream counter is +-- seeded to the current epoch in microseconds so sequence numbers are +-- time-ordered and survive restarts (mirrors the PostgreSQL setval seed). +CREATE TABLE IF NOT EXISTS seq_counters ( + name TEXT PRIMARY KEY, + value INTEGER NOT NULL DEFAULT 0 +); + +INSERT OR IGNORE INTO seq_counters (name, value) + VALUES ('stream', CAST(strftime('%s','now') AS INTEGER) * 1000000); + +-- Seed settings (mirror PostgreSQL defaults). +-- Recorded catalog version. Upserts rather than INSERT OR IGNORE: this schema is +-- re-applied by `extenddb migrate`, and with IGNORE the recorded version would +-- never advance, so a migration could add objects and still leave the server +-- refusing to start on a version mismatch. Must stay in step with +-- `CATALOG_VERSION` above; they are checked against each other in a test. +INSERT INTO settings (key, value) VALUES ('catalog_version', '0.0.3') + ON CONFLICT(key) DO UPDATE SET value = excluded.value; +INSERT OR IGNORE INTO settings (key, value) VALUES ('control_plane_delay_seconds', '0.25'); +INSERT OR IGNORE INTO settings (key, value) VALUES ('index_propagation_delay_ms', '10'); diff --git a/crates/storage/src/backup_definition.rs b/crates/storage/src/backup_definition.rs new file mode 100644 index 000000000..dbf479111 --- /dev/null +++ b/crates/storage/src/backup_definition.rs @@ -0,0 +1,316 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! The table definition a backup records, shared by the SQL backends. +//! +//! RestoreTableFromBackup reproduces what the service reproduces: key schema, +//! attribute definitions, global and local secondary indexes, billing mode and +//! provisioned throughput, table class, and encryption settings. Streams, TTL, +//! tags, deletion protection, and point-in-time recovery settings are not +//! restored, matching the service. The key schema and attribute definitions +//! keep their own catalog columns; everything else lives here. +//! +//! Stored in the wire's own shape behind a version marker, not as a copy of +//! catalog rows. A backup outlives the schema that produced it, so freezing +//! physical column names into it would let a later catalog change silently +//! alter the meaning of backups already on disk. A reader that finds a version +//! it does not know refuses the restore rather than guessing. + +use extenddb_core::types::{ + BillingMode, CreateTableInput, GsiInput, KeySchemaElement, KeyType, LsiInput, + OnDemandThroughput, ProvisionedThroughput, +}; +use serde::{Deserialize, Serialize}; + +use crate::error::StorageError; + +/// Items a backup or a restore buffers before writing them out, at most. +pub const COPY_BATCH_ITEMS: usize = 500; + +/// Bytes of stored item JSON a backup or a restore buffers before writing +/// them out, at most (one item over the budget is still taken whole). Counted +/// on the stored text, not the DynamoDB item size, because the text is what +/// is held: a 400 KB item can be several MB of JSON (a list of booleans is +/// about seven times its DynamoDB size). +pub const COPY_BATCH_BYTES: usize = 4 * 1024 * 1024; + +/// Refuse to restore a table whose base key has more than one HASH or more +/// than one RANGE attribute. +/// +/// Multi-part base keys are an opt-in preview (`enable_multipart_keys`), and +/// the item read and write paths of the SQL backends address only the first +/// HASH and first RANGE attribute of a base table. A restore that laid such a +/// table out by its full key would produce rows those paths cannot find; one +/// that followed the read paths would collapse distinct items. Refusing is the +/// only result that is not silently wrong. +/// +/// # Errors +/// +/// [`StorageError::Unsupported`] for a multi-part base key. +pub fn ensure_single_part_base_key( + key_schema: &[KeySchemaElement], + backup_arn: &str, +) -> Result<(), StorageError> { + let hashes = key_schema + .iter() + .filter(|k| k.key_type == KeyType::Hash) + .count(); + let ranges = key_schema + .iter() + .filter(|k| k.key_type == KeyType::Range) + .count(); + if hashes <= 1 && ranges <= 1 { + return Ok(()); + } + Err(StorageError::Unsupported(format!( + "backup {backup_arn} is of a table with a multi-part key ({hashes} HASH, {ranges} \ + RANGE attributes); restoring it is not supported" + ))) +} + +/// The only definition version this build writes and reads. +pub const BACKUP_DEFINITION_VERSION: u32 = 1; + +/// The parts of a table's definition a restore recreates, beyond its keys. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct BackupTableDefinition { + #[serde(rename = "Version")] + pub version: u32, + /// `PROVISIONED` or `PAY_PER_REQUEST`. + #[serde(rename = "BillingMode")] + pub billing_mode: String, + #[serde( + rename = "ProvisionedThroughput", + default, + skip_serializing_if = "Option::is_none" + )] + pub provisioned_throughput: Option, + #[serde(rename = "GlobalSecondaryIndexes", default)] + pub global_secondary_indexes: Vec, + #[serde(rename = "LocalSecondaryIndexes", default)] + pub local_secondary_indexes: Vec, + #[serde( + rename = "TableClass", + default, + skip_serializing_if = "Option::is_none" + )] + pub table_class: Option, + #[serde( + rename = "SSESpecification", + default, + skip_serializing_if = "Option::is_none" + )] + pub sse_specification: Option, + #[serde( + rename = "OnDemandThroughput", + default, + skip_serializing_if = "Option::is_none" + )] + pub on_demand_throughput: Option, + /// Names of the source table's vector indexes. Restore refuses a backup + /// that has any, rather than producing a table without them. + #[serde(rename = "VectorIndexNames", default)] + pub vector_index_names: Vec, +} + +impl BackupTableDefinition { + /// Parse a stored definition, refusing a version this build cannot read. + /// + /// # Errors + /// + /// [`StorageError::Unsupported`] for an unknown version, and + /// [`StorageError::Internal`] for a document that does not parse. + pub fn from_json(value: serde_json::Value, backup_arn: &str) -> Result { + let version = value.get("Version").and_then(serde_json::Value::as_u64); + if version != Some(u64::from(BACKUP_DEFINITION_VERSION)) { + return Err(StorageError::Unsupported(format!( + "backup {backup_arn} carries a table definition this build cannot read \ + (version {version:?})" + ))); + } + serde_json::from_value(value).map_err(|e| { + StorageError::Internal(format!("backup {backup_arn} table definition: {e}")) + }) + } + + /// Serialize for storage. + /// + /// # Errors + /// + /// [`StorageError::Internal`] if serialization fails. + pub fn to_json(&self) -> Result { + serde_json::to_value(self).map_err(|e| StorageError::Internal(e.to_string())) + } + + /// Refuse a definition the restore cannot reproduce in full. + /// + /// # Errors + /// + /// [`StorageError::Unsupported`] if the source table had vector indexes. + pub fn ensure_restorable(&self, backup_arn: &str) -> Result<(), StorageError> { + if self.vector_index_names.is_empty() { + return Ok(()); + } + Err(StorageError::Unsupported(format!( + "backup {backup_arn} has {} vector index(es); restoring a table with vector \ + indexes is not supported by this storage backend", + self.vector_index_names.len() + ))) + } + + /// Apply the definition to a restore target's `CreateTableInput`. + pub fn apply_to(self, input: &mut CreateTableInput) { + let on_demand = self.billing_mode == "PAY_PER_REQUEST"; + input.billing_mode = Some(if on_demand { + BillingMode::PayPerRequest + } else { + BillingMode::Provisioned + }); + input.provisioned_throughput = self.provisioned_throughput; + input.global_secondary_indexes = + (!self.global_secondary_indexes.is_empty()).then_some(self.global_secondary_indexes); + input.local_secondary_indexes = + (!self.local_secondary_indexes.is_empty()).then_some(self.local_secondary_indexes); + input.table_class = self.table_class; + input.sse_specification = self.sse_specification; + input.on_demand_throughput = self.on_demand_throughput; + } +} + +/// Throughput stored in a catalog `provisioned_throughput` column, which holds +/// either the request shape or the description shape depending on the writer. +/// Both carry the two capacity members under the same names. +#[must_use] +pub fn throughput_from_catalog(value: Option<&serde_json::Value>) -> Option { + let v = value?; + let read = v.get("ReadCapacityUnits")?.as_i64()?; + let write = v.get("WriteCapacityUnits")?.as_i64()?; + Some(ProvisionedThroughput { + read_capacity_units: read, + write_capacity_units: write, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use extenddb_core::types::{KeySchemaElement, KeyType, Projection, ProjectionType}; + + fn sample() -> BackupTableDefinition { + BackupTableDefinition { + version: BACKUP_DEFINITION_VERSION, + billing_mode: "PROVISIONED".to_owned(), + provisioned_throughput: Some(ProvisionedThroughput { + read_capacity_units: 7, + write_capacity_units: 9, + }), + global_secondary_indexes: vec![GsiInput { + index_name: "g".to_owned(), + key_schema: vec![KeySchemaElement { + attribute_name: "gpk".to_owned(), + key_type: KeyType::Hash, + }], + projection: Projection { + projection_type: ProjectionType::KeysOnly, + non_key_attributes: None, + }, + provisioned_throughput: Some(ProvisionedThroughput { + read_capacity_units: 3, + write_capacity_units: 4, + }), + }], + local_secondary_indexes: Vec::new(), + table_class: Some("STANDARD_INFREQUENT_ACCESS".to_owned()), + sse_specification: None, + on_demand_throughput: None, + vector_index_names: Vec::new(), + } + } + + #[test] + fn refuses_multi_part_base_keys() { + let k = |n: &str, t| KeySchemaElement { + attribute_name: n.to_owned(), + key_type: t, + }; + assert!(ensure_single_part_base_key(&[k("a", KeyType::Hash)], "arn").is_ok()); + assert!( + ensure_single_part_base_key(&[k("a", KeyType::Hash), k("b", KeyType::Range)], "arn") + .is_ok() + ); + for ks in [ + vec![k("a", KeyType::Hash), k("b", KeyType::Hash)], + vec![ + k("a", KeyType::Hash), + k("b", KeyType::Range), + k("c", KeyType::Range), + ], + ] { + assert!(matches!( + ensure_single_part_base_key(&ks, "arn"), + Err(StorageError::Unsupported(_)) + )); + } + } + + #[test] + fn round_trips_through_json() { + let def = sample(); + let back = BackupTableDefinition::from_json(def.to_json().unwrap(), "arn").unwrap(); + assert_eq!(back, def); + } + + #[test] + fn refuses_an_unknown_version() { + let mut json = sample().to_json().unwrap(); + json["Version"] = serde_json::json!(2); + let err = BackupTableDefinition::from_json(json, "arn").unwrap_err(); + assert!(matches!(err, StorageError::Unsupported(_)), "{err:?}"); + } + + #[test] + fn refuses_vector_indexes() { + let mut def = sample(); + def.vector_index_names.push("v".to_owned()); + assert!(matches!( + def.ensure_restorable("arn"), + Err(StorageError::Unsupported(_)) + )); + } + + #[test] + fn applies_to_a_create_input() { + let mut input = CreateTableInput::default(); + sample().apply_to(&mut input); + assert_eq!(input.billing_mode, Some(BillingMode::Provisioned)); + assert_eq!( + input + .provisioned_throughput + .as_ref() + .map(|p| p.read_capacity_units), + Some(7) + ); + assert_eq!( + input.global_secondary_indexes.as_ref().map(Vec::len), + Some(1) + ); + assert!(input.local_secondary_indexes.is_none()); + assert_eq!( + input.table_class.as_deref(), + Some("STANDARD_INFREQUENT_ACCESS") + ); + } + + #[test] + fn reads_both_catalog_throughput_shapes() { + let input = serde_json::json!({"ReadCapacityUnits": 5, "WriteCapacityUnits": 6}); + let desc = serde_json::json!({ + "ReadCapacityUnits": 5, "WriteCapacityUnits": 6, "NumberOfDecreasesToday": 0 + }); + for v in [input, desc] { + let pt = throughput_from_catalog(Some(&v)).unwrap(); + assert_eq!((pt.read_capacity_units, pt.write_capacity_units), (5, 6)); + } + assert!(throughput_from_catalog(None).is_none()); + } +} diff --git a/crates/storage/src/lib.rs b/crates/storage/src/lib.rs index 607c1fafc..7159de963 100755 --- a/crates/storage/src/lib.rs +++ b/crates/storage/src/lib.rs @@ -9,6 +9,7 @@ pub mod authorization_store; pub mod backend; +pub mod backup_definition; pub mod bootstrapper; pub mod config; pub mod diagnostics; diff --git a/docs/getting-started.md b/docs/getting-started.md index ffbbc9b3f..7318252e9 100755 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -199,7 +199,7 @@ You should see all checks pass: --- Checking catalog connection... OK: Connected to catalog. --- Checking catalog version... - OK: Catalog version 0.0.3 + OK: Catalog version 0.0.4 --- Checking data database... OK: Connected to data database 'extenddb_catalog_data'. --- Enumerating tables... @@ -215,7 +215,7 @@ extenddb runs as a daemon (background process) and logs to syslog. On startup it ```bash ./target/release/extenddb serve --config extenddb.toml -# extenddb 0.1.13 (catalog 0.0.3) starting on 127.0.0.1:18443 +# extenddb 0.1.13 (catalog 0.0.4) starting on 127.0.0.1:18443 # storage: postgres (postgresql://extenddb:***@localhost:5432/extenddb_catalog) ``` @@ -1294,7 +1294,7 @@ Each runner requires its tools to be installed. The runner checks prerequisites ```bash ./target/release/extenddb version # extenddb 0.1.13 -# catalog 0.0.3 (postgres) +# catalog 0.0.4 (postgres) # commit abc1234 # built 2026-04-17T12:00:00Z ``` diff --git a/docs/manuals/01-architecture-guide.md b/docs/manuals/01-architecture-guide.md index bd40427c9..43a1d07e6 100755 --- a/docs/manuals/01-architecture-guide.md +++ b/docs/manuals/01-architecture-guide.md @@ -172,7 +172,7 @@ extenddb uses a dual-database architecture: - **Catalog database** (e.g., `extenddb_catalog`): Stores table metadata, account/user/group/role/policy definitions, access keys, settings, stream metadata, and metrics. Shared across all accounts. - **Data database** (e.g., `extenddb_catalog_data`): Stores user items, GSI/LSI data, and stream records. Each table gets its own PostgreSQL table. -The catalog version is 0.0.3 on PostgreSQL and SQLite, stored in the `settings` table under the key `catalog_version` and checked at startup. +The catalog version is 0.0.4 on PostgreSQL and SQLite, stored in the `settings` table under the key `catalog_version` and checked at startup. The MongoDB backend tracks its own catalog version, 0.0.2. A mismatch between the compiled-in version and the stored one prevents the server from starting; run `extenddb migrate` to upgrade. diff --git a/docs/manuals/04-quickstart-setup-guide.md b/docs/manuals/04-quickstart-setup-guide.md index b3d8c53f1..cb5926a09 100755 --- a/docs/manuals/04-quickstart-setup-guide.md +++ b/docs/manuals/04-quickstart-setup-guide.md @@ -130,7 +130,7 @@ Check the version: ```bash ./target/release/extenddb version # extenddb 0.1.13 -# catalog 0.0.3 (postgres) +# catalog 0.0.4 (postgres) # commit abc1234 # built 2026-04-17T12:00:00Z ``` @@ -171,7 +171,7 @@ Expected output: --- Checking catalog connection... OK: Connected to catalog. --- Checking catalog version... - OK: Catalog version 0.0.3 + OK: Catalog version 0.0.4 --- Checking data database... OK: Connected to data database 'extenddb_catalog_data'. --- Enumerating tables... diff --git a/docs/manuals/07-upgrade-manual.md b/docs/manuals/07-upgrade-manual.md index fe2bf5eb2..2cc7c0b99 100755 --- a/docs/manuals/07-upgrade-manual.md +++ b/docs/manuals/07-upgrade-manual.md @@ -4,9 +4,9 @@ ## Current Status -Catalog 0.0.3 is current. The 0.0.2 to 0.0.3 upgrade is the first in-place catalog upgrade ExtendDB has, and **every existing PostgreSQL deployment must run it**, including deployments that never use vector indexes: the server refuses to start against a catalog version it was not built for. +Catalog 0.0.4 is current. **Every existing PostgreSQL and SQLite deployment must run `extenddb migrate`** to reach it: the server refuses to start against a catalog version it was not built for. -See [Catalog 0.0.3](#catalog-003-current) below for what changes and the exact sequence. +See [Catalog 0.0.4](#catalog-004-current) below for what changes and the exact sequence. ## How Catalog Upgrades Work @@ -16,7 +16,8 @@ Migrations are SQL files in `crates/storage-postgres/migrations/`, applied in fi ``` 001_schema.sql ← the complete initial schema -002_vector_indexes.sql ← vector index metadata, catalog 0.0.3 +002_vector_indexes.sql ← vector index metadata, catalog 0.0.3 +003_backup_definitions.sql ← backup table definitions, catalog 0.0.4 ``` The `schema_history` table tracks which files have been applied. When `extenddb migrate` runs, it: @@ -181,7 +182,25 @@ psql -d extenddb_catalog -f catalog_backup_YYYYMMDD.sql ## Version History -### Catalog 0.0.3 (Current) +### Catalog 0.0.4 (Current) + +Adds the table definition a backup records: + +- New `backup_definitions` table: one row per backup, holding the source table's global and local secondary indexes, billing mode and provisioned throughput, table class, and encryption settings. RestoreTableFromBackup recreates them. A backup taken before this upgrade has no row and restores as before, with its keys and items but no secondary indexes. + +Upgrade sequence, on PostgreSQL and SQLite alike: + +```bash +extenddb stop --config extenddb.toml +extenddb migrate --yes --config extenddb.toml +extenddb serve --config extenddb.toml +``` + +Run `extenddb migrate` without `--yes` first to see what is pending; it reports `catalog 0.0.3 -> 0.0.4` and changes nothing. + +The upgrade is not reversible in place: a 0.0.3 binary refuses to start against a 0.0.4 catalog. Roll back by restoring the catalog backup taken before the upgrade, as described above. Backups live in the catalog, so that restore also removes every backup created after the upgrade. + +### Catalog 0.0.3 Adds vector index metadata: diff --git a/docs/manuals/08-install-linux.md b/docs/manuals/08-install-linux.md index 1512201bd..7414ab006 100755 --- a/docs/manuals/08-install-linux.md +++ b/docs/manuals/08-install-linux.md @@ -118,7 +118,7 @@ Expected: ``` === extenddb verify === ... - OK: Catalog version 0.0.3 + OK: Catalog version 0.0.4 ... === HEALTHY: All checks passed === ``` diff --git a/docs/manuals/09-install-macos.md b/docs/manuals/09-install-macos.md index 4e2d87e8d..feec0a409 100755 --- a/docs/manuals/09-install-macos.md +++ b/docs/manuals/09-install-macos.md @@ -96,7 +96,7 @@ Expected: ``` === extenddb verify === ... - OK: Catalog version 0.0.3 + OK: Catalog version 0.0.4 ... === HEALTHY: All checks passed === ``` diff --git a/tests/test_backup_restore_fidelity.py b/tests/test_backup_restore_fidelity.py index 9d00652a2..57895649c 100644 --- a/tests/test_backup_restore_fidelity.py +++ b/tests/test_backup_restore_fidelity.py @@ -3,9 +3,11 @@ """Backup and restore reproduce the table they were taken from. -Every backend, every key shape. A restore must reach ACTIVE with the source's -key schema, attribute definitions, and exactly the source's items, attribute -for attribute. Comparing whole items rather than counts is what catches a +Every backend; hash-only keys and composite keys with S, N, and B sort keys, +plus secondary indexes and provisioned throughput. A restore must reach ACTIVE +with the source's key schema, attribute definitions, and exactly the source's +items, attribute for attribute. (Multi-part base keys are a preview that +restore refuses; that refusal is tested at the storage level.) Comparing whole items rather than counts is what catches a backend that restores the right number of rows under the wrong keys. Readiness assessment P0-2: on PostgreSQL every table with a sort key backed up @@ -49,15 +51,23 @@ def _create_backup(client, table_name: str) -> str: raise time.sleep(0.5) deadline = time.monotonic() + BACKUP_TIMEOUT_S - while True: - status = client.describe_backup(BackupArn=arn)["BackupDescription"][ - "BackupDetails" - ]["BackupStatus"] - if status == "AVAILABLE": - return arn - assert status == "CREATING", f"backup {arn} is {status}" - assert time.monotonic() < deadline, f"backup {arn} not AVAILABLE in time" - time.sleep(_poll_interval()) + try: + while True: + status = client.describe_backup(BackupArn=arn)["BackupDescription"][ + "BackupDetails" + ]["BackupStatus"] + if status == "AVAILABLE": + return arn + assert status == "CREATING", f"backup {arn} is {status}" + assert time.monotonic() < deadline, f"backup {arn} not AVAILABLE in time" + time.sleep(_poll_interval()) + except BaseException: + # The caller never learns the ARN, so the backup is removed here. + try: + client.delete_backup(BackupArn=arn) + except ClientError: + pass + raise def _scan_all(client, table_name: str) -> list[dict]: @@ -106,11 +116,21 @@ def sort_sets(av): def _drop(client, *names: str) -> None: + """Delete tables, waiting out ones still being created or restored.""" for name in names: - try: - client.delete_table(TableName=name) - except client.exceptions.ResourceNotFoundException: - continue + deadline = time.monotonic() + RESTORE_TIMEOUT_S + while True: + try: + client.delete_table(TableName=name) + break + except client.exceptions.ResourceNotFoundException: + break + except client.exceptions.ResourceInUseException: + # The service refuses to delete a table that is CREATING, + # which a restore target is until the restore finishes. + if time.monotonic() >= deadline: + raise + time.sleep(_poll_interval() * 25) wait_for_deleted(client, name) @@ -123,11 +143,19 @@ def _sort_value(kind: str, i: int) -> dict: return {"B": i.to_bytes(2, "big") + b"\x00\xff"} -def _item(i: int, sort_kind: str | None) -> dict: +def _hash_value(kind: str, n: int) -> dict: + if kind == "S": + return {"S": f"part-{n}"} + if kind == "N": + return {"N": str(n * 7 - 100)} + return {"B": b"\x00p" + n.to_bytes(2, "big")} + + +def _item(i: int, sort_kind: str | None, hash_kind: str = "S") -> dict: # With a sort key, several items share a partition; without one, every # partition key must be distinct or the puts overwrite each other. item: dict = { - "pk": {"S": f"part-{i % 7}" if sort_kind else f"part-{i}"}, + "pk": _hash_value(hash_kind, i % 7 if sort_kind else i), "str": {"S": f"value-{i}"}, "num": {"N": str(i)}, "nested": {"M": {"list": {"L": [{"N": "1"}, {"S": "two"}, {"BOOL": i % 2 == 0}]}}}, @@ -141,25 +169,71 @@ def _item(i: int, sort_kind: str | None) -> dict: return item -def _round_trip(client, sort_kind: str | None) -> None: +def _query_all(client, **kwargs) -> list[dict]: + items: list[dict] = [] + while True: + resp = client.query(**kwargs) + items += resp["Items"] + if "LastEvaluatedKey" not in resp: + return items + kwargs["ExclusiveStartKey"] = resp["LastEvaluatedKey"] + + +def _index_items(client, table: str, index: str, hash_attr: str, values: list) -> list[str]: + """Canonical items an index serves, across every listed partition. + + Values are strings for an S hash key, or typed attribute values. + """ + out: list[str] = [] + for v in values: + out += [ + _canonical(i) + for i in _query_all( + client, + TableName=table, + IndexName=index, + KeyConditionExpression="#h = :v", + ExpressionAttributeNames={"#h": hash_attr}, + ExpressionAttributeValues={":v": v if isinstance(v, dict) else {"S": v}}, + ) + ] + return sorted(out) + + +def _cleanup(client, backup_arn: str | None, *tables: str) -> None: + """Delete tables and the backup; each step runs even if an earlier one fails.""" + if tables: + try: + _drop(client, tables[0]) + finally: + _cleanup(client, backup_arn, *tables[1:]) + return + if backup_arn: + try: + client.delete_backup(BackupArn=backup_arn) + except client.exceptions.BackupNotFoundException: + pass + + +def _round_trip(client, sort_kind: str | None, hash_kind: str = "S") -> None: source = f"restore-fid-{uuid.uuid4().hex[:10]}" restored = f"{source}-r" key_schema = [{"AttributeName": "pk", "KeyType": "HASH"}] - attr_defs = [{"AttributeName": "pk", "AttributeType": "S"}] + attr_defs = [{"AttributeName": "pk", "AttributeType": hash_kind}] if sort_kind: key_schema.append({"AttributeName": "sk", "KeyType": "RANGE"}) attr_defs.append({"AttributeName": "sk", "AttributeType": sort_kind}) - client.create_table( - TableName=source, - KeySchema=key_schema, - AttributeDefinitions=attr_defs, - BillingMode="PAY_PER_REQUEST", - ) backup_arn = None try: + client.create_table( + TableName=source, + KeySchema=key_schema, + AttributeDefinitions=attr_defs, + BillingMode="PAY_PER_REQUEST", + ) wait_for_active(client, source) for i in range(ITEMS_PER_TABLE): - client.put_item(TableName=source, Item=_item(i, sort_kind)) + client.put_item(TableName=source, Item=_item(i, sort_kind, hash_kind)) # Compare against what the source serves, not what was sent: numbers # come back canonicalized (`-42.0` reads as `-42`), here as in DynamoDB. expected = _scan_all(client, source) @@ -175,6 +249,7 @@ def _round_trip(client, sort_kind: str | None) -> None: assert sorted(table["AttributeDefinitions"], key=lambda a: a["AttributeName"]) == sorted( attr_defs, key=lambda a: a["AttributeName"] ) + assert table["BillingModeSummary"]["BillingMode"] == "PAY_PER_REQUEST" got = sorted(_canonical(i) for i in _scan_all(client, restored)) want = sorted(_canonical(i) for i in expected) @@ -191,26 +266,216 @@ def _round_trip(client, sort_kind: str | None) -> None: assert _canonical(fetched.get("Item", {})) == _canonical(item) # The restored table takes writes like any other. - client.put_item(TableName=restored, Item=_item(ITEMS_PER_TABLE, sort_kind)) + client.put_item( + TableName=restored, Item=_item(ITEMS_PER_TABLE, sort_kind, hash_kind) + ) finally: - # Each cleanup step runs even if an earlier one fails. - try: - _drop(client, restored) - finally: - try: - _drop(client, source) - finally: - if backup_arn: - try: - client.delete_backup(BackupArn=backup_arn) - except client.exceptions.BackupNotFoundException: - pass + _cleanup(client, backup_arn, restored, source) -def test_restore_hash_only_table(dynamodb_client): - _round_trip(dynamodb_client, None) +@pytest.mark.parametrize("hash_kind", ["S", "N", "B"]) +def test_restore_hash_only_table(dynamodb_client, hash_kind): + _round_trip(dynamodb_client, None, hash_kind) @pytest.mark.parametrize("sort_kind", ["S", "N", "B"]) def test_restore_composite_key_table(dynamodb_client, sort_kind): _round_trip(dynamodb_client, sort_kind) + + +def _gsi(name: str, hash_attr: str, range_attr: str | None, projection: dict, **extra) -> dict: + ks = [{"AttributeName": hash_attr, "KeyType": "HASH"}] + if range_attr: + ks.append({"AttributeName": range_attr, "KeyType": "RANGE"}) + return {"IndexName": name, "KeySchema": ks, "Projection": projection, **extra} + + +def _index_by_name(indexes: list[dict]) -> dict: + return {i["IndexName"]: i for i in indexes} + + +def test_restore_preserves_indexes_and_provisioned_throughput(dynamodb_client): + """GSIs and LSIs, their key schemas, projections, and throughput, the + table's provisioned throughput, and every index's contents come back.""" + client = dynamodb_client + source = f"restore-idx-{uuid.uuid4().hex[:10]}" + restored = f"{source}-r" + gsi_tp = {"ReadCapacityUnits": 3, "WriteCapacityUnits": 4} + gsis = [ + _gsi("by_owner", "owner", "rank", {"ProjectionType": "ALL"}, ProvisionedThroughput=gsi_tp), + _gsi( + "by_kind", + "kind", + None, + {"ProjectionType": "INCLUDE", "NonKeyAttributes": ["note"]}, + ProvisionedThroughput=gsi_tp, + ), + _gsi("by_kind_keys", "kind", "rank", {"ProjectionType": "KEYS_ONLY"}, + ProvisionedThroughput=gsi_tp), + _gsi("by_code", "code", "tag", {"ProjectionType": "ALL"}, + ProvisionedThroughput=gsi_tp), + ] + lsis = [ + { + "IndexName": "by_rank", + "KeySchema": [ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "rank", "KeyType": "RANGE"}, + ], + "Projection": {"ProjectionType": "ALL"}, + } + ] + backup_arn = None + try: + client.create_table( + TableName=source, + KeySchema=[ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"}, + ], + AttributeDefinitions=[ + {"AttributeName": "pk", "AttributeType": "S"}, + {"AttributeName": "sk", "AttributeType": "S"}, + {"AttributeName": "owner", "AttributeType": "S"}, + {"AttributeName": "kind", "AttributeType": "S"}, + {"AttributeName": "rank", "AttributeType": "N"}, + {"AttributeName": "code", "AttributeType": "N"}, + {"AttributeName": "tag", "AttributeType": "B"}, + ], + ProvisionedThroughput={"ReadCapacityUnits": 7, "WriteCapacityUnits": 9}, + GlobalSecondaryIndexes=gsis, + LocalSecondaryIndexes=lsis, + ) + wait_for_active(client, source) + owners = [f"owner-{i}" for i in range(3)] + kinds = ["red", "blue"] + for i in range(ITEMS_PER_TABLE): + item = { + "pk": {"S": f"part-{i % 4}"}, + "sk": {"S": f"sort-{i:04d}"}, + # Unique, so every ordered index read is fully determined. + "rank": {"N": str(i * 3 - 50)}, + "code": {"N": str(i % 5 - 2)}, + "tag": {"B": bytes([255 - i, i])}, + "note": {"S": f"note-{i}"}, + "other": {"S": "not projected into by_kind"}, + } + # Sparse: some items lack one index key or the other. + if i % 3: + item["owner"] = {"S": owners[i % 3]} + if i % 4: + item["kind"] = {"S": kinds[i % 2]} + client.put_item(TableName=source, Item=item) + + # What each index serves on the source, read after the source's GSIs + # have caught up with the writes. + def index_view(table: str) -> dict: + return { + "by_owner": _index_items(client, table, "by_owner", "owner", owners), + "by_kind": _index_items(client, table, "by_kind", "kind", kinds), + "by_kind_keys": _index_items(client, table, "by_kind_keys", "kind", kinds), + "by_rank": _index_items(client, table, "by_rank", "pk", + [f"part-{p}" for p in range(4)]), + "by_code": _index_items(client, table, "by_code", "code", + [{"N": str(c)} for c in range(-2, 3)]), + } + + expected_sizes = { + "by_owner": sum(1 for i in range(ITEMS_PER_TABLE) if i % 3), + "by_kind": sum(1 for i in range(ITEMS_PER_TABLE) if i % 4), + "by_kind_keys": sum(1 for i in range(ITEMS_PER_TABLE) if i % 4), + "by_rank": ITEMS_PER_TABLE, + "by_code": ITEMS_PER_TABLE, + } + deadline = time.monotonic() + 120 + while True: + want = index_view(source) + sizes = {k: len(v) for k, v in want.items()} + if sizes == expected_sizes: + break + assert time.monotonic() < deadline, f"source GSIs did not converge: {sizes}" + time.sleep(_poll_interval() * 10) + + backup_arn = _create_backup(client, source) + client.restore_table_from_backup(TargetTableName=restored, BackupArn=backup_arn) + wait_for_active(client, restored, timeout=RESTORE_TIMEOUT_S) + + table = client.describe_table(TableName=restored)["Table"] + assert table["ProvisionedThroughput"]["ReadCapacityUnits"] == 7 + assert table["ProvisionedThroughput"]["WriteCapacityUnits"] == 9 + got_gsis = _index_by_name(table.get("GlobalSecondaryIndexes", [])) + assert sorted(got_gsis) == sorted(g["IndexName"] for g in gsis) + for g in gsis: + r = got_gsis[g["IndexName"]] + assert r["KeySchema"] == g["KeySchema"] + assert r["Projection"] == g["Projection"] + assert r["ProvisionedThroughput"]["ReadCapacityUnits"] == 3 + assert r["ProvisionedThroughput"]["WriteCapacityUnits"] == 4 + for g in gsis: + assert got_gsis[g["IndexName"]]["IndexStatus"] == "ACTIVE" + got_lsis = _index_by_name(table.get("LocalSecondaryIndexes", [])) + assert sorted(got_lsis) == ["by_rank"] + assert got_lsis["by_rank"]["KeySchema"] == lsis[0]["KeySchema"] + assert got_lsis["by_rank"]["Projection"] == lsis[0]["Projection"] + + # The service backfills GSIs after the table turns ACTIVE; poll. + deadline = time.monotonic() + RESTORE_TIMEOUT_S + while True: + got = index_view(restored) + if got == want: + break + assert time.monotonic() < deadline, { + k: (len(got[k]), len(want[k])) for k in want + } + time.sleep(_poll_interval() * 10) + + # Projections checked against their definitions, not just against the + # source table, so a projection bug both tables share still fails. + def attr_names(index: str, hash_attr: str, value: dict) -> set[frozenset]: + return { + frozenset(i) + for i in _query_all( + client, + TableName=restored, + IndexName=index, + KeyConditionExpression="#h = :v", + ExpressionAttributeNames={"#h": hash_attr}, + ExpressionAttributeValues={":v": value}, + ) + } + + assert attr_names("by_kind", "kind", {"S": "blue"}) == { + frozenset({"pk", "sk", "kind", "note"}) + } + assert attr_names("by_kind_keys", "kind", {"S": "blue"}) == { + frozenset({"pk", "sk", "kind", "rank"}) + } + all_attrs = attr_names("by_owner", "owner", {"S": owners[1]}) + assert all_attrs and all({"other", "note", "rank", "tag"} <= a for a in all_attrs) + + # Ordered, ranged reads on a numeric index sort key come back in the + # same order from both tables, in both directions. + for forward in (True, False): + def ranked(table: str) -> list[str]: + return [ + _canonical(i) + for i in _query_all( + client, + TableName=table, + IndexName="by_owner", + KeyConditionExpression="#o = :o AND #r BETWEEN :lo AND :hi", + ExpressionAttributeNames={"#o": "owner", "#r": "rank"}, + ExpressionAttributeValues={ + ":o": {"S": owners[1]}, + ":lo": {"N": "-20"}, + ":hi": {"N": "100"}, + }, + ScanIndexForward=forward, + ) + ] + + src_order = ranked(source) + assert src_order + assert ranked(restored) == src_order + finally: + _cleanup(client, backup_arn, restored, source) diff --git a/tests/test_cli_vector_catalog_migration.py b/tests/test_cli_vector_catalog_migration.py index 1b260e680..e50f9f175 100644 --- a/tests/test_cli_vector_catalog_migration.py +++ b/tests/test_cli_vector_catalog_migration.py @@ -1,7 +1,10 @@ # Copyright 2026 ExtendDB contributors # SPDX-License-Identifier: Apache-2.0 -"""Catalog migration tests for the vector index schema (catalog 0.0.2 -> 0.0.3). +"""Catalog migration tests against a live PostgreSQL deployment. + +The vector index schema (catalog 0.0.3) and the backup table definitions +(catalog 0.0.4) are exercised through an upgrade from the pre-vector shape. These exercise the migration against a live deployment rather than asserting the SQL text: init one, roll its catalog back to the pre-vector shape, and check that @@ -29,6 +32,14 @@ ) VECTOR_MIGRATION = "002_vector_indexes.sql" +BACKUP_DEFINITIONS_MIGRATION = "003_backup_definitions.sql" + +# The version the binary expects, written by the last catalog migration. Every +# current-version assertion reads it from here so the next bump is one edit. +CURRENT_CATALOG_VERSION = "0.0.4" + +# A version no release has reached, for the symmetric version-gate test. +FUTURE_CATALOG_VERSION = "0.0.5" def _catalog_conn(cli_env): @@ -67,6 +78,12 @@ def _backups_has_vector_column(cli_env): ) +def _backup_definitions_table_exists(cli_env): + return _catalog_query( + cli_env, "SELECT to_regclass('public.backup_definitions') IS NOT NULL" + ) + + def _init(cli_env): result = _run_extenddb( "init", *_init_args(cli_env), @@ -80,17 +97,19 @@ def _roll_catalog_back_to_pre_vector(cli_env): """Turn a freshly initialised catalog into the shape 0.0.2 deployments have. Reproduces an upgrade rather than a fresh install, which is the case that - matters: a fresh install applies both migrations in order and reaches the same + matters: a fresh install applies the migrations in order and reaches the same end state trivially. """ conn = _catalog_conn(cli_env) try: conn.autocommit = True with conn.cursor() as cur: + cur.execute("DROP TABLE IF EXISTS backup_definitions") cur.execute("DROP TABLE IF EXISTS vector_indexes") cur.execute("ALTER TABLE backups DROP COLUMN IF EXISTS vector_indexes") cur.execute( - "DELETE FROM schema_history WHERE filename = %s", (VECTOR_MIGRATION,) + "DELETE FROM schema_history WHERE filename IN (%s, %s)", + (VECTOR_MIGRATION, BACKUP_DEFINITIONS_MIGRATION), ) cur.execute("UPDATE settings SET value = '0.0.2' WHERE key = 'catalog_version'") finally: @@ -98,10 +117,10 @@ def _roll_catalog_back_to_pre_vector(cli_env): class TestVectorCatalogMigration: - """Catalog version 0.0.3: the vector index table and the backup snapshot.""" + """Migrations from the pre-vector catalog shape to the current version.""" - def test_init_creates_the_vector_catalog_at_0_0_3(self, cli_env): - """A fresh init applies both catalog migrations and records the version. + def test_init_creates_the_catalog_at_the_current_version(self, cli_env): + """A fresh init applies every catalog migration and records the version. The version is what the binary checks at startup, so a migration that creates the table without moving the version, or the reverse, would leave @@ -109,9 +128,10 @@ def test_init_creates_the_vector_catalog_at_0_0_3(self, cli_env): """ _init(cli_env) - assert _catalog_version(cli_env) == "0.0.3" + assert _catalog_version(cli_env) == CURRENT_CATALOG_VERSION assert _vector_table_exists(cli_env) is True assert _backups_has_vector_column(cli_env) is True + assert _backup_definitions_table_exists(cli_env) is True conn = _catalog_conn(cli_env) try: @@ -124,6 +144,7 @@ def test_init_creates_the_vector_catalog_at_0_0_3(self, cli_env): # deployment that is already current. The statements are idempotent too, # so a replay after a crash between applying and recording is harmless. assert VECTOR_MIGRATION in tracked, tracked + assert BACKUP_DEFINITIONS_MIGRATION in tracked, tracked def test_migrate_upgrades_a_pre_vector_deployment(self, cli_env): """A 0.0.2 deployment is refused, upgraded by migrate, and then serves.""" @@ -148,7 +169,7 @@ def test_migrate_upgrades_a_pre_vector_deployment(self, cli_env): ) from None combined = refused_serve.stdout + refused_serve.stderr assert refused_serve.returncode != 0, combined - assert "0.0.3" in combined and "0.0.2" in combined, combined + assert CURRENT_CATALOG_VERSION in combined and "0.0.2" in combined, combined # Without --yes, migrate reports the pending upgrade and changes nothing. pending = _run_extenddb( @@ -156,7 +177,7 @@ def test_migrate_upgrades_a_pre_vector_deployment(self, cli_env): ) pending_output = pending.stdout + pending.stderr assert pending.returncode != 0, pending_output - assert "0.0.2 -> 0.0.3" in pending_output, pending_output + assert f"0.0.2 -> {CURRENT_CATALOG_VERSION}" in pending_output, pending_output assert _vector_table_exists(cli_env) is False assert _catalog_version(cli_env) == "0.0.2" @@ -164,7 +185,7 @@ def test_migrate_upgrades_a_pre_vector_deployment(self, cli_env): "migrate", "--yes", *_pg_args(), config=cli_env["config_path"], check=False ) assert applied.returncode == 0, applied.stdout + applied.stderr - assert _catalog_version(cli_env) == "0.0.3" + assert _catalog_version(cli_env) == CURRENT_CATALOG_VERSION assert _vector_table_exists(cli_env) is True assert _backups_has_vector_column(cli_env) is True @@ -235,7 +256,7 @@ def test_migrate_survives_an_applied_but_unrecorded_migration(self, cli_env): assert f"Migration {VECTOR_MIGRATION} failed" not in output, output # The replay leaves the same end state, and the ledger is repaired. - assert _catalog_version(cli_env) == "0.0.3" + assert _catalog_version(cli_env) == CURRENT_CATALOG_VERSION assert _vector_table_exists(cli_env) is True assert _backups_has_vector_column(cli_env) is True conn = _catalog_conn(cli_env) @@ -270,7 +291,8 @@ def test_serve_refuses_a_catalog_newer_than_the_binary(self, cli_env): conn.autocommit = True with conn.cursor() as cur: cur.execute( - "UPDATE settings SET value = '0.0.4' WHERE key = 'catalog_version'" + "UPDATE settings SET value = %s WHERE key = 'catalog_version'", + (FUTURE_CATALOG_VERSION,), ) finally: conn.close() @@ -285,8 +307,11 @@ def test_serve_refuses_a_catalog_newer_than_the_binary(self, cli_env): except subprocess.TimeoutExpired: _run_extenddb("stop", config=cli_env["config_path"], check=False) raise AssertionError( - "serve started against a 0.0.4 catalog; the version gate is not symmetric" + f"serve started against a {FUTURE_CATALOG_VERSION} catalog; " + "the version gate is not symmetric" ) from None combined = refused.stdout + refused.stderr assert refused.returncode != 0, combined - assert "0.0.3" in combined and "0.0.4" in combined, combined + assert CURRENT_CATALOG_VERSION in combined and FUTURE_CATALOG_VERSION in combined, ( + combined + ) From ac657bd436e8d049a56d23e34ea402a78190b8a7 Mon Sep 17 00:00:00 2001 From: Scott Robinson Date: Tue, 6 Oct 2026 00:13:32 +0000 Subject: [PATCH 3/4] fix(backup): report RestoreSummary from DescribeTable, refuse DeleteBackup during a restore RestoreSummary was only fabricated on the RestoreTableFromBackup response; DescribeTable never reported it. And DeleteBackup during a restore deleted the backup and failed the restore, where the service refuses with BackupInUseException. Catalog 0.0.4 (same migration as the previous commit; nothing extra to migrate) adds `table_restores`: one row per restored table with the source backup ARN and the restore time. Not a foreign key to `backups`, because the summary keeps naming a backup that is deleted later, as on the service. - PostgreSQL and SQLite DescribeTable report RestoreSummary for a restored table, RestoreInProgress true while it is CREATING and false once ACTIVE. The restore response and later DescribeTable calls carry the same RestoreDateTime. Tables restored before the upgrade have no row and no summary. - DeleteBackup returns BackupInUseException (HTTP 400) while a table still CREATING names the backup. On PostgreSQL the backup row is locked FOR UPDATE by DeleteBackup and FOR SHARE by the restore while it records itself; on SQLite both run under the write lock. Once the restore is done the backup can be deleted. - New DynamoDbError::BackupInUseException and StorageError::BackupInUse. Tests: restore_summary_and_backup_in_use on both backends (summary absent on an ordinary table, in progress while CREATING, done when ACTIVE, delete refused then allowed, summary kept after the backup is deleted), and the wire fidelity test checks RestoreSummary on the response and on DescribeTable. By hand, DeleteBackup during a 60,000-item (SQLite) and 100,000-item (PostgreSQL) restore returned BackupInUseException 400 and the restore completed. Signed-off-by: Scott Robinson --- crates/core/src/error/mod.rs | 6 + crates/engine/src/backup.rs | 16 ++- crates/engine/src/create_table.rs | 1 + .../migrations/003_backup_definitions.sql | 16 ++- crates/storage-postgres/src/backup_engine.rs | 62 +++++++- crates/storage-postgres/src/table_helpers.rs | 24 ++++ .../storage-postgres/tests/backup_restore.rs | 85 +++++++++++ crates/storage-sqlite/src/backup.rs | 136 +++++++++++++++++- crates/storage-sqlite/src/schema.rs | 10 ++ crates/storage-sqlite/src/table_helpers.rs | 31 ++++ crates/storage/src/error.rs | 4 + docs/manuals/07-upgrade-manual.md | 1 + tests/test_backup_restore_fidelity.py | 8 ++ 13 files changed, 393 insertions(+), 7 deletions(-) diff --git a/crates/core/src/error/mod.rs b/crates/core/src/error/mod.rs index c0a28c0b3..e950fe84d 100755 --- a/crates/core/src/error/mod.rs +++ b/crates/core/src/error/mod.rs @@ -24,6 +24,8 @@ pub enum DynamoDbError { #[error("{0}")] BackupNotFoundException(String), #[error("{0}")] + BackupInUseException(String), + #[error("{0}")] ResourceInUseException(String), /// A per-table or per-account limit was exceeded. /// @@ -114,6 +116,7 @@ impl DynamoDbError { Self::ValidationException(_) | Self::ResourceNotFoundException(_) | Self::BackupNotFoundException(_) + | Self::BackupInUseException(_) | Self::ResourceInUseException(_) | Self::LimitExceededException(_) | Self::ConditionalCheckFailedException(..) @@ -155,6 +158,7 @@ impl DynamoDbError { Self::ValidationException(_) => "ValidationException", Self::ResourceNotFoundException(_) => "ResourceNotFoundException", Self::BackupNotFoundException(_) => "BackupNotFoundException", + Self::BackupInUseException(_) => "BackupInUseException", Self::ResourceInUseException(_) => "ResourceInUseException", Self::LimitExceededException(_) => "LimitExceededException", Self::ConditionalCheckFailedException(..) => "ConditionalCheckFailedException", @@ -221,6 +225,7 @@ impl DynamoDbError { Self::ValidationException(m) | Self::ResourceNotFoundException(m) | Self::BackupNotFoundException(m) + | Self::BackupInUseException(m) | Self::ResourceInUseException(m) | Self::LimitExceededException(m) | Self::ConditionalCheckFailedException(m, _) @@ -317,6 +322,7 @@ mod tests { (DynamoDbError::ResourceInUseException(String::new()), 400), (DynamoDbError::ResourceNotFoundException(String::new()), 400), (DynamoDbError::BackupNotFoundException(String::new()), 400), + (DynamoDbError::BackupInUseException(String::new()), 400), (DynamoDbError::SerializationException(String::new()), 400), (DynamoDbError::ServiceUnavailable(String::new()), 503), (DynamoDbError::ThrottlingException(String::new()), 400), diff --git a/crates/engine/src/backup.rs b/crates/engine/src/backup.rs index 36d42d098..447480f1e 100755 --- a/crates/engine/src/backup.rs +++ b/crates/engine/src/backup.rs @@ -157,16 +157,21 @@ pub(crate) async fn handle_restore_table_from_backup( // The restore response's TableDescription reports where the data came // from and that the restore is under way: SourceBackupArn and // RestoreInProgress: true, pinned by the ground-truth runs of 2026-08-24 - // (us-east-1 and eu-west-2). Set here rather than in each backend because - // the summary is response metadata about this call, not table state the - // backends persist. + // (us-east-1 and eu-west-2). The service returns the table CREATING with + // the restore in progress; the backends report CREATING here too, whatever + // the copy has reached. The time is the backend's own record where it + // keeps one, so the response and later DescribeTable calls agree. let now = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .map(|d| d.as_secs_f64()) .unwrap_or_default(); + let restore_date_time = desc + .restore_summary + .as_ref() + .map_or(now, |r| r.restore_date_time); desc.restore_summary = Some(extenddb_core::types::RestoreSummary { source_backup_arn: Some(backup_arn.clone()), - restore_date_time: now, + restore_date_time, restore_in_progress: true, }); @@ -296,6 +301,9 @@ fn storage_err_to_dynamo(e: extenddb_storage::error::StorageError) -> DynamoDbEr extenddb_storage::error::StorageError::LimitExceeded(msg) => { DynamoDbError::LimitExceededException(msg) } + extenddb_storage::error::StorageError::BackupInUse(msg) => { + DynamoDbError::BackupInUseException(msg) + } other => { tracing::error!(internal_error = %other, "backup storage error"); DynamoDbError::InternalServerError("Internal server error".to_owned()) diff --git a/crates/engine/src/create_table.rs b/crates/engine/src/create_table.rs index 257cd2299..019147f3d 100755 --- a/crates/engine/src/create_table.rs +++ b/crates/engine/src/create_table.rs @@ -134,6 +134,7 @@ pub(crate) fn storage_err_to_dynamo(e: extenddb_storage::error::StorageError) -> // state: the documented delete-table sentence and the measured // phase-dependent vector refusal both arrive through this one arm. StorageError::IndexesInUse(msg) => DynamoDbError::ResourceInUseException(msg), + StorageError::BackupInUse(msg) => DynamoDbError::BackupInUseException(msg), StorageError::LimitExceeded(msg) => DynamoDbError::LimitExceededException(msg), // Retryable by definition, so it maps like Connection: a 503 the SDKs // retry, rather than a 500 they surface. StorageError::Transient(msg) => { diff --git a/crates/storage-postgres/migrations/003_backup_definitions.sql b/crates/storage-postgres/migrations/003_backup_definitions.sql index 654537494..0756c3cc2 100644 --- a/crates/storage-postgres/migrations/003_backup_definitions.sql +++ b/crates/storage-postgres/migrations/003_backup_definitions.sql @@ -1,6 +1,7 @@ -- Copyright 2026 ExtendDB contributors -- SPDX-License-Identifier: Apache-2.0 --- Migration 003: record a backup's table definition (catalog version 0.0.4). +-- Migration 003: record a backup's table definition, and which backup a +-- restored table came from (catalog version 0.0.4). -- -- A backup kept the source table's key schema, attribute definitions, and -- billing mode, and nothing else, so a restored table came back without its @@ -22,6 +23,19 @@ CREATE TABLE IF NOT EXISTS backup_definitions ( definition JSONB NOT NULL ); +-- One row per table created by RestoreTableFromBackup: the backup it came +-- from and when, reported by DescribeTable as RestoreSummary, and used to +-- refuse DeleteBackup while the restore is still running. Not a foreign key +-- to backups: the backup may be deleted after the restore, and the summary +-- keeps naming it, as on the service. +CREATE TABLE IF NOT EXISTS table_restores ( + table_id TEXT PRIMARY KEY REFERENCES tables(table_id) ON DELETE CASCADE, + source_backup_arn TEXT NOT NULL, + restore_date_time TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX IF NOT EXISTS idx_table_restores_backup ON table_restores (source_backup_arn); + UPDATE settings SET value = '0.0.4' WHERE key = 'catalog_version'; COMMIT; diff --git a/crates/storage-postgres/src/backup_engine.rs b/crates/storage-postgres/src/backup_engine.rs index c20bfbc84..24538a634 100755 --- a/crates/storage-postgres/src/backup_engine.rs +++ b/crates/storage-postgres/src/backup_engine.rs @@ -1016,6 +1016,31 @@ impl BackupEngine for PostgresEngine { // than relying on the lookup above, so the statements are correct // on their own terms. let mut tx = self.pool.begin().await.map_err(db_err)?; + // Refuse while a restore from this backup is still running, as the + // service does. The backup row is locked first; a restore records + // itself in `table_restores` while holding the same row FOR SHARE, + // so one of the two always sees the other. + sqlx::query( + "SELECT 1 FROM backups WHERE backup_arn = $1 AND account_id = $2 FOR UPDATE", + ) + .bind(&backup_arn) + .bind(&account_id) + .execute(&mut *tx) + .await + .map_err(db_err)?; + let restoring: Option = sqlx::query_scalar( + "SELECT t.table_name FROM table_restores r JOIN tables t ON t.table_id = r.table_id \ + WHERE r.source_backup_arn = $1 AND t.table_status = 'CREATING' LIMIT 1", + ) + .bind(&backup_arn) + .fetch_optional(&mut *tx) + .await + .map_err(db_err)?; + if let Some(table) = restoring { + return Err(StorageError::BackupInUse(format!( + "Backup is being used to restore table {table}: {backup_arn}" + ))); + } sqlx::query( "DELETE FROM backup_items WHERE backup_arn = $1 AND EXISTS (\ SELECT 1 FROM backups b WHERE b.backup_arn = $1 AND b.account_id = $2)", @@ -1210,7 +1235,36 @@ impl BackupEngine for PostgresEngine { // From here on a failure must not leave the target behind: it is // CREATING with no scheduled transition, so nothing else would // move it on, and the name would stay taken. - let copied = self.copy_backup_items(&desc, &backup_arn).await; + let copied = async { + // Provenance first: DescribeTable reports it as RestoreSummary, + // and DeleteBackup refuses while a table still CREATING names + // this backup. + let mut tx = self.pool.begin().await.map_err(db_err)?; + let available: Option = sqlx::query_scalar( + "SELECT 1 FROM backups WHERE backup_arn = $1 \ + AND backup_status = 'AVAILABLE' FOR SHARE", + ) + .bind(&backup_arn) + .fetch_optional(&mut *tx) + .await + .map_err(db_err)?; + if available.is_none() { + return Err(StorageError::Validation(format!( + "Backup not found: {backup_arn}" + ))); + } + sqlx::query( + "INSERT INTO table_restores (table_id, source_backup_arn) VALUES ($1, $2)", + ) + .bind(&desc.table_id) + .bind(&backup_arn) + .execute(&mut *tx) + .await + .map_err(db_err)?; + tx.commit().await.map_err(db_err)?; + self.copy_backup_items(&desc, &backup_arn).await + } + .await; if let Err(e) = copied { match self.abort_restore(&desc.table_id).await { // Not ours to remove: it was already claimed for removal @@ -1239,6 +1293,12 @@ impl BackupEngine for PostgresEngine { // Return CREATING — the API response shows the initial status, // but the table is already ACTIVE by the time the caller polls. + // The summary carries the recorded restore time, so the response + // and later DescribeTable calls agree. + let mut desc = desc; + desc.restore_summary = self + .restore_summary(&desc.table_id, &desc.table_status) + .await?; Ok(desc) }) } diff --git a/crates/storage-postgres/src/table_helpers.rs b/crates/storage-postgres/src/table_helpers.rs index d75049ba0..5bbf78373 100755 --- a/crates/storage-postgres/src/table_helpers.rs +++ b/crates/storage-postgres/src/table_helpers.rs @@ -242,7 +242,9 @@ impl PostgresEngine { .map_err(|e| StorageError::Internal(e.to_string()))?; let table_name_owned = row.table_name.clone(); + let table_id = row.table_id.clone(); let mut desc = self.build_table_description_from_row(account_id, row, index_rows)?; + desc.restore_summary = self.restore_summary(&table_id, &desc.table_status).await?; desc.vector_indexes = vector_index_descriptions( &self.region, account_id, @@ -435,4 +437,26 @@ impl PostgresEngine { ..Default::default() }) } + + /// The RestoreSummary of a table created by RestoreTableFromBackup, or + /// `None` for any other table. In progress until the table is ACTIVE. + pub(crate) async fn restore_summary( + &self, + table_id: &str, + status: &extenddb_core::types::TableStatus, + ) -> Result, StorageError> { + let row: Option<(String, f64)> = sqlx::query_as( + "SELECT source_backup_arn, EXTRACT(EPOCH FROM restore_date_time)::FLOAT8 \ + FROM table_restores WHERE table_id = $1", + ) + .bind(table_id) + .fetch_optional(&self.pool) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + Ok(row.map(|(arn, at)| extenddb_core::types::RestoreSummary { + source_backup_arn: Some(arn), + restore_date_time: at, + restore_in_progress: *status == extenddb_core::types::TableStatus::Creating, + })) + } } diff --git a/crates/storage-postgres/tests/backup_restore.rs b/crates/storage-postgres/tests/backup_restore.rs index 0ae45af1f..fa16624a5 100644 --- a/crates/storage-postgres/tests/backup_restore.rs +++ b/crates/storage-postgres/tests/backup_restore.rs @@ -1152,3 +1152,88 @@ async fn delete_table_refuses_a_restore_in_progress() { ); s.cleanup().await; } + +/// DescribeTable reports where a restored table came from, in progress +/// while it is CREATING and done once ACTIVE; DeleteBackup is refused while a +/// restore from the backup is still running, and allowed afterwards, with +/// the summary still naming the deleted backup. +#[tokio::test] +async fn restore_summary_and_backup_in_use() { + if base_conn().is_none() { + eprintln!("SKIP restore_summary_and_backup_in_use: no PostgreSQL"); + return; + } + let s = scratch().await; + indexed_source(&s, "rs_src").await; + let backup = s + .engine + .create_backup(ACCOUNT, "rs_src", "b") + .await + .expect("backup"); + let desc = s + .engine + .restore_table_from_backup(ACCOUNT, "rs_dst", &backup.backup_arn) + .await + .expect("restore"); + let describe = |name: &'static str| { + let engine = &s.engine; + async move { + engine + .describe_table( + ACCOUNT, + DescribeTableInput { + table_name: name.to_owned(), + }, + ) + .await + .expect("describe") + } + }; + let done = describe("rs_dst").await.restore_summary.expect("summary"); + assert_eq!( + done.source_backup_arn.as_deref(), + Some(backup.backup_arn.as_str()) + ); + assert!(!done.restore_in_progress); + assert!(done.restore_date_time > 0.0); + assert!(describe("rs_src").await.restore_summary.is_none()); + + // While the restore is still running: in progress, and the backup in use. + sqlx::query( + "UPDATE tables SET table_status = 'CREATING', status_transition_at = NULL \ + WHERE table_id = $1", + ) + .bind(&desc.table_id) + .execute(&s.catalog) + .await + .expect("back to in progress"); + assert!( + describe("rs_dst") + .await + .restore_summary + .expect("summary") + .restore_in_progress + ); + let err = s + .engine + .delete_backup(ACCOUNT, &backup.backup_arn) + .await + .expect_err("in use"); + assert!(matches!(err, StorageError::BackupInUse(_)), "{err:?}"); + + sqlx::query("UPDATE tables SET table_status = 'ACTIVE' WHERE table_id = $1") + .bind(&desc.table_id) + .execute(&s.catalog) + .await + .expect("finished"); + s.engine + .delete_backup(ACCOUNT, &backup.backup_arn) + .await + .expect("deletable once the restore is done"); + let after = describe("rs_dst").await.restore_summary.expect("summary"); + assert_eq!( + after.source_backup_arn.as_deref(), + Some(backup.backup_arn.as_str()) + ); + s.cleanup().await; +} diff --git a/crates/storage-sqlite/src/backup.rs b/crates/storage-sqlite/src/backup.rs index a4e3fc52a..fb89c945e 100644 --- a/crates/storage-sqlite/src/backup.rs +++ b/crates/storage-sqlite/src/backup.rs @@ -734,6 +734,22 @@ impl BackupEngine for SqliteEngine { .begin_with("BEGIN IMMEDIATE") .await .map_err(db_err)?; + // Refuse while a restore from this backup is still running, as the + // service does. Under the write lock, which the restore also takes + // to record itself, so the check cannot miss one. + let restoring: Option = sqlx::query_scalar( + "SELECT t.table_name FROM table_restores r JOIN tables t ON t.table_id = r.table_id \ + WHERE r.source_backup_arn = ? AND t.table_status = 'CREATING' LIMIT 1", + ) + .bind(&backup_arn) + .fetch_optional(&mut *tx) + .await + .map_err(db_err)?; + if let Some(table) = restoring { + return Err(StorageError::BackupInUse(format!( + "Backup is being used to restore table {table}: {backup_arn}" + ))); + } sqlx::query( "DELETE FROM backup_items WHERE backup_arn = ?1 AND EXISTS (\ SELECT 1 FROM backups b WHERE b.backup_arn = ?1 AND b.account_id = ?2)", @@ -851,7 +867,43 @@ impl BackupEngine for SqliteEngine { // last one; until then the target is CREATING and refuses every // data-plane request. A failure part-way is cleaned up here, a // crash part-way by the startup sweep. - let copied = self.copy_backup_items(&desc, &backup_arn).await; + let copied = async { + // Provenance first: DescribeTable reports it as RestoreSummary, + // and DeleteBackup refuses while a table still CREATING names + // this backup. + { + let _writer = self.write_lock.lock().await; + let mut tx = self + .pool + .begin_with("BEGIN IMMEDIATE") + .await + .map_err(db_err)?; + let available: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM backups WHERE backup_arn = ? \ + AND backup_status = 'AVAILABLE')", + ) + .bind(&backup_arn) + .fetch_one(&mut *tx) + .await + .map_err(db_err)?; + if !available { + return Err(StorageError::Validation(format!( + "Backup not found: {backup_arn}" + ))); + } + sqlx::query( + "INSERT INTO table_restores (table_id, source_backup_arn) VALUES (?, ?)", + ) + .bind(&desc.table_id) + .bind(&backup_arn) + .execute(&mut *tx) + .await + .map_err(db_err)?; + tx.commit().await.map_err(db_err)?; + } + self.copy_backup_items(&desc, &backup_arn).await + } + .await; if let Err(e) = copied { tracing::error!( "restore of {backup_arn} into {target_table_name} failed, \ @@ -867,6 +919,12 @@ impl BackupEngine for SqliteEngine { return Err(e); } + // The summary carries the recorded restore time, so the response + // and later DescribeTable calls agree. + let mut desc = desc; + desc.restore_summary = self + .restore_summary(&desc.table_id, &desc.table_status) + .await?; Ok(desc) }) } @@ -1608,4 +1666,80 @@ mod restore_tests { "{err:?}" ); } + + #[tokio::test] + async fn restore_summary_and_backup_in_use() { + use extenddb_storage::TableEngine; + let engine = engine().await; + indexed_table(&engine, "src").await; + let backup = engine + .create_backup(ACCOUNT, "src", "bkp") + .await + .expect("backup"); + let desc = engine + .restore_table_from_backup(ACCOUNT, "dst", &backup.backup_arn) + .await + .expect("restore"); + let describe = |name: &'static str| { + let engine = &engine; + async move { + engine + .describe_table( + ACCOUNT, + extenddb_core::types::DescribeTableInput { + table_name: name.to_owned(), + }, + ) + .await + .expect("describe") + } + }; + let done = describe("dst").await.restore_summary.expect("summary"); + assert_eq!( + done.source_backup_arn.as_deref(), + Some(backup.backup_arn.as_str()) + ); + assert!(!done.restore_in_progress); + assert!(done.restore_date_time > 0.0); + assert!(describe("src").await.restore_summary.is_none()); + + sqlx::query( + "UPDATE tables SET table_status = 'CREATING', status_transition_at = NULL \ + WHERE table_id = ?", + ) + .bind(&desc.table_id) + .execute(&engine.pool) + .await + .expect("in progress"); + assert!( + describe("dst") + .await + .restore_summary + .expect("summary") + .restore_in_progress + ); + let err = engine + .delete_backup(ACCOUNT, &backup.backup_arn) + .await + .expect_err("in use"); + assert!( + matches!(err, extenddb_storage::error::StorageError::BackupInUse(_)), + "{err:?}" + ); + + sqlx::query("UPDATE tables SET table_status = 'ACTIVE' WHERE table_id = ?") + .bind(&desc.table_id) + .execute(&engine.pool) + .await + .expect("finished"); + engine + .delete_backup(ACCOUNT, &backup.backup_arn) + .await + .expect("deletable once the restore is done"); + let after = describe("dst").await.restore_summary.expect("summary"); + assert_eq!( + after.source_backup_arn.as_deref(), + Some(backup.backup_arn.as_str()) + ); + } } diff --git a/crates/storage-sqlite/src/schema.rs b/crates/storage-sqlite/src/schema.rs index a558b777f..a8c13297f 100644 --- a/crates/storage-sqlite/src/schema.rs +++ b/crates/storage-sqlite/src/schema.rs @@ -385,6 +385,16 @@ CREATE TABLE IF NOT EXISTS backup_definitions ( definition TEXT NOT NULL ); +-- Which backup a restored table came from, and when (catalog 0.0.4): reported +-- by DescribeTable as RestoreSummary, and used to refuse DeleteBackup while the +-- restore is running. Not a foreign key to backups, which may be deleted later. +CREATE TABLE IF NOT EXISTS table_restores ( + table_id TEXT PRIMARY KEY REFERENCES tables(table_id) ON DELETE CASCADE, + source_backup_arn TEXT NOT NULL, + restore_date_time TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')) +); +CREATE INDEX IF NOT EXISTS idx_table_restores_backup ON table_restores (source_backup_arn); + -- Continuous backups / PITR status. CREATE TABLE IF NOT EXISTS continuous_backups ( account_id TEXT NOT NULL, diff --git a/crates/storage-sqlite/src/table_helpers.rs b/crates/storage-sqlite/src/table_helpers.rs index b8adfbf1e..7554565ca 100644 --- a/crates/storage-sqlite/src/table_helpers.rs +++ b/crates/storage-sqlite/src/table_helpers.rs @@ -150,7 +150,9 @@ impl SqliteEngine { .map_err(|e| StorageError::Internal(e.to_string()))?; let table_name_owned = row.table_name.clone(); + let table_id = row.table_id.clone(); let mut desc = self.build_table_description_from_row(account_id, row, index_rows)?; + desc.restore_summary = self.restore_summary(&table_id, &desc.table_status).await?; let catalog_rows = vector_rows .into_iter() .map(VectorIndexRow::into_catalog_row) @@ -388,3 +390,32 @@ fn zero_throughput() -> ProvisionedThroughputDescription { last_decrease_date_time: None, } } + +impl SqliteEngine { + /// The RestoreSummary of a table created by RestoreTableFromBackup, or + /// `None` for any other table. In progress until the table is ACTIVE. + pub(crate) async fn restore_summary( + &self, + table_id: &str, + status: &extenddb_core::types::TableStatus, + ) -> Result, StorageError> { + let row: Option<(String, String)> = sqlx::query_as( + "SELECT source_backup_arn, restore_date_time FROM table_restores WHERE table_id = ?", + ) + .bind(table_id) + .fetch_optional(&self.pool) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + Ok(row.map(|(arn, at)| { + #[allow(clippy::cast_precision_loss)] + let at = crate::sqlite_util::parse_timestamp(&at) + .map(|t| t.unix_timestamp_nanos() as f64 / 1e9) + .unwrap_or(0.0); + extenddb_core::types::RestoreSummary { + source_backup_arn: Some(arn), + restore_date_time: at, + restore_in_progress: *status == extenddb_core::types::TableStatus::Creating, + } + })) + } +} diff --git a/crates/storage/src/error.rs b/crates/storage/src/error.rs index b98593c09..563dfcba6 100755 --- a/crates/storage/src/error.rs +++ b/crates/storage/src/error.rs @@ -51,6 +51,10 @@ pub enum StorageError { /// measured 2026-08-13). #[error("{0}")] LimitExceeded(String), + /// A backup an in-progress restore is still reading. Maps to + /// `BackupInUseException`. + #[error("{0}")] + BackupInUse(String), /// A failure that is expected to succeed on retry: I/O errors, pool /// timeouts, SQLITE_BUSY / SQLITE_LOCKED. Exists so queue workers can tell /// "this row can never be applied" (drop it, or the whole queue stalls) diff --git a/docs/manuals/07-upgrade-manual.md b/docs/manuals/07-upgrade-manual.md index 2cc7c0b99..5f4689b9a 100755 --- a/docs/manuals/07-upgrade-manual.md +++ b/docs/manuals/07-upgrade-manual.md @@ -187,6 +187,7 @@ psql -d extenddb_catalog -f catalog_backup_YYYYMMDD.sql Adds the table definition a backup records: - New `backup_definitions` table: one row per backup, holding the source table's global and local secondary indexes, billing mode and provisioned throughput, table class, and encryption settings. RestoreTableFromBackup recreates them. A backup taken before this upgrade has no row and restores as before, with its keys and items but no secondary indexes. +- New `table_restores` table: one row per table created by RestoreTableFromBackup, naming the backup and the restore time. DescribeTable reports it as `RestoreSummary`, and DeleteBackup is refused with `BackupInUseException` while a restore from the backup is still running. Tables restored before this upgrade have no row and report no `RestoreSummary`. Upgrade sequence, on PostgreSQL and SQLite alike: diff --git a/tests/test_backup_restore_fidelity.py b/tests/test_backup_restore_fidelity.py index 57895649c..c0e25edc3 100644 --- a/tests/test_backup_restore_fidelity.py +++ b/tests/test_backup_restore_fidelity.py @@ -242,9 +242,17 @@ def _round_trip(client, sort_kind: str | None, hash_kind: str = "S") -> None: backup_arn = _create_backup(client, source) resp = client.restore_table_from_backup(TargetTableName=restored, BackupArn=backup_arn) assert resp["TableDescription"]["TableName"] == restored + started = resp["TableDescription"]["RestoreSummary"] + assert started["SourceBackupArn"] == backup_arn + assert started["RestoreInProgress"] is True wait_for_active(client, restored, timeout=RESTORE_TIMEOUT_S) table = client.describe_table(TableName=restored)["Table"] + # DescribeTable keeps reporting where the table came from, finished. + summary = table["RestoreSummary"] + assert summary["SourceBackupArn"] == backup_arn + assert summary["RestoreInProgress"] is False + assert summary["RestoreDateTime"] == started["RestoreDateTime"] assert table["KeySchema"] == key_schema assert sorted(table["AttributeDefinitions"], key=lambda a: a["AttributeName"]) == sorted( attr_defs, key=lambda a: a["AttributeName"] From 205bc8948d480f1d0542200b4888138a6734c832 Mon Sep 17 00:00:00 2001 From: Scott Robinson Date: Tue, 6 Oct 2026 01:53:13 +0000 Subject: [PATCH 4/4] fix(mongodb): report RestoreSummary and refuse DeleteBackup during a restore The wire fidelity test from the previous commit checks RestoreSummary on every backend, and failed on MongoDB, which never reported it. A restore now records `restore_source_backup_arn` and `restore_date_time` on the target's catalog document before the copy. DescribeTable reports them as RestoreSummary, in progress while the table is CREATING (on this backend that includes the index backfill that runs after the restore call returns), and DeleteBackup returns BackupInUseException while a CREATING table names the backup. Verified against a mongo:7 replica set: test_backup_restore_fidelity.py 7/7, test_mongodb_backup_restore.py and test_streams.py pass (26 passed, 1 xfailed); DeleteBackup during a 30,000-item restore returned BackupInUseException 400 and the restore completed. Signed-off-by: Scott Robinson --- crates/storage-mongodb/src/backup_engine.rs | 50 ++++++++++++++++++++- crates/storage-mongodb/src/table_engine.rs | 15 +++++++ 2 files changed, 64 insertions(+), 1 deletion(-) diff --git a/crates/storage-mongodb/src/backup_engine.rs b/crates/storage-mongodb/src/backup_engine.rs index 5476c7cc0..ade31c126 100644 --- a/crates/storage-mongodb/src/backup_engine.rs +++ b/crates/storage-mongodb/src/backup_engine.rs @@ -502,6 +502,30 @@ impl BackupEngine for MongoEngine { StorageError::Validation(format!("Backup not found: {backup_arn}")) })?; + // Refuse while a restore from this backup is still running, as the + // service does. + let restoring = self + .catalog_db + .collection::("tables") + .find_one(doc! { + "_id.account_id": &account_id, + "restore_source_backup_arn": &backup_arn, + "table_status": "CREATING", + }) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + if let Some(t) = restoring { + let name = t + .get_document("_id") + .ok() + .and_then(|id| id.get_str("table_name").ok()) + .unwrap_or("?") + .to_owned(); + return Err(StorageError::BackupInUse(format!( + "Backup is being used to restore table {name}: {backup_arn}" + ))); + } + // Drop the backup collection. If backup_id is absent (e.g., a // pre-`$out` backup on an old catalog) we skip — nothing to drop // at the collection level in that case. @@ -641,10 +665,34 @@ impl BackupEngine for MongoEngine { // Create the table with the ACTIVE transition deferred: it enters // CREATING with no scheduled flip, so the table cannot become // ACTIVE until we schedule it below, after the data copy completes. - let desc = self + let mut desc = self .create_table_impl(&account_id, create_input, true) .await?; + // Provenance, recorded before the copy: DescribeTable reports it as + // RestoreSummary, and DeleteBackup refuses while a table still + // CREATING names this backup. + let restore_at = bson::DateTime::now(); + self.catalog_db + .collection::("tables") + .update_one( + doc! { "_id": { "account_id": &account_id, "table_name": &target_table_name } }, + doc! { "$set": { + "restore_source_backup_arn": &backup_arn, + "restore_date_time": restore_at, + } }, + ) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + #[allow(clippy::cast_precision_loss)] + { + desc.restore_summary = Some(extenddb_core::types::RestoreSummary { + source_backup_arn: Some(backup_arn.clone()), + restore_date_time: restore_at.timestamp_millis() as f64 / 1000.0, + restore_in_progress: true, + }); + } + // Restore items from the backup collection using server-side `$out`. // The backup collection was written by `create_backup` in the same // document shape as the source data collection, so this is a diff --git a/crates/storage-mongodb/src/table_engine.rs b/crates/storage-mongodb/src/table_engine.rs index 87a5df208..0ef661d10 100644 --- a/crates/storage-mongodb/src/table_engine.rs +++ b/crates/storage-mongodb/src/table_engine.rs @@ -1456,6 +1456,20 @@ impl MongoEngine { let on_demand_throughput: Option = doc .get("on_demand_throughput") .and_then(|b| bson::from_bson(b.clone()).ok()); + // Set on a table created by RestoreTableFromBackup; in progress until + // the table is ACTIVE (the restore's index backfill runs while it is + // CREATING). + #[allow(clippy::cast_precision_loss)] + let restore_summary = doc.get_str("restore_source_backup_arn").ok().map(|arn| { + extenddb_core::types::RestoreSummary { + source_backup_arn: Some(arn.to_owned()), + restore_date_time: doc + .get_datetime("restore_date_time") + .map(|d| d.timestamp_millis() as f64 / 1000.0) + .unwrap_or(0.0), + restore_in_progress: table_status == TableStatus::Creating, + } + }); Ok(TableDescription { table_name, @@ -1479,6 +1493,7 @@ impl MongoEngine { sse_description, table_class_summary, on_demand_throughput, + restore_summary, // Fields for features this backend does not implement, vector // indexes today, take their defaults. Adding one to // TableDescription then does not break this build.