Squashed commit of the following:

commit c836f8f
Author: TimPansino <TimPansino@users.noreply.github.com>
Date:   Thu Jul 27 19:54:35 2023 +0000

    [Mega-Linter] Apply linters fixes

commit 02a55a1
Author: Tim Pansino <timpansino@gmail.com>
Date:   Thu Jul 27 12:46:46 2023 -0700

    Add collection group instrumentation

commit ab1f4ff
Author: Tim Pansino <timpansino@gmail.com>
Date:   Thu Jul 27 12:00:33 2023 -0700

    Better parallelization safeguards

commit fa5f39a
Author: TimPansino <TimPansino@users.noreply.github.com>
Date:   Wed Jul 26 22:59:11 2023 +0000

    [Mega-Linter] Apply linters fixes

commit 9d411e0
Author: Tim Pansino <timpansino@gmail.com>
Date:   Wed Jul 26 15:57:39 2023 -0700

    Clean out unnecessary code

commit cb550ba
Author: Tim Pansino <timpansino@gmail.com>
Date:   Wed Jul 26 14:27:01 2023 -0700

    Allow better parallelization in firestore tests

newrelic/hooks/datastore_firestore.py

@@ -53,7 +53,7 @@ def _wrapper(wrapped, instance, args, kwargs):
         trace = DatastoreTrace(product="Firestore", target=target_, operation=method_name)
         wrapped = async_wrapper(wrapped, trace)
         return wrapped(*args, **kwargs)
-    
+
     class_ = getattr(module, class_name)
     if class_ is not None:
         if hasattr(class_, method_name):
@@ -92,7 +92,11 @@ def instrument_google_cloud_firestore_v1_collection(module):
         for method in ("add", "get"):
             if hasattr(class_, method):
                 wrap_datastore_trace(
-                    module, "CollectionReference.%s" % method, product="Firestore", target=_get_object_id, operation=method
+                    module,
+                    "CollectionReference.%s" % method,
+                    product="Firestore",
+                    target=_get_object_id,
+                    operation=method,
                 )
 
         for method in ("stream", "list_documents"):
@@ -120,7 +124,11 @@ def instrument_google_cloud_firestore_v1_document(module):
         for method in ("create", "delete", "get", "set", "update"):
             if hasattr(class_, method):
                 wrap_datastore_trace(
-                    module, "DocumentReference.%s" % method, product="Firestore", target=_get_object_id, operation=method
+                    module,
+                    "DocumentReference.%s" % method,
+                    product="Firestore",
+                    target=_get_object_id,
+                    operation=method,
                 )
 
         for method in ("collections",):
@@ -155,6 +163,12 @@ def instrument_google_cloud_firestore_v1_query(module):
             if hasattr(class_, method):
                 wrap_generator_method(module, "Query", method, target=_get_parent_id)
 
+    if hasattr(module, "CollectionGroup"):
+        class_ = module.CollectionGroup
+        for method in ("get_partitions",):
+            if hasattr(class_, method):
+                wrap_generator_method(module, "CollectionGroup", method, target=_get_parent_id)
+
 
 def instrument_google_cloud_firestore_v1_async_query(module):
     if hasattr(module, "AsyncQuery"):
@@ -176,7 +190,11 @@ def instrument_google_cloud_firestore_v1_aggregation(module):
         for method in ("get",):
             if hasattr(class_, method):
                 wrap_datastore_trace(
-                    module, "AggregationQuery.%s" % method, product="Firestore", target=_get_collection_ref_id, operation=method
+                    module,
+                    "AggregationQuery.%s" % method,
+                    product="Firestore",
+                    target=_get_collection_ref_id,
+                    operation=method,
                 )
 
         for method in ("stream",):

tests/datastore_firestore/conftest.py

@@ -15,45 +15,52 @@
 import uuid
 
 import pytest
-
 from google.cloud.firestore import Client, AsyncClient
-
-from newrelic.api.time_trace import current_trace
-from newrelic.api.datastore_trace import DatastoreTrace
 from testing_support.db_settings import firestore_settings
 from testing_support.fixture.event_loop import event_loop as loop  # noqa: F401; pylint: disable=W0611
