Merge pull request googleapis#136 from jgeewax/koss-collection-4

Add Query.order_by method.

firestore/google/cloud/firestore/collection.py

@@ -20,6 +20,8 @@
 
 from google.protobuf.wrappers_pb2 import Int32Value
 
+_DEFAULT_DIRECTION = 'asc'
+
 
 class Query(object):
  """Firestore Query.
@@ -33,20 +35,51 @@ class Query(object):
  :type limit: int
  :param limit: (optional) Maximum number of results to return.
 
+ :type field_orders: sequence of
+ :class:`~google.cloud.firestore.collection.FieldOrder`
+ :param field_orders: (Optional) Sequence of fields to control order of
+ results.
+
  :raises ValueError: Path must be a valid Collection path (not a
  Document path).
  """
 
  def __init__(self,
  client,
  ref_path,
- limit=None):
+ limit=None,
+ field_orders=()):
  if not ref_path.is_collection:
  raise ValueError('Invalid collection path: %s' % (ref_path,))
 
  self._client = client
  self._path = ref_path
  self._limit = limit
+ self._field_orders = field_orders
+
+ def order_by(self, field, direction=_DEFAULT_DIRECTION):
+ """Modify the query to add an order clause on a specific field.
+
+ Successive order_by calls will further refine the ordering of
+ return results.
+
+ :type field: str
+ :param field: The name of a document field (property) on which to
+ order the query results.
+
+ :type direction: str
+ :param direction: (optional) One of 'asc' (default) or 'desc' to
+ set the ordering direction to ascending or
+ descending, respectively.
+
+ :rtype: :class:`~google.cloud.firestore.collection.Query`
+ :return: An ordered ``Query``.
+ """
+ return Query(
+ self._client,
+ self._path,
+ limit=self._limit,
+ field_orders=self._field_orders + (FieldOrder(field, direction),))
 
  def limit(self, count):
  """Limit a query to return a fixed number of results.
@@ -83,7 +116,8 @@ def get(self):
  query = query_pb2.Query(
  kind=[query_pb2.KindExpression(name=self._path.kind)],
  filter=complete_filter,
- limit=Int32Value(value=self._limit))
+ limit=Int32Value(value=self._limit),
+ order=[order.to_proto() for order in self._field_orders])
 
  response = self._client._api.run_query(
  self._client.project,
@@ -149,3 +183,55 @@ def new_document(self):
  :returns: A ``DocumentRef`` to the location of a new document.
  """
  return self.document(self._client.auto_id())
+
+
+class FieldOrder(object):
+ """Query order-by field.
+
+ :type field_name: str
+ :param field_name : The name of a document field (property) on which to
+ order query results.
+
+ :type direction: str
+ :param direction: (optional) One of 'asc' (default) or 'desc' to
+ set the ordering direction to ascending or
+ descending, respectively.
+
+ :raises ValueError: Direction must be valid or None.
+ """
+
+ def __init__(self, field_name, direction=_DEFAULT_DIRECTION):
+ self._field_name = field_name
+ self._direction_enum = _direction_to_enum(direction)
+
+ def to_proto(self):
+ """Convert FieldOrder to the protobuf PropertyReference.
+
+ Intended to be used as a submessage in a Query protobuf.
+
+ :rtype: :class:`~google.firestore.v1alpha1.query_pb2.PropertyOrder`
+ :returns: A protobuf PropertyOrder.
+ """
+ return query_pb2.PropertyOrder(
+ property=query_pb2.PropertyReference(name=self._field_name),
+ direction=self._direction_enum)
+
+
+def _direction_to_enum(direction):
+ """Encode direction string into direction enum.
+
+ :type direction: str
+ :param direction: One of 'asc' or 'desc' to set the ordering direction to
+ ascending or descending, respectively.
+
+ :rtype: :class:`~enums.Direction`
+ :returns: A Query direction enum value.
+ """
+ if direction == 'asc':
+ return enums.Direction.ASCENDING
+ elif direction == 'desc':
+ return enums.Direction.DESCENDING
+ else:
+ raise ValueError('Invalid order_by direction (%s) - '
+ 'expect one of \'asc\' or \'desc\'.' %
+ (direction,))

firestore/unit_tests/test_collection.py

@@ -15,6 +15,21 @@
 import unittest
 
 
+def _make_credentials():
+ import google.auth.credentials
+ import mock
+
+ class _CredentialsWithScopes(
+ google.auth.credentials.Credentials,
+ google.auth.credentials.Scoped):
+ pass
+
+ credentials = mock.Mock(spec=_CredentialsWithScopes)
+ # Return self when being scoped.
+ credentials.with_scopes.return_value = credentials
+ return credentials
+
+
 class TestCollectionRef(unittest.TestCase):
 
  @staticmethod
@@ -78,6 +93,8 @@ def test_new_document(self):
 
 class TestQuery(unittest.TestCase):
 
+ PROJECT = 'project'
+
  @staticmethod
  def _get_target_class():
  from google.cloud.firestore.collection import Query
@@ -87,11 +104,16 @@ def _get_target_class():
  def _make_one(self, *args, **kwargs):
  import mock
  from google.cloud.firestore.client import Client
- from google.cloud.gapic.firestore.v1alpha1.datastore_api import (
- DatastoreApi)
 
