Firestore: add support for CollectionGroup queries.

firestore/google/cloud/firestore_v1/client.py

@@ -26,6 +26,7 @@
 from google.cloud.client import ClientWithProject
 
 from google.cloud.firestore_v1 import _helpers
+from google.cloud.firestore_v1 import query
 from google.cloud.firestore_v1 import types
 from google.cloud.firestore_v1.batch import WriteBatch
 from google.cloud.firestore_v1.collection import CollectionReference
@@ -179,6 +180,31 @@ def collection(self, *collection_path):
 
         return CollectionReference(*path, client=self)
 
+    def collection_group(self, collection_id):
+        """
+        Creates and returns a new Query that includes all documents in the
+        database that are contained in a collection or subcollection with the
+        given collection_id.
+
+        .. code-block:: python
+
+            >>> query = firestore.collection_group('mygroup')
+
+        @param {string} collectionId Identifies the collections to query over.
+        Every collection or subcollection with this ID as the last segment of its
+        path will be included. Cannot contain a slash.
+        @returns {Query} The created Query.
+        """
+        if "/" in collection_id:
+            raise ValueError(
+                "Invalid collection_id "
+                + collection_id
+                + ". Collection IDs must not contain '/'."
+            )
+
+        collection = self.collection(collection_id)
+        return query.Query(collection, all_descendants=True)
+
     def document(self, *document_path):
         """Get a reference to a document in a collection.
 
@@ -215,6 +241,13 @@ def document(self, *document_path):
         else:
             path = document_path
 
+        # DocumentReference takes a relative path. Strip the database string if present.
+        base_path = self._database_string + "/documents/"
+        joined_path = _helpers.DOCUMENT_PATH_DELIMITER.join(path)
+        if joined_path.startswith(base_path):
+            joined_path = joined_path[len(base_path) :]
+        path = joined_path.split(_helpers.DOCUMENT_PATH_DELIMITER)
+
         return DocumentReference(*path, client=self)
 
     @staticmethod

firestore/google/cloud/firestore_v1/query.py

@@ -111,6 +111,10 @@ class Query(object):
             any matching documents will be included in the result set.
             When the query is formed, the document values
             will be used in the order given by ``orders``.
+        all_descendants (Optional[bool]): When false, selects only collections
+            that are immediate children of the `parent` specified in the
+            containing `RunQueryRequest`. When true, selects all descendant
+            collections.
     """
 
     ASCENDING = "ASCENDING"
@@ -128,6 +132,7 @@ def __init__(
         offset=None,
         start_at=None,
         end_at=None,
+        all_descendants=False,
     ):
         self._parent = parent
         self._projection = projection
@@ -137,6 +142,7 @@ def __init__(
         self._offset = offset
         self._start_at = start_at
         self._end_at = end_at
+        self._all_descendants = all_descendants
 
     def __eq__(self, other):
         if not isinstance(other, self.__class__):
@@ -150,6 +156,7 @@ def __eq__(self, other):
             and self._offset == other._offset
             and self._start_at == other._start_at
             and self._end_at == other._end_at
+            and self._all_descendants == other._all_descendants
         )
 
     @property
@@ -203,6 +210,7 @@ def select(self, field_paths):
             offset=self._offset,
             start_at=self._start_at,
             end_at=self._end_at,
+            all_descendants=self._all_descendants,
         )
 
     def where(self, field_path, op_string, value):
@@ -270,6 +278,7 @@ def where(self, field_path, op_string, value):
             offset=self._offset,
             start_at=self._start_at,
             end_at=self._end_at,
+            all_descendants=self._all_descendants,
         )
 
     @staticmethod
@@ -321,6 +330,7 @@ def order_by(self, field_path, direction=ASCENDING):
             offset=self._offset,
             start_at=self._start_at,
             end_at=self._end_at,
+            all_descendants=self._all_descendants,
         )
 
     def limit(self, count):
@@ -346,6 +356,7 @@ def limit(self, count):
             offset=self._offset,
             start_at=self._start_at,
             end_at=self._end_at,
+            all_descendants=self._all_descendants,
         )
 
     def offset(self, num_to_skip):
@@ -372,6 +383,7 @@ def offset(self, num_to_skip):
             offset=num_to_skip,
             start_at=self._start_at,
             end_at=self._end_at,
+            all_descendants=self._all_descendants,
         )
 
     def _cursor_helper(self, document_fields, before, start):
