diff --git a/firestore/google/cloud/firestore/collection.py b/firestore/google/cloud/firestore/collection.py index 4632017db750..99e39e31d507 100644 --- a/firestore/google/cloud/firestore/collection.py +++ b/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,6 +35,11 @@ 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). """ @@ -40,13 +47,39 @@ class Query(object): 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,)) diff --git a/firestore/unit_tests/test_collection.py b/firestore/unit_tests/test_collection.py index dd4d16d13200..4187777ab7e6 100644 --- a/firestore/unit_tests/test_collection.py +++ b/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')