Repository: getmoto/moto

Issue:
DynamoDB query with brackets in KeyConditionExpression causes: Invalid function name; function:
See the attached testcase

I have found that when mocking DynamoDB and running a query, if I put brackets in my KeyConditionExpression then it raises an error

My KeyConditionExpression is `(#audit_object_id = :object_id) AND (#audit_timestamp < :before_timestamp)`

Example traceback (there is no more detail)

    Traceback (most recent call last):
      File "C:\Users\david.pottage\git\labs.platform-hub.audit-trail-api\src\tests\test_demo_moto_query_bug.py", line 140, in test_query_with_brackets
        query_results = self.table.query(**self.query_args)
    botocore.exceptions.ClientError: An error occurred (ValidationException) when calling the Query operation: Invalid KeyConditionExpression: Invalid function name; function:

If I omit brackets from my query then it works fine.

Seen on moto version 4.1.6 with boto3 1.26.99



Discussion:
I thought I could attach my example code to the ticket. Looks like that won't work. Instead here it is pasted inline:

    import boto3
    import json
    import unittest
    from datetime import datetime, timedelta
    from moto import mock_dynamodb

    # Reported to the moto project on github as:
    #   https://github.com/getmoto/moto/issues/6150

    class DynamoDBQueryTests(unittest.TestCase):

        # Do not decorate the class, instead create & remove the mock via setUp and tearDown
        mock_dynamodb = mock_dynamodb()

        def setUp(self) -> None:
            self.mock_dynamodb.start()
            kwargs = {}
            boto3.setup_default_session()
            dynamodb = boto3.resource('dynamodb', **kwargs)

            self.query_args = {
                # NB the KeyConditionExpression will be added by each test
                "ScanIndexForward": False,
                "ExpressionAttributeNames": {
                    "#audit_object_id": "object_id",
                    "#audit_timestamp": "timestamp",
                    "#audit_user": "user"
                },
                "ExpressionAttributeValues": {
                    ":object_id": "test.com",
                    ":before_timestamp": 1676246400_000000,
                    ":user": "test@test.com"
                },
                "FilterExpression": "#audit_user = :user",
                "Limit": 20,
            }

            table = dynamodb.create_table(
                TableName="test-table",
                TableClass="STANDARD_INFREQUENT_ACCESS",
                BillingMode="PAY_PER_REQUEST",
                KeySchema=[
                    {"AttributeName": "object_id", "KeyType": "HASH"},
                    {"AttributeName": "timestamp", "KeyType": "RANGE"}
                ],
                AttributeDefinitions=[
                    {"AttributeName": "object_id", "AttributeType": "S"},
                    {"AttributeName": "user", "AttributeType": "S"},
                    {"AttributeName": "timestamp", "AttributeType": "N"}
                ],
                GlobalSecondaryIndexes=[
                    {
                        "IndexName": "gsi_user",
                        "Projection": {"ProjectionType": "ALL"},
                        "KeySchema": [
                            {"AttributeName": "user", "KeyType": "HASH"},
                            {"AttributeName": "timestamp", "KeyType": "RANGE"}
                        ]
                    }
                ]
            )
            table.wait_until_exists()

            self.dynamodb = dynamodb
            self.table = table
            return

        def tearDown(self) -> None:
            dynamodb = self.dynamodb
            table = dynamodb.Table('test-table')
            table.delete()
            table.wait_until_not_exists()
            self.mock_dynamodb.stop()
            return

        def _store_fake_records(self, record_count=1, start_timestamp=1676246400_000000, timestamp_increment=-1_000_000):
            # Put some records in the mocked DynamoDB. Returns the records that a future query on them should return.
            ttl = datetime.now() + timedelta(days=365)
            timestamp_seq = [start_timestamp +
                            (i*timestamp_increment) for i in range(record_count)]

            for i in range(record_count):
                self.table.put_item(Item={
                    'object_id': 'test.com',
                    'object_type': 'domain',
                    'system': 'test_system',
                    'user': 'test@test.com',
                    'category': 'classification',
                    'timestamp': timestamp_seq[i],
                    'details': json.dumps({'comment': 'some details', 'domain': 'test.com'}),
                    'ttl': int(ttl.timestamp()),
                })

            # All done
            return

        @unittest.expectedFailure
        def test_query_with_brackets(self):
            # Query that should work but does not
            events = self._store_fake_records(25)

            # With brackets
            self.query_args["KeyConditionExpression"] = "(#audit_object_id = :object_id) AND (#audit_timestamp < :before_timestamp)"

            try:
                query_results = self.table.query(**self.query_args)
            except:
                self.fail("query call raised and exception")

            self.assertEqual(query_results['Count'], 20, 'Did not get the expected 20 results back')

        def test_query_without_brackets(self):
            # Same query but without brackets in the KeyConditionExpression - Works
            events = self._store_fake_records(25)

            # Without brackets
            self.query_args["KeyConditionExpression"] = "#audit_object_id = :object_id AND #audit_timestamp < :before_timestamp"

            try:
                query_results = self.table.query(**self.query_args)
            except:
                self.fail("query call raised and exception")

            self.assertEqual(query_results['Count'], 20, 'Did not get the expected 20 results back')


    if __name__ == '__main__':
        unittest.main()

Thanks for providing the test case @davidpottage - looks like our validation is too strict. Marking it as a bug.

Resolving patch:
diff --git a/moto/dynamodb/parsing/key_condition_expression.py b/moto/dynamodb/parsing/key_condition_expression.py
--- a/moto/dynamodb/parsing/key_condition_expression.py
+++ b/moto/dynamodb/parsing/key_condition_expression.py
@@ -160,8 +160,10 @@ def parse_expression(
         if crnt_char == "(":
             # hashkey = :id and begins_with(      sortkey,     :sk)
             #                              ^ --> ^
+            # (hash_key = :id) and (sortkey = :sk)
+            #                     ^
             if current_stage in [EXPRESSION_STAGES.INITIAL_STAGE]:
-                if current_phrase != "begins_with":
+                if current_phrase not in ["begins_with", ""]:
                     raise MockValidationException(
                         f"Invalid KeyConditionExpression: Invalid function name; function: {current_phrase}"
                     )


Test patch:
diff --git a/tests/test_dynamodb/models/test_key_condition_expression_parser.py b/tests/test_dynamodb/models/test_key_condition_expression_parser.py
--- a/tests/test_dynamodb/models/test_key_condition_expression_parser.py
+++ b/tests/test_dynamodb/models/test_key_condition_expression_parser.py
@@ -179,6 +179,24 @@ def test_reverse_keys(self, expr):
         )
         desired_hash_key.should.equal("pk")
 
+    @pytest.mark.parametrize(
+        "expr",
+        [
+            "(job_id = :id) and start_date = :sk",
+            "job_id = :id and (start_date = :sk)",
+            "(job_id = :id) and (start_date = :sk)",
+        ],
+    )
+    def test_brackets(self, expr):
+        eav = {":id": "pk", ":sk1": "19", ":sk2": "21"}
+        desired_hash_key, comparison, range_values = parse_expression(
+            expression_attribute_values=eav,
+            key_condition_expression=expr,
+            schema=self.schema,
+            expression_attribute_names=dict(),
+        )
+        desired_hash_key.should.equal("pk")
+
 
 class TestNamesAndValues:
     schema = [{"AttributeName": "job_id", "KeyType": "HASH"}]