ProjectionExpression works with table.scan()

moto/dynamodb2/models.py

@@ -724,7 +724,7 @@ class Table(BaseModel):
                 if idx_col_set.issubset(set(hash_set.attrs)):
                     yield hash_set
 
-    def scan(self, filters, limit, exclusive_start_key, filter_expression=None, index_name=None):
+    def scan(self, filters, limit, exclusive_start_key, filter_expression=None, index_name=None, projection_expression=None):
         results = []
         scanned_count = 0
         all_indexes = self.all_indexes()
@@ -763,6 +763,14 @@ class Table(BaseModel):
             if passes_all_conditions:
                 results.append(item)
 
+        if projection_expression:
+            expressions = [x.strip() for x in projection_expression.split(',')]
+            results = copy.deepcopy(results)
+            for result in results:
+                for attr in list(result.attrs):
+                    if attr not in expressions:
+                        result.attrs.pop(attr)
+
         results, last_evaluated_key = self._trim_results(results, limit,
                                                          exclusive_start_key, index_name)
         return results, scanned_count, last_evaluated_key
@@ -962,7 +970,7 @@ class DynamoDBBackend(BaseBackend):
         return table.query(hash_key, range_comparison, range_values, limit,
                            exclusive_start_key, scan_index_forward, projection_expression, index_name, filter_expression, **filter_kwargs)
 
-    def scan(self, table_name, filters, limit, exclusive_start_key, filter_expression, expr_names, expr_values, index_name):
+    def scan(self, table_name, filters, limit, exclusive_start_key, filter_expression, expr_names, expr_values, index_name, projection_expression):
         table = self.tables.get(table_name)
         if not table:
             return None, None, None
@@ -977,7 +985,7 @@ class DynamoDBBackend(BaseBackend):
         else:
             filter_expression = Op(None, None)  # Will always eval to true
 
-        return table.scan(scan_filters, limit, exclusive_start_key, filter_expression, index_name)
+        return table.scan(scan_filters, limit, exclusive_start_key, filter_expression, index_name, projection_expression)
 
     def update_item(self, table_name, key, update_expression, attribute_updates, expression_attribute_names,
                     expression_attribute_values, expected=None):

moto/dynamodb2/responses.py

@@ -558,7 +558,7 @@ class DynamoHandler(BaseResponse):
         filter_expression = self.body.get('FilterExpression')
         expression_attribute_values = self.body.get('ExpressionAttributeValues', {})
         expression_attribute_names = self.body.get('ExpressionAttributeNames', {})
-
+        projection_expression = self.body.get('ProjectionExpression')
         exclusive_start_key = self.body.get('ExclusiveStartKey')
         limit = self.body.get("Limit")
         index_name = self.body.get('IndexName')
@@ -570,7 +570,8 @@ class DynamoHandler(BaseResponse):
                                                                                   filter_expression,
                                                                                   expression_attribute_names,
                                                                                   expression_attribute_values,
-                                                                                  index_name)
+                                                                                  index_name,
+                                                                                  projection_expression)
         except InvalidIndexNameError as err:
             er = 'com.amazonaws.dynamodb.v20111205#ValidationException'
             return self.error(er, str(err))

tests/test_dynamodb2/test_dynamodb.py

@@ -452,6 +452,90 @@ def test_basic_projection_expressions():
     assert 'body' in results['Items'][1]
     assert 'forum_name' in results['Items'][1]
 
+@mock_dynamodb2
+def test_basic_projection_expressions_using_scan():
+    dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
+
+    # Create the DynamoDB table.
+    table = dynamodb.create_table(
+        TableName='users',
+        KeySchema=[
+            {
+                'AttributeName': 'forum_name',
+                'KeyType': 'HASH'
+            },
+            {
+                'AttributeName': 'subject',
+                'KeyType': 'RANGE'
+            },
+        ],
+        AttributeDefinitions=[
+            {
+                'AttributeName': 'forum_name',
+                'AttributeType': 'S'
+            },
+            {
+                'AttributeName': 'subject',
+                'AttributeType': 'S'
+            },
+        ],
+        ProvisionedThroughput={
+            'ReadCapacityUnits': 5,
+            'WriteCapacityUnits': 5
+        }
+    )
+    table = dynamodb.Table('users')
+
+    table.put_item(Item={
+        'forum_name': 'the-key',
+        'subject': '123',
+        'body': 'some test message'
+    })
+
+    table.put_item(Item={
+        'forum_name': 'not-the-key',
+        'subject': '123',
+        'body': 'some other test message'
+    })
+    # Test a query returning all items
+    results = table.scan(
+        FilterExpression=Key('forum_name').eq(
+            'the-key'),
+        ProjectionExpression='body, subject'
+    )
+
+    assert 'body' in results['Items'][0]
+    assert results['Items'][0]['body'] == 'some test message'
+    assert 'subject' in results['Items'][0]
+
+    table.put_item(Item={
+        'forum_name': 'the-key',
+        'subject': '1234',
+        'body': 'yet another test message'
+    })
+
+    results = table.scan(
+        FilterExpression=Key('forum_name').eq(
+            'the-key'),
+        ProjectionExpression='body'
+    )
+
+    assert 'body' in results['Items'][0]
+    assert 'subject' not in results['Items'][0]
+    assert results['Items'][0]['body'] == 'some test message'
+    assert 'body' in results['Items'][1]
+    assert 'subject' not in results['Items'][1]
+    assert results['Items'][1]['body'] == 'yet another test message'
+
+    # The projection expression should not remove data from storage
+    results = table.query(
+        KeyConditionExpression=Key('forum_name').eq(
+            'the-key'),
+    )
+    assert 'subject' in results['Items'][0]
+    assert 'body' in results['Items'][1]
+    assert 'forum_name' in results['Items'][1]
+
 
 @mock_dynamodb2
 def test_basic_projection_expressions_with_attr_expression_names():