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:
Mike Grima 2019-04-19 10:04:04 -07:00 committed by GitHub
commit 734a39b3e4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 84 additions and 12 deletions

View File

@ -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(

View File

@ -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'},
})