From d3e473204bb2d840d6a73ec1b5de897e11e193ee Mon Sep 17 00:00:00 2001
From: Tim Pansino <timpansino@gmail.com>
Date: Thu, 27 Jul 2023 13:20:37 -0700
Subject: [PATCH] Add async collection group instrumentation

---
 newrelic/hooks/datastore_firestore.py         |  6 ++
 tests/datastore_firestore/test_async_query.py | 81 +++++++++++++++++++
 2 files changed, 87 insertions(+)

diff --git a/newrelic/hooks/datastore_firestore.py b/newrelic/hooks/datastore_firestore.py
index ea92b4e5ec..9f28e7caf3 100644
--- a/newrelic/hooks/datastore_firestore.py
+++ b/newrelic/hooks/datastore_firestore.py
@@ -183,6 +183,12 @@ def instrument_google_cloud_firestore_v1_async_query(module):
             if hasattr(class_, method):
                 wrap_async_generator_method(module, "AsyncQuery", method, target=_get_parent_id)
 
+    if hasattr(module, "AsyncCollectionGroup"):
+        class_ = module.AsyncCollectionGroup
+        for method in ("get_partitions",):
+            if hasattr(class_, method):
+                wrap_async_generator_method(module, "AsyncCollectionGroup", method, target=_get_parent_id)
+
 
 def instrument_google_cloud_firestore_v1_aggregation(module):
     if hasattr(module, "AggregationQuery"):
diff --git a/tests/datastore_firestore/test_async_query.py b/tests/datastore_firestore/test_async_query.py
index 3319e73a53..d00449a122 100644
--- a/tests/datastore_firestore/test_async_query.py
+++ b/tests/datastore_firestore/test_async_query.py
@@ -26,6 +26,11 @@ def sample_data(collection):
     for x in range(1, 6):
         collection.add({"x": x})
 
+    subcollection_doc = collection.document("subcollection")
+    subcollection_doc.set({})
+    subcollection_doc.collection("subcollection1").add({})
+
+
 # ===== AsyncQuery =====
 
 async def _exercise_async_query(async_collection):
@@ -103,3 +108,79 @@ def _test():
 def test_firestore_async_aggregation_query_generators(async_collection, assert_trace_for_async_generator):
     async_aggregation_query = async_collection.select("x").where(field_path="x", op_string="<=", value=3).count()
     assert_trace_for_async_generator(async_aggregation_query.stream)
+
+
+# ===== CollectionGroup =====
+
+
+@pytest.fixture()
+def patch_partition_queries(monkeypatch, async_client, collection, sample_data):
+    """
+    Partitioning is not implemented in the Firestore emulator.
+
+    Ordinarily this method would return a coroutine that returns an async_generator of Cursor objects. 
+    Each Cursor must point at a valid document path. To test this, we can patch the RPC to return 1 Cursor 
+    which is pointed at any document available. The get_partitions will take that and make 2 QueryPartition 
+    objects out of it, which should be enough to ensure we can exercise the generator's tracing.
+    """
+    from google.cloud.firestore_v1.types.document import Value
+    from google.cloud.firestore_v1.types.query import Cursor
+
+    subcollection = collection.document("subcollection").collection("subcollection1")
+    documents = [d for d in subcollection.list_documents()]
+
+    async def mock_partition_query(*args, **kwargs):
+        async def _mock_partition_query():
+            yield Cursor(before=False, values=[Value(reference_value=documents[0].path)])
+        return _mock_partition_query()
+
+    monkeypatch.setattr(async_client._firestore_api, "partition_query", mock_partition_query)
+    yield
+
+
+async def _exercise_async_collection_group(async_client, async_collection):
+    async_collection_group = async_client.collection_group(async_collection.id)
+    assert len(await async_collection_group.get())
+    assert len([d async for d in async_collection_group.stream()])
+
+    partitions = [p async for p in async_collection_group.get_partitions(1)]
+    assert len(partitions) == 2
+    documents = []
+    while partitions:
+        documents.extend(await partitions.pop().query().get())
+    assert len(documents) == 6
+
+
+def test_firestore_async_collection_group(loop, async_client, async_collection, patch_partition_queries):
+    _test_scoped_metrics = [
+        ("Datastore/statement/Firestore/%s/get" % async_collection.id, 3),
+        ("Datastore/statement/Firestore/%s/stream" % async_collection.id, 1),
+        ("Datastore/statement/Firestore/%s/get_partitions" % async_collection.id, 1),
+    ]
+
+    _test_rollup_metrics = [
+        ("Datastore/operation/Firestore/get", 3),
+        ("Datastore/operation/Firestore/stream", 1),
+        ("Datastore/operation/Firestore/get_partitions", 1),
+        ("Datastore/all", 5),
+        ("Datastore/allOther", 5),
+    ]
+
+    @validate_database_duration()
+    @validate_transaction_metrics(
+        "test_firestore_async_collection_group",
+        scoped_metrics=_test_scoped_metrics,
+        rollup_metrics=_test_rollup_metrics,
+        background_task=True,
+    )
+    @background_task(name="test_firestore_async_collection_group")
+    def _test():
+        loop.run_until_complete(_exercise_async_collection_group(async_client, async_collection))
+
+    _test()
+
+
+@background_task()
+def test_firestore_async_collection_group_generators(async_client, async_collection, assert_trace_for_async_generator, patch_partition_queries):
+    async_collection_group = async_client.collection_group(async_collection.id)
+    assert_trace_for_async_generator(async_collection_group.get_partitions, 1)