Merge pull request googleapis#134 from jgeewax/koss-collection-2

Add Query class and tests.

firestore/google/cloud/firestore/collection.py

@@ -14,48 +14,133 @@
 
 """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.
 
  :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.
 
  :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.

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)

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'):