From 9072153474c9e4afbe271ab398a4c9fa7ec9aacd Mon Sep 17 00:00:00 2001 From: Mike Fogel Date: Tue, 14 Jul 2020 09:42:13 -0300 Subject: [PATCH] Fix dynamodb2 KEYS_ONLY Indexes (#3125) KEYS_ONLY indexes include table keys. --- moto/dynamodb2/models/__init__.py | 30 ++++++++++++++++++++------- tests/test_dynamodb2/test_dynamodb.py | 23 +++++++++++++++----- 2 files changed, 40 insertions(+), 13 deletions(-) diff --git a/moto/dynamodb2/models/__init__.py b/moto/dynamodb2/models/__init__.py index eafa2743a..233c4001f 100644 --- a/moto/dynamodb2/models/__init__.py +++ b/moto/dynamodb2/models/__init__.py @@ -283,17 +283,18 @@ class SecondaryIndex(BaseModel): if self.projection: if self.projection.get("ProjectionType", None) == "KEYS_ONLY": allowed_attributes = ",".join( - [key["AttributeName"] for key in self.schema] + self.table_key_attrs + [key["AttributeName"] for key in self.schema] ) item.filter(allowed_attributes) return item class LocalSecondaryIndex(SecondaryIndex): - def __init__(self, index_name, schema, projection): + def __init__(self, index_name, schema, projection, table_key_attrs): self.name = index_name self.schema = schema self.projection = projection + self.table_key_attrs = table_key_attrs def describe(self): return { @@ -303,21 +304,29 @@ class LocalSecondaryIndex(SecondaryIndex): } @staticmethod - def create(dct): + def create(dct, table_key_attrs): return LocalSecondaryIndex( index_name=dct["IndexName"], schema=dct["KeySchema"], projection=dct["Projection"], + table_key_attrs=table_key_attrs, ) class GlobalSecondaryIndex(SecondaryIndex): def __init__( - self, index_name, schema, projection, status="ACTIVE", throughput=None + self, + index_name, + schema, + projection, + table_key_attrs, + status="ACTIVE", + throughput=None, ): self.name = index_name self.schema = schema self.projection = projection + self.table_key_attrs = table_key_attrs self.status = status self.throughput = throughput or { "ReadCapacityUnits": 0, @@ -334,11 +343,12 @@ class GlobalSecondaryIndex(SecondaryIndex): } @staticmethod - def create(dct): + def create(dct, table_key_attrs): return GlobalSecondaryIndex( index_name=dct["IndexName"], schema=dct["KeySchema"], projection=dct["Projection"], + table_key_attrs=table_key_attrs, throughput=dct.get("ProvisionedThroughput", None), ) @@ -374,16 +384,20 @@ class Table(BaseModel): else: self.range_key_attr = elem["AttributeName"] self.range_key_type = elem["KeyType"] + self.table_key_attrs = [ + key for key in (self.hash_key_attr, self.range_key_attr) if key + ] if throughput is None: self.throughput = {"WriteCapacityUnits": 10, "ReadCapacityUnits": 10} else: self.throughput = throughput self.throughput["NumberOfDecreasesToday"] = 0 self.indexes = [ - LocalSecondaryIndex.create(i) for i in (indexes if indexes else []) + LocalSecondaryIndex.create(i, self.table_key_attrs) + for i in (indexes if indexes else []) ] self.global_indexes = [ - GlobalSecondaryIndex.create(i) + GlobalSecondaryIndex.create(i, self.table_key_attrs) for i in (global_indexes if global_indexes else []) ] self.created_at = datetime.datetime.utcnow() @@ -1015,7 +1029,7 @@ class DynamoDBBackend(BaseBackend): ) gsis_by_name[gsi_to_create["IndexName"]] = GlobalSecondaryIndex.create( - gsi_to_create + gsi_to_create, table.table_key_attrs, ) # in python 3.6, dict.values() returns a dict_values object, but we expect it to be a list in other diff --git a/tests/test_dynamodb2/test_dynamodb.py b/tests/test_dynamodb2/test_dynamodb.py index 2dfb8fd2d..d56fd3f11 100644 --- a/tests/test_dynamodb2/test_dynamodb.py +++ b/tests/test_dynamodb2/test_dynamodb.py @@ -5358,14 +5358,23 @@ def test_gsi_projection_type_keys_only(): IndexName="GSI-K1", )["Items"] items.should.have.length_of(1) - # Item should only include GSI Keys, as per the ProjectionType - items[0].should.equal({"gsiK1PartitionKey": "gsi-pk", "gsiK1SortKey": "gsi-sk"}) + # Item should only include GSI Keys and Table Keys, as per the ProjectionType + items[0].should.equal( + { + "gsiK1PartitionKey": "gsi-pk", + "gsiK1SortKey": "gsi-sk", + "partitionKey": "pk-1", + } + ) @mock_dynamodb2 def test_lsi_projection_type_keys_only(): table_schema = { - "KeySchema": [{"AttributeName": "partitionKey", "KeyType": "HASH"}], + "KeySchema": [ + {"AttributeName": "partitionKey", "KeyType": "HASH"}, + {"AttributeName": "sortKey", "KeyType": "RANGE"}, + ], "LocalSecondaryIndexes": [ { "IndexName": "LSI", @@ -5378,12 +5387,14 @@ def test_lsi_projection_type_keys_only(): ], "AttributeDefinitions": [ {"AttributeName": "partitionKey", "AttributeType": "S"}, + {"AttributeName": "sortKey", "AttributeType": "S"}, {"AttributeName": "lsiK1SortKey", "AttributeType": "S"}, ], } item = { "partitionKey": "pk-1", + "sortKey": "sk-1", "lsiK1SortKey": "lsi-sk", "someAttribute": "lore ipsum", } @@ -5399,5 +5410,7 @@ def test_lsi_projection_type_keys_only(): KeyConditionExpression=Key("partitionKey").eq("pk-1"), IndexName="LSI", )["Items"] items.should.have.length_of(1) - # Item should only include GSI Keys, as per the ProjectionType - items[0].should.equal({"partitionKey": "pk-1", "lsiK1SortKey": "lsi-sk"}) + # Item should only include GSI Keys and Table Keys, as per the ProjectionType + items[0].should.equal( + {"partitionKey": "pk-1", "sortKey": "sk-1", "lsiK1SortKey": "lsi-sk"} + )