From 9e3b23758af52b315beaf8bf6277e4c07c3e5c77 Mon Sep 17 00:00:00 2001
From: Garrett <garrettheel@users.noreply.github.com>
Date: Wed, 25 Nov 2020 15:28:05 -0500
Subject: [PATCH] [dynamodb2] Support include projection on indexes (#3498)

* [dynamodb2] Support include projection on indexes

* linter
---
 moto/dynamodb2/models/__init__.py     | 16 ++++++--
 tests/test_dynamodb2/test_dynamodb.py | 55 +++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 4 deletions(-)

diff --git a/moto/dynamodb2/models/__init__.py b/moto/dynamodb2/models/__init__.py
index 18b0b918..7218fe0c 100644
--- a/moto/dynamodb2/models/__init__.py
+++ b/moto/dynamodb2/models/__init__.py
@@ -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
 
 
diff --git a/tests/test_dynamodb2/test_dynamodb.py b/tests/test_dynamodb2/test_dynamodb.py
index 3571239e..0e0fcb08 100644
--- a/tests/test_dynamodb2/test_dynamodb.py
+++ b/tests/test_dynamodb2/test_dynamodb.py
@@ -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 = {