Repository: getmoto/moto

Issue:
DynamoDB2 queries on GSI's fail to raise exceptions on wrong key names
The following test case succeeds when run directly against AWS dynamodb, but fails when run using moto. It appears that moto is not validating the names of key attributes when querying an index.

- tested with lastest from `master`
- this is a keys-only index but I think that's irrelevant to the test
- this doesn't explicity test using `begins_with` or `between` but based on [the implementation here](https://github.com/spulec/moto/blob/9894e1785a610f1dd8c9bfce1fa416692b3e0c2f/moto/dynamodb2/responses.py#L503-L515) it seems that tests for those cases might be valuable

Thanks again for all your work on this awesome project, let me know if you have questions plz

```python
@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 assert_raises(ClientError) as ex:
        table.query(
            KeyConditionExpression="gsiK1PartitionKey = :pk AND wrongName = :sk",
            ExpressionAttributeValues={':pk': 'gsi-pk', ':sk': 'gsi-sk'},
            IndexName="GSI-K1",
        )["Items"]
    ex.exception.response["Error"]["Code"].should.equal("ValidationException")
    ex.exception.response["ResponseMetadata"]["HTTPStatusCode"].should.equal(400)
    ex.exception.response["Error"]["Message"].should.equal('Query condition missed key schema element: gsiK1SortKey')

    # check using wrong name for partition key throws exception
    with assert_raises(ClientError) as ex:
        table.query(
            KeyConditionExpression="wrongName = :pk AND gsiK1SortKey = :sk",
            ExpressionAttributeValues={':pk': 'gsi-pk', ':sk': 'gsi-sk'},
            IndexName="GSI-K1",
        )["Items"]
    ex.exception.response["Error"]["Code"].should.equal("ValidationException")
    ex.exception.response["ResponseMetadata"]["HTTPStatusCode"].should.equal(400)
    ex.exception.response["Error"]["Message"].should.equal('Query condition missed key schema element: gsiK1PartitionKey')
```


Discussion:
Thanks for raising this, and providing the test case @mfogel - much appreciated.

Marking this as an enhancement.

Resolving patch:
diff --git a/moto/dynamodb2/models/__init__.py b/moto/dynamodb2/models/__init__.py
--- a/moto/dynamodb2/models/__init__.py
+++ b/moto/dynamodb2/models/__init__.py
@@ -757,6 +757,7 @@ def query(
                 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(
diff --git a/moto/dynamodb2/responses.py b/moto/dynamodb2/responses.py
--- a/moto/dynamodb2/responses.py
+++ b/moto/dynamodb2/responses.py
@@ -550,29 +550,50 @@ def query(self):
                 )
                 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)
+                    (
+                        (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 @@ def query(self):
                         ),
                     )
                 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


Test patch:
diff --git a/tests/test_dynamodb2/test_dynamodb_exceptions.py b/tests/test_dynamodb2/test_dynamodb_exceptions.py
new file mode 100644
--- /dev/null
+++ b/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"
+    )