From 6aa2472545830932408039be5bb39b035446a58d Mon Sep 17 00:00:00 2001 From: Kyle Cannon Date: Fri, 30 Oct 2026 11:05:09 -0700 Subject: [PATCH] fix(dynamodb): return LastEvaluatedKey when Limit ends the page DynamoDB returns a LastEvaluatedKey whenever Limit stops a Query or Scan, even if nothing is left, and the next page just comes back empty. rustack only returned it when more items followed, so code that pages with Limit stopped a page early compared to AWS. LocalStack behaves the same way as this change. Query, Scan and GSI queries now share one paging helper for it. Returning a key on the last page means the follow-up request always carries a start key, which turned up three places that didn't honour it: - queries on tables without a sort key returned the same item forever - `sk = :v` queries ignored the start key and looped the same way - a scan whose start item had been deleted started over from the top All three now move past the start key. --- crates/rustack-dynamodb-core/src/provider.rs | 178 +++++++++++--- crates/rustack-dynamodb-core/src/storage.rs | 235 +++++++++++++------ 2 files changed, 310 insertions(+), 103 deletions(-) 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());