DynamoDB - Raise exceptions when query is missing GSI keys (#4379)

moto/dynamodb2/models/__init__.py

@@ -757,6 +757,7 @@ class Table(CloudFormationModel):
                 for item in list(self.all_items())
                 if isinstance(item, Item) and item.hash_key == hash_key
             ]
+
         if range_comparison:
             if index_name and not index_range_key:
                 raise ValueError(

moto/dynamodb2/responses.py

@@ -550,29 +550,50 @@ class DynamoHandler(BaseResponse):
                 )
                 hash_key_regex = r"(^|[\s(]){0}\b".format(hash_key_var)
                 i, hash_key_expression = next(
+                    (
                         (i, e)
                         for i, e in enumerate(expressions)
                         if re.search(hash_key_regex, e)
+                    ),
+                    (None, None),
+                )
+                if hash_key_expression is None:
+                    return self.error(
+                        "ValidationException",
+                        "Query condition missed key schema element: {}".format(
+                            hash_key_var
+                        ),
                     )
                 hash_key_expression = hash_key_expression.strip("()")
                 expressions.pop(i)
 
                 # TODO implement more than one range expression and OR operators
                 range_key_expression = expressions[0].strip("()")
-                range_key_expression_components = range_key_expression.split()
+                # Split expression, and account for all kinds of whitespacing around commas and brackets
+                range_key_expression_components = re.split(
+                    r"\s*\(\s*|\s*,\s*|\s", range_key_expression
+                )
+                # Skip whitespace
+                range_key_expression_components = [
+                    c for c in range_key_expression_components if c
+                ]
                 range_comparison = range_key_expression_components[1]
 
                 if " and " in range_key_expression.lower():
                     range_comparison = "BETWEEN"
+                    # [range_key, between, x, and, y]
                     range_values = [
                         value_alias_map[range_key_expression_components[2]],
                         value_alias_map[range_key_expression_components[4]],
                     ]
+                    supplied_range_key = range_key_expression_components[0]
                 elif "begins_with" in range_key_expression:
                     range_comparison = "BEGINS_WITH"
+                    # [begins_with, range_key, x]
                     range_values = [
                         value_alias_map[range_key_expression_components[-1]]
                     ]
+                    supplied_range_key = range_key_expression_components[1]
                 elif "begins_with" in range_key_expression.lower():
                     function_used = range_key_expression[
                         range_key_expression.lower().index("begins_with") : len(
@@ -586,7 +607,23 @@ class DynamoHandler(BaseResponse):
                         ),
                     )
                 else:
+                    # [range_key, =, x]
                     range_values = [value_alias_map[range_key_expression_components[2]]]
+                    supplied_range_key = range_key_expression_components[0]
+
+                supplied_range_key = expression_attribute_names.get(
+                    supplied_range_key, supplied_range_key
+                )
+                range_keys = [
+                    k["AttributeName"] for k in index if k["KeyType"] == "RANGE"
+                ]
+                if supplied_range_key not in range_keys:
+                    return self.error(
+                        "ValidationException",
+                        "Query condition missed key schema element: {}".format(
+                            range_keys[0]
+                        ),
+                    )
             else:
                 hash_key_expression = key_condition_expression.strip("()")
                 range_comparison = None

tests/test_dynamodb2/test_dynamodb_exceptions.py

@@ -0,0 +1,98 @@
+import boto3
+import pytest
+import sure  # noqa
+
+from botocore.exceptions import ClientError
+from moto import mock_dynamodb2
+
+
+@mock_dynamodb2
+def test_query_gsi_with_wrong_key_attribute_names_throws_exception():
+    table_schema = {
+        "KeySchema": [{"AttributeName": "partitionKey", "KeyType": "HASH"}],
+        "GlobalSecondaryIndexes": [
+            {
+                "IndexName": "GSI-K1",
+                "KeySchema": [
+                    {"AttributeName": "gsiK1PartitionKey", "KeyType": "HASH"},
+                    {"AttributeName": "gsiK1SortKey", "KeyType": "RANGE"},
+                ],
+                "Projection": {"ProjectionType": "KEYS_ONLY",},
+            }
+        ],
+        "AttributeDefinitions": [
+            {"AttributeName": "partitionKey", "AttributeType": "S"},
+            {"AttributeName": "gsiK1PartitionKey", "AttributeType": "S"},
+            {"AttributeName": "gsiK1SortKey", "AttributeType": "S"},
+        ],
+    }
+
+    item = {
+        "partitionKey": "pk-1",
+        "gsiK1PartitionKey": "gsi-pk",
+        "gsiK1SortKey": "gsi-sk",
+        "someAttribute": "lore ipsum",
+    }
+
+    dynamodb = boto3.resource("dynamodb", region_name="us-east-1")
+    dynamodb.create_table(
+        TableName="test-table", BillingMode="PAY_PER_REQUEST", **table_schema
+    )
+    table = dynamodb.Table("test-table")
+    table.put_item(Item=item)
+
+    # check using wrong name for sort key throws exception
+    with pytest.raises(ClientError) as exc:
+        table.query(
+            KeyConditionExpression="gsiK1PartitionKey = :pk AND wrongName = :sk",
+            ExpressionAttributeValues={":pk": "gsi-pk", ":sk": "gsi-sk"},
+            IndexName="GSI-K1",
+        )["Items"]
+    err = exc.value.response["Error"]
+    err["Code"].should.equal("ValidationException")
+    err["Message"].should.equal(
+        "Query condition missed key schema element: gsiK1SortKey"
+    )
+
+    # check using wrong name for partition key throws exception
+    with pytest.raises(ClientError) as exc:
+        table.query(
+            KeyConditionExpression="wrongName = :pk AND gsiK1SortKey = :sk",
+            ExpressionAttributeValues={":pk": "gsi-pk", ":sk": "gsi-sk"},
+            IndexName="GSI-K1",
+        )["Items"]
+    err = exc.value.response["Error"]
+    err["Code"].should.equal("ValidationException")
+    err["Message"].should.equal(
+        "Query condition missed key schema element: gsiK1PartitionKey"
+    )
+
+    # verify same behaviour for begins_with
+    with pytest.raises(ClientError) as exc:
+        table.query(
+            KeyConditionExpression="gsiK1PartitionKey = :pk AND begins_with ( wrongName , :sk )",
+            ExpressionAttributeValues={":pk": "gsi-pk", ":sk": "gsi-sk"},
+            IndexName="GSI-K1",
+        )["Items"]
+    err = exc.value.response["Error"]
+    err["Code"].should.equal("ValidationException")
+    err["Message"].should.equal(
+        "Query condition missed key schema element: gsiK1SortKey"
+    )
+
+    # verify same behaviour for between
+    with pytest.raises(ClientError) as exc:
+        table.query(
+            KeyConditionExpression="gsiK1PartitionKey = :pk AND wrongName BETWEEN :sk1 and :sk2",
+            ExpressionAttributeValues={
+                ":pk": "gsi-pk",
+                ":sk1": "gsi-sk",
+                ":sk2": "gsi-sk2",
+            },
+            IndexName="GSI-K1",
+        )["Items"]
+    err = exc.value.response["Error"]
+    err["Code"].should.equal("ValidationException")
+    err["Message"].should.equal(
+        "Query condition missed key schema element: gsiK1SortKey"
+    )