Merge pull request #2162 from cm-iwata/fix_dynamo_query
fix #2161 DyanmoDB query method using GSI with range key
This commit is contained in:
		
						commit
						734a39b3e4
					
				@ -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(
 | 
			
		||||
 | 
			
		||||
@ -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'},
 | 
			
		||||
    })
 | 
			
		||||
 | 
			
		||||
		Loading…
	
	
			
			x
			
			
		
	
		Reference in New Issue
	
	Block a user