-from testing_support.fixtures import collector_agent_registration_fixture, collector_available_fixture  # noqa: F401; pylint: disable=W0611
+from testing_support.fixtures import (  # noqa: F401; pylint: disable=W0611
+    collector_agent_registration_fixture,
+    collector_available_fixture,
+)
+
+from newrelic.api.datastore_trace import DatastoreTrace
+from newrelic.api.time_trace import current_trace
 
 DB_SETTINGS = firestore_settings()[0]
 FIRESTORE_HOST = DB_SETTINGS["host"]
 FIRESTORE_PORT = DB_SETTINGS["port"]
 
 _default_settings = {
-    'transaction_tracer.explain_threshold': 0.0,
-    'transaction_tracer.transaction_threshold': 0.0,
-    'transaction_tracer.stack_trace_threshold': 0.0,
-    'debug.log_data_collector_payloads': True,
-    'debug.record_transaction_failure': True,
-    'debug.log_explain_plan_queries': True
+    "transaction_tracer.explain_threshold": 0.0,
+    "transaction_tracer.transaction_threshold": 0.0,
+    "transaction_tracer.stack_trace_threshold": 0.0,
+    "debug.log_data_collector_payloads": True,
+    "debug.record_transaction_failure": True,
+    "debug.log_explain_plan_queries": True,
 }
 
 collector_agent_registration = collector_agent_registration_fixture(
-        app_name='Python Agent Test (datastore_firestore)',
-        default_settings=_default_settings,
-        linked_applications=['Python Agent Test (datastore)'])
+    app_name="Python Agent Test (datastore_firestore)",
+    default_settings=_default_settings,
+    linked_applications=["Python Agent Test (datastore)"],
+)
 
 
 @pytest.fixture(scope="session")
 def client():
     os.environ["FIRESTORE_EMULATOR_HOST"] = "%s:%d" % (FIRESTORE_HOST, FIRESTORE_PORT)
     client = Client()
-    client.collection("healthcheck").document("healthcheck").set({}, retry=None, timeout=5)  # Ensure connection is available
+    client.collection("healthcheck").document("healthcheck").set(
+        {}, retry=None, timeout=5
+    )  # Ensure connection is available
     return client
 
 
 @pytest.fixture(scope="function")
 def collection(client):
-    yield client.collection("firestore_collection_" + str(uuid.uuid4()))
+    collection_ = client.collection("firestore_collection_" + str(uuid.uuid4()))
+    yield collection_
+    client.recursive_delete(collection_)
 
 
 @pytest.fixture(scope="session")
@@ -70,12 +77,6 @@ def async_collection(async_client, collection):
     yield async_client.collection(collection.id)
 
 
-@pytest.fixture(scope="function", autouse=True)
-def reset_firestore(client):
-    for coll in client.collections():
-        client.recursive_delete(coll)
-
-
 @pytest.fixture(scope="session")
 def assert_trace_for_generator():
     def _assert_trace_for_generator(generator_func, *args, **kwargs):

tests/datastore_firestore/test_batching.py

@@ -12,14 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
-from newrelic.api.background_task import background_task
 from testing_support.validators.validate_database_duration import (
     validate_database_duration,
 )
+from testing_support.validators.validate_transaction_metrics import (
+    validate_transaction_metrics,
+)
+
+from newrelic.api.background_task import background_task
 
 # ===== WriteBatch =====
 
+
 def _exercise_write_batch(client, collection):
     docs = [collection.document(str(x)) for x in range(1, 4)]
     batch = client.batch()
@@ -38,6 +42,7 @@ def test_firestore_write_batch(client, collection):
         ("Datastore/all", 1),
         ("Datastore/allOther", 1),
     ]
+
     @validate_database_duration()
     @validate_transaction_metrics(
         "test_firestore_write_batch",
@@ -51,11 +56,13 @@ def _test():
 
     _test()
 
+
 # ===== BulkWriteBatch =====
 
+
 def _exercise_bulk_write_batch(client, collection):
     from google.cloud.firestore_v1.bulk_batch import BulkWriteBatch
-    
+
     docs = [collection.document(str(x)) for x in range(1, 4)]
     batch = BulkWriteBatch(client)
     for doc in docs:
@@ -73,6 +80,7 @@ def test_firestore_bulk_write_batch(client, collection):
         ("Datastore/all", 1),
         ("Datastore/allOther", 1),
     ]
