[dynamodb2] Support include projection on indexes (#3498)

* [dynamodb2] Support include projection on indexes

* linter
This commit is contained in:
Garrett 2020-11-25 15:28:05 -05:00 committed by GitHub
parent d58d3e2c2e
commit 9e3b23758a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 67 additions and 4 deletions

View File

@ -292,11 +292,19 @@ class SecondaryIndex(BaseModel):
:return:
"""
if self.projection:
if self.projection.get("ProjectionType", None) == "KEYS_ONLY":
allowed_attributes = ",".join(
self.table_key_attrs + [key["AttributeName"] for key in self.schema]
projection_type = self.projection.get("ProjectionType", None)
key_attributes = self.table_key_attrs + [
key["AttributeName"] for key in self.schema
]
if projection_type == "KEYS_ONLY":
item.filter(",".join(key_attributes))
elif projection_type == "INCLUDE":
allowed_attributes = key_attributes + self.projection.get(
"NonKeyAttributes", []
)
item.filter(allowed_attributes)
item.filter(",".join(allowed_attributes))
# ALL is handled implicitly by not filtering
return item

View File

@ -5523,6 +5523,61 @@ def test_gsi_projection_type_keys_only():
)
@mock_dynamodb2
def test_gsi_projection_type_include():
table_schema = {
"KeySchema": [{"AttributeName": "partitionKey", "KeyType": "HASH"}],
"GlobalSecondaryIndexes": [
{
"IndexName": "GSI-INC",
"KeySchema": [
{"AttributeName": "gsiK1PartitionKey", "KeyType": "HASH"},
{"AttributeName": "gsiK1SortKey", "KeyType": "RANGE"},
],
"Projection": {
"ProjectionType": "INCLUDE",
"NonKeyAttributes": ["projectedAttribute"],
},
}
],
"AttributeDefinitions": [
{"AttributeName": "partitionKey", "AttributeType": "S"},
{"AttributeName": "gsiK1PartitionKey", "AttributeType": "S"},
{"AttributeName": "gsiK1SortKey", "AttributeType": "S"},
],
}
item = {
"partitionKey": "pk-1",
"gsiK1PartitionKey": "gsi-pk",
"gsiK1SortKey": "gsi-sk",
"projectedAttribute": "lore ipsum",
"nonProjectedAttribute": "dolor sit amet",
}
dynamodb = boto3.resource("dynamodb", region_name="us-east-1")
dynamodb.create_table(
TableName="test-table", BillingMode="PAY_PER_REQUEST", **table_schema
)
table = dynamodb.Table("test-table")
table.put_item(Item=item)
items = table.query(
KeyConditionExpression=Key("gsiK1PartitionKey").eq("gsi-pk"),
IndexName="GSI-INC",
)["Items"]
items.should.have.length_of(1)
# Item should only include keys and additionally projected attributes only
items[0].should.equal(
{
"gsiK1PartitionKey": "gsi-pk",
"gsiK1SortKey": "gsi-sk",
"partitionKey": "pk-1",
"projectedAttribute": "lore ipsum",
}
)
@mock_dynamodb2
def test_lsi_projection_type_keys_only():
table_schema = {