diff --git a/firestore/google/cloud/firestore/collection.py b/firestore/google/cloud/firestore/collection.py index f0c03c9c8501..4632017db750 100644 --- a/firestore/google/cloud/firestore/collection.py +++ b/firestore/google/cloud/firestore/collection.py @@ -14,20 +14,106 @@ """Firestore Collection and Query classes.""" +from google.cloud.gapic.firestore.v1alpha1 import enums + +from google.firestore.v1alpha1 import query_pb2 + +from google.protobuf.wrappers_pb2 import Int32Value -class CollectionRef(object): - """Reference to a collection location in a Firestore database. + +class Query(object): + """Firestore Query. :type client: :class:`~google.cloud.firestore.Client` :param client: Firestore client context. - :type path: :class:`~google.cloud.firestore.Path` - :param path: Path location of this collection. + :type ref_path: :class: `~google.cloud.firestore._path.Path` + :param ref_path: Path to a collection to be queried. + + :type limit: int + :param limit: (optional) Maximum number of results to return. + + :raises ValueError: Path must be a valid Collection path (not a + Document path). """ - def __init__(self, client, path): + def __init__(self, + client, + ref_path, + limit=None): + if not ref_path.is_collection: + raise ValueError('Invalid collection path: %s' % (ref_path,)) + self._client = client - self._path = path + self._path = ref_path + self._limit = limit + + def limit(self, count): + """Limit a query to return a fixed number of results. + + This function returns a new (immutable) instance of the ``Query`` + (rather than modify the existing ``Query``) to impose the limit. + + :type count: int + :param count: Maximum number of documents to return when :meth:`get` + is called. + + :rtype: :class:`~google.cloud.firestore.collection.Query` + :return: A limited :class:`Query`. + """ + return Query(self._client, self._path, limit=count) + + def get(self): + """Read the documents in the collection indicated by this query. + + :rtype: list of :class:`~google.cloud.firestore.DocumentResult` + :returns: A sequence of Documents that fulfill the query + conditions from a collection. + """ + import google.cloud.firestore._values as values + from google.cloud.firestore.document import DocumentResult + from google.cloud.firestore._path import Path + + complete_filter = query_pb2.Filter( + property_filter=query_pb2.PropertyFilter( + property=query_pb2.PropertyReference(name='__key__'), + op=enums.Operator.HAS_PARENT, + value=values.encode_value(self._path.parent))) + + query = query_pb2.Query( + kind=[query_pb2.KindExpression(name=self._path.kind)], + filter=complete_filter, + limit=Int32Value(value=self._limit)) + + response = self._client._api.run_query( + self._client.project, + self._client.database_id, + partition_id=None, + read_options=None, + query=query, + gql_query=None, + property_mask=None) + + result = [ + DocumentResult( + self._client, + Path.from_key_proto(entity_result.entity.key), + values.decode_dict(entity_result.entity.properties)) + for entity_result in response.batch.entity_results + ] + + return result + + +class CollectionRef(Query): + """Reference to a collection location in a Firestore database. + + :type client: :class:`~google.cloud.firestore.Client` + :param client: Firestore client context. + + :type path: :class:`~google.cloud.firestore._path.Path` + :param path: Path location of this collection. + """ def document(self, doc_id): """Reference to a child document in the collection. @@ -35,13 +121,12 @@ def document(self, doc_id): :type doc_id: int or str :param doc_id: The (unique) document identifier. - :rtype: :class:`~google.cloud.firestore.Document` + :rtype: :class:`~google.cloud.firestore.DocumentResult` :returns: A reference to the child document. """ - import google.cloud.firestore.document as document + from google.cloud.firestore.document import DocumentRef - return document.DocumentRef(self._client, - self._path.child(doc_id)) + return DocumentRef(self._client, self._path.child(doc_id)) def add(self, value): """Create a new document and save it in this colleciton. @@ -49,13 +134,13 @@ def add(self, value): :type value: dict :param value: Dictionary of values to save to the document. - :rtype: :class:`~google.cloud.firestore.Document` - :returns: A ``Document`` that was saved to the database. + :rtype: :class:`~google.cloud.firestore.DocumentResult` + :returns: A ``DocumentResult`` that was saved to the database. """ return self.new_document().set(value) def new_document(self): - """Create a reference to a new ``Document`` in this collection. + """Create a reference to a new Document in this collection. New documents are created with uniquely created client-side identifiers. diff --git a/firestore/unit_tests/test_collection.py b/firestore/unit_tests/test_collection.py index 095db880deec..dd4d16d13200 100644 --- a/firestore/unit_tests/test_collection.py +++ b/firestore/unit_tests/test_collection.py @@ -69,8 +69,55 @@ def test_new_document(self): DatastoreApi) client = Client() - client._api = mock.MagicMock(spec=DatastoreApi) + client._api = mock.Mock(spec=DatastoreApi) col_ref = self._make_one(client, Path('my-collection')) doc_ref = col_ref.new_document() self.assertIsInstance(doc_ref, DocumentRef) self.assertTrue(len(doc_ref.id) >= 20) + + +class TestQuery(unittest.TestCase): + + @staticmethod + def _get_target_class(): + from google.cloud.firestore.collection import Query + + return Query + + 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) + return self._get_target_class()(client, *args, **kwargs) + + def test_constructor(self): + from google.cloud.firestore._path import Path + + query = self._make_one(Path('my-collection')) + self.assertIsInstance(query, self._get_target_class()) + self.assertIsNone(query._limit) + + def test_constructor_illegal_path(self): + from google.cloud.firestore._path import Path + + with self.assertRaisesRegexp(ValueError, r'Invalid.*path'): + self._make_one(Path('my-collection', 'my-document')) + + def test_limit(self): + from google.cloud.firestore._path import Path + + query = self._make_one(Path('my-collection')).limit(123) + + self.assertEqual(query._limit, 123) + + def test_get(self): + from google.cloud.firestore._path import Path + + query = self._make_one(Path('my-collection')) + docs = query.get() + self.assertEqual(query._client._api.run_query.call_count, 1) + self.assertIsInstance(docs, list) diff --git a/firestore/unit_tests/test_document.py b/firestore/unit_tests/test_document.py index 1d9f31e6ae0b..327e62b6f2e9 100644 --- a/firestore/unit_tests/test_document.py +++ b/firestore/unit_tests/test_document.py @@ -107,7 +107,8 @@ def _make_one(self, *args, **kwargs): def test_constructor(self): doc_ref = self._make_one(self.client, self.doc_path) - self.assertIsNotNone(doc_ref) + self.assertIsInstance(doc_ref, self._get_target_class()) + self.assertEqual(doc_ref._path, self.doc_path) def test_invalid_path(self): with self.assertRaisesRegexp(ValueError, 'Invalid'):