add basic projection expressions

tests/test_dynamodb2/test_dynamodb.py

@@ -9,6 +9,7 @@ from moto import mock_dynamodb2, mock_dynamodb2_deprecated
 from moto.dynamodb2 import dynamodb_backend2
 from boto.exception import JSONResponseError
 from botocore.exceptions import ClientError
+from boto3.dynamodb.conditions import Key
 from tests.helpers import requires_boto_gte
 import tests.backport_assert_raises
 from nose.tools import assert_raises
@@ -228,3 +229,126 @@ def test_scan_returns_consumed_capacity():
     assert 'ConsumedCapacity' in response
     assert 'CapacityUnits' in response['ConsumedCapacity']
     assert response['ConsumedCapacity']['TableName'] == name
+
+@requires_boto_gte("2.9")
+@mock_dynamodb2
+def test_query_returns_consumed_capacity():
+    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'
+    })
+
+    results = table.query(
+        KeyConditionExpression=Key('forum_name').eq(
+            'the-key')
+    )
+
+    assert 'ConsumedCapacity' in results
+    assert 'CapacityUnits' in results['ConsumedCapacity']
+    assert results['ConsumedCapacity']['CapacityUnits'] == 1
+
+@mock_dynamodb2
+def test_projection_expressions():
+    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.query(
+        KeyConditionExpression=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.query(
+        KeyConditionExpression=Key('forum_name').eq(
+            'the-key'),
+        ProjectionExpression='body'
+    )
+
+    assert 'body' in results['Items'][0]
+    assert results['Items'][0]['body'] == 'some test message'
+    assert 'body' in results['Items'][1]
+    assert results['Items'][1]['body'] == 'yet another test message'