From 4e0d5883073b3c12c27e888b978dba530708035a Mon Sep 17 00:00:00 2001
From: Bert Blommers <info@bertblommers.nl>
Date: Fri, 3 Jul 2020 14:20:04 +0100
Subject: [PATCH] DynamoDB - Allow ProjectionType to be set for LSIs

---
 moto/dynamodb2/models/__init__.py     | 36 ++++++++++++-----------
 tests/test_dynamodb2/test_dynamodb.py | 41 +++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 17 deletions(-)

diff --git a/moto/dynamodb2/models/__init__.py b/moto/dynamodb2/models/__init__.py
index 7e288bb9..eafa2743 100644
--- a/moto/dynamodb2/models/__init__.py
+++ b/moto/dynamodb2/models/__init__.py
@@ -272,7 +272,24 @@ class StreamShard(BaseModel):
         return [i.to_json() for i in self.items[start:end]]
 
 
-class LocalSecondaryIndex(BaseModel):
+class SecondaryIndex(BaseModel):
+    def project(self, item):
+        """
+        Enforces the ProjectionType of this Index (LSI/GSI)
+        Removes any non-wanted attributes from the item
+        :param item:
+        :return:
+        """
+        if self.projection:
+            if self.projection.get("ProjectionType", None) == "KEYS_ONLY":
+                allowed_attributes = ",".join(
+                    [key["AttributeName"] for key in self.schema]
+                )
+                item.filter(allowed_attributes)
+        return item
+
+
+class LocalSecondaryIndex(SecondaryIndex):
     def __init__(self, index_name, schema, projection):
         self.name = index_name
         self.schema = schema
@@ -294,7 +311,7 @@ class LocalSecondaryIndex(BaseModel):
         )
 
 
-class GlobalSecondaryIndex(BaseModel):
+class GlobalSecondaryIndex(SecondaryIndex):
     def __init__(
         self, index_name, schema, projection, status="ACTIVE", throughput=None
     ):
@@ -331,21 +348,6 @@ class GlobalSecondaryIndex(BaseModel):
         self.projection = u.get("Projection", self.projection)
         self.throughput = u.get("ProvisionedThroughput", self.throughput)
 
-    def project(self, item):
-        """
-        Enforces the ProjectionType of this GSI
-        Removes any non-wanted attributes from the item
-        :param item:
-        :return:
-        """
-        if self.projection:
-            if self.projection.get("ProjectionType", None) == "KEYS_ONLY":
-                allowed_attributes = ",".join(
-                    [key["AttributeName"] for key in self.schema]
-                )
-                item.filter(allowed_attributes)
-        return item
-
 
 class Table(BaseModel):
     def __init__(
diff --git a/tests/test_dynamodb2/test_dynamodb.py b/tests/test_dynamodb2/test_dynamodb.py
index cf1548e0..2dfb8fd2 100644
--- a/tests/test_dynamodb2/test_dynamodb.py
+++ b/tests/test_dynamodb2/test_dynamodb.py
@@ -5360,3 +5360,44 @@ def test_gsi_projection_type_keys_only():
     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"})
+
+
+@mock_dynamodb2
+def test_lsi_projection_type_keys_only():
+    table_schema = {
+        "KeySchema": [{"AttributeName": "partitionKey", "KeyType": "HASH"}],
+        "LocalSecondaryIndexes": [
+            {
+                "IndexName": "LSI",
+                "KeySchema": [
+                    {"AttributeName": "partitionKey", "KeyType": "HASH"},
+                    {"AttributeName": "lsiK1SortKey", "KeyType": "RANGE"},
+                ],
+                "Projection": {"ProjectionType": "KEYS_ONLY",},
+            }
+        ],
+        "AttributeDefinitions": [
+            {"AttributeName": "partitionKey", "AttributeType": "S"},
+            {"AttributeName": "lsiK1SortKey", "AttributeType": "S"},
+        ],
+    }
+
+    item = {
+        "partitionKey": "pk-1",
+        "lsiK1SortKey": "lsi-sk",
+        "someAttribute": "lore ipsum",
+    }
+
+    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("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"})