fix: match real DynamoDB on Scan and Query ExclusiveStartKey malformations

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)
+    )
+}

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

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