fix(query): classify and order key conditions on a multi-attribute GSI the way the service does

crates/engine/src/legacy_filter.rs

@@ -22,6 +22,12 @@ pub fn desugar_key_conditions(
     conditions: &HashMap<String, Condition>,
     key_schema: &[(String, bool)], // Vec of (attr_name, is_hash)
 ) -> Result<(KeyCondition, ExpressionMaps), DynamoDbError> {
+    let hash_count = key_schema.iter().filter(|(_, is_hash)| *is_hash).count();
+    let range_count = key_schema.len() - hash_count;
+    if hash_count > 1 || range_count > 1 {
+        return desugar_multipart_key_conditions(conditions, key_schema);
+    }
+
     // Find the hash key condition
     let hash_attr = key_schema
         .iter()
@@ -162,6 +168,163 @@ fn desugar_sk_condition(
     }
 }
 
+/// Desugar legacy `KeyConditions` on a key schema with several HASH or several
+/// RANGE attributes into the same `KeyCondition` shape the expression parser
+/// produces, so `KeyCondition::resolve_multipart` classifies both forms by one
+/// rule (measured on the service 2026-10-21: `KeyConditions` on such an index
+/// answer the same refusals as the equivalent `KeyConditionExpression`).
+fn desugar_multipart_key_conditions(
+    conditions: &HashMap<String, Condition>,
+    key_schema: &[(String, bool)],
+) -> Result<(KeyCondition, ExpressionMaps), DynamoDbError> {
+    let mut values = HashMap::new();
+    let mut equalities: Vec<(Vec<PathElement>, Expr)> = Vec::new();
+    let mut non_equalities: Vec<(String, SortKeyCondition)> = Vec::new();
+
+    // Key-schema attributes first, in key-schema order, then any other name,
+    // so placeholder numbering and clause order are deterministic.
+    let mut ordered: Vec<&String> = key_schema
+        .iter()
+        .map(|(name, _)| name)
+        .filter(|name| conditions.contains_key(*name))
+        .collect();
+    let mut others: Vec<&String> = conditions
+        .keys()
+        .filter(|name| !key_schema.iter().any(|(k, _)| k == *name))
+        .collect();
+    others.sort();
+    ordered.extend(others);
+
+    for (i, attr) in ordered.into_iter().enumerate() {
+        let cond = &conditions[attr];
+        if cond.comparison_operator == "EQ" {
+            if cond.attribute_value_list.len() != 1 {
+                return Err(DynamoDbError::ValidationException(
+                    "Query key condition must have exactly one value for EQ".to_owned(),
+                ));
+            }
+            let placeholder = format!("_kc_eq{i}");
+            values.insert(placeholder.clone(), cond.attribute_value_list[0].clone());
+            equalities.push((
+                vec![PathElement::Attribute(attr.clone())],
+                Expr::Placeholder(placeholder),
+            ));
+        } else {
+            let mut own_values = HashMap::new();
+            let sk = desugar_sk_condition(cond, attr, &mut own_values)?;
+            // Rename this condition's placeholders so two non-equalities
+            // cannot collide on `_kc_sk`.
+            let sk = rename_placeholders(sk, &own_values, &format!("_kc_ne{i}_"), &mut values);
+            non_equalities.push((attr.clone(), sk));
+        }
+    }
+
+    let hash_attrs: Vec<&str> = key_schema
+        .iter()
+        .filter(|(_, is_hash)| *is_hash)
+        .map(|(name, _)| name.as_str())
+        .collect();
+    let range_attrs: Vec<&str> = key_schema
+        .iter()
+        .filter(|(_, is_hash)| !*is_hash)
+        .map(|(name, _)| name.as_str())
+        .collect();
+
+    // Two non-equalities: the service names the RANGE attributes before the
+    // last conditioned one that lack an equality (measured: `h EQ, a GT,
+    // b GT` is `RANGE key attributes a must have equality conditions
+    // specified in the query because a condition is present on key attribute
+    // b`).
+    if non_equalities.len() > 1 {
+        let conditioned: Vec<usize> = conditions
+            .keys()
+            .filter_map(|name| range_attrs.iter().position(|r| r == name))
+            .collect();
+        if let Some(&last) = conditioned.iter().max() {
+            let without_eq: Vec<&str> = range_attrs[..last]
+                .iter()
+                .enumerate()
+                .filter(|(_, attr)| {
+                    conditions
+                        .get(**attr)
+                        .is_none_or(|c| c.comparison_operator != "EQ")
+                })
+                .map(|(_, attr)| *attr)
+                .collect();
+            if !without_eq.is_empty() {
+                return Err(DynamoDbError::ValidationException(format!(
+                    "RANGE key attributes {} must have equality conditions specified in the \
+                     query because a condition is present on key attribute {}",
+                    without_eq.join(", "),
+                    range_attrs[last]
+                )));
+            }
+        }
+        return Err(DynamoDbError::ValidationException(
+            "Query key condition not supported".to_owned(),
+        ));
+    }
+
+    let Some((pk_path, pk_value)) = equalities.first().cloned() else {
+        // No equality at all: every HASH attribute lacks one.
+        return Err(DynamoDbError::ValidationException(format!(
+            "HASH key attributes {} must have equality conditions specified in the query",
+            hash_attrs.join(", ")
+        )));
+    };
+    let extra_pk_conditions = equalities.into_iter().skip(1).collect();
+    let sk_condition = non_equalities.into_iter().next().map(|(_, sk)| sk);
+
+    Ok((
+        KeyCondition {
+            pk_path,
+            pk_value,
+            sk_condition,
+            extra_pk_conditions,
+            extra_sk_conditions: Vec::new(),
+        },
+        ExpressionMaps::new(HashMap::new(), values),
+    ))
+}
+
+/// Move `own_values` into `values` under `prefix`-ed names and rewrite the
+/// condition's placeholders to match.
+fn rename_placeholders(
+    sk: SortKeyCondition,
+    own_values: &HashMap<String, AttributeValue>,
+    prefix: &str,
+    values: &mut HashMap<String, AttributeValue>,
+) -> SortKeyCondition {
+    let mut rename = |expr: Expr| -> Expr {
+        match expr {
+            Expr::Placeholder(name) => {
+                let new_name = format!("{prefix}{name}");
+                if let Some(v) = own_values.get(&name) {
+                    values.insert(new_name.clone(), v.clone());
+                }
+                Expr::Placeholder(new_name)
+            }
+            other => other,
+        }
+    };
+    match sk {
+        SortKeyCondition::Compare { path, op, value } => SortKeyCondition::Compare {
+            path,
+            op,
+            value: rename(value),
+        },
+        SortKeyCondition::Between { path, low, high } => SortKeyCondition::Between {
+            path,
+            low: rename(low),
+            high: rename(high),
+        },
+        SortKeyCondition::BeginsWith { path, prefix: p } => SortKeyCondition::BeginsWith {
+            path,
+            prefix: rename(p),
+        },
+    }
+}
+
 /// Desugar legacy `QueryFilter` or `ScanFilter` into a condition expression AST.
 ///
 /// Each entry maps an attribute name to a Condition with `ComparisonOperator` and values.

