From 3de3cf2bc41429dd5e30f4df040074b807f83d88 Mon Sep 17 00:00:00 2001 From: Palagiri Subba Reddy Date: Tue, 7 Jul 2026 23:07:06 +0530 Subject: [PATCH] fix: return proper error for invalid dynamo query instead of panic --- lore-aws/src/dynamodb.rs | 24 +++++++++--------------- 1 file changed, 9 insertions(+), 15 deletions(-) diff --git a/lore-aws/src/dynamodb.rs b/lore-aws/src/dynamodb.rs index a41b9eb..5f74c4f 100644 --- a/lore-aws/src/dynamodb.rs +++ b/lore-aws/src/dynamodb.rs @@ -69,15 +69,6 @@ const QUERY_ROWS: [f64; 10] = [ pub const METRICS_ACCUMULATION_OPERATION_LATENCY_METRIC_NAME: &str = "accumulation_operation_duration"; - -/// Temporary environment variable to allow enabling/disabling query pagination. -/// TODO(jcohen): Remove this when we're confident in the pagination support. -static ENABLE_QUERY_PAGINATION: LazyLock = - LazyLock::new(|| match std::env::var("LORE_ENABLE_QUERY_PAGINATION") { - Ok(v) => v.eq_ignore_ascii_case("true"), - Err(_) => true, // Default to enabled if the env var isn't present - }); - pub trait DynamoDbQuery { fn index_name(&self) -> Option { None @@ -575,6 +566,14 @@ impl DynamoDbImpl { where T: DynamoDbQuery + 'static, { + if matches!(query.select(), Some(Select::Count)) && query.limit().is_some() { + return Err(crate::aws_error::AwsError::AwsSdkError( + aws_sdk_dynamodb::error::SdkError::construction_failure( + "Combining Select::Count and a limit is disallowed for paginated queries", + ), + )); + } + let base_labels = { let mut labels = self .instruments @@ -651,12 +650,7 @@ impl DynamoDbImpl { // the desired number of rows, there's no need to continue fetching more // rows. // 3. *Unless* the query is also a count query (in which case we need to fetch - // until there's no more data to collect to ensure we get the correct count). - // TODO(jcohen): We should probably disallow the combination of `Select::Count` - // and a limit to avoid the degenerate case of fetching a full count one row at - // a time. - if *ENABLE_QUERY_PAGINATION - && let Some(last_evaluated_key) = r.last_evaluated_key + if let Some(last_evaluated_key) = r.last_evaluated_key && !last_evaluated_key.is_empty() && (query.limit().is_none() || matches!(query.select(), Some(Select::Count))