diff --git a/crates/rustack-dynamodb-core/src/provider.rs b/crates/rustack-dynamodb-core/src/provider.rs index a1ff0a0a..92d580b6 100644 --- a/crates/rustack-dynamodb-core/src/provider.rs +++ b/crates/rustack-dynamodb-core/src/provider.rs @@ -53,7 +53,7 @@ use crate::{ state::{DynamoDBServiceState, DynamoDBTable}, storage::{ KeyAttribute, KeySchema, PrimaryKey, SortKeyCondition, SortableAttributeValue, - TableStorage, calculate_item_size, extract_primary_key, partition_key_segment, + TableStorage, calculate_item_size, extract_primary_key, partition_key_segment, take_page, }, stream::{ ChangeEvent, ChangeEventName, NoopStreamEmitter, NoopStreamLifecycle, StreamEmitter, @@ -1662,7 +1662,12 @@ impl RustackDynamoDB { } else { let start_pk = extract_primary_key(&table.key_schema, &input.exclusive_start_key) .map_err(storage_error_to_dynamodb)?; - start_pk.sort_key + // Items of a table without a sort key are stored under the sentinel. + Some( + start_pk + .sort_key + .unwrap_or(SortableAttributeValue::Sentinel), + ) }; let (items, last_key_sort) = table.storage.query( @@ -4840,9 +4845,6 @@ fn build_last_evaluated_key( key } -/// Maximum response size for a single Query/Scan response (1 MB). -const MAX_RESPONSE_BYTES: u64 = 1_048_576; - /// Result type for GSI query operations. type GsiQueryResult = ( Vec>, @@ -4891,8 +4893,7 @@ fn query_gsi( ); } - // Apply limit and 1 MB size cap. - let (selected, has_more) = gsi_apply_limit(matching, limit); + let (selected, has_more) = take_page(matching, limit, calculate_item_size); // Build last_evaluated_key containing both GSI keys and table primary keys. let last_evaluated_key = if has_more { @@ -5011,30 +5012,6 @@ fn gsi_apply_pagination( } } -/// Apply count limit and 1 MB size cap, returning `(selected_items, has_more)`. -fn gsi_apply_limit( - items: Vec>, - limit: Option, -) -> (Vec>, bool) { - let effective_limit = limit.unwrap_or(usize::MAX); - let mut selected: Vec> = Vec::new(); - let mut cumulative_size: u64 = 0; - - for item in items { - if selected.len() >= effective_limit { - return (selected, true); - } - let item_size = calculate_item_size(&item); - if !selected.is_empty() && cumulative_size + item_size > MAX_RESPONSE_BYTES { - return (selected, true); - } - cumulative_size += item_size; - selected.push(item); - } - - (selected, false) -} - /// Build the `last_evaluated_key` for a GSI query response. /// /// DynamoDB includes both the GSI keys and the table's primary keys in the @@ -5087,8 +5064,8 @@ mod tests { DescribeTableInput, UpdateContinuousBackupsInput, UpdateItemInput, }, types::{ - AttributeDefinition, KeySchemaElement, KeyType, PutRequest, ScalarAttributeType, - WriteRequest, + AttributeDefinition, GlobalSecondaryIndex, KeySchemaElement, KeyType, Projection, + ProjectionType, PutRequest, ScalarAttributeType, WriteRequest, }, }; @@ -5622,4 +5599,139 @@ mod tests { )]) ); } + + #[test] + fn test_should_return_last_key_when_limit_ends_on_last_gsi_item() { + let provider = RustackDynamoDB::new(DynamoDBConfig::default()); + provider + .handle_create_table(CreateTableInput { + table_name: "GsiTable".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, + }, + AttributeDefinition { + attribute_name: "gpk".to_owned(), + attribute_type: ScalarAttributeType::S, + }, + AttributeDefinition { + attribute_name: "gsk".to_owned(), + attribute_type: ScalarAttributeType::N, + }, + ], + global_secondary_indexes: vec![GlobalSecondaryIndex { + index_name: "ByGroup".to_owned(), + key_schema: vec![ + KeySchemaElement { + attribute_name: "gpk".to_owned(), + key_type: KeyType::Hash, + }, + KeySchemaElement { + attribute_name: "gsk".to_owned(), + key_type: KeyType::Range, + }, + ], + projection: Projection { + projection_type: Some(ProjectionType::All), + non_key_attributes: Vec::new(), + }, + provisioned_throughput: None, + }], + billing_mode: Some(BillingMode::PayPerRequest), + ..Default::default() + }) + .unwrap(); + for i in 1..=4 { + provider + .handle_put_item(PutItemInput { + table_name: "GsiTable".to_owned(), + item: HashMap::from([ + ("pk".to_owned(), AttributeValue::S(format!("item{i}"))), + ("gpk".to_owned(), AttributeValue::S("group".to_owned())), + ("gsk".to_owned(), AttributeValue::N(i.to_string())), + ]), + ..Default::default() + }) + .unwrap(); + } + let query = |exclusive_start_key| { + provider + .handle_query(QueryInput { + table_name: "GsiTable".to_owned(), + index_name: Some("ByGroup".to_owned()), + key_condition_expression: Some("gpk = :g".to_owned()), + expression_attribute_values: HashMap::from([( + ":g".to_owned(), + AttributeValue::S("group".to_owned()), + )]), + limit: Some(2), + exclusive_start_key, + ..Default::default() + }) + .unwrap() + }; + + let first = query(HashMap::new()); + assert_eq!(first.count, 2); + assert!(!first.last_evaluated_key.is_empty()); + + // Limit stops this page on the last item: DynamoDB still returns a key. + let second = query(first.last_evaluated_key); + assert_eq!(second.count, 2); + assert_eq!( + second.last_evaluated_key, + HashMap::from([ + ("gpk".to_owned(), AttributeValue::S("group".to_owned())), + ("gsk".to_owned(), AttributeValue::N("4".to_owned())), + ("pk".to_owned(), AttributeValue::S("item4".to_owned())), + ]) + ); + + let third = query(second.last_evaluated_key); + assert_eq!(third.count, 0); + assert!(third.last_evaluated_key.is_empty()); + } + + #[test] + fn test_should_end_query_on_partition_key_only_table() { + let provider = setup_provider_with_table(); + provider + .handle_put_item(PutItemInput { + table_name: "TestTable".to_owned(), + item: HashMap::from([("pk".to_owned(), AttributeValue::S("k1".to_owned()))]), + ..Default::default() + }) + .unwrap(); + let query = |exclusive_start_key| { + provider + .handle_query(QueryInput { + table_name: "TestTable".to_owned(), + key_condition_expression: Some("pk = :p".to_owned()), + expression_attribute_values: HashMap::from([( + ":p".to_owned(), + AttributeValue::S("k1".to_owned()), + )]), + limit: Some(1), + exclusive_start_key, + ..Default::default() + }) + .unwrap() + }; + + let first = query(HashMap::new()); + assert_eq!(first.count, 1); + assert_eq!( + first.last_evaluated_key, + HashMap::from([("pk".to_owned(), AttributeValue::S("k1".to_owned()))]) + ); + + let second = query(first.last_evaluated_key); + assert_eq!(second.count, 0); + assert!(second.last_evaluated_key.is_empty()); + } } diff --git a/crates/rustack-dynamodb-core/src/storage.rs b/crates/rustack-dynamodb-core/src/storage.rs index 0db38a02..637be1c4 100644 --- a/crates/rustack-dynamodb-core/src/storage.rs +++ b/crates/rustack-dynamodb-core/src/storage.rs @@ -589,8 +589,11 @@ impl TableStorage { /// Queries items in a single partition with optional sort key conditions. /// /// Returns a tuple of (items, last_evaluated_key). The last evaluated key - /// is `Some` when the result was truncated by the count limit or the 1 MB - /// response size cap, indicating that more items are available for pagination. + /// is `Some` whenever the count limit or the 1 MB response size cap stopped + /// the read, even when no items remain; the page after it is then empty. + /// + /// `exclusive_start_key` is the sort key of the last item already returned, + /// which is [`SortableAttributeValue::Sentinel`] on a table without a sort key. #[must_use] pub fn query( &self, @@ -600,9 +603,6 @@ impl TableStorage { limit: Option, exclusive_start_key: Option<&SortableAttributeValue>, ) -> (Vec>, Option) { - /// DynamoDB caps a single Query response at 1 MB. - const MAX_RESPONSE_BYTES: u64 = 1_048_576; - let Some(partition) = self.data.get(partition_key) else { return (Vec::new(), None); }; @@ -615,29 +615,11 @@ impl TableStorage { exclusive_start_key, ); - // Apply the 1 MB size cap and optional count limit. - let effective_limit = limit.unwrap_or(usize::MAX); - let mut selected: Vec<&StoredItem> = Vec::new(); - let mut cumulative_size: u64 = 0; - let mut size_limited = false; - - for item in all_items.iter().take(effective_limit) { - let item_size = calculate_item_size(&item.attributes); - // Always include at least one item even if it exceeds 1 MB on its - // own, matching DynamoDB behaviour. - if !selected.is_empty() && cumulative_size + item_size > MAX_RESPONSE_BYTES { - size_limited = true; - break; - } - cumulative_size += item_size; - selected.push(item); - } - - let count_limited = - !size_limited && selected.len() == effective_limit && effective_limit < all_items.len(); - let has_more = size_limited || count_limited; + let (selected, stopped) = take_page(all_items, limit, |item| { + calculate_item_size(&item.attributes) + }); - let last_key = if has_more { + let last_key = if stopped { build_last_evaluated_key( &selected, Some(selected.len()), @@ -659,8 +641,8 @@ impl TableStorage { /// Scans all items in the table. /// /// Returns a tuple of (items, last_evaluated_key). The last evaluated key - /// is `Some` when the result was truncated by the limit or by the 1 MB - /// response size cap. + /// is `Some` whenever the limit or the 1 MB response size cap stopped the + /// read, even when no items remain; the page after it is then empty. /// /// When `segment` and `total_segments` are provided, only items whose /// partition key hashes into the given segment are included (parallel scan). @@ -672,9 +654,6 @@ impl TableStorage { segment: Option, total_segments: Option, ) -> (Vec>, Option) { - /// DynamoDB caps a single Scan/Query response at 1 MB. - const MAX_RESPONSE_BYTES: u64 = 1_048_576; - // Collect all items sorted deterministically by partition key then sort key. let mut all_items: Vec<(AttributeValue, SortableAttributeValue, &StoredItem)> = Vec::new(); @@ -700,46 +679,26 @@ impl TableStorage { } } - // Apply exclusive_start_key: skip items up to and including the start key. - let start_idx = if let Some(start_key) = exclusive_start_key { - let start_sk = start_key + // Apply exclusive_start_key: skip every item up to and including its place in scan order, + // so the scan resumes correctly even if the item it names has since been deleted. + let start_idx = exclusive_start_key.map_or(0, |start_key| { + let start_partition = format!("{}", start_key.partition_key); + let start_sort = start_key .sort_key .as_ref() - .map_or(&SortableAttributeValue::Sentinel, |sk| sk); - - all_items - .iter() - .position(|(pk, sk, _)| *pk == start_key.partition_key && *sk == *start_sk) - .map_or(0, |idx| idx + 1) - } else { - 0 - }; - - let effective_limit = limit.unwrap_or(usize::MAX); - - // Select items respecting both the count limit and the 1 MB size cap. - let mut selected: Vec<&(AttributeValue, SortableAttributeValue, &StoredItem)> = Vec::new(); - let mut cumulative_size: u64 = 0; - let mut size_limited = false; - - for entry in all_items.iter().skip(start_idx).take(effective_limit) { - let item_size = calculate_item_size(&entry.2.attributes); - // Always include at least one item even if it exceeds 1 MB on its - // own, matching DynamoDB behaviour. - if !selected.is_empty() && cumulative_size + item_size > MAX_RESPONSE_BYTES { - size_limited = true; - break; - } - cumulative_size += item_size; - selected.push(entry); - } + .unwrap_or(&SortableAttributeValue::Sentinel); + all_items.partition_point(|(pk, sk, _)| match format!("{pk}").cmp(&start_partition) { + Ordering::Less => true, + Ordering::Equal => sk <= start_sort, + Ordering::Greater => false, + }) + }); - let count_limited = !size_limited - && selected.len() == effective_limit - && start_idx + effective_limit < all_items.len(); - let has_more = size_limited || count_limited; + let (selected, stopped) = take_page(all_items.iter().skip(start_idx), limit, |entry| { + calculate_item_size(&entry.2.attributes) + }); - let last_key = if has_more { + let last_key = if stopped { selected.last().map(|(pk, sk, _)| { PrimaryKey { partition_key: pk.clone(), @@ -771,6 +730,38 @@ impl TableStorage { } } +// --------------------------------------------------------------------------- +// Paging +// --------------------------------------------------------------------------- + +/// DynamoDB caps a single Query or Scan response at 1 MB. +const MAX_RESPONSE_BYTES: u64 = 1_048_576; + +/// Takes items in order until `limit` or the 1 MB response cap stops the read. +/// +/// Returns the page and whether the read was stopped. DynamoDB returns a +/// `LastEvaluatedKey` whenever it was, even when no items remain, so the page +/// after it is empty. At least one item is taken even if it alone exceeds 1 MB. +pub(crate) fn take_page( + items: impl IntoIterator, + limit: Option, + item_size: impl Fn(&T) -> u64, +) -> (Vec, bool) { + let limit = limit.unwrap_or(usize::MAX); + let mut page = Vec::new(); + let mut page_size: u64 = 0; + for item in items.into_iter().take(limit) { + let size = item_size(&item); + if !page.is_empty() && page_size.saturating_add(size) > MAX_RESPONSE_BYTES { + return (page, true); + } + page_size = page_size.saturating_add(size); + page.push(item); + } + let limited = page.len() == limit; + (page, limited) +} + // --------------------------------------------------------------------------- // Parallel scan helpers // --------------------------------------------------------------------------- @@ -857,8 +848,19 @@ fn collect_with_condition<'a>( ) -> Vec<&'a StoredItem> { match condition { SortKeyCondition::Eq(val) => { - // Exact match - at most one item. - partition.get(val).into_iter().collect() + // Exact match - at most one item, and none once paging has passed it. + let after_start = exclusive_start_key.is_none_or(|start| { + if scan_forward { + val > start + } else { + val < start + } + }); + partition + .get(val) + .filter(|_| after_start) + .into_iter() + .collect() } SortKeyCondition::Lt(val) => collect_range( partition, @@ -1606,6 +1608,99 @@ mod tests { } } + #[test] + fn test_should_return_last_key_when_limit_ends_on_last_item() { + let storage = TableStorage::new(numeric_sort_key_schema()); + + for i in 1..=4 { + let item = make_item(&[ + ("pk", AttributeValue::S("p1".to_owned())), + ("sk", AttributeValue::N(i.to_string())), + ]); + storage.put_item(item).ok(); + } + + let pk = AttributeValue::S("p1".to_owned()); + + let (items, last_key) = storage.query(&pk, None, true, Some(2), None); + assert_eq!(items.len(), 2); + let start_sk = last_key.as_ref().and_then(|k| k.sort_key.as_ref()); + + let (items2, last_key2) = storage.query(&pk, None, true, Some(2), start_sk); + assert_eq!(items2.len(), 2); + assert!(last_key2.is_some()); + let start_sk2 = last_key2.as_ref().and_then(|k| k.sort_key.as_ref()); + + let (items3, last_key3) = storage.query(&pk, None, true, Some(2), start_sk2); + assert_eq!(items3.len(), 0); + assert!(last_key3.is_none()); + + let (scanned, scan_key) = storage.scan(Some(4), None, None, None); + assert_eq!(scanned.len(), 4); + assert!(scan_key.is_some()); + + // A Limit the items never reach does not stop the read. + let (items, last_key) = storage.query(&pk, None, true, Some(5), None); + assert_eq!(items.len(), 4); + assert!(last_key.is_none()); + } + + #[test] + fn test_should_end_equality_query_after_its_item() { + let storage = TableStorage::new(numeric_sort_key_schema()); + for i in 1..=3 { + let item = make_item(&[ + ("pk", AttributeValue::S("p1".to_owned())), + ("sk", AttributeValue::N(i.to_string())), + ]); + storage.put_item(item).ok(); + } + let pk = AttributeValue::S("p1".to_owned()); + let equals_two = SortKeyCondition::Eq(SortableAttributeValue::N("2".to_owned())); + + for scan_forward in [true, false] { + let (items, last_key) = + storage.query(&pk, Some(&equals_two), scan_forward, Some(1), None); + assert_eq!(items.len(), 1); + let start = last_key.as_ref().and_then(|key| key.sort_key.as_ref()); + assert!(start.is_some()); + + let (items, last_key) = + storage.query(&pk, Some(&equals_two), scan_forward, Some(1), start); + assert_eq!(items.len(), 0); + assert!(last_key.is_none()); + } + } + + #[test] + fn test_should_resume_scan_after_deleted_start_key() { + let storage = TableStorage::new(numeric_sort_key_schema()); + for i in 1..=4 { + let item = make_item(&[ + ("pk", AttributeValue::S("p1".to_owned())), + ("sk", AttributeValue::N(i.to_string())), + ]); + storage.put_item(item).ok(); + } + + let (items, last_key) = storage.scan(Some(2), None, None, None); + assert_eq!(items.len(), 2); + let Some(last_key) = last_key else { + panic!("Limit stopped the scan, so it returns a key"); + }; + assert!(storage.delete_item(&last_key).is_some()); + + let (items, _) = storage.scan(Some(2), Some(&last_key), None, None); + let sort_keys: Vec<_> = items.iter().filter_map(|item| item.get("sk")).collect(); + assert_eq!( + sort_keys, + [ + &AttributeValue::N("3".to_owned()), + &AttributeValue::N("4".to_owned()) + ] + ); + } + #[test] fn test_should_handle_composite_key_put_get_delete() { let storage = TableStorage::new(composite_key_schema());