@@ -418,6 +430,7 @@ def _cursor_helper(self, document_fields, before, start):
             "orders": self._orders,
             "limit": self._limit,
             "offset": self._offset,
+            "all_descendants": self._all_descendants,
         }
         if start:
             query_kwargs["start_at"] = cursor_pair
@@ -679,7 +692,7 @@ def _to_protobuf(self):
             "select": projection,
             "from": [
                 query_pb2.StructuredQuery.CollectionSelector(
-                    collection_id=self._parent.id
+                    collection_id=self._parent.id, all_descendants=self._all_descendants
                 )
             ],
             "where": self._filters_pb(),
@@ -739,9 +752,14 @@ def stream(self, transaction=None):
         )
 
         for response in response_iterator:
-            snapshot = _query_response_to_snapshot(
-                response, self._parent, expected_prefix
-            )
+            if self._all_descendants:
+                snapshot = _collection_group_query_response_to_snapshot(
+                    response, self._parent
+                )
+            else:
+                snapshot = _query_response_to_snapshot(
+                    response, self._parent, expected_prefix
+                )
             if snapshot is not None:
                 yield snapshot
 
@@ -968,3 +986,32 @@ def _query_response_to_snapshot(response_pb, collection, expected_prefix):
         update_time=response_pb.document.update_time,
     )
     return snapshot
+
+
+def _collection_group_query_response_to_snapshot(response_pb, collection):
+    """Parse a query response protobuf to a document snapshot.
+
+    Args:
+        response_pb (google.cloud.proto.firestore.v1.\
+            firestore_pb2.RunQueryResponse): A
+        collection (~.firestore_v1.collection.CollectionReference): A
+            reference to the collection that initiated the query.
+
+    Returns:
+        Optional[~.firestore.document.DocumentSnapshot]: A
+        snapshot of the data returned in the query. If ``response_pb.document``
+        is not set, the snapshot will be :data:`None`.
+    """
+    if not response_pb.HasField("document"):
+        return None
+    reference = collection._client.document(response_pb.document.name)
+    data = _helpers.decode_dict(response_pb.document.fields, collection._client)
+    snapshot = document.DocumentSnapshot(
+        reference,
+        data,
+        exists=True,
+        read_time=response_pb.read_time,
+        create_time=response_pb.document.create_time,
+        update_time=response_pb.document.update_time,
+    )
+    return snapshot

firestore/tests/system.py

@@ -634,6 +634,120 @@ def test_query_unary(client, cleanup):
     assert math.isnan(data1[field_name])
 
 
