From f40cc42442fd5583de321a4af9e0164405fb2d2f Mon Sep 17 00:00:00 2001 From: Lee Hannigan Date: Sun, 25 Oct 2026 02:42:15 +0000 Subject: [PATCH] fix(streams): resolve the shard against its stream inside the GetRecords read GetRecords called validate_shard before get_stream_records, which added two catalog queries to every poll on top of the account ownership guard inside the read. get_stream_records now takes the iterator's stream ARN and resolves the shard against the table that ARN names, scoped to the caller's account, in the same query as the records read (SQLite), or in one catalog lookup plus one data query (PostgreSQL, MongoDB), or one catalog lookup plus the shard row and records read issued together (Cassandra). Whether the shard belongs to the stream the ARN's label names is decided in one function, extenddb_storage::util::shard_in_stream, from the table's current label. A shard of another table, another account, or an earlier stream label on the same table gets the stream not-found answer validate_shard gave. The authorization decision is untouched: crates/server does not change. Storage statements per GetRecords call, measured on an idle shard and on a shard with records: SQLite 4 to 1, PostgreSQL 5 to 2, MongoDB 5 to 2, Cassandra 4 to 3 (two sequential round trips instead of four; idle p50 2.1 ms to 1.6 ms over loopback). GetRecords refuses a Limit under 1 with "1 validation error detected: Value at 'limit' failed to satisfy constraint: Member must have value greater than or equal to 1" and over 1000 with "One or more parameter values were invalid: Limit should be less than or equal to 1000" before the read (both measured on the service 2026-10-25; before, 0 reached storage and 1001 was clamped). MongoDB streams the page one record per document, so a page over 16 MB reads instead of answering InternalServerError. Every streams path answers one stream not-found sentence, without a trailing period, as measured on DescribeStream. validate_shard stays for GetShardIterator. A GetRecords token with no stream ARN is answered as expired by the engine, as the authorization layer already does; in a dev-mode build, where authorization does not classify the token, such a token with a real shard id read before and is expired now. Tests: seven SQLite storage tests, three unit tests on shard_in_stream, one engine test on the Limit bounds, storage-level ownership tests on PostgreSQL, Cassandra and MongoDB (same-label other table, same-named other account, an earlier label on the same shards; a 40 by 450 KB page on MongoDB, inserted out of order so the read's sort is what the assertion sees), and three wire tests under tests/python. The MongoDB storage tests run in the backend's CI job against its container. Closes #438 --- .github/workflows/integration.yml | 2 +- crates/engine/src/streams.rs | 110 ++++- crates/storage-cassandra/src/stream_engine.rs | 196 ++++----- .../tests/direct/stream_ownership.rs | 268 +++++++++++++ .../storage-cassandra/tests/direct/streams.rs | 32 +- .../tests/direct_integration.rs | 2 + crates/storage-mongodb/src/stream_engine.rs | 157 ++++---- .../src/stream_records_live_tests.rs | 326 +++++++++++++++ crates/storage-postgres/src/stream_engine.rs | 111 +++--- .../tests/stream_shard_ownership.rs | 377 ++++++++++++++++++ crates/storage-sqlite/src/lib.rs | 2 + crates/storage-sqlite/src/stream.rs | 121 +++--- crates/storage-sqlite/src/stream_tests.rs | 298 ++++++++++++++ crates/storage/src/lib.rs | 14 +- crates/storage/src/util/arn.rs | 14 + crates/storage/src/util/mod.rs | 4 +- crates/storage/src/util/shard_stream.rs | 63 +++ devtools/run-mongodb-tests | 11 + tests/python/helpers.py | 7 + tests/python/test_stream_shard_ownership.py | 239 +++++++++++ 20 files changed, 2045 insertions(+), 309 deletions(-) create mode 100644 crates/storage-cassandra/tests/direct/stream_ownership.rs create mode 100644 crates/storage-mongodb/src/stream_records_live_tests.rs create mode 100644 crates/storage-postgres/tests/stream_shard_ownership.rs create mode 100644 crates/storage-sqlite/src/stream_tests.rs create mode 100644 crates/storage/src/util/shard_stream.rs create mode 100644 tests/python/test_stream_shard_ownership.py diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 8a4818f75..5ec819421 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -310,7 +310,7 @@ jobs: run: >- cargo test --release -p extenddb-storage-postgres --test vector_control_plane --test put_create_race --test twi_conflict - --test stream_shard_retention --test backup_restore + --test stream_shard_retention --test stream_shard_ownership --test backup_restore --test gsi_online_build --test gsi_stale_writer # The daemonized server logs to syslog; dump it so server-side failures diff --git a/crates/engine/src/streams.rs b/crates/engine/src/streams.rs index 83b297253..4d08d930c 100755 --- a/crates/engine/src/streams.rs +++ b/crates/engine/src/streams.rs @@ -164,6 +164,16 @@ pub async fn handle_get_shard_iterator( /// Shard iterator expiration: 15 minutes (900 seconds), matching real `DynamoDB`. const SHARD_ITERATOR_EXPIRY_SECS: u64 = 900; +/// The model's constraint sentence for a member value, as the service renders +/// it for the streams API (no value quoted; measured 2026-10-24 on `shardId`, +/// `exclusiveStartShardId`, `sequenceNumber`, `limit` and +/// `exclusiveStartStreamArn`). +fn constraint(member: &str, rule: &str) -> String { + format!( + "1 validation error detected: Value at '{member}' failed to satisfy constraint: Member must {rule}" + ) +} + /// The message `DynamoDB` returns for an expired shard iterator (measured live, /// `live-r2/iter-live.json`: an iterator 990 seconds old returns /// `ExpiredIteratorException` with exactly this text). Shared so the @@ -262,13 +272,16 @@ pub fn classify_shard_iterator(token: &str) -> ShardIteratorClass { /// Handle `GetRecords`. /// -/// Decodes the shard iterator, checks expiration, reads records, and returns +/// Decodes the shard iterator, checks expiration, reads records from the +/// shard as one of the shards of the stream the iterator names, and returns /// a new iterator. /// /// # Errors /// /// Returns [`DynamoDbError::ExpiredIteratorException`] if the iterator is older -/// than 15 minutes. +/// than 15 minutes or carries no stream ARN. +/// Returns [`DynamoDbError::ResourceNotFoundException`] if the stream does not +/// exist for the caller's account or the shard is not one of its shards. /// Returns [`DynamoDbError::ValidationException`] on invalid iterator. pub async fn handle_get_records( body: Value, @@ -293,18 +306,21 @@ pub async fn handle_get_records( } } - // Confirm the shard belongs to the stream named in the token before - // reading. The request is authorized against that stream ARN, so this - // stops a forged token that pairs an authorized stream ARN with another - // stream's shard id from reading the other stream's records. - if let Some(ref stream_arn) = it.stream_arn { - ctx.storage - .validate_shard(&ctx.account_id, stream_arn, &it.shard_id) - .await - .map_err(storage_to_dynamo)?; - } + // The request was authorized against the stream ARN carried in the token. + // The storage read below resolves the shard against that stream (table + // name and stream label, within the caller's account) in the same round + // trip as the records, so a forged token that pairs an authorized stream + // ARN with another stream's shard id reads nothing and gets the stream + // not-found answer. A token with no stream ARN was issued by a build before + // the ARN was carried; the authorization layer answers it as expired, and + // the engine gives the same answer when called directly. + let Some(stream_arn) = it.stream_arn.as_deref() else { + return Err(DynamoDbError::ExpiredIteratorException( + SHARD_ITERATOR_EXPIRED_MESSAGE.to_owned(), + )); + }; - let limit = input.limit.unwrap_or(1000).min(1000); + let limit = get_records_limit(input.limit)?; // All iterator types are resolved to AFTER_SEQUENCE_NUMBER at // GetShardIterator time. Empty seq means "read from beginning". @@ -318,6 +334,7 @@ pub async fn handle_get_records( .storage .get_stream_records( &ctx.account_id, + stream_arn, &it.shard_id, after_sequence.as_deref(), limit, @@ -338,7 +355,7 @@ pub async fn handle_get_records( "AFTER_SEQUENCE_NUMBER", &next_seq, now, - it.stream_arn.as_deref().unwrap_or(""), + stream_arn, )) }; @@ -349,6 +366,34 @@ pub async fn handle_get_records( serialize_output(&output) } +/// `GetRecords` with a `Limit` over 1000 (measured 2026-10-25 with 1001). +pub const GET_RECORDS_LIMIT_TOO_LARGE: &str = + "One or more parameter values were invalid: Limit should be less than or equal to 1000"; + +/// The largest `Limit` `GetRecords` accepts. +const GET_RECORDS_MAX_LIMIT: i64 = 1000; + +/// The page size a `GetRecords` call asked for, or the service's refusal. +/// +/// Measured 2026-10-25: `Limit` 0 and -1 get the model's constraint sentence, +/// 1001 the parameter sentence; 1 and 1000 are accepted. The bounds are +/// checked before storage runs, so the records read never carries a limit +/// under 1, which the backends read as a shard that is not the stream's (no +/// row) or refuse outright. +fn get_records_limit(limit: Option) -> Result { + match limit { + Some(l) if l < 1 => Err(DynamoDbError::ValidationException(constraint( + "limit", + "have value greater than or equal to 1", + ))), + Some(l) if l > GET_RECORDS_MAX_LIMIT => Err(DynamoDbError::ValidationException( + GET_RECORDS_LIMIT_TOO_LARGE.to_owned(), + )), + Some(l) => Ok(l), + None => Ok(GET_RECORDS_MAX_LIMIT), + } +} + fn storage_to_dynamo(e: StorageError) -> DynamoDbError { match e { StorageError::Validation(msg) => DynamoDbError::ValidationException(msg), @@ -415,6 +460,43 @@ mod tests { assert!(super::decode_shard_iterator("!!!not base64!!!").is_none()); } + #[test] + fn stream_not_found_maps_to_resource_not_found_with_the_arn() { + // `get_stream_records` answers a shard that is not the iterator's + // stream's with the storage not-found error; the wire answer is the + // ResourceNotFoundException `validate_shard` produced for GetRecords + // before the check moved into the read. + let arn = "arn:aws:dynamodb:us-east-1:1:table/T/stream/2026-10-20T16:17:40.188"; + match super::storage_to_dynamo(extenddb_storage::util::stream_not_found(arn)) { + super::DynamoDbError::ResourceNotFoundException(msg) => assert_eq!( + msg, + format!("Requested resource not found: Stream: {arn} not found") + ), + other => panic!("expected ResourceNotFoundException, got {other:?}"), + } + } + + #[test] + fn get_records_limit_bounds_carry_the_measured_sentences() { + let refused = |l: i64| match super::get_records_limit(Some(l)) { + Err(super::DynamoDbError::ValidationException(msg)) => msg, + other => panic!("expected ValidationException for Limit {l}, got {other:?}"), + }; + for l in [0, -1] { + assert_eq!( + refused(l), + "1 validation error detected: Value at 'limit' failed to satisfy constraint: Member must have value greater than or equal to 1" + ); + } + assert_eq!( + refused(1001), + "One or more parameter values were invalid: Limit should be less than or equal to 1000" + ); + assert_eq!(super::get_records_limit(Some(1)).unwrap(), 1); + assert_eq!(super::get_records_limit(Some(1000)).unwrap(), 1000); + assert_eq!(super::get_records_limit(None).unwrap(), 1000); + } + fn b64(s: &str) -> String { base64::Engine::encode(&base64::engine::general_purpose::STANDARD, s) } diff --git a/crates/storage-cassandra/src/stream_engine.rs b/crates/storage-cassandra/src/stream_engine.rs index e3c6088fa..17b8b0049 100755 --- a/crates/storage-cassandra/src/stream_engine.rs +++ b/crates/storage-cassandra/src/stream_engine.rs @@ -10,7 +10,7 @@ use extenddb_core::types::{ }; use extenddb_storage::StreamEngine; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{parse_stream_arn, stream_arn}; +use extenddb_storage::util::{parse_stream_arn, shard_in_stream, stream_arn}; use futures::future::BoxFuture; use crate::CassandraEngine; @@ -60,44 +60,55 @@ impl CassandraEngine { Ok(self.account_keyspace(&account_id)) } - /// Like `account_keyspace_for_shard`, but also enforces that the shard's - /// table belongs to `caller_account_id`. Returns `Validation("Invalid - /// ShardIterator")` for both "shard not found" and "shard belongs to a - /// different account" — no information leakage about other accounts. - async fn account_keyspace_for_shard_owned_by( + /// The table id and current stream label of table `table_name` within + /// `account_id` (a primary-key read of the catalog), or `None` when the + /// account has no table of that name. + async fn stream_table_row( &self, - shard_id: &str, - caller_account_id: &str, - ) -> Result { - let without_prefix = shard_id - .strip_prefix("shardId-") - .ok_or_else(|| StorageError::Validation("Invalid ShardIterator".to_owned()))?; - let table_id = without_prefix - .rsplit_once('-') - .map(|x| x.0) - .ok_or_else(|| StorageError::Validation("Invalid ShardIterator".to_owned()))?; - + account_id: &str, + table_name: &str, + ) -> Result)>, StorageError> { let catalog_keyspace = self.catalog_keyspace(); - let query = format!("SELECT account_id FROM {catalog_keyspace}.tables WHERE table_id = ?"); + let query = format!( + "SELECT table_id, stream_label FROM {catalog_keyspace}.tables \ + WHERE account_id = ? AND table_name = ?" + ); let result = self .session - .query_with_values(&query, cdrs_tokio::query_values!(table_id)) + .query_with_values(&query, cdrs_tokio::query_values!(account_id, table_name)) .await .map_err(|e| StorageError::Internal(e.to_string()))?; - - let table_account_id: Option = result + let rows = result .response_body() .map_err(|e| StorageError::Internal(e.to_string()))? .into_rows() - .unwrap_or_default() - .into_iter() + .unwrap_or_default(); + rows.into_iter() .next() - .and_then(|row| row.get_r_by_name("account_id").ok()); + .map(|row| { + let table_id: String = row + .get_r_by_name("table_id") + .map_err(|e| StorageError::Internal(e.to_string()))?; + let label: Option = row.get_by_name("stream_label").ok().flatten(); + Ok((table_id, label)) + }) + .transpose() + } - // Collapse "not found" and "wrong account" into the same error. - match table_account_id { - Some(ref owner) if owner == caller_account_id => Ok(self.account_keyspace(owner)), - _ => Err(StorageError::Validation("Invalid ShardIterator".to_owned())), + /// The table id of the stream `stream_arn` names within `account_id`, or + /// the stream not-found error when no table of that account carries that + /// name and stream label. + async fn stream_table_id( + &self, + account_id: &str, + stream_arn: &str, + ) -> Result { + let (table_name, stream_label) = parse_stream_arn(stream_arn)?; + match self.stream_table_row(account_id, &table_name).await? { + Some((table_id, label)) if label.as_deref() == Some(stream_label.as_str()) => { + Ok(table_id) + } + _ => Err(extenddb_storage::util::stream_not_found(stream_arn)), } } } @@ -125,23 +136,40 @@ impl StreamEngine for CassandraEngine { fn get_stream_records( &self, account_id: &str, + stream_arn: &str, shard_id: &str, after_sequence: Option<&str>, limit: i64, ) -> BoxFuture<'_, Result<(Vec, Option), StorageError>> { let account_id = account_id.to_string(); + let stream_arn = stream_arn.to_string(); let shard_id = shard_id.to_string(); let after_sequence = after_sequence.map(str::to_string); Box::pin(async move { - // Ownership guard: resolve the shard's owning account from the - // catalog and reject the request if it doesn't match the caller. - // Mirrors the PostgreSQL two-step check. Both "shard not found" and - // "shard belongs to a different account" collapse to the same - // `Invalid ShardIterator` error — consistent with real DynamoDB - // Streams behaviour (no information leakage about other accounts). - let keyspace = self - .account_keyspace_for_shard_owned_by(&shard_id, &account_id) - .await?; + // Cassandra has no joins, so the stream resolves in two sequential + // round trips: the table the ARN names (catalog keyspace, by + // primary key: its id and current stream label), then the shard + // row and the records, issued together and awaited together on the + // account keyspace. The shard is the table's when its row carries + // that table id, and the stream's when `shard_in_stream` says so + // from the table's label and the shard id; otherwise (another + // table, another account, or an earlier stream label on the same + // table) the answer is the stream not-found error and the records + // read are discarded. The keyspace follows from the caller's + // account, so no table id is parsed out of the shard id. + let (table_name, stream_label) = parse_stream_arn(&stream_arn)?; + let Some((table_id, table_label)) = + self.stream_table_row(&account_id, &table_name).await? + else { + return Err(extenddb_storage::util::stream_not_found(&stream_arn)); + }; + let keyspace = self.account_keyspace_for(&account_id); + + let shard_query = + format!("SELECT table_id FROM {keyspace}.stream_shards WHERE shard_id = ?"); + let shard_read = self + .session + .query_with_values(&shard_query, cdrs_tokio::query_values!(shard_id.as_str())); let query = if let Some(ref after) = after_sequence { format!( @@ -156,22 +184,39 @@ impl StreamEngine for CassandraEngine { LIMIT {limit}" ) }; + let records_read = async { + if let Some(ref after) = after_sequence { + self.session + .query_with_values( + &query, + cdrs_tokio::query_values!(shard_id.as_str(), after.as_str()), + ) + .await + } else { + self.session + .query_with_values(&query, cdrs_tokio::query_values!(shard_id.as_str())) + .await + } + }; - let result = if let Some(ref after) = after_sequence { - self.session - .query_with_values( - &query, - cdrs_tokio::query_values!(shard_id.as_str(), after.as_str()), - ) - .await - } else { - self.session - .query_with_values(&query, cdrs_tokio::query_values!(shard_id.as_str())) - .await + let (shard_result, result) = futures::join!(shard_read, records_read); + let shard_table_id: Option = shard_result + .map_err(|e| StorageError::Internal(e.to_string()))? + .response_body() + .map_err(|e| StorageError::Internal(e.to_string()))? + .into_rows() + .unwrap_or_default() + .into_iter() + .next() + .and_then(|row| row.get_r_by_name("table_id").ok()); + if shard_table_id.as_deref() != Some(table_id.as_str()) + || !shard_in_stream(&table_id, &shard_id, table_label.as_deref(), &stream_label) + { + return Err(extenddb_storage::util::stream_not_found(&stream_arn)); } - .map_err(|e| StorageError::Internal(e.to_string()))?; let rows = result + .map_err(|e| StorageError::Internal(e.to_string()))? .response_body() .map_err(|e| StorageError::Internal(e.to_string()))? .into_rows() @@ -233,11 +278,10 @@ impl StreamEngine for CassandraEngine { .into_rows() .unwrap_or_default(); - let row = rows.into_iter().next().ok_or_else(|| { - StorageError::TableNotFound(format!( - "Requested resource not found: Stream: {stream_arn_str} not found." - )) - })?; + let row = rows + .into_iter() + .next() + .ok_or_else(|| extenddb_storage::util::stream_not_found(&stream_arn_str))?; let ks_json: String = row .get_r_by_name("key_schema") @@ -493,47 +537,9 @@ impl StreamEngine for CassandraEngine { let stream_arn = stream_arn.to_string(); let shard_id = shard_id.to_string(); Box::pin(async move { - let (table_name, stream_label) = parse_stream_arn(&stream_arn)?; - let catalog_keyspace = self.catalog_keyspace(); + let table_id = self.stream_table_id(&account_id, &stream_arn).await?; let account_keyspace = self.account_keyspace_for(&account_id); - let query = format!( - "SELECT table_id FROM {catalog_keyspace}.tables \ - WHERE account_id = ? AND table_name = ? AND stream_label = ? \ - ALLOW FILTERING" - ); - let result = self - .session - .query_with_values( - &query, - cdrs_tokio::query_values!( - account_id.as_str(), - table_name.as_str(), - stream_label.as_str() - ), - ) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - - let rows = result - .response_body() - .map_err(|e| StorageError::Internal(e.to_string()))? - .into_rows() - .unwrap_or_default(); - - let table_id: String = rows - .into_iter() - .next() - .ok_or_else(|| { - StorageError::TableNotFound(format!( - "Requested resource not found: Stream: {stream_arn} not found." - )) - }) - .and_then(|row| { - row.get_r_by_name("table_id") - .map_err(|e| StorageError::Internal(e.to_string())) - })?; - let shard_query = format!( "SELECT shard_id FROM {account_keyspace}.stream_shards \ WHERE shard_id = ? AND table_id = ? ALLOW FILTERING" @@ -554,9 +560,7 @@ impl StreamEngine for CassandraEngine { .is_some_and(|r| !r.is_empty()); if !found { - return Err(StorageError::TableNotFound(format!( - "Requested resource not found: Stream: {stream_arn} not found." - ))); + return Err(extenddb_storage::util::stream_not_found(&stream_arn)); } Ok(()) }) diff --git a/crates/storage-cassandra/tests/direct/stream_ownership.rs b/crates/storage-cassandra/tests/direct/stream_ownership.rs new file mode 100644 index 000000000..c141e3c7f --- /dev/null +++ b/crates/storage-cassandra/tests/direct/stream_ownership.rs @@ -0,0 +1,268 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! `get_stream_records` reads a shard only under the stream its iterator names. +//! +//! The catalog read is keyed on the caller's account and the table name the +//! ARN carries, the shard read on that table's id, and the stream label is +//! decided in Rust from the table's current label. Each test here removes one +//! of those from the picture (same label on two tables, same name in two +//! accounts, a new label on the same shards) so that one predicate alone +//! decides the answer. + +use std::collections::BTreeMap; +use std::sync::Arc; + +use extenddb_core::types::{AttributeValue, StreamViewType, TableKeyInfo}; +use extenddb_storage::error::StorageError; +use extenddb_storage::util::stream_arn; +use extenddb_storage::{DataEngine, StreamCapture, StreamEngine, StreamRecordsResult}; +use extenddb_storage_cassandra::CassandraEngine; + +use crate::helpers::setup_engine; +use crate::streams::{setup_stream_table, shard_for}; + +const REGION: &str = "us-east-1"; +const LABEL: &str = "2026-01-01T00:00:00.000"; +const LATER_LABEL: &str = "2026-02-02T00:00:00.000"; + +/// Pin a table's stream label so two tables can share one. +async fn relabel(engine: &CassandraEngine, key_info: &TableKeyInfo, label: &str) { + let catalog_keyspace = engine.catalog_keyspace(); + let query = format!( + "UPDATE {catalog_keyspace}.tables SET stream_label = ? \ + WHERE account_id = ? AND table_name = ?" + ); + engine + .session_arc() + .query_with_values( + &query, + cdrs_tokio::query_values!( + label, + key_info.account_id.as_str(), + key_info.table_name.as_str() + ), + ) + .await + .expect("relabel"); +} + +/// A second stream-enabled table in an account `setup_stream_table` created. +async fn stream_table_in( + engine: &CassandraEngine, + account_id: &str, + table_name: &str, +) -> TableKeyInfo { + use extenddb_core::types::{ + AttributeDefinition, CreateTableInput, KeySchemaElement, KeyType, ScalarAttributeType, + StreamSpecification, + }; + use extenddb_storage::TableEngine; + let key_schema = vec![KeySchemaElement { + attribute_name: "id".to_string(), + key_type: KeyType::Hash, + }]; + let attribute_definitions = vec![AttributeDefinition { + attribute_name: "id".to_string(), + attribute_type: ScalarAttributeType::S, + }]; + let desc = engine + .create_table( + account_id, + CreateTableInput { + table_name: table_name.to_string(), + key_schema: key_schema.clone(), + attribute_definitions: attribute_definitions.clone(), + stream_specification: Some(StreamSpecification { + stream_enabled: true, + stream_view_type: Some(StreamViewType::KeysOnly), + }), + ..Default::default() + }, + ) + .await + .expect("create the second table"); + TableKeyInfo { + vector_indexes: Vec::new(), + table_name: table_name.to_string(), + account_id: account_id.to_string(), + table_id: desc.table_id, + key_schema: key_schema.clone(), + base_key_schema: key_schema, + attribute_definitions, + has_lsi: false, + global_secondary_indexes: Vec::new(), + local_secondary_indexes: Vec::new(), + stream_specification: desc.stream_specification, + } +} + +/// One item written under `pk`, which lands in `shard_for(pk, table_id)`. +async fn put(engine: &CassandraEngine, key_info: &TableKeyInfo, pk: &str, val: u32) { + let mut item = BTreeMap::new(); + item.insert("id".to_string(), AttributeValue::S(pk.to_string())); + item.insert("val".to_string(), AttributeValue::N(val.to_string())); + let capture = StreamCapture { + view_type: StreamViewType::KeysOnly, + user_identity: None, + region: Arc::from(REGION), + }; + engine + .put_item( + key_info, + item, + false, + None, + &Default::default(), + Some(&capture), + ) + .await + .expect("put_item"); +} + +fn arn_of(key_info: &TableKeyInfo, label: &str) -> String { + stream_arn(REGION, &key_info.account_id, &key_info.table_name, label) +} + +fn assert_stream_not_found(result: StreamRecordsResult, arn: &str) { + let StorageError::TableNotFound(expected) = extenddb_storage::util::stream_not_found(arn) + else { + unreachable!("stream_not_found is a TableNotFound error") + }; + match result { + Err(StorageError::TableNotFound(msg)) => assert_eq!(msg, expected), + other => panic!("expected the stream not-found error, got {other:?}"), + } +} + +/// Two tables of one account share a stream label: a shard of table `b` +/// under table `a`'s ARN is refused, so the table name alone decides. +#[tokio::test] +async fn test_another_tables_shard_with_the_same_label_is_not_found() { + if crate::helpers::skip_without_cassandra() { + return; + } + let engine = setup_engine().await; + let (a, _) = setup_stream_table(&engine, "OwnA").await; + let b = stream_table_in(&engine, &a.account_id, "OwnB").await; + relabel(&engine, &a, LABEL).await; + relabel(&engine, &b, LABEL).await; + put(&engine, &a, "pk-a", 1).await; + put(&engine, &b, "pk-b", 1).await; + let shard_a = shard_for("pk-a", &a.table_id); + let shard_b = shard_for("pk-b", &b.table_id); + let arn_a = arn_of(&a, LABEL); + let arn_b = arn_of(&b, LABEL); + + // Both directions, so whichever table a lookup without the name would + // pick, one of them reads a shard that is not its own. + assert_stream_not_found( + engine + .get_stream_records(&a.account_id, &arn_a, &shard_b, None, 10) + .await, + &arn_a, + ); + assert_stream_not_found( + engine + .get_stream_records(&a.account_id, &arn_b, &shard_a, None, 10) + .await, + &arn_b, + ); + assert_stream_not_found( + engine + .get_stream_records( + &a.account_id, + &arn_a, + "shardId-forged-000000000000", + None, + 10, + ) + .await, + &arn_a, + ); + // The control: table A's own shard reads under A's ARN. + let (records, _) = engine + .get_stream_records(&a.account_id, &arn_a, &shard_a, None, 10) + .await + .expect("own shard reads"); + assert_eq!(records.len(), 1); +} + +/// Two accounts own a table of the same name and label: the caller's ARN +/// resolves in the caller's account only, so the other account's shard is +/// refused under either spelling of the account in the ARN. +#[tokio::test] +async fn test_another_accounts_same_named_stream_is_not_found() { + if crate::helpers::skip_without_cassandra() { + return; + } + let engine = setup_engine().await; + // `setup_stream_table` creates a fresh account per call, so these two + // tables of one name belong to two accounts. + let (own, _) = setup_stream_table(&engine, "SameName").await; + let (other, _) = setup_stream_table(&engine, "SameName").await; + assert_ne!(own.account_id, other.account_id); + relabel(&engine, &own, LABEL).await; + relabel(&engine, &other, LABEL).await; + put(&engine, &own, "pk-o", 1).await; + put(&engine, &other, "pk-o", 1).await; + let own_shard = shard_for("pk-o", &own.table_id); + let other_shard = shard_for("pk-o", &other.table_id); + + // From both callers, so whichever table a lookup without the account + // would pick, one caller reads the other account's shard. + for (caller, foreign_shard, own_shard) in [ + (&own, &other_shard, &own_shard), + (&other, &own_shard, &other_shard), + ] { + for arn in [ + stream_arn(REGION, &own.account_id, "SameName", LABEL), + stream_arn(REGION, &other.account_id, "SameName", LABEL), + ] { + assert_stream_not_found( + engine + .get_stream_records(&caller.account_id, &arn, foreign_shard, None, 10) + .await, + &arn, + ); + let (records, _) = engine + .get_stream_records(&caller.account_id, &arn, own_shard, None, 10) + .await + .expect("own shard reads"); + assert_eq!(records.len(), 1); + } + } +} + +/// The table takes a new stream label while keeping its shard rows: the old +/// ARN's iterator is refused and the new ARN reads every record. +#[tokio::test] +async fn test_an_earlier_label_on_the_same_shards_is_not_found() { + if crate::helpers::skip_without_cassandra() { + return; + } + let engine = setup_engine().await; + let (t, _) = setup_stream_table(&engine, "Relabel").await; + relabel(&engine, &t, LABEL).await; + put(&engine, &t, "pk-r", 1).await; + let shard = shard_for("pk-r", &t.table_id); + let old_arn = arn_of(&t, LABEL); + engine + .get_stream_records(&t.account_id, &old_arn, &shard, None, 10) + .await + .expect("current label reads"); + + relabel(&engine, &t, LATER_LABEL).await; + put(&engine, &t, "pk-r", 2).await; + assert_stream_not_found( + engine + .get_stream_records(&t.account_id, &old_arn, &shard, None, 10) + .await, + &old_arn, + ); + let (records, _) = engine + .get_stream_records(&t.account_id, &arn_of(&t, LATER_LABEL), &shard, None, 10) + .await + .expect("new label reads"); + assert_eq!(records.len(), 2); +} diff --git a/crates/storage-cassandra/tests/direct/streams.rs b/crates/storage-cassandra/tests/direct/streams.rs index 865774d3a..16c86540e 100644 --- a/crates/storage-cassandra/tests/direct/streams.rs +++ b/crates/storage-cassandra/tests/direct/streams.rs @@ -22,7 +22,7 @@ fn capture(view_type: StreamViewType) -> StreamCapture { } /// Create a table with streams enabled and return its key_info + stream_label. -async fn setup_stream_table( +pub(crate) async fn setup_stream_table( engine: &CassandraEngine, table_name: &str, ) -> (extenddb_core::types::TableKeyInfo, String) { @@ -139,7 +139,7 @@ async fn fetch_stream_records( .collect() } -fn shard_for(pk: &str, table_id: &str) -> String { +pub(crate) fn shard_for(pk: &str, table_id: &str) -> String { extenddb_storage_cassandra::stream_util::assign_shard_id(pk, table_id) } @@ -492,7 +492,13 @@ async fn test_get_stream_records_returns_written_records() { use extenddb_storage::StreamEngine; let engine = setup_engine().await; - let (key_info, _) = setup_stream_table(&engine, "GetStreamRecords").await; + let (key_info, stream_label) = setup_stream_table(&engine, "GetStreamRecords").await; + let arn = extenddb_storage::util::stream_arn( + "us-east-1", + &key_info.account_id, + &key_info.table_name, + &stream_label, + ); for i in 0..3u32 { let mut item = BTreeMap::new(); @@ -513,7 +519,7 @@ async fn test_get_stream_records_returns_written_records() { let shard_id = shard_for("pk-read", &key_info.table_id); let (records, last_seq) = engine - .get_stream_records(&key_info.account_id, &shard_id, None, 10) + .get_stream_records(&key_info.account_id, &arn, &shard_id, None, 10) .await .unwrap(); @@ -533,7 +539,13 @@ async fn test_get_stream_records_after_sequence_paginates() { use extenddb_storage::StreamEngine; let engine = setup_engine().await; - let (key_info, _) = setup_stream_table(&engine, "GetStreamRecordsPaginate").await; + let (key_info, stream_label) = setup_stream_table(&engine, "GetStreamRecordsPaginate").await; + let arn = extenddb_storage::util::stream_arn( + "us-east-1", + &key_info.account_id, + &key_info.table_name, + &stream_label, + ); for i in 0..4u32 { let mut item = BTreeMap::new(); @@ -554,13 +566,19 @@ async fn test_get_stream_records_after_sequence_paginates() { let shard_id = shard_for("pk-page", &key_info.table_id); let (first_page, last_seq) = engine - .get_stream_records(&key_info.account_id, &shard_id, None, 2) + .get_stream_records(&key_info.account_id, &arn, &shard_id, None, 2) .await .unwrap(); assert_eq!(first_page.len(), 2); let (second_page, _) = engine - .get_stream_records(&key_info.account_id, &shard_id, last_seq.as_deref(), 10) + .get_stream_records( + &key_info.account_id, + &arn, + &shard_id, + last_seq.as_deref(), + 10, + ) .await .unwrap(); assert_eq!(second_page.len(), 2); diff --git a/crates/storage-cassandra/tests/direct_integration.rs b/crates/storage-cassandra/tests/direct_integration.rs index 43eb28b81..c3cae92f9 100644 --- a/crates/storage-cassandra/tests/direct_integration.rs +++ b/crates/storage-cassandra/tests/direct_integration.rs @@ -45,6 +45,8 @@ mod roles; mod scan; #[path = "direct/settings_store.rs"] mod settings_store; +#[path = "direct/stream_ownership.rs"] +mod stream_ownership; #[path = "direct/streams.rs"] mod streams; #[path = "direct/table_engine.rs"] diff --git a/crates/storage-mongodb/src/stream_engine.rs b/crates/storage-mongodb/src/stream_engine.rs index e976a667a..d4609e65d 100644 --- a/crates/storage-mongodb/src/stream_engine.rs +++ b/crates/storage-mongodb/src/stream_engine.rs @@ -22,7 +22,7 @@ use extenddb_core::types::{ }; use extenddb_storage::StreamEngine; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{parse_stream_arn, stream_arn}; +use extenddb_storage::util::{parse_stream_arn, shard_in_stream, stream_arn}; use extenddb_storage::{StreamListResult, StreamRecordsResult}; use crate::MongoEngine; @@ -285,90 +285,99 @@ impl StreamEngine for MongoEngine { fn get_stream_records( &self, account_id: &str, + stream_arn_val: &str, shard_id: &str, after_sequence: Option<&str>, limit: i64, ) -> BoxFuture<'_, StreamRecordsResult> { let account_id = account_id.to_owned(); + let stream_arn_val = stream_arn_val.to_owned(); let shard_id = shard_id.to_owned(); let after_sequence = after_sequence.map(std::borrow::ToOwned::to_owned); Box::pin(async move { - // Ownership guard: only return records if the shard's backing table - // belongs to the calling account. `stream_shards`/`stream_records` - // live in the data database while the `tables` catalog (which - // carries account_id inside its compound `_id`) lives in the catalog - // database, so ownership is resolved in two steps across the two - // databases: shard_id -> table_id (data db), then table_id + - // account_id (catalog db). Mirrors the postgres backend. - let shards_coll = self.data_db.collection::("stream_shards"); - let shard_doc = shards_coll - .find_one(doc! { "shard_id": &shard_id }) + let (table_name, stream_label) = parse_stream_arn(&stream_arn_val)?; + // The `tables` catalog lives in the catalog database and + // `stream_shards`/`stream_records` in the data database, so the + // stream resolves in two round trips: the table the ARN names + // (catalog: its id and current stream label), then the shard and + // its records in one aggregation (data). The shard document drives + // the aggregation with the records joined by `$lookup` and + // unwound, so each record is its own output document and the page + // is bounded by `limit` rather than by the 16 MB document cap (the + // server coalesces the `$lookup` with the `$unwind` and never + // materializes the array). A shard of the named table with no + // records returns one document without `record_data`, and a shard + // that is not that table's (another table, or another account) + // returns no document. Whether the shard belongs to the stream the + // ARN's label names is decided in Rust from the table's label and + // the shard id (`shard_in_stream`). + let table_doc = self + .catalog_db + .collection::("tables") + .find_one(doc! { + "_id": { "account_id": &account_id, "table_name": &table_name }, + }) .await .map_err(|e| StorageError::Internal(e.to_string()))?; - let owned = match shard_doc.as_ref().and_then(|d| d.get_str("table_id").ok()) { - // Look the table up by its globally-unique `table_id` (a - // top-level field), then compare the owning account_id read out - // of the compound `_id` subdocument in Rust. account_id lives - // only inside the embedded `_id` doc; resolving ownership in - // Rust after a single table_id lookup keeps the check explicit - // and avoids depending on query-time behaviour of a partial - // `_id.account_id` path. - Some(table_id) => self - .catalog_db - .collection::("tables") - .find_one(doc! { "table_id": table_id }) - .await - .map_err(|e| StorageError::Internal(e.to_string()))? - .as_ref() - .and_then(|t| t.get_document("_id").ok()) - .and_then(|id| id.get_str("account_id").ok()) - .is_some_and(|owner| owner == account_id), - None => false, + let Some(table_doc) = table_doc else { + return Err(extenddb_storage::util::stream_not_found(&stream_arn_val)); }; - // A shard iterator that resolves to a shard the caller does not own - // (catalog step fails) or one that does not exist (data step fails) - // is rejected identically. Real DynamoDB returns - // `ValidationException: Invalid ShardIterator` for a GetRecords - // iterator it did not issue, and does not distinguish "exists but not - // yours" from "does not exist" — so neither do we. - if !owned { - return Err(StorageError::Validation("Invalid ShardIterator".to_owned())); + let table_id = table_doc + .get_str("table_id") + .map_err(|e| StorageError::Internal(e.to_string()))? + .to_owned(); + let table_label = table_doc.get_str("stream_label").ok(); + + let mut record_match = doc! { "$eq": ["$shard_id", "$$shard"] }; + if let Some(ref after) = after_sequence { + record_match = doc! { "$and": [ + record_match, + { "$gt": ["$sequence_number", after] }, + ] }; } - - let records_coll = self.data_db.collection::("stream_records"); - - let filter = if let Some(ref after) = after_sequence { - doc! { - "shard_id": &shard_id, - "sequence_number": { "$gt": after }, - } - } else { - doc! { "shard_id": &shard_id } - }; - - let opts = FindOptions::builder() - .sort(doc! { "sequence_number": 1 }) - .limit(limit) - .build(); - - let cursor = records_coll - .find(filter) - .with_options(opts) + let pipeline = vec![ + doc! { "$match": { "shard_id": &shard_id, "table_id": &table_id } }, + doc! { "$limit": 1 }, + doc! { "$lookup": { + "from": "stream_records", + "let": { "shard": "$shard_id" }, + "pipeline": [ + { "$match": { "$expr": record_match } }, + { "$sort": { "sequence_number": 1 } }, + { "$limit": limit }, + { "$project": { "_id": 0, "record_data": 1 } }, + ], + "as": "records", + } }, + doc! { "$unwind": { "path": "$records", "preserveNullAndEmptyArrays": true } }, + doc! { "$project": { + "_id": 0, + "shard_id": 1, + "record_data": "$records.record_data", + } }, + ]; + let cursor = self + .data_db + .collection::("stream_shards") + .aggregate(pipeline) .await .map_err(|e| StorageError::Internal(e.to_string()))?; - let docs: Vec = cursor .try_collect() .await .map_err(|e| StorageError::Internal(e.to_string()))?; + let Some(row_shard_id) = docs.first().and_then(|d| d.get_str("shard_id").ok()) else { + return Err(extenddb_storage::util::stream_not_found(&stream_arn_val)); + }; + if !shard_in_stream(&table_id, row_shard_id, table_label, &stream_label) { + return Err(extenddb_storage::util::stream_not_found(&stream_arn_val)); + } let records: Vec = docs .into_iter() - .map(|d| { - let record_bson = d - .get("record_data") - .ok_or_else(|| StorageError::Internal("Missing record_data".to_owned()))?; - let json_val: serde_json::Value = bson::from_bson(record_bson.clone()) + .filter_map(|mut d| d.remove("record_data")) + .map(|record_bson| { + let json_val: serde_json::Value = bson::from_bson(record_bson) .map_err(|e| StorageError::Internal(e.to_string()))?; serde_json::from_value(json_val) .map_err(|e| StorageError::Internal(e.to_string())) @@ -400,11 +409,7 @@ impl StreamEngine for MongoEngine { }) .await .map_err(|e| StorageError::Internal(e.to_string()))? - .ok_or_else(|| { - StorageError::TableNotFound(format!( - "Requested resource not found: Stream: {stream_arn_val} not found." - )) - })?; + .ok_or_else(|| extenddb_storage::util::stream_not_found(&stream_arn_val))?; let key_schema = table_doc .get("key_schema") @@ -709,9 +714,7 @@ impl StreamEngine for MongoEngine { .map_err(|e| StorageError::Internal(e.to_string()))?; let Some(table_doc) = table_doc else { - return Err(StorageError::TableNotFound(format!( - "Requested resource not found: Stream: {stream_arn_val} not found." - ))); + return Err(extenddb_storage::util::stream_not_found(&stream_arn_val)); }; let table_id = table_doc.get_str("table_id").unwrap_or_default(); @@ -723,9 +726,7 @@ impl StreamEngine for MongoEngine { .map_err(|e| StorageError::Internal(e.to_string()))?; if exists.is_none() { - return Err(StorageError::TableNotFound(format!( - "Requested resource not found: Stream: {stream_arn_val} not found." - ))); + return Err(extenddb_storage::util::stream_not_found(&stream_arn_val)); } Ok(()) }) @@ -808,3 +809,9 @@ mod tests { assert!(s2 < s3); } } + +/// Live `get_stream_records` tests against a MongoDB server; see the module +/// documentation. +#[cfg(test)] +#[path = "stream_records_live_tests.rs"] +mod records_live_tests; diff --git a/crates/storage-mongodb/src/stream_records_live_tests.rs b/crates/storage-mongodb/src/stream_records_live_tests.rs new file mode 100644 index 000000000..240168d46 --- /dev/null +++ b/crates/storage-mongodb/src/stream_records_live_tests.rs @@ -0,0 +1,326 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! Live MongoDB tests for `get_stream_records`: the shard is resolved against +//! the stream its iterator names, and a page is streamed record by record so +//! its size is bounded by `Limit`, not by the 16 MB document cap. +//! +//! They read a MongoDB URI from `EXTENDDB_TEST_MONGODB_URI` (a replica set is +//! not needed here) and are `#[ignore]`d so `cargo test --lib` passes without +//! one; run them with +//! `cargo test -p extenddb-storage-mongodb --lib -- --ignored records_live`. +//! Each test uses its own catalog and data databases and drops them at the +//! end. Rows are inserted directly in the shape the engine writes them. + +use std::sync::Arc; + +use extenddb_storage::StreamEngine; +use extenddb_storage::error::StorageError; +use extenddb_storage::util::stream_arn; +use mongodb::bson::{self, Document, doc}; + +use crate::MongoEngine; + +const REGION: &str = "us-east-1"; +const ACCOUNT: &str = "111111111111"; +const OTHER_ACCOUNT: &str = "222222222222"; +const LABEL: &str = "2026-01-01T00:00:00.000"; +const LATER_LABEL: &str = "2026-02-02T00:00:00.000"; + +fn uri() -> String { + std::env::var("EXTENDDB_TEST_MONGODB_URI").expect( + "set EXTENDDB_TEST_MONGODB_URI to a MongoDB server, e.g. \ + mongodb://127.0.0.1:27771/?directConnection=true", + ) +} + +struct Fixture { + engine: MongoEngine, +} + +impl Fixture { + async fn new() -> Self { + let tag = uuid::Uuid::new_v4().simple().to_string(); + let client = mongodb::Client::with_uri_str(uri()).await.expect("client"); + let slots = crate::backup_engine::restore_server_slot_count(10); + let engine = MongoEngine { + client: client.clone(), + catalog_db: client.database(&format!("s{}_cat", &tag[..12])), + data_db: client.database(&format!("s{}_data", &tag[..12])), + region: REGION.to_owned(), + gsi_cache: Arc::default(), + tx_read_concern: mongodb::options::ReadConcern::snapshot(), + backup_liveness: crate::backup_engine::BackupLiveness::default(), + restore_server_slot_count: slots, + restore_server_slots: Arc::new(tokio::sync::Semaphore::new(slots)), + restore_account_slots: crate::backup_engine::RestoreAccountSlots::default(), + }; + Self { engine } + } + + /// A stream-enabled table with four shards, as `create_table` and + /// `init_stream_shards` write them. Returns the table id. + async fn stream_table(&self, account: &str, name: &str, label: &str) -> String { + let table_id = format!("{name}-{account}-id"); + self.engine + .catalog_db + .collection::("tables") + .insert_one(doc! { + "_id": { "account_id": account, "table_name": name }, + "table_id": &table_id, + "stream_label": label, + }) + .await + .expect("table"); + for i in 0..4 { + self.engine + .data_db + .collection::("stream_shards") + .insert_one(doc! { + "shard_id": super::build_shard_id(&table_id, i), + "table_id": &table_id, + "starting_sequence_number": format!("{:021}", 0), + }) + .await + .expect("shard"); + } + table_id + } + + async fn write_records(&self, shard_id: &str, table_id: &str, seqs: &[u64], blob: usize) { + for seq in seqs { + let record_bson = bson::to_bson(&record_json(*seq, blob)).expect("record json to bson"); + self.engine + .data_db + .collection::("stream_records") + .insert_one(doc! { + "sequence_number": format!("{seq:021}"), + "shard_id": shard_id, + "table_id": table_id, + "event_name": "INSERT", + "record_data": record_bson, + "created_at": bson::DateTime::now(), + }) + .await + .expect("record"); + } + } + + async fn relabel(&self, account: &str, name: &str, label: &str) { + self.engine + .catalog_db + .collection::("tables") + .update_one( + doc! { "_id": { "account_id": account, "table_name": name } }, + doc! { "$set": { "stream_label": label } }, + ) + .await + .expect("relabel"); + } + + async fn finish(self) { + self.engine.catalog_db.drop().await.expect("drop catalog"); + self.engine.data_db.drop().await.expect("drop data"); + } +} + +/// A stream record whose new image carries a string of `blob` bytes. +fn record_json(seq: u64, blob: usize) -> serde_json::Value { + serde_json::json!({ + "eventID": format!("e{seq}"), + "eventName": "INSERT", + "eventVersion": "1.1", + "eventSource": "aws:dynamodb", + "awsRegion": REGION, + "dynamodb": { + "ApproximateCreationDateTime": 1_700_000_000, + "Keys": { "pk": { "S": format!("k{seq}") } }, + "NewImage": { + "pk": { "S": format!("k{seq}") }, + "blob": { "S": "x".repeat(blob) }, + }, + "SequenceNumber": format!("{seq:021}"), + "SizeBytes": blob, + "StreamViewType": "NEW_IMAGE" + } + }) +} + +fn assert_stream_not_found(result: extenddb_storage::StreamRecordsResult, arn: &str) { + let StorageError::TableNotFound(expected) = extenddb_storage::util::stream_not_found(arn) + else { + unreachable!("stream_not_found is a TableNotFound error") + }; + match result { + Err(StorageError::TableNotFound(msg)) => assert_eq!(msg, expected), + other => panic!("expected the stream not-found error, got {other:?}"), + } +} + +/// 40 records of about 450 KB each on one shard, read with `Limit` 1000: the +/// page is about 18 MB, past the 16 MB BSON document cap, so it can only be +/// read if each record is its own output document. +#[tokio::test] +#[ignore = "needs a MongoDB server in EXTENDDB_TEST_MONGODB_URI"] +async fn records_live_a_page_over_16_mb_is_read_in_full() { + let f = Fixture::new().await; + let table_id = f.stream_table(ACCOUNT, "big", LABEL).await; + let arn = stream_arn(REGION, ACCOUNT, "big", LABEL); + let shard_id = super::build_shard_id(&table_id, 0); + // Inserted in descending order so the read's sort, not insertion order, + // produces the sequence order the page is asserted in. + let seqs: Vec = (1..=40).rev().collect(); + f.write_records(&shard_id, &table_id, &seqs, 450_000).await; + + let (records, last) = f + .engine + .get_stream_records(ACCOUNT, &arn, &shard_id, None, 1000) + .await + .expect("a page over 16 MB reads"); + assert_eq!(records.len(), 40); + let got: Vec<&str> = records + .iter() + .map(|r| r.dynamodb.sequence_number.as_str()) + .collect(); + let want: Vec = (1..=40).map(|i| format!("{i:021}")).collect(); + assert_eq!(got, want.iter().map(String::as_str).collect::>()); + assert_eq!(last.as_deref(), Some(format!("{:021}", 40).as_str())); + + // The idle page after the last record, and a bounded page in the middle. + let (rest, last) = f + .engine + .get_stream_records(ACCOUNT, &arn, &shard_id, last.as_deref(), 1000) + .await + .expect("idle page"); + assert!(rest.is_empty()); + assert!(last.is_none()); + let (page, _) = f + .engine + .get_stream_records(ACCOUNT, &arn, &shard_id, Some(&format!("{:021}", 10)), 5) + .await + .expect("bounded page"); + assert_eq!(page.len(), 5); + assert_eq!(page[0].dynamodb.sequence_number, format!("{:021}", 11)); + f.finish().await; +} + +/// The table takes a new stream label while keeping its shards: the old +/// ARN's iterator is refused and the new ARN reads. +#[tokio::test] +#[ignore = "needs a MongoDB server in EXTENDDB_TEST_MONGODB_URI"] +async fn records_live_an_earlier_label_on_the_same_shards_is_not_found() { + let f = Fixture::new().await; + let table_id = f.stream_table(ACCOUNT, "t", LABEL).await; + let shard_id = super::build_shard_id(&table_id, 0); + f.write_records(&shard_id, &table_id, &[1], 10).await; + let old_arn = stream_arn(REGION, ACCOUNT, "t", LABEL); + f.engine + .get_stream_records(ACCOUNT, &old_arn, &shard_id, None, 10) + .await + .expect("current label reads"); + + f.relabel(ACCOUNT, "t", LATER_LABEL).await; + f.write_records(&shard_id, &table_id, &[2], 10).await; + assert_stream_not_found( + f.engine + .get_stream_records(ACCOUNT, &old_arn, &shard_id, None, 10) + .await, + &old_arn, + ); + let new_arn = stream_arn(REGION, ACCOUNT, "t", LATER_LABEL); + let (records, _) = f + .engine + .get_stream_records(ACCOUNT, &new_arn, &shard_id, None, 10) + .await + .expect("new label reads"); + assert_eq!(records.len(), 2); + f.finish().await; +} + +/// Two accounts own a table of the same name and label: the caller's ARN +/// resolves in the caller's account only, so the other account's shard is +/// refused under the ARN the caller spells. +#[tokio::test] +#[ignore = "needs a MongoDB server in EXTENDDB_TEST_MONGODB_URI"] +async fn records_live_another_accounts_same_named_stream_is_not_found() { + let f = Fixture::new().await; + let own_id = f.stream_table(ACCOUNT, "t", LABEL).await; + let other_id = f.stream_table(OTHER_ACCOUNT, "t", LABEL).await; + let own_shard = super::build_shard_id(&own_id, 0); + let other_shard = super::build_shard_id(&other_id, 0); + f.write_records(&own_shard, &own_id, &[1], 10).await; + f.write_records(&other_shard, &other_id, &[1], 10).await; + + // From both callers, so whichever table a lookup without the account + // would pick, one caller reads the other account's shard. + for (caller, foreign_shard, own) in [ + (ACCOUNT, &other_shard, &own_shard), + (OTHER_ACCOUNT, &own_shard, &other_shard), + ] { + for arn in [ + stream_arn(REGION, ACCOUNT, "t", LABEL), + stream_arn(REGION, OTHER_ACCOUNT, "t", LABEL), + ] { + assert_stream_not_found( + f.engine + .get_stream_records(caller, &arn, foreign_shard, None, 10) + .await, + &arn, + ); + // The control: the caller's own shard reads under either spelling + // of the account, since storage binds the caller's account and + // parses only the table name and label out of the ARN. + let (records, _) = f + .engine + .get_stream_records(caller, &arn, own, None, 10) + .await + .expect("own shard reads"); + assert_eq!(records.len(), 1); + } + } + f.finish().await; +} + +/// Two tables of one account share a stream label: a shard of table `b` +/// under table `a`'s ARN is refused, so the table name alone decides. +#[tokio::test] +#[ignore = "needs a MongoDB server in EXTENDDB_TEST_MONGODB_URI"] +async fn records_live_another_tables_shard_with_the_same_label_is_not_found() { + let f = Fixture::new().await; + let a_id = f.stream_table(ACCOUNT, "a", LABEL).await; + let b_id = f.stream_table(ACCOUNT, "b", LABEL).await; + f.write_records(&super::build_shard_id(&a_id, 0), &a_id, &[1], 10) + .await; + f.write_records(&super::build_shard_id(&b_id, 0), &b_id, &[1, 2], 10) + .await; + let arn_a = stream_arn(REGION, ACCOUNT, "a", LABEL); + let arn_b = stream_arn(REGION, ACCOUNT, "b", LABEL); + + // Both directions, so whichever table a lookup without the name would + // pick, one of them reads a shard that is not its own. + assert_stream_not_found( + f.engine + .get_stream_records(ACCOUNT, &arn_a, &super::build_shard_id(&b_id, 0), None, 10) + .await, + &arn_a, + ); + assert_stream_not_found( + f.engine + .get_stream_records(ACCOUNT, &arn_b, &super::build_shard_id(&a_id, 0), None, 10) + .await, + &arn_b, + ); + assert_stream_not_found( + f.engine + .get_stream_records(ACCOUNT, &arn_a, "shardId-forged-000000000000", None, 10) + .await, + &arn_a, + ); + let (records, _) = f + .engine + .get_stream_records(ACCOUNT, &arn_a, &super::build_shard_id(&a_id, 0), None, 10) + .await + .expect("own shard reads"); + assert_eq!(records.len(), 1); + f.finish().await; +} diff --git a/crates/storage-postgres/src/stream_engine.rs b/crates/storage-postgres/src/stream_engine.rs index 6eee09658..0fb17674f 100755 --- a/crates/storage-postgres/src/stream_engine.rs +++ b/crates/storage-postgres/src/stream_engine.rs @@ -9,7 +9,7 @@ use extenddb_core::types::{ }; use extenddb_storage::StreamEngine; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{new_stream_label, parse_stream_arn, stream_arn}; +use extenddb_storage::util::{new_stream_label, parse_stream_arn, shard_in_stream, stream_arn}; use futures::future::BoxFuture; use sqlx::PgPool; @@ -195,58 +195,53 @@ impl StreamEngine for PostgresEngine { fn get_stream_records( &self, account_id: &str, + stream_arn: &str, shard_id: &str, after_sequence: Option<&str>, limit: i64, ) -> BoxFuture<'_, Result<(Vec, Option), StorageError>> { let account_id = account_id.to_string(); + let stream_arn = stream_arn.to_string(); let shard_id = shard_id.to_string(); let after_sequence = after_sequence.map(std::string::ToString::to_string); Box::pin(async move { - // Ownership guard: only return records if the shard's backing table - // belongs to the calling account. `stream_shards`/`stream_records` - // rows live in the data database while the `tables` catalog (which - // carries `account_id`) lives in the catalog database, so ownership - // is resolved in two steps across the two pools: shard -> table_id - // (data pool), then table_id + account_id (catalog pool). A shard - // iterator for a table the caller does not own is rejected as an - // invalid shard iterator (below). - let shard_table_id: Option<(String,)> = - sqlx::query_as("SELECT table_id FROM stream_shards WHERE shard_id = $1") - .bind(&shard_id) - .fetch_optional(&self.data_pool) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - let owned = match &shard_table_id { - Some((table_id,)) => sqlx::query_as::<_, (i32,)>( - "SELECT 1 FROM tables WHERE table_id = $1 AND account_id = $2", - ) - .bind(table_id) - .bind(account_id) - .fetch_optional(&self.pool) - .await - .map_err(|e| StorageError::Internal(e.to_string()))? - .is_some(), - None => false, + let (table_name, stream_label) = parse_stream_arn(&stream_arn)?; + // The `tables` catalog lives in the catalog database and + // `stream_shards`/`stream_records` in the data database, which may + // be a different server, so the stream resolves in two round + // trips: the table the ARN names (catalog pool: its id and current + // stream label), then the shard and its records in one query + // (data pool). The shard row drives the data query with the + // records LEFT JOINed: a shard of the named table with no records + // returns one row with a NULL record, and a shard that is not that + // table's (another table, or another account) returns no rows. + // Whether the shard belongs to the stream the ARN's label names is + // decided in Rust from the table's label and the shard id + // (`shard_in_stream`). + let table: Option<(String, Option)> = sqlx::query_as( + "SELECT table_id, stream_label FROM tables \ + WHERE account_id = $1 AND table_name = $2", + ) + .bind(&account_id) + .bind(&table_name) + .fetch_optional(&self.pool) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + let Some((table_id, table_label)) = table else { + return Err(extenddb_storage::util::stream_not_found(&stream_arn)); }; - // A shard iterator resolves to a shard the caller does not own (step - // 2 fails) or one that does not exist (step 1 fails). Real DynamoDB - // returns `ValidationException: Invalid ShardIterator` for a - // GetRecords iterator it did not issue, and does NOT distinguish - // "exists but not yours" from "does not exist" — so neither do we - // (both collapse here). Verified against DynamoDB - // Streams (us-east-1). - if !owned { - return Err(StorageError::Validation("Invalid ShardIterator".to_owned())); - } - let rows: Vec<(serde_json::Value,)> = if let Some(after) = after_sequence { + let rows: Vec<(String, Option)> = if let Some(after) = after_sequence + { sqlx::query_as( - "SELECT record_data FROM stream_records \ - WHERE shard_id = $1 AND sequence_number > $2 \ - ORDER BY sequence_number LIMIT $3", + "SELECT s.shard_id, r.record_data FROM stream_shards s \ + LEFT JOIN stream_records r \ + ON r.shard_id = s.shard_id AND r.sequence_number > $3 \ + WHERE s.shard_id = $1 AND s.table_id = $2 \ + ORDER BY r.sequence_number LIMIT $4", ) .bind(&shard_id) + .bind(&table_id) .bind(&after) .bind(limit) .fetch_all(&self.data_pool) @@ -254,20 +249,34 @@ impl StreamEngine for PostgresEngine { .map_err(|e| StorageError::Internal(e.to_string()))? } else { sqlx::query_as( - "SELECT record_data FROM stream_records \ - WHERE shard_id = $1 \ - ORDER BY sequence_number LIMIT $2", + "SELECT s.shard_id, r.record_data FROM stream_shards s \ + LEFT JOIN stream_records r ON r.shard_id = s.shard_id \ + WHERE s.shard_id = $1 AND s.table_id = $2 \ + ORDER BY r.sequence_number LIMIT $3", ) .bind(&shard_id) + .bind(&table_id) .bind(limit) .fetch_all(&self.data_pool) .await .map_err(|e| StorageError::Internal(e.to_string()))? }; + let Some((row_shard_id, _)) = rows.first() else { + return Err(extenddb_storage::util::stream_not_found(&stream_arn)); + }; + if !shard_in_stream( + &table_id, + row_shard_id, + table_label.as_deref(), + &stream_label, + ) { + return Err(extenddb_storage::util::stream_not_found(&stream_arn)); + } let records: Vec = rows .into_iter() - .map(|(data,)| { + .filter_map(|(_, data)| data) + .map(|data| { serde_json::from_value(data).map_err(|e| StorageError::Internal(e.to_string())) }) .collect::, _>>()?; @@ -302,11 +311,7 @@ impl StreamEngine for PostgresEngine { .map_err(|e| StorageError::Internal(e.to_string()))?; let (ks_json, _ad_json, stream_spec_json, table_status, table_id) = - row.ok_or_else(|| { - StorageError::TableNotFound(format!( - "Requested resource not found: Stream: {stream_arn} not found." - )) - })?; + row.ok_or_else(|| extenddb_storage::util::stream_not_found(&stream_arn))?; let key_schema = serde_json::from_value(ks_json) .map_err(|e| StorageError::Internal(e.to_string()))?; @@ -575,9 +580,7 @@ impl StreamEngine for PostgresEngine { .map_err(|e| StorageError::Internal(e.to_string()))?; let Some(table_id) = table_id else { - return Err(StorageError::TableNotFound(format!( - "Requested resource not found: Stream: {stream_arn} not found." - ))); + return Err(extenddb_storage::util::stream_not_found(&stream_arn)); }; let exists: Option<(i32,)> = @@ -589,9 +592,7 @@ impl StreamEngine for PostgresEngine { .map_err(|e| StorageError::Internal(e.to_string()))?; if exists.is_none() { - return Err(StorageError::TableNotFound(format!( - "Requested resource not found: Stream: {stream_arn} not found." - ))); + return Err(extenddb_storage::util::stream_not_found(&stream_arn)); } Ok(()) }) diff --git a/crates/storage-postgres/tests/stream_shard_ownership.rs b/crates/storage-postgres/tests/stream_shard_ownership.rs new file mode 100644 index 000000000..c2b2003b0 --- /dev/null +++ b/crates/storage-postgres/tests/stream_shard_ownership.rs @@ -0,0 +1,377 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! `get_stream_records` reads a shard only under the stream its iterator names. +//! +//! The catalog read is keyed on the caller's account and the table name the +//! ARN carries, the data read on that table's id, and the stream label is +//! decided in Rust from the table's current label. Each test here removes one +//! of those from the picture (same label on two tables, same name in two +//! accounts, a new label on the same shards) so that one predicate alone +//! decides the answer. +//! +//! Needs `EXTENDDB_TEST_PG_CONNECTION_STRING` (host-only, e.g. +//! `postgres://user:pass@127.0.0.1:5432`); skips otherwise. Each test builds +//! its own scratch database and drops it afterwards. + +use extenddb_core::types::{ + AttributeDefinition, BillingMode, CreateTableInput, KeySchemaElement, KeyType, + ScalarAttributeType, StreamSpecification, StreamViewType, +}; +use extenddb_storage::error::StorageError; +use extenddb_storage::util::stream_arn; +use extenddb_storage::{StreamEngine, StreamRecordsResult, TableEngine}; +use extenddb_storage_postgres::{PostgresConfig, PostgresEngine}; +use sqlx::PgPool; +use sqlx::postgres::PgPoolOptions; + +const ACCOUNT: &str = "111111111111"; +const OTHER_ACCOUNT: &str = "222222222222"; +const REGION: &str = "us-east-1"; +const LABEL: &str = "2026-01-01T00:00:00.000"; +const LATER_LABEL: &str = "2026-02-02T00:00:00.000"; + +struct Scratch { + engine: PostgresEngine, + db: PgPool, + admin: PgPool, + db_name: String, +} + +impl Scratch { + async fn cleanup(self) { + let Scratch { + engine, + db, + admin, + db_name, + } = self; + drop(engine); + db.close().await; + sqlx::query(&format!( + "DROP DATABASE IF EXISTS \"{db_name}\" WITH (FORCE)" + )) + .execute(&admin) + .await + .expect("drop the scratch database"); + admin.close().await; + } + + /// A stream-enabled table whose label is pinned to `label`. Returns its + /// table id and the id of its first shard. + async fn stream_table(&self, account: &str, name: &str, label: &str) -> (String, String) { + let desc = self + .engine + .create_table(account, stream_table(name)) + .await + .expect("create a stream-enabled table"); + sqlx::query("UPDATE tables SET stream_label = $1 WHERE table_id = $2") + .bind(label) + .bind(&desc.table_id) + .execute(&self.db) + .await + .expect("pin the stream label"); + let shard: String = sqlx::query_scalar( + "SELECT shard_id FROM stream_shards WHERE table_id = $1 ORDER BY shard_id LIMIT 1", + ) + .bind(&desc.table_id) + .fetch_one(&self.db) + .await + .expect("one shard id"); + (desc.table_id, shard) + } + + async fn write_records(&self, shard_id: &str, table_id: &str, seqs: &[u64]) { + for seq in seqs { + sqlx::query( + "INSERT INTO stream_records (shard_id, sequence_number, table_id, event_name, \ + record_data, created_at) VALUES ($1, $2, $3, 'INSERT', $4::jsonb, NOW())", + ) + .bind(shard_id) + .bind(format!("{seq:021}")) + .bind(table_id) + .bind(record_json(*seq)) + .execute(&self.db) + .await + .expect("seed a record"); + } + } + + async fn relabel(&self, table_id: &str, label: &str) { + sqlx::query("UPDATE tables SET stream_label = $1 WHERE table_id = $2") + .bind(label) + .bind(table_id) + .execute(&self.db) + .await + .expect("relabel"); + } +} + +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()) +} + +/// Every `.sql` file under `migrations/` and `data_migrations/`, each +/// directory in filename order. +fn shipped_migrations() -> Vec { + let root = std::path::Path::new(env!("CARGO_MANIFEST_DIR")); + let mut out = Vec::new(); + for dir in ["migrations", "data_migrations"] { + let mut files: Vec<_> = std::fs::read_dir(root.join(dir)) + .expect("read the migrations directory") + .map(|e| e.expect("directory entry").path()) + .filter(|p| p.extension().is_some_and(|e| e == "sql")) + .collect(); + files.sort(); + for f in files { + out.push(std::fs::read_to_string(&f).expect("read a migration file")); + } + } + out +} + +async fn scratch() -> Scratch { + let base = base_conn().expect("caller checks base_conn() first"); + let db_name = format!("eddb_sown_{}", uuid::Uuid::new_v4().simple())[..24].to_owned(); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect(&format!("{base}/postgres")) + .await + .expect("connect to the postgres maintenance database"); + sqlx::query(&format!("CREATE DATABASE \"{db_name}\"")) + .execute(&admin) + .await + .expect("create the scratch database"); + let url = format!("{base}/{db_name}"); + let db = PgPoolOptions::new() + .max_connections(2) + .connect(&url) + .await + .expect("connect to the scratch database"); + // One scratch database stands in for both the catalog and the data + // database; the catalog schema's foreign key from `stream_shards` to + // `tables` cannot exist in a real deployment, so it is dropped. + for sql in shipped_migrations() { + sqlx::raw_sql(&sql) + .execute(&db) + .await + .expect("apply a shipped migration"); + } + sqlx::query("ALTER TABLE stream_shards DROP CONSTRAINT IF EXISTS stream_shards_table_id_fkey") + .execute(&db) + .await + .expect("match the two-database layout"); + sqlx::query("UPDATE settings SET value = '0' WHERE key = 'control_plane_delay_seconds'") + .execute(&db) + .await + .expect("pin the control-plane delay to zero"); + for account in [ACCOUNT, OTHER_ACCOUNT] { + sqlx::query("INSERT INTO accounts (account_id, account_name) VALUES ($1, $2)") + .bind(account) + .bind(format!("acct-{account}")) + .execute(&db) + .await + .expect("seed the account row"); + } + let engine = PostgresEngine::new( + &PostgresConfig { + connection_string: url, + pool_size: 10, + max_item_size_bytes: 400_000, + }, + REGION, + ) + .await + .expect("open a PostgresEngine on the scratch database"); + Scratch { + engine, + db, + admin, + db_name, + } +} + +fn stream_table(name: &str) -> CreateTableInput { + CreateTableInput { + table_name: name.to_owned(), + key_schema: vec![KeySchemaElement { + attribute_name: "pk".to_owned(), + key_type: KeyType::Hash, + }], + attribute_definitions: vec![AttributeDefinition { + attribute_name: "pk".to_owned(), + attribute_type: ScalarAttributeType::S, + }], + billing_mode: Some(BillingMode::PayPerRequest), + stream_specification: Some(StreamSpecification { + stream_enabled: true, + stream_view_type: Some(StreamViewType::KeysOnly), + }), + ..Default::default() + } +} + +fn record_json(seq: u64) -> String { + serde_json::json!({ + "eventID": format!("e{seq}"), + "eventName": "INSERT", + "eventVersion": "1.1", + "eventSource": "aws:dynamodb", + "awsRegion": REGION, + "dynamodb": { + "ApproximateCreationDateTime": 1_700_000_000, + "Keys": { "pk": { "S": format!("k{seq}") } }, + "SequenceNumber": format!("{seq:021}"), + "SizeBytes": 10, + "StreamViewType": "KEYS_ONLY" + } + }) + .to_string() +} + +fn assert_stream_not_found(result: StreamRecordsResult, arn: &str) { + let StorageError::TableNotFound(expected) = extenddb_storage::util::stream_not_found(arn) + else { + unreachable!("stream_not_found is a TableNotFound error") + }; + match result { + Err(StorageError::TableNotFound(msg)) => assert_eq!(msg, expected), + other => panic!("expected the stream not-found error, got {other:?}"), + } +} + +/// Two tables of one account share a stream label: a shard of table `b` +/// under table `a`'s ARN is refused, so the table name alone decides. +#[tokio::test] +async fn another_tables_shard_with_the_same_label_is_not_found() { + if base_conn().is_none() { + eprintln!("SKIP: EXTENDDB_TEST_PG_CONNECTION_STRING is not set"); + return; + } + let s = scratch().await; + let (a_id, shard_a) = s.stream_table(ACCOUNT, "a", LABEL).await; + let (b_id, shard_b) = s.stream_table(ACCOUNT, "b", LABEL).await; + s.write_records(&shard_a, &a_id, &[1]).await; + // Inserted out of order so the read's ORDER BY, not insertion order, + // produces the sequence order the control asserts. + s.write_records(&shard_b, &b_id, &[2, 1]).await; + let arn_a = stream_arn(REGION, ACCOUNT, "a", LABEL); + let arn_b = stream_arn(REGION, ACCOUNT, "b", LABEL); + + // Both directions, so whichever table a lookup without the name would + // pick, one of them reads a shard that is not its own. + assert_stream_not_found( + s.engine + .get_stream_records(ACCOUNT, &arn_a, &shard_b, None, 10) + .await, + &arn_a, + ); + assert_stream_not_found( + s.engine + .get_stream_records(ACCOUNT, &arn_b, &shard_a, None, 10) + .await, + &arn_b, + ); + assert_stream_not_found( + s.engine + .get_stream_records(ACCOUNT, &arn_a, "shardId-forged-0000000000000000", None, 10) + .await, + &arn_a, + ); + let (records, _) = s + .engine + .get_stream_records(ACCOUNT, &arn_a, &shard_a, None, 10) + .await + .expect("own shard reads"); + assert_eq!(records.len(), 1); + let (records, last) = s + .engine + .get_stream_records(ACCOUNT, &arn_b, &shard_b, None, 10) + .await + .expect("own shard reads in sequence order"); + assert_eq!(records[0].dynamodb.sequence_number, format!("{:021}", 1)); + assert_eq!(records[1].dynamodb.sequence_number, format!("{:021}", 2)); + assert_eq!(last.as_deref(), Some(format!("{:021}", 2).as_str())); + s.cleanup().await; +} + +/// Two accounts own a table of the same name and label: the caller's ARN +/// resolves in the caller's account only, so the other account's shard is +/// refused under either spelling of the account in the ARN. +#[tokio::test] +async fn another_accounts_same_named_stream_is_not_found() { + if base_conn().is_none() { + eprintln!("SKIP: EXTENDDB_TEST_PG_CONNECTION_STRING is not set"); + return; + } + let s = scratch().await; + let (own_id, own_shard) = s.stream_table(ACCOUNT, "t", LABEL).await; + let (other_id, other_shard) = s.stream_table(OTHER_ACCOUNT, "t", LABEL).await; + s.write_records(&own_shard, &own_id, &[1]).await; + s.write_records(&other_shard, &other_id, &[1]).await; + + // From both callers, so whichever table a lookup without the account + // would pick, one caller reads the other account's shard. + for (caller, foreign_shard, own) in [ + (ACCOUNT, &other_shard, &own_shard), + (OTHER_ACCOUNT, &own_shard, &other_shard), + ] { + for arn in [ + stream_arn(REGION, ACCOUNT, "t", LABEL), + stream_arn(REGION, OTHER_ACCOUNT, "t", LABEL), + ] { + assert_stream_not_found( + s.engine + .get_stream_records(caller, &arn, foreign_shard, None, 10) + .await, + &arn, + ); + // The control: the caller's own shard reads under either spelling + // of the account, since storage binds the caller's account and + // parses only the table name and label out of the ARN. + let (records, _) = s + .engine + .get_stream_records(caller, &arn, own, None, 10) + .await + .expect("own shard reads"); + assert_eq!(records.len(), 1); + } + } + s.cleanup().await; +} + +/// The table takes a new stream label while keeping its shard rows: the old +/// ARN's iterator is refused and the new ARN reads every record. +#[tokio::test] +async fn an_earlier_label_on_the_same_shards_is_not_found() { + if base_conn().is_none() { + eprintln!("SKIP: EXTENDDB_TEST_PG_CONNECTION_STRING is not set"); + return; + } + let s = scratch().await; + let (table_id, shard) = s.stream_table(ACCOUNT, "t", LABEL).await; + s.write_records(&shard, &table_id, &[1]).await; + let old_arn = stream_arn(REGION, ACCOUNT, "t", LABEL); + s.engine + .get_stream_records(ACCOUNT, &old_arn, &shard, None, 10) + .await + .expect("current label reads"); + + s.relabel(&table_id, LATER_LABEL).await; + s.write_records(&shard, &table_id, &[2]).await; + assert_stream_not_found( + s.engine + .get_stream_records(ACCOUNT, &old_arn, &shard, None, 10) + .await, + &old_arn, + ); + let new_arn = stream_arn(REGION, ACCOUNT, "t", LATER_LABEL); + let (records, _) = s + .engine + .get_stream_records(ACCOUNT, &new_arn, &shard, None, 10) + .await + .expect("new label reads"); + assert_eq!(records.len(), 2); + s.cleanup().await; +} diff --git a/crates/storage-sqlite/src/lib.rs b/crates/storage-sqlite/src/lib.rs index 3953fc350..c0f25a816 100644 --- a/crates/storage-sqlite/src/lib.rs +++ b/crates/storage-sqlite/src/lib.rs @@ -40,6 +40,8 @@ mod serve_lock; mod sqlite_util; mod store; mod stream; +#[cfg(test)] +mod stream_tests; mod table_engine; mod table_helpers; mod update_table; diff --git a/crates/storage-sqlite/src/stream.rs b/crates/storage-sqlite/src/stream.rs index 62943f678..27df489ae 100644 --- a/crates/storage-sqlite/src/stream.rs +++ b/crates/storage-sqlite/src/stream.rs @@ -14,7 +14,7 @@ use extenddb_core::types::{ StreamSummary, StreamViewType, }; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{parse_stream_arn, stream_arn}; +use extenddb_storage::util::{parse_stream_arn, shard_in_stream, stream_arn}; use extenddb_storage::{StreamEngine, StreamListResult, StreamRecordsResult}; use futures::future::BoxFuture; @@ -112,62 +112,78 @@ impl StreamEngine for SqliteEngine { fn get_stream_records( &self, account_id: &str, + stream_arn: &str, shard_id: &str, after_sequence: Option<&str>, limit: i64, ) -> BoxFuture<'_, StreamRecordsResult> { let account_id = account_id.to_owned(); + let stream_arn = stream_arn.to_owned(); let shard_id = shard_id.to_owned(); let after = after_sequence.map(str::to_owned); Box::pin(async move { - // Ownership guard: only return records if the shard's backing table - // belongs to the calling account. SQLite keeps shards and the table - // catalog in the same database, so ownership resolves in one join. - // A shard iterator resolves to a shard the caller does not own or - // one that does not exist. Real DynamoDB returns - // `ValidationException: Invalid ShardIterator` for a GetRecords - // iterator it did not issue, and does NOT distinguish "exists but - // not yours" from "does not exist" — so neither do we (both - // collapse here). Verified against DynamoDB Streams (us-east-1). - let owned: Option<(i32,)> = sqlx::query_as( - "SELECT 1 FROM stream_shards s \ - JOIN tables t ON t.table_id = s.table_id \ - WHERE s.shard_id = ? AND t.account_id = ?", - ) - .bind(&shard_id) - .bind(&account_id) - .fetch_optional(&self.pool) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - if owned.is_none() { - return Err(StorageError::Validation("Invalid ShardIterator".to_owned())); - } - - let rows: Vec<(String,)> = if let Some(after) = after { - sqlx::query_as( - "SELECT record_data FROM stream_records \ - WHERE shard_id = ? AND sequence_number > ? ORDER BY sequence_number LIMIT ?", - ) - .bind(&shard_id) - .bind(&after) - .bind(limit) - .fetch_all(&self.pool) - .await - } else { - sqlx::query_as( - "SELECT record_data FROM stream_records \ - WHERE shard_id = ? ORDER BY sequence_number LIMIT ?", - ) - .bind(&shard_id) - .bind(limit) - .fetch_all(&self.pool) - .await + let (table_name, stream_label) = parse_stream_arn(&stream_arn)?; + // One query resolves the shard against the table the iterator's + // ARN names (scoped to the caller's account) and reads its + // records. The shard row is the driving row and the records are + // LEFT JOINed, so a shard of the named table with no records + // returns one row with a NULL record, and a shard that is not that + // table's (another table, or another account) returns no rows at + // all. The sequence bound sits in the join condition so the LEFT + // JOIN keeps its one row for an idle shard. Whether the shard + // belongs to the stream the ARN's label names is decided in Rust + // from the table's label and the shard id (`shard_in_stream`). + let rows: Vec<(String, Option, String, Option)> = + if let Some(after) = after { + sqlx::query_as( + "SELECT t.table_id, t.stream_label, s.shard_id, r.record_data \ + FROM stream_shards s \ + JOIN tables t ON t.table_id = s.table_id \ + LEFT JOIN stream_records r \ + ON r.shard_id = s.shard_id AND r.sequence_number > ? \ + WHERE s.shard_id = ? AND t.account_id = ? AND t.table_name = ? \ + ORDER BY r.sequence_number LIMIT ?", + ) + .bind(&after) + .bind(&shard_id) + .bind(&account_id) + .bind(&table_name) + .bind(limit) + .fetch_all(&self.pool) + .await + } else { + sqlx::query_as( + "SELECT t.table_id, t.stream_label, s.shard_id, r.record_data \ + FROM stream_shards s \ + JOIN tables t ON t.table_id = s.table_id \ + LEFT JOIN stream_records r ON r.shard_id = s.shard_id \ + WHERE s.shard_id = ? AND t.account_id = ? AND t.table_name = ? \ + ORDER BY r.sequence_number LIMIT ?", + ) + .bind(&shard_id) + .bind(&account_id) + .bind(&table_name) + .bind(limit) + .fetch_all(&self.pool) + .await + } + .map_err(|e| StorageError::Internal(e.to_string()))?; + let Some((table_id, table_label, row_shard_id, _)) = rows.first() else { + return Err(extenddb_storage::util::stream_not_found(&stream_arn)); + }; + if !shard_in_stream( + table_id, + row_shard_id, + table_label.as_deref(), + &stream_label, + ) { + return Err(extenddb_storage::util::stream_not_found(&stream_arn)); } - .map_err(|e| StorageError::Internal(e.to_string()))?; let records: Vec = rows .into_iter() - .map(|(d,)| { + .filter_map(|(_, _, _, d)| d) + .map(|d| { serde_json::from_str(&d).map_err(|e| StorageError::Internal(e.to_string())) }) .collect::>()?; @@ -198,11 +214,8 @@ impl StreamEngine for SqliteEngine { .await .map_err(|e| StorageError::Internal(e.to_string()))?; - let (ks_json, stream_spec_json, table_status, table_id) = row.ok_or_else(|| { - StorageError::TableNotFound(format!( - "Requested resource not found: Stream: {stream_arn} not found." - )) - })?; + let (ks_json, stream_spec_json, table_status, table_id) = + row.ok_or_else(|| extenddb_storage::util::stream_not_found(&stream_arn))?; let key_schema = serde_json::from_str(&ks_json) .map_err(|e| StorageError::Internal(e.to_string()))?; @@ -462,9 +475,7 @@ impl StreamEngine for SqliteEngine { .map_err(|e| StorageError::Internal(e.to_string()))?; let Some(table_id) = table_id else { - return Err(StorageError::TableNotFound(format!( - "Requested resource not found: Stream: {stream_arn} not found." - ))); + return Err(extenddb_storage::util::stream_not_found(&stream_arn)); }; let exists: Option<(i64,)> = sqlx::query_as("SELECT 1 FROM stream_shards WHERE shard_id = ? AND table_id = ?") @@ -474,9 +485,7 @@ impl StreamEngine for SqliteEngine { .await .map_err(|e| StorageError::Internal(e.to_string()))?; if exists.is_none() { - return Err(StorageError::TableNotFound(format!( - "Requested resource not found: Stream: {stream_arn} not found." - ))); + return Err(extenddb_storage::util::stream_not_found(&stream_arn)); } Ok(()) }) diff --git a/crates/storage-sqlite/src/stream_tests.rs b/crates/storage-sqlite/src/stream_tests.rs new file mode 100644 index 000000000..da83ceaac --- /dev/null +++ b/crates/storage-sqlite/src/stream_tests.rs @@ -0,0 +1,298 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! `get_stream_records` resolves the shard against the stream the iterator +//! names in the same query as the records read. + +use extenddb_storage::StreamEngine; +use extenddb_storage::error::StorageError; +use extenddb_storage::util::stream_arn; + +use crate::store::SqliteEngine; + +const REGION: &str = "us-east-1"; +const ACCOUNT: &str = "111111111111"; +const OTHER_ACCOUNT: &str = "222222222222"; + +async fn engine() -> SqliteEngine { + let engine = SqliteEngine::new(":memory:", 1, REGION, 409_600) + .await + .expect("engine"); + crate::schema::apply(&engine.pool).await.expect("schema"); + for account in [ACCOUNT, OTHER_ACCOUNT] { + sqlx::query("INSERT INTO accounts (account_id, account_name) VALUES (?, ?)") + .bind(account) + .bind(account) + .execute(&engine.pool) + .await + .expect("account"); + } + engine +} + +/// A stream-enabled table with four shards. Returns its table id and the +/// shard ids, which follow `init_stream_shards`. +async fn stream_table(engine: &SqliteEngine, account: &str, name: &str, label: &str) -> String { + let table_id = format!("{name}-{account}-id"); + sqlx::query( + "INSERT INTO tables (account_id, table_name, key_schema, attribute_definitions, \ + table_status, table_arn, table_id, stream_label) \ + VALUES (?, ?, '[]', '[]', 'ACTIVE', ?, ?, ?)", + ) + .bind(account) + .bind(name) + .bind(extenddb_storage::util::table_arn(REGION, account, name)) + .bind(&table_id) + .bind(label) + .execute(&engine.pool) + .await + .expect("table"); + for i in 0..4u32 { + sqlx::query( + "INSERT INTO stream_shards (shard_id, table_id, starting_sequence_number) \ + VALUES (?, ?, ?)", + ) + .bind(format!("shardId-{table_id}-{i:016}")) + .bind(&table_id) + .bind(format!("{:021}", 0)) + .execute(&engine.pool) + .await + .expect("shard"); + } + table_id +} + +fn shard(table_id: &str, i: u32) -> String { + format!("shardId-{table_id}-{i:016}") +} + +fn record_json(seq: u64) -> String { + serde_json::json!({ + "eventID": format!("e{seq}"), + "eventName": "INSERT", + "eventVersion": "1.1", + "eventSource": "aws:dynamodb", + "awsRegion": REGION, + "dynamodb": { + "ApproximateCreationDateTime": 1_700_000_000, + "Keys": { "pk": { "S": format!("k{seq}") } }, + "SequenceNumber": format!("{seq:021}"), + "SizeBytes": 10, + "StreamViewType": "KEYS_ONLY" + } + }) + .to_string() +} + +async fn write_records(engine: &SqliteEngine, shard_id: &str, table_id: &str, seqs: &[u64]) { + for seq in seqs { + sqlx::query( + "INSERT INTO stream_records (sequence_number, shard_id, table_id, event_name, \ + record_data) VALUES (?, ?, ?, 'Insert', ?)", + ) + .bind(format!("{seq:021}")) + .bind(shard_id) + .bind(table_id) + .bind(record_json(*seq)) + .execute(&engine.pool) + .await + .expect("record"); + } +} + +fn assert_stream_not_found(result: extenddb_storage::StreamRecordsResult, arn: &str) { + let StorageError::TableNotFound(expected) = extenddb_storage::util::stream_not_found(arn) + else { + unreachable!("stream_not_found is a TableNotFound error") + }; + match result { + Err(StorageError::TableNotFound(msg)) => assert_eq!(msg, expected), + other => panic!("expected the stream not-found error, got {other:?}"), + } +} + +#[tokio::test] +async fn own_shard_reads_its_records_and_paginates() { + let engine = engine().await; + let table_id = stream_table(&engine, ACCOUNT, "t", "2026-01-01T00:00:00.000").await; + let arn = stream_arn(REGION, ACCOUNT, "t", "2026-01-01T00:00:00.000"); + let shard_id = shard(&table_id, 0); + // Inserted out of order so the read's ORDER BY, not insertion order, + // produces the sequence order the pages are asserted in. + write_records(&engine, &shard_id, &table_id, &[3, 1, 2]).await; + + let (page, last) = engine + .get_stream_records(ACCOUNT, &arn, &shard_id, None, 2) + .await + .expect("first page"); + assert_eq!(page.len(), 2); + assert_eq!(page[0].dynamodb.sequence_number, format!("{:021}", 1)); + assert_eq!(page[1].dynamodb.sequence_number, format!("{:021}", 2)); + assert_eq!(last.as_deref(), Some(format!("{:021}", 2).as_str())); + + let (rest, last) = engine + .get_stream_records(ACCOUNT, &arn, &shard_id, last.as_deref(), 10) + .await + .expect("second page"); + assert_eq!(rest.len(), 1); + assert_eq!(last.as_deref(), Some(format!("{:021}", 3).as_str())); +} + +#[tokio::test] +async fn idle_shard_of_the_stream_reads_nothing_without_error() { + let engine = engine().await; + let table_id = stream_table(&engine, ACCOUNT, "t", "2026-01-01T00:00:00.000").await; + let arn = stream_arn(REGION, ACCOUNT, "t", "2026-01-01T00:00:00.000"); + let shard_id = shard(&table_id, 1); + + for after in [None, Some(format!("{:021}", 5))] { + let (records, last) = engine + .get_stream_records(ACCOUNT, &arn, &shard_id, after.as_deref(), 10) + .await + .expect("idle shard reads"); + assert!(records.is_empty()); + assert!(last.is_none()); + } +} + +#[tokio::test] +async fn another_streams_shard_under_an_allowed_arn_is_not_found() { + let engine = engine().await; + // Both tables carry the same label, so the table name alone decides. + let a_id = stream_table(&engine, ACCOUNT, "a", "2026-01-01T00:00:00.000").await; + let b_id = stream_table(&engine, ACCOUNT, "b", "2026-01-01T00:00:00.000").await; + let arn_a = stream_arn(REGION, ACCOUNT, "a", "2026-01-01T00:00:00.000"); + let arn_b = stream_arn(REGION, ACCOUNT, "b", "2026-01-01T00:00:00.000"); + write_records(&engine, &shard(&a_id, 0), &a_id, &[1]).await; + write_records(&engine, &shard(&b_id, 0), &b_id, &[1, 2]).await; + + // Stream A's ARN with stream B's shard: nothing is read, and the answer is + // the one `validate_shard` gives for a shard that is not the stream's. + // Both directions, so whichever table a lookup without the name would + // pick, one of them reads a shard that is not its own. + assert_stream_not_found( + engine + .get_stream_records(ACCOUNT, &arn_a, &shard(&b_id, 0), None, 10) + .await, + &arn_a, + ); + assert_stream_not_found( + engine + .get_stream_records(ACCOUNT, &arn_b, &shard(&a_id, 0), None, 10) + .await, + &arn_b, + ); + // A shard id that exists for no table. + assert_stream_not_found( + engine + .get_stream_records(ACCOUNT, &arn_a, "shardId-forged-0000000000000000", None, 10) + .await, + &arn_a, + ); + // The control: stream A's own shard reads. + let (records, _) = engine + .get_stream_records(ACCOUNT, &arn_a, &shard(&a_id, 0), None, 10) + .await + .expect("own shard"); + assert_eq!(records.len(), 1); +} + +#[tokio::test] +async fn another_accounts_stream_is_not_found_under_its_own_arn() { + let engine = engine().await; + let other_id = stream_table(&engine, OTHER_ACCOUNT, "t", "2026-01-01T00:00:00.000").await; + write_records(&engine, &shard(&other_id, 0), &other_id, &[1]).await; + // The caller spells the other account's stream ARN exactly; the table + // lookup is scoped to the caller's account, so it resolves to nothing. + let arn = stream_arn(REGION, OTHER_ACCOUNT, "t", "2026-01-01T00:00:00.000"); + assert_stream_not_found( + engine + .get_stream_records(ACCOUNT, &arn, &shard(&other_id, 0), None, 10) + .await, + &arn, + ); +} + +#[tokio::test] +async fn another_accounts_same_named_stream_is_not_found() { + let engine = engine().await; + // Both accounts own a table `t` with the same label, so the account alone + // decides which table the ARN names. + let own_id = stream_table(&engine, ACCOUNT, "t", "2026-01-01T00:00:00.000").await; + let other_id = stream_table(&engine, OTHER_ACCOUNT, "t", "2026-01-01T00:00:00.000").await; + write_records(&engine, &shard(&own_id, 0), &own_id, &[1]).await; + write_records(&engine, &shard(&other_id, 0), &other_id, &[1]).await; + + // From both callers, so whichever table a lookup without the account + // would pick, one caller reads the other account's shard. + for (caller, foreign, own) in [ + (ACCOUNT, &other_id, &own_id), + (OTHER_ACCOUNT, &own_id, &other_id), + ] { + for arn in [ + stream_arn(REGION, ACCOUNT, "t", "2026-01-01T00:00:00.000"), + stream_arn(REGION, OTHER_ACCOUNT, "t", "2026-01-01T00:00:00.000"), + ] { + assert_stream_not_found( + engine + .get_stream_records(caller, &arn, &shard(foreign, 0), None, 10) + .await, + &arn, + ); + // The control: the caller's own shard reads under either spelling + // of the account, since storage binds the caller's account and + // parses only the table name and label out of the ARN. + let (records, _) = engine + .get_stream_records(caller, &arn, &shard(own, 0), None, 10) + .await + .expect("own shard reads"); + assert_eq!(records.len(), 1); + } + } +} + +#[tokio::test] +async fn an_earlier_stream_label_on_the_same_shards_is_not_found() { + let engine = engine().await; + let table_id = stream_table(&engine, ACCOUNT, "t", "2026-01-01T00:00:00.000").await; + let old_arn = stream_arn(REGION, ACCOUNT, "t", "2026-01-01T00:00:00.000"); + let shard_id = shard(&table_id, 0); + write_records(&engine, &shard_id, &table_id, &[1]).await; + engine + .get_stream_records(ACCOUNT, &old_arn, &shard_id, None, 10) + .await + .expect("current label reads"); + + // The table keeps its shard rows and takes a new stream label, as a + // stream re-enabled on a table whose label was cleared does. + sqlx::query("UPDATE tables SET stream_label = ? WHERE table_id = ?") + .bind("2026-02-02T00:00:00.000") + .bind(&table_id) + .execute(&engine.pool) + .await + .expect("relabel"); + write_records(&engine, &shard_id, &table_id, &[2]).await; + + assert_stream_not_found( + engine + .get_stream_records(ACCOUNT, &old_arn, &shard_id, None, 10) + .await, + &old_arn, + ); + let new_arn = stream_arn(REGION, ACCOUNT, "t", "2026-02-02T00:00:00.000"); + let (records, _) = engine + .get_stream_records(ACCOUNT, &new_arn, &shard_id, None, 10) + .await + .expect("new label reads"); + assert_eq!(records.len(), 2); +} + +#[tokio::test] +async fn a_malformed_stream_arn_is_a_validation_error() { + let engine = engine().await; + let table_id = stream_table(&engine, ACCOUNT, "t", "2026-01-01T00:00:00.000").await; + let err = engine + .get_stream_records(ACCOUNT, "not-an-arn", &shard(&table_id, 0), None, 10) + .await + .expect_err("malformed ARN"); + assert!(matches!(err, StorageError::Validation(_)), "{err:?}"); +} diff --git a/crates/storage/src/lib.rs b/crates/storage/src/lib.rs index c760583a6..9f4e86d56 100755 --- a/crates/storage/src/lib.rs +++ b/crates/storage/src/lib.rs @@ -603,13 +603,19 @@ pub trait StreamEngine: Send + Sync { /// Read stream records from a shard starting after a sequence number. /// - /// `account_id` is the authenticated caller's account. Implementations MUST - /// return records only for shards whose backing table belongs to that - /// account; a shard whose table belongs to a different account must yield no - /// records, so a shard iterator only reads its owning account's stream data. + /// `stream_arn` is the stream the shard iterator was issued for, and the + /// resource the request was authorized against. Implementations MUST + /// resolve the shard against the table and stream label that ARN names, + /// scoped to `account_id`, in the same round trip as the records read, and + /// return `StorageError::TableNotFound` with the message + /// `Requested resource not found: Stream: not found.` when the + /// stream does not exist for that account or the shard is not one of that + /// stream's shards. A shard of another table, another account, or an + /// earlier stream label on the same table therefore yields no records. fn get_stream_records( &self, account_id: &str, + stream_arn: &str, shard_id: &str, after_sequence: Option<&str>, limit: i64, diff --git a/crates/storage/src/util/arn.rs b/crates/storage/src/util/arn.rs index 083e62c35..00609f008 100755 --- a/crates/storage/src/util/arn.rs +++ b/crates/storage/src/util/arn.rs @@ -90,6 +90,20 @@ pub fn parse_stream_arn(arn: &str) -> Result<(String, String), StorageError> { } } +/// The error a stream operation returns when the stream ARN names no stream of +/// the caller's account, or a shard that is not one of that stream's shards. +/// `describe_stream`, `validate_shard` and `get_stream_records` on every +/// backend answer with this same text, so a shard iterator that pairs an +/// authorized stream ARN with another stream's shard gets the same answer +/// whichever path checks it. Measured on `DescribeStream` 2026-10-24: no +/// trailing period. +#[must_use] +pub fn stream_not_found(stream_arn: &str) -> StorageError { + StorageError::TableNotFound(format!( + "Requested resource not found: Stream: {stream_arn} not found" + )) +} + #[cfg(test)] mod stream_label_tests { use super::format_stream_label; diff --git a/crates/storage/src/util/mod.rs b/crates/storage/src/util/mod.rs index a762738e7..68b4dae64 100755 --- a/crates/storage/src/util/mod.rs +++ b/crates/storage/src/util/mod.rs @@ -10,10 +10,11 @@ mod arn; mod backup_id; mod catalog_version; mod key; +mod shard_stream; pub use arn::{ ARN_PARTITION, format_stream_label, index_arn, new_stream_label, parse_stream_arn, stream_arn, - table_arn, + stream_not_found, table_arn, }; pub use backup_id::{ BackupArnError, BackupArnParts, MAX_BACKUP_ARN_LEN, backup_arn, backup_creation_date_time, @@ -27,6 +28,7 @@ pub use key::{ merge_attribute_definitions, parse_sk, pk_to_text, recover_sort_key_definitions, sk_column, sk_column_n, sk_info, }; +pub use shard_stream::shard_in_stream; /// Whether an UpdateTable's GSI actions must be refused because the GSIs named /// in `building` are `CREATING`: any Create, or a Delete of a GSI that is not diff --git a/crates/storage/src/util/shard_stream.rs b/crates/storage/src/util/shard_stream.rs new file mode 100644 index 000000000..93982c78b --- /dev/null +++ b/crates/storage/src/util/shard_stream.rs @@ -0,0 +1,63 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! Which stream of its table a shard belongs to. +//! +//! `GetRecords` carries a stream ARN in its iterator and reads one shard. Every +//! backend keys its catalog read on the caller's account and the table name +//! the ARN carries, keys its shard read on that table's id, and then asks the +//! one function here whether the shard row belongs to the stream the ARN's +//! label names. The rule lives in one place so a change to how streams are +//! identified touches one function and its test. + +/// Whether shard `shard_id`, a row of table `table_id` whose current stream +/// label is `table_label` (`None` when the table never had a stream), belongs +/// to the stream labelled `arn_label`. +/// +/// Today a table has one stream and every shard row of the table belongs to +/// it, so the shard is the stream's exactly when the ARN's label is the +/// table's current label. `table_id` and `shard_id` are part of the contract +/// for a scheme that names the stream inside the shard id; this rule does not +/// read them. +#[must_use] +pub fn shard_in_stream( + table_id: &str, + shard_id: &str, + table_label: Option<&str>, + arn_label: &str, +) -> bool { + let _ = (table_id, shard_id); + table_label == Some(arn_label) +} + +#[cfg(test)] +mod tests { + use super::shard_in_stream; + + const TABLE: &str = "0f3c1c1e-7b1d-4e3b-9a6e-4a2d3e5f6a7b"; + const LABEL: &str = "2026-01-01T00:00:00.000"; + const LATER: &str = "2026-02-02T00:00:00.000"; + + fn shard(i: u32) -> String { + format!("shardId-{TABLE}-{i:016}") + } + + #[test] + fn a_shard_belongs_to_its_tables_current_stream() { + assert!(shard_in_stream(TABLE, &shard(0), Some(LABEL), LABEL)); + assert!(shard_in_stream(TABLE, &shard(3), Some(LABEL), LABEL)); + } + + #[test] + fn an_earlier_label_on_the_same_shards_is_not_the_current_stream() { + // The table took a new label while keeping its shard rows: the old + // ARN no longer names a stream of the table. + assert!(!shard_in_stream(TABLE, &shard(0), Some(LATER), LABEL)); + assert!(shard_in_stream(TABLE, &shard(0), Some(LATER), LATER)); + } + + #[test] + fn a_table_without_a_stream_owns_no_stream_shard() { + assert!(!shard_in_stream(TABLE, &shard(0), None, LABEL)); + } +} diff --git a/devtools/run-mongodb-tests b/devtools/run-mongodb-tests index 14ecb9867..a877b84cc 100755 --- a/devtools/run-mongodb-tests +++ b/devtools/run-mongodb-tests @@ -249,6 +249,17 @@ fi # --------------------------------------------------------------------------- echo "" +if [[ " ${RUN_TESTS_ARGS[*]} " == *" --rust"* ]]; then + # The storage-level MongoDB tests that need a live server (the stream + # records page over 16 MB and the shard ownership predicates) are + # #[ignore]d in the crate and read their server from this variable; this + # container is the one place in CI that has one. + echo "=== Running MongoDB storage-level live tests (records_live) ===" + echo "" + EXTENDDB_TEST_MONGODB_URI="mongodb://localhost:$MONGO_PORT/?replicaSet=rs0&directConnection=true" \ + cargo test -p extenddb-storage-mongodb --lib -- --ignored records_live + echo "" +fi echo "=== Running test suite via devtools/run-tests --backend mongodb ===" echo " passthrough args: ${RUN_TESTS_ARGS[*]}" echo "" diff --git a/tests/python/helpers.py b/tests/python/helpers.py index f4ed85479..a12816cb4 100755 --- a/tests/python/helpers.py +++ b/tests/python/helpers.py @@ -71,3 +71,10 @@ def wait_for_deleted(client, table_name: str, timeout: float = 60.0) -> None: def unique_name(prefix: str = "test") -> str: """Generate a unique name for test isolation.""" return f"{prefix}-{uuid.uuid4().hex[:12]}" + + +def stream_not_found_message(stream_arn: str) -> str: + """The sentence a streams call answers for a stream ARN that names no stream + of the caller's account, or a shard that is not that stream's (measured on + DescribeStream 2026-10-24: no trailing period).""" + return f"Requested resource not found: Stream: {stream_arn} not found" diff --git a/tests/python/test_stream_shard_ownership.py b/tests/python/test_stream_shard_ownership.py new file mode 100644 index 000000000..54a03cbd6 --- /dev/null +++ b/tests/python/test_stream_shard_ownership.py @@ -0,0 +1,239 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""GetRecords reads only from a shard of the stream its iterator names. + +The shard iterator carries the stream ARN the request is authorized against. +The records read resolves the shard against that stream, so an iterator that +pairs an allowed stream ARN with another stream's shard reads nothing, and an +iterator reads only while its ARN is the table's current stream. + +These tests require an ExtendDB endpoint (DYNAMODB_ENDPOINT): Amazon DynamoDB +serves streams on a separate endpoint and its iterator tokens are opaque, so +a token naming another stream's shard cannot be built against it. +""" + +from __future__ import annotations + +import base64 +import os +import time + +import boto3 +import botocore.config +import pytest +from botocore.exceptions import ClientError + +from helpers import stream_not_found_message, unique_name, wait_for_active, wait_for_deleted + +pytestmark = pytest.mark.skipif( + not os.environ.get("DYNAMODB_ENDPOINT", "").strip(), + reason="Streams tests require extenddb (DYNAMODB_ENDPOINT)", +) + +STREAM_SPEC = {"StreamEnabled": True, "StreamViewType": "NEW_AND_OLD_IMAGES"} + + +def _streams_client(endpoint_url, parameter_validation: bool): + kwargs: dict = { + "service_name": "dynamodbstreams", + "region_name": os.environ.get("AWS_DEFAULT_REGION", "us-east-1"), + "config": botocore.config.Config( + retries={"max_attempts": 0}, parameter_validation=parameter_validation + ), + } + if endpoint_url: + kwargs["endpoint_url"] = endpoint_url + if endpoint_url.startswith("https://"): + ca_cert = os.environ.get("EXTENDDB_CA_CERT", "") + kwargs["verify"] = ca_cert if ca_cert else False + return boto3.client(**kwargs) + + +@pytest.fixture(scope="module") +def streams_client(endpoint_url): + """DynamoDB Streams client pointing at the extenddb endpoint.""" + return _streams_client(endpoint_url, parameter_validation=True) + + +@pytest.fixture(scope="module") +def raw_streams_client(endpoint_url): + """The same client with the SDK's parameter validation off, so a request + the SDK would refuse locally reaches the server.""" + return _streams_client(endpoint_url, parameter_validation=False) + + +def _create_stream_table(dynamodb_client) -> tuple[str, str]: + name = unique_name("shardown") + resp = dynamodb_client.create_table( + TableName=name, + AttributeDefinitions=[{"AttributeName": "pk", "AttributeType": "S"}], + KeySchema=[{"AttributeName": "pk", "KeyType": "HASH"}], + BillingMode="PAY_PER_REQUEST", + StreamSpecification=STREAM_SPEC, + ) + wait_for_active(dynamodb_client, name) + return name, resp["TableDescription"]["LatestStreamArn"] + + +def _shards(streams_client, stream_arn: str) -> list[str]: + desc = streams_client.describe_stream(StreamArn=stream_arn) + return [s["ShardId"] for s in desc["StreamDescription"]["Shards"]] + + +def _first_shard(streams_client, stream_arn: str) -> str: + return _shards(streams_client, stream_arn)[0] + + +def _trim_horizon(streams_client, stream_arn: str, shard_id: str) -> str: + return streams_client.get_shard_iterator( + StreamArn=stream_arn, ShardId=shard_id, ShardIteratorType="TRIM_HORIZON" + )["ShardIterator"] + + +class TestShardBelongsToIteratorStream: + def test_allowed_stream_arn_with_another_streams_shard_reads_nothing( + self, dynamodb_client, streams_client + ): + """A token naming stream A's ARN and one of stream B's shards reads + nothing: ResourceNotFoundException for stream A, the same answer as a + shard that does not exist. Stream B's own iterator reads its records.""" + name_a, arn_a = _create_stream_table(dynamodb_client) + name_b, arn_b = _create_stream_table(dynamodb_client) + try: + for i in range(8): + dynamodb_client.put_item(TableName=name_b, Item={"pk": {"S": f"k{i}"}}) + shard_b = _first_shard(streams_client, arn_b) + now = int(time.time()) + raw = f"{shard_b}|AFTER_SEQUENCE_NUMBER||{now}|{arn_a}" + forged = base64.standard_b64encode(raw.encode()).decode() + + with pytest.raises(ClientError) as exc: + streams_client.get_records(ShardIterator=forged) + err = exc.value.response["Error"] + assert err["Code"] == "ResourceNotFoundException", err + assert err["Message"] == stream_not_found_message(arn_a), err + + raw = f"shardId-forged-0000000000000000|AFTER_SEQUENCE_NUMBER||{now}|{arn_a}" + forged = base64.standard_b64encode(raw.encode()).decode() + with pytest.raises(ClientError) as exc: + streams_client.get_records(ShardIterator=forged) + err = exc.value.response["Error"] + assert err["Code"] == "ResourceNotFoundException", err + assert err["Message"] == stream_not_found_message(arn_a), err + + # Control: the shard reads under its own stream's iterator, and + # stream A's own shard reads (nothing was written to A). + resp = streams_client.get_records( + ShardIterator=_trim_horizon(streams_client, arn_b, shard_b) + ) + assert resp["NextShardIterator"] + resp = streams_client.get_records( + ShardIterator=_trim_horizon( + streams_client, arn_a, _first_shard(streams_client, arn_a) + ) + ) + assert resp["Records"] == [] + assert resp["NextShardIterator"] + finally: + for name in (name_a, name_b): + dynamodb_client.delete_table(TableName=name) + for name in (name_a, name_b): + wait_for_deleted(dynamodb_client, name) + + def test_iterator_reads_only_while_its_arn_is_the_tables_stream( + self, dynamodb_client, streams_client + ): + """An iterator issued before the stream was disabled and enabled again + reads nothing that was written to the stream that replaced it. + + While the table keeps its stream ARN across the cycle, the old + iterators read the same records a fresh iterator on each shard reads. + When the re-enabled stream has a new ARN (as on the service, measured + 2026-10-25), the old stream is closed with its own shard ids: an old + iterator answers 200, returns none of the keys written after the + re-enable, and its shard id is not one of the new stream's.""" + name, old_arn = _create_stream_table(dynamodb_client) + try: + old_shards = _shards(streams_client, old_arn) + old_its = {s: _trim_horizon(streams_client, old_arn, s) for s in old_shards} + dynamodb_client.update_table( + TableName=name, StreamSpecification={"StreamEnabled": False} + ) + wait_for_active(dynamodb_client, name) + dynamodb_client.update_table(TableName=name, StreamSpecification=STREAM_SPEC) + wait_for_active(dynamodb_client, name) + new_arn = dynamodb_client.describe_table(TableName=name)["Table"]["LatestStreamArn"] + new_keys = {f"n{i}" for i in range(8)} + for key in sorted(new_keys): + dynamodb_client.put_item(TableName=name, Item={"pk": {"S": key}}) + time.sleep(1) + + new_shards = _shards(streams_client, new_arn) + new_records = [ + r + for s in new_shards + for r in streams_client.get_records( + ShardIterator=_trim_horizon(streams_client, new_arn, s) + )["Records"] + ] + assert {r["dynamodb"]["Keys"]["pk"]["S"] for r in new_records} == new_keys + + old_pages = { + s: streams_client.get_records(ShardIterator=it) for s, it in old_its.items() + } + if new_arn == old_arn: + # The table kept its stream: the old iterators read, across the + # four shards, exactly what the fresh iterators read. + old_seqs = sorted( + r["dynamodb"]["SequenceNumber"] + for p in old_pages.values() + for r in p["Records"] + ) + assert old_seqs == sorted(r["dynamodb"]["SequenceNumber"] for r in new_records) + else: + # The old stream was replaced: its shards are not the new + # stream's, and its iterators return none of the new keys. + assert not set(old_shards) & set(new_shards) + for page in old_pages.values(): + old_keys = {r["dynamodb"]["Keys"]["pk"]["S"] for r in page["Records"]} + assert not old_keys & new_keys + finally: + dynamodb_client.delete_table(TableName=name) + wait_for_deleted(dynamodb_client, name) + +class TestGetRecordsLimitBounds: + def test_limit_outside_1_to_1000_is_refused_before_the_read( + self, dynamodb_client, streams_client, raw_streams_client + ): + """Limit 0 and -1 get the model's constraint sentence and Limit 1001 + the parameter sentence (both measured 2026-10-25); 1 and 1000 read. + The bounds are checked before storage runs, so a Limit of 0 is never + mistaken for a shard that is not the stream's.""" + name, arn = _create_stream_table(dynamodb_client) + try: + dynamodb_client.put_item(TableName=name, Item={"pk": {"S": "k0"}}) + shard = _first_shard(streams_client, arn) + it = _trim_horizon(streams_client, arn, shard) + for limit in (0, -1): + with pytest.raises(ClientError) as exc: + raw_streams_client.get_records(ShardIterator=it, Limit=limit) + err = exc.value.response["Error"] + assert err["Code"] == "ValidationException", err + assert err["Message"] == ( + "1 validation error detected: Value at 'limit' failed to satisfy " + "constraint: Member must have value greater than or equal to 1" + ), err + with pytest.raises(ClientError) as exc: + raw_streams_client.get_records(ShardIterator=it, Limit=1001) + err = exc.value.response["Error"] + assert err["Code"] == "ValidationException", err + assert err["Message"] == ( + "One or more parameter values were invalid: " + "Limit should be less than or equal to 1000" + ), err + for limit in (1, 1000): + resp = raw_streams_client.get_records(ShardIterator=it, Limit=limit) + assert resp["NextShardIterator"] + finally: + dynamodb_client.delete_table(TableName=name) + wait_for_deleted(dynamodb_client, name)