- client = Client()
- client._api = mock.MagicMock(spec=DatastoreApi)
+ credentials = _make_credentials()
+ client = Client(project=self.PROJECT, credentials=credentials)
+
+ # Mock the API to return an empty list of results.
+ client._api = mock.Mock(spec=['run_query'])
+ batch = mock.Mock(entity_results=[], spec=['entity_results'])
+ client._api.run_query.return_value = mock.Mock(
+ batch=batch, spec=['batch'])
+
  return self._get_target_class()(client, *args, **kwargs)
 
  def test_constructor(self):
@@ -114,10 +136,142 @@ def test_limit(self):
 
  self.assertEqual(query._limit, 123)
 
+ @staticmethod
+ def _make_query_pb(collection, **kwargs):
+ from google.cloud.gapic.firestore.v1alpha1 import enums
+ from google.firestore.v1alpha1.entity_pb2 import Value
+ from google.firestore.v1alpha1 import query_pb2
+ from google.protobuf.struct_pb2 import NULL_VALUE
+ from google.protobuf.wrappers_pb2 import Int32Value
+
+ return query_pb2.Query(
+ kind=[query_pb2.KindExpression(name=collection)],
+ filter=query_pb2.Filter(
+ property_filter=query_pb2.PropertyFilter(
+ property=query_pb2.PropertyReference(name='__key__'),
+ op=enums.Operator.HAS_PARENT,
+ value=Value(null_value=NULL_VALUE),
+ ),
+ ),
+ limit=Int32Value(value=None),
+ **kwargs
+ )
+
  def test_get(self):
  from google.cloud.firestore._path import Path
 
- query = self._make_one(Path('my-collection'))
+ collection = 'my-collection'
+ query = self._make_one(Path(collection))
  docs = query.get()
- self.assertEqual(query._client._api.run_query.call_count, 1)
- self.assertIsInstance(docs, list)
+ self.assertEqual(docs, [])
+
+ mock_query = query._client._api.run_query
+ expected = self._make_query_pb(collection)
+ mock_query.assert_called_once_with(
+ self.PROJECT,
+ None,
+ partition_id=None,
+ read_options=None,
+ query=expected,
+ gql_query=None,
+ property_mask=None,
+ )
+
+ def test_order_by(self):
+ from google.cloud.firestore._path import Path
+ from google.firestore.v1alpha1 import query_pb2
+ from google.cloud.gapic.firestore.v1alpha1 import enums
+
+ collection = 'my-collection'
+ query = self._make_one(Path(collection))
+ prop_name = 'my-property'
+ new_query = query.order_by(prop_name)
+ docs = new_query.get()
+ self.assertEqual(docs, [])
+
+ mock_query = query._client._api.run_query
+ order_pb = query_pb2.PropertyOrder(
+ property=query_pb2.PropertyReference(name=prop_name),
+ direction=enums.Direction.ASCENDING,
+ )
+ expected = self._make_query_pb(collection, order=[order_pb])
+ mock_query.assert_called_once_with(
+ self.PROJECT,
+ None,
+ partition_id=None,
+ read_options=None,
+ query=expected,
+ gql_query=None,
+ property_mask=None,
+ )
+
+ def test_order_by_chain(self):
+ from google.cloud.firestore._path import Path
+ from google.firestore.v1alpha1 import query_pb2
+ from google.cloud.gapic.firestore.v1alpha1 import enums
+
+ collection = 'my-collection'
+ query1 = self._make_one(Path(collection))
+ prop_name1 = 'my-property'
+ query2 = query1.order_by(prop_name1)
+ prop_name2 = 'another-prop'
+ query3 = query2.order_by(prop_name2, 'desc')
+ docs = query3.get()
+ self.assertEqual(docs, [])
+
+ mock_query = query1._client._api.run_query
+ order_pb1 = query_pb2.PropertyOrder(
+ property=query_pb2.PropertyReference(name=prop_name1),
+ direction=enums.Direction.ASCENDING,
+ )
+ order_pb2 = query_pb2.PropertyOrder(
+ property=query_pb2.PropertyReference(name=prop_name2),
+ direction=enums.Direction.DESCENDING,
+ )
+ expected = self._make_query_pb(
+ collection, order=[order_pb1, order_pb2])
+ mock_query.assert_called_once_with(
+ self.PROJECT,
+ None,
+ partition_id=None,
+ read_options=None,
+ query=expected,
+ gql_query=None,
+ property_mask=None,
+ )
+
+
+class TestFieldOrder(unittest.TestCase):
+ @staticmethod
+ def _get_target_class():
+ from google.cloud.firestore.collection import FieldOrder
+
+ return FieldOrder
+
+ def _make_one(self, *args, **kwargs):
+ return self._get_target_class()(*args, **kwargs)
+
+ def test_constructor(self):
+ from google.cloud.gapic.firestore.v1alpha1 import enums
+
+ order = self._make_one('my-property')
+ self.assertEqual(order._field_name, 'my-property')
+ self.assertEqual(order._direction_enum, enums.Direction.ASCENDING)
+
+ def test_to_proto(self):
+ from google.cloud.gapic.firestore.v1alpha1 import enums
+ from google.firestore.v1alpha1 import query_pb2
+
+ order = self._make_one('my-property')
+ expected = query_pb2.PropertyOrder(
+ property=query_pb2.PropertyReference(
+ name='my-property'
+ ),
+ direction=enums.Direction.ASCENDING
+ )
+
+ self.assertEqual(order.to_proto(), expected)
+
+ def test_invalid_direction(self):
+ with self.assertRaisesRegexp(ValueError, r'Invalid.*direction'):
+ self._make_one('my-property', 'descending')