crates/engine/src/query.rs

@@ -270,50 +270,67 @@ pub async fn handle_query(
         &maps
     };
 
-    // Correct PK/SK assignment when both clauses are equality comparisons.
-    // The parser can't distinguish PK from SK without the key schema.
     let pk_attr = &query_key_info.key_schema[0].attribute_name;
-    key_condition.resolve_pk_sk(pk_attr, &effective_maps.names)?;
-
-    // Validate that the partition key is actually referenced in the condition.
-    let pk_resolved = resolve_path_attr_name(&key_condition.pk_path, &effective_maps.names);
-    if pk_resolved.as_deref() != Some(pk_attr.as_str()) {
-        return Err(DynamoDbError::ValidationException(format!(
-            "Query condition missed key schema element: {pk_attr}"
-        )));
-    }
-
-    // For multi-part key schemas (GSIs with >1 HASH attribute), reclassify
-    // the parsed conditions so all HASH attributes go to pk_path/extra_pk_conditions
-    // and the RANGE condition stays as sk_condition.
-    if extenddb_core::types::is_multipart_key_schema(&query_key_info.key_schema) {
-        let hash_elements = extenddb_core::types::hash_key_elements(&query_key_info.key_schema);
-        let hash_attrs: Vec<&str> = hash_elements
-            .iter()
-            .map(|ks| ks.attribute_name.as_str())
-            .collect();
-        key_condition.resolve_multipart(&hash_attrs, &effective_maps.names)?;
-
-        // Validate all HASH attributes are present in the KeyConditionExpression.
-        let provided_count = 1 + key_condition.extra_pk_conditions.len();
-        if provided_count != hash_attrs.len() {
-            // Find the first missing HASH attribute for the error message.
-            let missing = hash_attrs
+    // Assign every clause to its key attribute through the classification
+    // authorization also uses to find dynamodb:LeadingKeys, so both read the
+    // same clause as the partition key. On a multi-attribute key schema (several
+    // HASH or several RANGE attributes) every clause is assigned by name, every
+    // HASH attribute needs an equality and the RANGE conditions must cover a
+    // leading run of the RANGE attributes with equalities before the last one.
+    // A wrong-typed value is reported where the service reports it: after those
+    // structural rules and before the non-key rule on a multi-attribute key
+    // schema, and ahead of the missing-key check on a single-attribute one
+    // (measured 2026-10-21: `pk = :n` on an S key with a second clause on a
+    // non-key attribute reports the type, not the missing key).
+    //
+    // On a multi-attribute key schema the classifier refuses a clause on an
+    // attribute outside the key schema (`Query condition missed key schema
+    // element`), so the type check reads only the queried key schema's
+    // definitions there. For an index query the table's list also holds the
+    // base table's key and the other indexes' key attributes. Measured
+    // 2026-10-30 on a GSI with HASH `GSI2PK` and RANGE `GSI2SK`: `GSI2PK = :h
+    // AND PK = :k` answers `Query condition missed key schema element: GSI2SK`
+    // whether `:k` has the table's type for `PK` or not, so the service's type
+    // check does not see an attribute outside the queried key schema. The
+    // single-attribute path has no non-key rule yet: a clause on another
+    // attribute reaches storage, which applies it to the first RANGE
+    // attribute, so it keeps checking against the table's definitions, which
+    // refuse a value of the wrong type before storage does.
+    let multipart = extenddb_core::types::is_multipart_key_schema(&query_key_info.key_schema);
+    let key_schema_definitions: Vec<extenddb_core::types::AttributeDefinition> = query_key_info
+        .attribute_definitions
+        .iter()
+        .filter(|def| {
+            query_key_info
+                .key_schema
                 .iter()
-                .find(|attr| {
-                    let pk_name =
-                        resolve_path_attr_name(&key_condition.pk_path, &effective_maps.names);
-                    if pk_name.as_deref() == Some(*attr) {
-                        return false;
-                    }
-                    !key_condition.extra_pk_conditions.iter().any(|(path, _)| {
-                        resolve_path_attr_name(path, &effective_maps.names).as_deref()
-                            == Some(*attr)
-                    })
-                })
-                .unwrap_or(&hash_attrs[0]);
+                .any(|ks| ks.attribute_name == def.attribute_name)
+        })
+        .cloned()
+        .collect();
+    let type_definitions = if multipart {
+        &key_schema_definitions
+    } else {
+        &query_key_info.attribute_definitions
+    };
+    key_condition.resolve_key_schema_with(
+        &query_key_info.key_schema,
+        &effective_maps.names,
+        |path, value| check_key_value_type(path, value, type_definitions, effective_maps),
+    )?;
+    refuse_begins_with_on_number(
+        &key_condition,
+        &query_key_info.attribute_definitions,
+        effective_maps,
+    )?;
+    // On a single-attribute key schema, validate that the partition key is
+    // actually referenced in the condition. The classifier has already required
+    // an equality on every HASH attribute of a multi-attribute one.
+    if !multipart {
+        let pk_resolved = resolve_path_attr_name(&key_condition.pk_path, &effective_maps.names);
+        if pk_resolved.as_deref() != Some(pk_attr.as_str()) {
             return Err(DynamoDbError::ValidationException(format!(
-                "Query condition missed key schema element: {missing}"
+                "Query condition missed key schema element: {pk_attr}"
             )));
         }
     }
@@ -574,6 +591,85 @@ fn resolve_path_attr_name(
     }
 }
 
+/// Refuse a key condition value whose type differs from the attribute
+/// definition of the attribute it names.
+///
+/// Measured on the service 2026-10-21: an S key compared with an N value,
+/// an N sort key compared with an S value in `=`, `>`, `begins_with` or a
+/// BETWEEN bound, and the same on an index, all answer
+/// `One or more parameter values were invalid: Condition parameter type does
+/// not match schema type`, on single-attribute and multi-attribute key
+/// schemas alike. A clause on an attribute without a definition in
+/// `attribute_definitions` is not checked here. On a multi-attribute key
+/// schema the caller passes the queried key schema's definitions only, and the
+/// classifier reports a clause on any other attribute.
+fn check_key_value_type(
+    path: &[PathElement],
+    expr: &extenddb_core::expression::Expr,
+    attribute_definitions: &[extenddb_core::types::AttributeDefinition],
+    maps: &ExpressionMaps,
+) -> Result<(), DynamoDbError> {
+    use extenddb_core::expression::Expr;
+    use extenddb_core::types::{AttributeValue, ScalarAttributeType};
+
+    let Some(attr) = resolve_path_attr_name(path, &maps.names) else {
+        return Ok(());
+    };
+    let Some(def) = attribute_definitions
+        .iter()
+        .find(|d| d.attribute_name == attr)
+    else {
+        return Ok(());
+    };
+    let Expr::Placeholder(name) = expr else {
+        return Ok(());
+    };
+    let Some(value) = maps.values.get(name) else {
+        return Ok(());
+    };
+    let matches = matches!(
+        (def.attribute_type, value),
+        (ScalarAttributeType::S, AttributeValue::S(_))
+            | (ScalarAttributeType::N, AttributeValue::N(_))
+            | (ScalarAttributeType::B, AttributeValue::B(_))
+    );
+    if matches {
+        Ok(())
+    } else {
+        Err(DynamoDbError::ValidationException(
+            "One or more parameter values were invalid: Condition parameter type does not match \
+             schema type"
+                .to_owned(),
+        ))
+    }
+}
+
+/// Refuse `begins_with` on a number key attribute (measured 2026-10-21:
+/// `Invalid KeyConditionExpression: Incorrect operand type for operator or
+/// function; operator or function: begins_with, operand type: N`).
+fn refuse_begins_with_on_number(
+    key_condition: &extenddb_core::expression::KeyCondition,
+    attribute_definitions: &[extenddb_core::types::AttributeDefinition],
+    maps: &ExpressionMaps,
+) -> Result<(), DynamoDbError> {
+    use extenddb_core::expression::SortKeyCondition;
+    use extenddb_core::types::ScalarAttributeType;
+
+    if let Some(SortKeyCondition::BeginsWith { path, .. }) = &key_condition.sk_condition
+        && let Some(attr) = resolve_path_attr_name(path, &maps.names)
+        && attribute_definitions
+            .iter()
+            .any(|d| d.attribute_name == attr && d.attribute_type == ScalarAttributeType::N)
+    {
+        return Err(DynamoDbError::ValidationException(
+            "Invalid KeyConditionExpression: Incorrect operand type for operator or function; \
+             operator or function: begins_with, operand type: N"
+                .to_owned(),
+        ));
+    }
+    Ok(())
+}
+
 /// Validate that all `#name` references in an expression AST are defined.
 fn validate_name_refs_in_expr(
     expr: &extenddb_core::expression::Expr,