feat(datastore): Add support for projection and distinctOn (#49)

datastore/gcloud/rest/datastore/__init__.py

@@ -20,6 +20,7 @@
 from gcloud.rest.datastore.key import Key
 from gcloud.rest.datastore.key import PathElement
 from gcloud.rest.datastore.lat_lng import LatLng
+from gcloud.rest.datastore.projection import Projection
 from gcloud.rest.datastore.property_order import PropertyOrder
 from gcloud.rest.datastore.query import GQLQuery
 from gcloud.rest.datastore.query import Query
@@ -30,6 +31,6 @@
 __all__ = ['__version__', 'CompositeFilter', 'CompositeFilterOperator',
            'Consistency', 'Datastore', 'DatastoreOperation', 'Direction',
            'Entity', 'EntityResult', 'Filter', 'GQLQuery', 'Key', 'LatLng',
-           'Mode', 'MoreResultsType', 'Operation', 'PathElement',
+           'Mode', 'MoreResultsType', 'Operation', 'PathElement', 'Projection',
            'PropertyFilter', 'PropertyFilterOperator', 'PropertyOrder',
            'Query', 'QueryResultBatch', 'ResultType', 'SCOPES', 'Value']

datastore/gcloud/rest/datastore/entity.py

@@ -44,7 +44,7 @@ def to_repr(self):
 class EntityResult(object):
     entity_kind = Entity
 
-    def __init__(self, entity, version, cursor=''):
+    def __init__(self, entity, version='', cursor=''):
         # type: (Entity, str, str) -> None
         self.entity = entity
         self.version = version
@@ -66,15 +66,17 @@ def __repr__(self):
     @classmethod
     def from_repr(cls, data):
         # type: (Dict[str, Any]) -> EntityResult
-        return cls(cls.entity_kind.from_repr(data['entity']), data['version'],
+        return cls(cls.entity_kind.from_repr(data['entity']),
+                   data.get('version', ''),
                    data.get('cursor', ''))
 
     def to_repr(self):
         # type: () -> Dict[str, Any]
         data = {
             'entity': self.entity.to_repr(),
-            'version': self.version,
         }
+        if self.version:
+            data['version'] = self.version
         if self.cursor:
             data['cursor'] = self.cursor
 

datastore/gcloud/rest/datastore/projection.py

@@ -0,0 +1,32 @@
+from typing import Any  # pylint: disable=unused-import
+from typing import Dict  # pylint: disable=unused-import
+
+# https://cloud.google.com/datastore/docs/reference/data/rest/v1/projects/runQuery#Projection
+
+
+class Projection(object):
+    def __init__(self, prop):
+        # type: (str) -> None
+        self.prop = prop
+
+    def __eq__(self, other):
+        # type: (Any) -> bool
+        if not isinstance(other, Projection):
+            return False
+
+        return bool(self.prop == other.prop)
+
+    def __repr__(self):
+        # type: () -> str
+        return str(self.to_repr())
+
+    @classmethod
+    def from_repr(cls, data):
+        # type: (Dict[str, Any]) -> Projection
+        return cls(prop=data['property']['name'])
+
+    def to_repr(self):
+        # type: () -> Dict[str, Any]
+        return {
+            'property': {'name': self.prop},
+        }

datastore/gcloud/rest/datastore/query.py

@@ -7,6 +7,7 @@
 from gcloud.rest.datastore.constants import ResultType
 from gcloud.rest.datastore.entity import EntityResult
 from gcloud.rest.datastore.filter import Filter
+from gcloud.rest.datastore.projection import Projection
 from gcloud.rest.datastore.property_order import PropertyOrder
 from gcloud.rest.datastore.value import Value
 
@@ -31,17 +32,19 @@ def to_repr(self):
 
 # https://cloud.google.com/datastore/docs/reference/data/rest/v1/projects/runQuery#Query
 class Query(BaseQuery):
+    # pylint: disable=too-many-instance-attributes
     json_key = 'query'
 
-    # TODO: support `projection` and `distinctOn`
     def __init__(self,
                  kind='',            # type: str
                  query_filter=None,  # type: Optional[Filter]
                  order=None,         # type: Optional[List[PropertyOrder]]
                  start_cursor='',    # type: str
                  end_cursor='',      # type: str
                  offset=0,           # type: int
-                 limit=0             # type: int
+                 limit=0,            # type: int
+                 projection=None,    # type: Optional[List[Projection]]
+                 distinct_on=None    # type: Optional[List[str]]
                  ):
         # type: (...) -> None
         self.kind = kind
@@ -51,6 +54,8 @@ def __init__(self,
         self.end_cursor = end_cursor
         self.offset = offset
         self.limit = limit
+        self.projection = projection or []
+        self.distinct_on = distinct_on or []
 
     def __eq__(self, other):
         # type: (Any) -> bool
@@ -70,13 +75,17 @@ def from_repr(cls, data):
         end_cursor = data.get('endCursor') or ''
         offset = int(data.get('offset') or 0)
         limit = int(data.get('limit') or 0)
+        projection = [Projection.from_repr(p)
+                      for p in data.get('projection', [])]
+        distinct_on = [d['name'] for d in data.get('distinct_on', [])]
 
         filter_ = data.get('filter')
         query_filter = Filter.from_repr(filter_) if filter_ else None
 
         return cls(kind=kind, query_filter=query_filter, order=orders,
                    start_cursor=start_cursor, end_cursor=end_cursor,
-                   offset=offset, limit=limit)
+                   offset=offset, limit=limit,
+                   projection=projection, distinct_on=distinct_on)
 
     def to_repr(self):
         # type: () -> Dict[str, Any]
@@ -93,6 +102,10 @@ def to_repr(self):
             data['offset'] = self.offset
         if self.limit:
             data['limit'] = self.limit
+        if self.projection:
+            data['projection'] = [p.to_repr() for p in self.projection]
+        if self.distinct_on:
+            data['distinctOn'] = [{'name': d} for d in self.distinct_on]
         return data
 
 

datastore/tests/integration/query_test.py

@@ -8,6 +8,7 @@
 from gcloud.rest.datastore import Key
 from gcloud.rest.datastore import Operation
 from gcloud.rest.datastore import PathElement
+from gcloud.rest.datastore import Projection
 from gcloud.rest.datastore import PropertyFilter
 from gcloud.rest.datastore import PropertyFilterOperator
 from gcloud.rest.datastore import Query
@@ -77,3 +78,65 @@ def test_gql_query(creds, kind, project):
 
         after = ds.runQuery(query, session=s)
         assert len(after.entity_results) == num_results + 3
+
+
+def test_query_with_value_projection(creds, kind, project):
+    # type: (str, str, str) -> None
+    with requests.Session() as s:
+        ds = Datastore(project=project, service_file=creds, session=s)
+        # setup test data
+        ds.insert(Key(project, [PathElement(kind)]), {'value': 30}, s)
+        projection = [Projection.from_repr({'property': {'name': 'value'}})]
+
+        query = Query(kind=kind, limit=1,
+                      projection=projection)
+        result = ds.runQuery(query, session=s)
+        assert result.entity_result_type.value == 'PROJECTION'
+        # clean up test data
+        ds.delete(result.entity_results[0].entity.key, s)
+
+
+def test_query_with_key_projection(creds, kind, project):
+    # type: (str, str, str) -> None
+    with requests.Session() as s:
+        ds = Datastore(project=project, service_file=creds, session=s)
+        # setup test data
+        ds.insert(Key(project, [PathElement(kind)]), {'value': 30}, s)
+        property_filter = PropertyFilter(
+            prop='value', operator=PropertyFilterOperator.EQUAL,
+            value=Value(30))
+        projection = [Projection.from_repr({'property': {'name': '__key__'}})]
+
+        query = Query(kind=kind, query_filter=Filter(property_filter), limit=1,
+                      projection=projection)
+        result = ds.runQuery(query, session=s)
+        assert result.entity_results[0].entity.properties == {}
+        assert result.entity_result_type.value == 'KEY_ONLY'
+        # clean up test data
+        ds.delete(result.entity_results[0].entity.key, s)
+
+
+def test_query_with_distinct_on(creds, kind, project):
+    # type: (str, str, str) -> None
+    keys1 = [Key(project, [PathElement(kind)])
+             for i in range(3)]  # pylint: disable=unused-variable
+    keys2 = [Key(project, [PathElement(kind)])
+             for i in range(3)]  # pylint: disable=unused-variable
+    with requests.Session() as s:
+        ds = Datastore(project=project, service_file=creds, session=s)
+
+        # setup test data
+        allocatedKeys1 = ds.allocateIds(keys1, session=s)
+        allocatedKeys2 = ds.allocateIds(keys2, session=s)
+        for key1 in allocatedKeys1:
+            ds.insert(key1, {'dist_value': 11}, s)
+        for key2 in allocatedKeys2:
+            ds.insert(key2, {'dist_value': 22}, s)
+        query = Query(kind=kind, limit=10, distinct_on=['dist_value'])
+        result = ds.runQuery(query, session=s)
+        assert len(result.entity_results) == 2
+        # clean up test data
+        for key1 in allocatedKeys1:
+            ds.delete(key1, s)
+        for key2 in allocatedKeys2:
+            ds.delete(key2, s)