+def test_collection_group_queries(client, cleanup):
+    collection_group = "b" + unique_resource_id("-")
+
+    doc_paths = [
+        "abc/123/" + collection_group + "/cg-doc1",
+        "abc/123/" + collection_group + "/cg-doc2",
+        collection_group + "/cg-doc3",
+        collection_group + "/cg-doc4",
+        "def/456/" + collection_group + "/cg-doc5",
+        collection_group + "/virtual-doc/nested-coll/not-cg-doc",
+        "x" + collection_group + "/not-cg-doc",
+        collection_group + "x/not-cg-doc",
+        "abc/123/" + collection_group + "x/not-cg-doc",
+        "abc/123/x" + collection_group + "/not-cg-doc",
+        "abc/" + collection_group,
+    ]
+
+    batch = client.batch()
+    for doc_path in doc_paths:
+        doc_ref = client.document(doc_path)
+        batch.set(doc_ref, {"x": 1})
+
+    batch.commit()
+
+    query = client.collection_group(collection_group)
+    snapshots = list(query.stream())
+    found = [snapshot.id for snapshot in snapshots]
+    expected = ["cg-doc1", "cg-doc2", "cg-doc3", "cg-doc4", "cg-doc5"]
+    assert found == expected
+
+
+def test_collection_group_queries_startat_endat(client, cleanup):
+    collection_group = "b" + unique_resource_id("-")
+
+    doc_paths = [
+        "a/a/" + collection_group + "/cg-doc1",
+        "a/b/a/b/" + collection_group + "/cg-doc2",
+        "a/b/" + collection_group + "/cg-doc3",
+        "a/b/c/d/" + collection_group + "/cg-doc4",
+        "a/c/" + collection_group + "/cg-doc5",
+        collection_group + "/cg-doc6",
+        "a/b/nope/nope",
+    ]
+
+    batch = client.batch()
+    for doc_path in doc_paths:
+        doc_ref = client.document(doc_path)
+        batch.set(doc_ref, {"x": doc_path})
+
+    batch.commit()
+
+    query = (
+        client.collection_group(collection_group)
+        .order_by("__name__")
+        .start_at([client.document("a/b")])
+        .end_at([client.document("a/b0")])
+    )
+    snapshots = list(query.stream())
+    found = set(snapshot.id for snapshot in snapshots)
+    assert found == set(["cg-doc2", "cg-doc3", "cg-doc4"])
+
+    query = (
+        client.collection_group(collection_group)
+        .order_by("__name__")
+        .start_after([client.document("a/b")])
+        .end_before([client.document("a/b/" + collection_group + "/cg-doc3")])
+    )
+    snapshots = list(query.stream())
+    found = set(snapshot.id for snapshot in snapshots)
+    assert found == set(["cg-doc2"])
+
+
+def test_collection_group_queries_filters(client, cleanup):
+    collection_group = "b" + unique_resource_id("-")
+
+    doc_paths = [
+        "a/a/" + collection_group + "/cg-doc1",
+        "a/b/a/b/" + collection_group + "/cg-doc2",
+        "a/b/" + collection_group + "/cg-doc3",
+        "a/b/c/d/" + collection_group + "/cg-doc4",
+        "a/c/" + collection_group + "/cg-doc5",
+        collection_group + "/cg-doc6",
+        "a/b/nope/nope",
+    ]
+
+    batch = client.batch()
+
+    for index, doc_path in enumerate(doc_paths):
+        doc_ref = client.document(doc_path)
+        batch.set(doc_ref, {"x": index})
+
+    batch.commit()
+
+    query = (
+        client.collection_group(collection_group)
+        .where("__name__", ">=", client.document("a/b"))
+        .where("__name__", "<=", client.document("a/b0"))
+    )
+    snapshots = list(query.stream())
+    found = set(snapshot.id for snapshot in snapshots)
+    assert found == set(["cg-doc2", "cg-doc3", "cg-doc4"])
+
+    query = (
+        client.collection_group(collection_group)
+        .where("__name__", ">", client.document("a/b"))
+        .where(
+            "__name__", "<", client.document("a/b/{}/cg-doc3".format(collection_group))
+        )
+    )
+    snapshots = list(query.stream())
+    found = set(snapshot.id for snapshot in snapshots)
+    assert found == set(["cg-doc2"])
+
+
 def test_get_all(client, cleanup):
     collection_name = "get-all" + unique_resource_id("-")
 

firestore/tests/unit/v1/test_client.py

@@ -130,6 +130,21 @@ def test_collection_factory_nested(self):
         self.assertIs(collection2._client, client)
         self.assertIsInstance(collection2, CollectionReference)
 
+    def test_collection_group(self):
+        client = self._make_default_one()
+        query = client.collection_group("collectionId").where("foo", "==", u"bar")
+
+        assert query._all_descendants
+        assert query._field_filters[0].field.field_path == "foo"
+        assert query._field_filters[0].value.string_value == u"bar"
+        assert query._field_filters[0].op == query._field_filters[0].EQUAL
+        assert query._parent.id == "collectionId"
+
+    def test_collection_group_no_slashes(self):
+        client = self._make_default_one()
+        with self.assertRaises(ValueError):
+            client.collection_group("foo/bar")
+
     def test_document_factory(self):
         from google.cloud.firestore_v1.document import DocumentReference
 
@@ -148,7 +163,20 @@ def test_document_factory(self):
         self.assertIs(document2._client, client)
         self.assertIsInstance(document2, DocumentReference)
 
-    def test_document_factory_nested(self):
+    def test_document_factory_w_absolute_path(self):
+        from google.cloud.firestore_v1.document import DocumentReference
+
+        parts = ("rooms", "roomA")
+        client = self._make_default_one()
+        doc_path = "/".join(parts)
+        to_match = client.document(doc_path)
+        document1 = client.document(to_match._document_path)
+
+        self.assertEqual(document1._path, parts)
+        self.assertIs(document1._client, client)
+        self.assertIsInstance(document1, DocumentReference)
+
+    def test_document_factory_w_nested_path(self):
         from google.cloud.firestore_v1.document import DocumentReference
 
         client = self._make_default_one()