+
     @validate_database_duration()
     @validate_transaction_metrics(
         "test_firestore_bulk_write_batch",

tests/datastore_firestore/test_client.py

@@ -12,29 +12,30 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import pytest
-
-from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
-from newrelic.api.background_task import background_task
 from testing_support.validators.validate_database_duration import (
     validate_database_duration,
 )
+from testing_support.validators.validate_transaction_metrics import (
+    validate_transaction_metrics,
+)
+
+from newrelic.api.background_task import background_task
 
 
 @pytest.fixture()
-def existing_document(collection, reset_firestore):
-    # reset_firestore must be run before, not after this fixture
+def sample_data(collection):
     doc = collection.document("document")
     doc.set({"x": 1})
     return doc
 
 
-def _exercise_client(client, collection, existing_document):
-    assert len([_ for _ in client.collections()]) >= 1
-    doc = [_ for _ in client.get_all([existing_document])][0]
+def _exercise_client(client, collection, sample_data):
+    assert len([_ for _ in client.collections()])
+    doc = [_ for _ in client.get_all([sample_data])][0]
     assert doc.to_dict()["x"] == 1
 
 
-def test_firestore_client(client, collection, existing_document):
+def test_firestore_client(client, collection, sample_data):
     _test_scoped_metrics = [
         ("Datastore/operation/Firestore/collections", 1),
         ("Datastore/operation/Firestore/get_all", 1),
@@ -54,15 +55,12 @@ def test_firestore_client(client, collection, existing_document):
     )
     @background_task(name="test_firestore_client")
     def _test():
-        _exercise_client(client, collection, existing_document)
+        _exercise_client(client, collection, sample_data)
 
     _test()
 
 
 @background_task()
-def test_firestore_client_generators(client, collection, assert_trace_for_generator):
-    doc = collection.document("test")
-    doc.set({})
-
+def test_firestore_client_generators(client, collection, sample_data, assert_trace_for_generator):
     assert_trace_for_generator(client.collections)
-    assert_trace_for_generator(client.get_all, [doc])
+    assert_trace_for_generator(client.get_all, [sample_data])

tests/datastore_firestore/test_collections.py

@@ -12,18 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
-from newrelic.api.background_task import background_task
 from testing_support.validators.validate_database_duration import (
     validate_database_duration,
 )
+from testing_support.validators.validate_transaction_metrics import (
+    validate_transaction_metrics,
+)
+
+from newrelic.api.background_task import background_task
 
 
 def _exercise_collections(collection):
     collection.document("DoesNotExist")
     collection.add({"capital": "Rome", "currency": "Euro", "language": "Italian"}, "Italy")
     collection.add({"capital": "Mexico City", "currency": "Peso", "language": "Spanish"}, "Mexico")
-    
+
     documents_get = collection.get()
     assert len(documents_get) == 2
     documents_stream = [_ for _ in collection.stream()]
@@ -67,6 +70,6 @@ def test_firestore_collections_generators(collection, assert_trace_for_generator
     collection.add({})
     collection.add({})
     assert len([_ for _ in collection.list_documents()]) == 2
-    
+
     assert_trace_for_generator(collection.stream)
     assert_trace_for_generator(collection.list_documents)

tests/datastore_firestore/test_documents.py

@@ -12,11 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
-from newrelic.api.background_task import background_task
 from testing_support.validators.validate_database_duration import (
     validate_database_duration,
 )
+from testing_support.validators.validate_transaction_metrics import (
+    validate_transaction_metrics,
+)
+
+from newrelic.api.background_task import background_task
 
 
 def _exercise_documents(collection):
@@ -78,5 +81,5 @@ def test_firestore_documents_generators(collection, assert_trace_for_generator):
     subcollection_doc.collection("collection1").add({})
     subcollection_doc.collection("collection2").add({})
     assert len([_ for _ in subcollection_doc.collections()]) == 2
-    
+
     assert_trace_for_generator(subcollection_doc.collections)

tests/datastore_firestore/test_query.py

@@ -13,22 +13,29 @@
 # limitations under the License.
 
 import pytest
-
-from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
-from newrelic.api.background_task import background_task
 from testing_support.validators.validate_database_duration import (
     validate_database_duration,
 )
+from testing_support.validators.validate_transaction_metrics import (
+    validate_transaction_metrics,
+)
+
+from newrelic.api.background_task import background_task
 
 
 @pytest.fixture(autouse=True)
-def sample_data(collection, reset_firestore):
-    # reset_firestore must be run before, not after this fixture
+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({})
+
+
 # ===== Query =====
 
+
 def _exercise_query(collection):
     query = collection.select("x").limit(10).order_by("x").where(field_path="x", op_string="<=", value=3)
     assert len(query.get()) == 3
@@ -47,7 +54,6 @@ def test_firestore_query(collection):
         ("Datastore/all", 2),
         ("Datastore/allOther", 2),
     ]
-    @validate_database_duration()
     @validate_transaction_metrics(
         "test_firestore_query",
         scoped_metrics=_test_scoped_metrics,
@@ -66,8 +72,15 @@ def test_firestore_query_generators(collection, assert_trace_for_generator):
     query = collection.select("x").where(field_path="x", op_string="<=", value=3)
     assert_trace_for_generator(query.stream)
 
+
+@validate_database_duration()
+@background_task()
+def test_firestore_query_db_duration(collection):
+    _exercise_query(collection)
+
 # ===== AggregationQuery =====
 
+
 def _exercise_aggregation_query(collection):
     aggregation_query = collection.select("x").where(field_path="x", op_string="<=", value=3).count()
     assert aggregation_query.get()[0][0].value == 3
@@ -86,7 +99,6 @@ def test_firestore_aggregation_query(collection):
         ("Datastore/all", 2),
         ("Datastore/allOther", 2),
     ]
-    @validate_database_duration()
     @validate_transaction_metrics(
         "test_firestore_aggregation_query",
         scoped_metrics=_test_scoped_metrics,
@@ -104,3 +116,9 @@ def _test():
 def test_firestore_aggregation_query_generators(collection, assert_trace_for_generator):
     aggregation_query = collection.select("x").where(field_path="x", op_string="<=", value=3).count()
     assert_trace_for_generator(aggregation_query.stream)
+
+
+@validate_database_duration()
+@background_task()
+def test_firestore_aggregation_query_db_duration(collection):
+    _exercise_aggregation_query(collection)