Merge pull request #2162 from cm-iwata/fix_dynamo_query

fix #2161 DyanmoDB query method using GSI with range key

moto/dynamodb2/models.py

@@ -570,6 +570,7 @@ class Table(BaseModel):
               exclusive_start_key, scan_index_forward, projection_expression,
               index_name=None, filter_expression=None, **filter_kwargs):
         results = []
+
         if index_name:
             all_indexes = (self.global_indexes or []) + (self.indexes or [])
             indexes_by_name = dict((i['IndexName'], i) for i in all_indexes)
@@ -586,24 +587,28 @@ class Table(BaseModel):
                 raise ValueError('Missing Hash Key. KeySchema: %s' %
                                  index['KeySchema'])
 
-            possible_results = []
-            for item in self.all_items():
-                if not isinstance(item, Item):
-                    continue
-                item_hash_key = item.attrs.get(index_hash_key['AttributeName'])
-                if item_hash_key and item_hash_key == hash_key:
-                    possible_results.append(item)
-        else:
-            possible_results = [item for item in list(self.all_items()) if isinstance(
-                item, Item) and item.hash_key == hash_key]
-
-        if index_name:
             try:
                 index_range_key = [key for key in index[
                     'KeySchema'] if key['KeyType'] == 'RANGE'][0]
             except IndexError:
                 index_range_key = None
 
+            possible_results = []
+            for item in self.all_items():
+                if not isinstance(item, Item):
+                    continue
+                item_hash_key = item.attrs.get(index_hash_key['AttributeName'])
+                if index_range_key is None:
+                    if item_hash_key and item_hash_key == hash_key:
+                        possible_results.append(item)
+                else:
+                    item_range_key = item.attrs.get(index_range_key['AttributeName'])
+                    if item_hash_key and item_hash_key == hash_key and item_range_key:
+                        possible_results.append(item)
+        else:
+            possible_results = [item 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(

tests/test_dynamodb2/test_dynamodb.py

@@ -1612,3 +1612,70 @@ def test_condition_expressions():
                 ':match': {'S': 'match2'}
             }
         )
+
+
+@mock_dynamodb2
+def test_query_gsi_with_range_key():
+    dynamodb = boto3.client('dynamodb', region_name='us-east-1')
+    dynamodb.create_table(
+        TableName='test',
+        KeySchema=[{'AttributeName': 'id', 'KeyType': 'HASH'}],
+        AttributeDefinitions=[
+            {'AttributeName': 'id', 'AttributeType': 'S'},
+            {'AttributeName': 'gsi_hash_key', 'AttributeType': 'S'},
+            {'AttributeName': 'gsi_range_key', 'AttributeType': 'S'}
+        ],
+        ProvisionedThroughput={'ReadCapacityUnits': 1, 'WriteCapacityUnits': 1},
+        GlobalSecondaryIndexes=[
+            {
+                'IndexName': 'test_gsi',
+                'KeySchema': [
+                    {
+                        'AttributeName': 'gsi_hash_key',
+                        'KeyType': 'HASH'
+                    },
+                    {
+                        'AttributeName': 'gsi_range_key',
+                        'KeyType': 'RANGE'
+                    },
+                ],
+                'Projection': {
+                    'ProjectionType': 'ALL',
+                },
+                'ProvisionedThroughput': {
+                    'ReadCapacityUnits': 1,
+                    'WriteCapacityUnits': 1
+                }
+            },
+        ]
+    )
+
+    dynamodb.put_item(
+        TableName='test',
+        Item={
+            'id': {'S': 'test1'},
+            'gsi_hash_key': {'S': 'key1'},
+            'gsi_range_key': {'S': 'range1'},
+        }
+    )
+    dynamodb.put_item(
+        TableName='test',
+        Item={
+            'id': {'S': 'test2'},
+            'gsi_hash_key': {'S': 'key1'},
+        }
+    )
+
+    res = dynamodb.query(TableName='test', IndexName='test_gsi',
+                         KeyConditionExpression='gsi_hash_key = :gsi_hash_key AND gsi_range_key = :gsi_range_key',
+                         ExpressionAttributeValues={
+                            ':gsi_hash_key': {'S': 'key1'},
+                            ':gsi_range_key': {'S': 'range1'}
+                         })
+    res.should.have.key("Count").equal(1)
+    res.should.have.key("Items")
+    res['Items'][0].should.equal({
+        'id': {'S': 'test1'},
+        'gsi_hash_key': {'S': 'key1'},
+        'gsi_range_key': {'S': 'range1'},
+    })