diff --git a/crates/engine/src/index_helpers.rs b/crates/engine/src/index_helpers.rs index 5ddf3dc..0dc333a 100755 --- a/crates/engine/src/index_helpers.rs +++ b/crates/engine/src/index_helpers.rs @@ -3,7 +3,10 @@ //! Shared helpers for secondary index operations in the engine layer. -use extenddb_core::types::{IndexInfo, Item, KeySchemaElement, ProjectionType}; +use extenddb_core::types::{ + AttributeDefinition, AttributeValue, IndexInfo, Item, KeySchemaElement, ProjectionType, + ScalarAttributeType, +}; /// Build the combined key schema for `LastEvaluatedKey` extraction. /// @@ -61,27 +64,101 @@ pub fn apply_index_projection( } } -/// Validate that an `ExclusiveStartKey` contains the required key elements for Scan. +/// Validate `ExclusiveStartKey` against the (table or index) key schema for +/// `Scan`. Base-table scans use the long DynamoDB error message; index scans +/// use the short one. /// -/// For base table scans, uses the long DynamoDB error message. -/// For index scans, uses the short message (matching real DynamoDB behavior). +/// # Errors +/// +/// Returns `ValidationException` if the start key has missing keys, extras, +/// or scalar type mismatches. pub fn validate_scan_exclusive_start_key( start_key: &Item, key_info: &extenddb_core::types::TableKeyInfo, index_info: Option<&IndexInfo>, ) -> Result<(), extenddb_core::error::DynamoDbError> { let required = combined_lek_key_schema(&key_info.key_schema, index_info); - for ks in &required { + let message = scan_invalid_start_key_message(index_info); + check_exclusive_start_key( + start_key, + &required, + &key_info.attribute_definitions, + message, + ) +} + +/// Validate `ExclusiveStartKey` for `Query`. Same rules as Scan; uses the +/// short DynamoDB error message in all cases. +/// +/// # Errors +/// +/// Returns `ValidationException` if the start key has missing keys, extras, +/// or scalar type mismatches. +pub fn validate_query_exclusive_start_key( + start_key: &Item, + key_info: &extenddb_core::types::TableKeyInfo, + index_info: Option<&IndexInfo>, +) -> Result<(), extenddb_core::error::DynamoDbError> { + let required = combined_lek_key_schema(&key_info.key_schema, index_info); + check_exclusive_start_key( + start_key, + &required, + &key_info.attribute_definitions, + QUERY_INVALID_START_KEY_MSG, + ) +} + +const QUERY_INVALID_START_KEY_MSG: &str = "The provided starting key is invalid"; +const SCAN_INVALID_START_KEY_MSG_BASE: &str = "The provided starting key is invalid: \ + The provided key element does not match the schema"; +const SCAN_INVALID_START_KEY_MSG_INDEX: &str = "The provided starting key is invalid"; + +fn scan_invalid_start_key_message(index_info: Option<&IndexInfo>) -> &'static str { + if index_info.is_some() { + SCAN_INVALID_START_KEY_MSG_INDEX + } else { + SCAN_INVALID_START_KEY_MSG_BASE + } +} + +/// Three rules: required-keys-present, no-extras, scalar-type-match. +fn check_exclusive_start_key( + start_key: &Item, + required: &[KeySchemaElement], + attribute_definitions: &[AttributeDefinition], + error_message: &str, +) -> Result<(), extenddb_core::error::DynamoDbError> { + let invalid = + || extenddb_core::error::DynamoDbError::ValidationException(error_message.to_owned()); + + for ks in required { if !start_key.contains_key(&ks.attribute_name) { - let msg = if index_info.is_some() { - "The provided starting key is invalid".to_owned() - } else { - "The provided starting key is invalid: The provided key element does not match the schema".to_owned() - }; - return Err(extenddb_core::error::DynamoDbError::ValidationException( - msg, - )); + return Err(invalid()); } } + if start_key.len() != required.len() { + return Err(invalid()); + } + for ks in required { + let declared = attribute_definitions + .iter() + .find(|ad| ad.attribute_name == ks.attribute_name) + .ok_or_else(invalid)?; + let supplied = start_key.get(&ks.attribute_name).ok_or_else(invalid)?; + if !attr_value_matches_scalar(supplied, declared.attribute_type) { + return Err(invalid()); + } + } + Ok(()) } + +/// Whether an `AttributeValue` matches the declared scalar type (S / N / B). +fn attr_value_matches_scalar(value: &AttributeValue, scalar: ScalarAttributeType) -> bool { + matches!( + (value, scalar), + (AttributeValue::S(_), ScalarAttributeType::S) + | (AttributeValue::N(_), ScalarAttributeType::N) + | (AttributeValue::B(_), ScalarAttributeType::B) + ) +} diff --git a/crates/engine/src/query.rs b/crates/engine/src/query.rs index e7431c8..be9ba56 100755 --- a/crates/engine/src/query.rs +++ b/crates/engine/src/query.rs @@ -20,7 +20,7 @@ use crate::OperationContext; use crate::capacity_helpers; use crate::create_table::storage_err_to_dynamo; use crate::expression_helpers::{build_expression_maps, parse_optional_filter}; -use crate::index_helpers::combined_lek_key_schema; +use crate::index_helpers::{combined_lek_key_schema, validate_query_exclusive_start_key}; use crate::legacy_filter::{desugar_filter, desugar_key_conditions}; use crate::read_helpers::apply_post_read; use crate::serialize_output; @@ -362,14 +362,7 @@ pub async fn handle_query( // Validate ExclusiveStartKey matches the key schema if let Some(ref start_key) = input.exclusive_start_key { - let required = combined_lek_key_schema(&key_info.key_schema, index_info.as_ref()); - for ks in &required { - if !start_key.contains_key(&ks.attribute_name) { - return Err(DynamoDbError::ValidationException( - "The provided starting key is invalid".to_owned(), - )); - } - } + validate_query_exclusive_start_key(start_key, &key_info, index_info.as_ref())?; } // Query storage diff --git a/tests/python/test_scan_edge_cases.py b/tests/python/test_scan_edge_cases.py index 9549e30..a8403f3 100755 --- a/tests/python/test_scan_edge_cases.py +++ b/tests/python/test_scan_edge_cases.py @@ -259,3 +259,196 @@ def test_query_malformed_exclusive_start_key_on_gsi(self, dynamodb_client): assert err["Code"] == "ValidationException" assert err["Message"] == "The provided starting key is invalid" dynamodb_client.delete_table(TableName=name) + + def test_scan_empty_exclusive_start_key(self, table_factory, dynamodb_client): + """Scan with an empty {} ExclusiveStartKey returns ValidationException.""" + name = table_factory() + with pytest.raises(ClientError) as exc_info: + dynamodb_client.scan(TableName=name, ExclusiveStartKey={}) + err = exc_info.value.response["Error"] + assert err["Code"] == "ValidationException" + assert err["Message"] == ( + "The provided starting key is invalid: " + "The provided key element does not match the schema" + ) + + def test_query_empty_exclusive_start_key(self, table_factory, dynamodb_client): + """Query with an empty {} ExclusiveStartKey returns ValidationException.""" + name = table_factory() + with pytest.raises(ClientError) as exc_info: + dynamodb_client.query( + TableName=name, + KeyConditionExpression="pk = :v", + ExpressionAttributeValues={":v": {"S": "x"}}, + ExclusiveStartKey={}, + ) + err = exc_info.value.response["Error"] + assert err["Code"] == "ValidationException" + assert err["Message"] == "The provided starting key is invalid" + + def test_scan_extra_attribute_in_exclusive_start_key(self, table_factory, dynamodb_client): + """Scan with a valid composite start key plus an extra non-key attr.""" + name = table_factory(range_key="sk") + dynamodb_client.put_item( + TableName=name, + Item={"pk": {"S": "P"}, "sk": {"S": "S"}}, + ) + with pytest.raises(ClientError) as exc_info: + dynamodb_client.scan( + TableName=name, + ExclusiveStartKey={ + "pk": {"S": "P"}, + "sk": {"S": "S"}, + "extra": {"S": "junk"}, + }, + ) + err = exc_info.value.response["Error"] + assert err["Code"] == "ValidationException" + assert err["Message"] == ( + "The provided starting key is invalid: " + "The provided key element does not match the schema" + ) + + def test_query_extra_attribute_in_exclusive_start_key(self, table_factory, dynamodb_client): + """Query with a valid composite start key plus an extra non-key attr.""" + name = table_factory(range_key="sk") + dynamodb_client.put_item( + TableName=name, + Item={"pk": {"S": "P"}, "sk": {"S": "S"}}, + ) + with pytest.raises(ClientError) as exc_info: + dynamodb_client.query( + TableName=name, + KeyConditionExpression="pk = :v", + ExpressionAttributeValues={":v": {"S": "P"}}, + ExclusiveStartKey={ + "pk": {"S": "P"}, + "sk": {"S": "S"}, + "extra": {"S": "junk"}, + }, + ) + err = exc_info.value.response["Error"] + assert err["Code"] == "ValidationException" + assert err["Message"] == "The provided starting key is invalid" + + def test_scan_wrong_scalar_type_in_exclusive_start_key( + self, table_factory, dynamodb_client + ): + """Scan with PK declared S but supplied as N raises ValidationException.""" + name = table_factory(range_key="sk") + dynamodb_client.put_item( + TableName=name, + Item={"pk": {"S": "P"}, "sk": {"S": "S"}}, + ) + with pytest.raises(ClientError) as exc_info: + dynamodb_client.scan( + TableName=name, + ExclusiveStartKey={ + "pk": {"N": "42"}, + "sk": {"S": "S"}, + }, + ) + err = exc_info.value.response["Error"] + assert err["Code"] == "ValidationException" + assert err["Message"] == ( + "The provided starting key is invalid: " + "The provided key element does not match the schema" + ) + + def test_query_wrong_scalar_type_in_exclusive_start_key( + self, table_factory, dynamodb_client + ): + """Query with PK declared S but supplied as N raises ValidationException.""" + name = table_factory(range_key="sk") + dynamodb_client.put_item( + TableName=name, + Item={"pk": {"S": "P"}, "sk": {"S": "S"}}, + ) + with pytest.raises(ClientError) as exc_info: + dynamodb_client.query( + TableName=name, + KeyConditionExpression="pk = :v", + ExpressionAttributeValues={":v": {"S": "P"}}, + ExclusiveStartKey={ + "pk": {"N": "42"}, + "sk": {"S": "S"}, + }, + ) + err = exc_info.value.response["Error"] + assert err["Code"] == "ValidationException" + assert err["Message"] == "The provided starting key is invalid" + + def test_scan_simple_table_extra_attribute_in_exclusive_start_key( + self, table_factory, dynamodb_client + ): + """Scan on a single-key table rejects an extra sk in the start key.""" + name = table_factory() # default: pk only + dynamodb_client.put_item( + TableName=name, + Item={"pk": {"S": "P"}}, + ) + with pytest.raises(ClientError) as exc_info: + dynamodb_client.scan( + TableName=name, + ExclusiveStartKey={"pk": {"S": "P"}, "sk": {"S": "S"}}, + ) + err = exc_info.value.response["Error"] + assert err["Code"] == "ValidationException" + assert err["Message"] == ( + "The provided starting key is invalid: " + "The provided key element does not match the schema" + ) + + def test_query_simple_table_extra_attribute_in_exclusive_start_key( + self, table_factory, dynamodb_client + ): + """Query on a single-key table rejects an extra sk in the start key.""" + name = table_factory() # default: pk only + dynamodb_client.put_item( + TableName=name, + Item={"pk": {"S": "P"}}, + ) + with pytest.raises(ClientError) as exc_info: + dynamodb_client.query( + TableName=name, + KeyConditionExpression="pk = :v", + ExpressionAttributeValues={":v": {"S": "P"}}, + ExclusiveStartKey={"pk": {"S": "P"}, "sk": {"S": "S"}}, + ) + err = exc_info.value.response["Error"] + assert err["Code"] == "ValidationException" + assert err["Message"] == "The provided starting key is invalid" + + def test_scan_full_composite_exclusive_start_key_is_accepted( + self, table_factory, dynamodb_client + ): + """Sanity: a complete, well-typed ExclusiveStartKey paginates without error.""" + name = table_factory(range_key="sk") + dynamodb_client.put_item( + TableName=name, + Item={"pk": {"S": "P"}, "sk": {"S": "S"}, "v": {"N": "1"}}, + ) + resp = dynamodb_client.scan( + TableName=name, + ExclusiveStartKey={"pk": {"S": "P"}, "sk": {"S": "S"}}, + ) + # Implementation detail: the storage layer is free to return zero or + # more items past the cursor. Contract under test is "no error". + assert "Items" in resp + + def test_query_full_composite_exclusive_start_key_is_accepted( + self, table_factory, dynamodb_client + ): + """Sanity: a complete, well-typed Query ExclusiveStartKey paginates without error.""" + name = table_factory(range_key="sk") + dynamodb_client.put_item( + TableName=name, + Item={"pk": {"S": "P"}, "sk": {"S": "S"}, "v": {"N": "1"}}, + ) + resp = dynamodb_client.query( + TableName=name, + KeyConditionExpression="pk = :v", + ExpressionAttributeValues={":v": {"S": "P"}}, + ExclusiveStartKey={"pk": {"S": "P"}, "sk": {"S": "S"}}, + ) + assert "Items" in resp