From db206e994b89de8b1a99bff4b905a320369b2bbf Mon Sep 17 00:00:00 2001
From: Bert Blommers <info@bertblommers.nl>
Date: Fri, 18 Oct 2019 09:58:09 +0100
Subject: [PATCH] #250 - DynamoDB - Add check for valid query
 keyconditionexpression

---
 moto/dynamodb2/responses.py           |  3 +++
 tests/test_dynamodb2/test_dynamodb.py | 25 +++++++++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/moto/dynamodb2/responses.py b/moto/dynamodb2/responses.py
index 333213c2a..b5e5e11a8 100644
--- a/moto/dynamodb2/responses.py
+++ b/moto/dynamodb2/responses.py
@@ -457,6 +457,9 @@ class DynamoHandler(BaseResponse):
                 range_comparison = None
                 range_values = []
 
+            if '=' not in hash_key_expression:
+                return self.error('com.amazonaws.dynamodb.v20111205#ValidationException',
+                                  'Query key condition not supported')
             hash_key_value_alias = hash_key_expression.split("=")[1].strip()
             # Temporary fix until we get proper KeyConditionExpression function
             hash_key = value_alias_map.get(hash_key_value_alias, {'S': hash_key_value_alias})
diff --git a/tests/test_dynamodb2/test_dynamodb.py b/tests/test_dynamodb2/test_dynamodb.py
index 9debe49d6..3d9914f14 100644
--- a/tests/test_dynamodb2/test_dynamodb.py
+++ b/tests/test_dynamodb2/test_dynamodb.py
@@ -2705,6 +2705,31 @@ def test_item_size_is_under_400KB():
                                     Item={'id': {'S': 'foo'}, 'itemlist': {'L': [{'M': {'item1': {'S': large_item}}}]}})
 
 
+@mock_dynamodb2
+# https://docs.aws.amazon.com/amazondynamodb/latest/APIReference/API_Query.html#DDB-Query-request-KeyConditionExpression
+def test_hash_key_cannot_use_begins_with_operations():
+    dynamodb = boto3.resource('dynamodb')
+    table = dynamodb.create_table(
+        TableName='test-table',
+        KeySchema=[{'AttributeName': 'key', 'KeyType': 'HASH'}],
+        AttributeDefinitions=[{'AttributeName': 'key', 'AttributeType': 'S'}],
+        ProvisionedThroughput={'ReadCapacityUnits': 1, 'WriteCapacityUnits': 1})
+
+    items = [{'key': 'prefix-$LATEST', 'value': '$LATEST'},
+             {'key': 'prefix-DEV', 'value': 'DEV'},
+             {'key': 'prefix-PROD', 'value': 'PROD'}]
+
+    with table.batch_writer() as batch:
+        for item in items:
+            batch.put_item(Item=item)
+
+    table = dynamodb.Table('test-table')
+    with assert_raises(ClientError) as ex:
+        table.query(KeyConditionExpression=Key('key').begins_with('prefix-'))
+    ex.exception.response['Error']['Code'].should.equal('ValidationException')
+    ex.exception.response['Error']['Message'].should.equal('Query key condition not supported')
+
+
 def assert_failure_due_to_item_size(func, **kwargs):
     with assert_raises(ClientError) as ex:
         func(**kwargs)