Client instrumentation

newrelic/config.py

@@ -2269,23 +2269,27 @@ def _process_module_builtin_defaults():
         "instrument_graphql_validate",
     )
 
+    _process_module_definition(
+        "google.cloud.firestore_v1.base_client",
+        "newrelic.hooks.datastore_firestore",
+        "instrument_google_cloud_firestore_v1_base_client",
+    )
+    _process_module_definition(
+        "google.cloud.firestore_v1.client",
+        "newrelic.hooks.datastore_firestore",
+        "instrument_google_cloud_firestore_v1_client",
+    )
     _process_module_definition(
         "google.cloud.firestore_v1.document",
         "newrelic.hooks.datastore_firestore",
         "instrument_google_cloud_firestore_v1_document",
     )
-
     _process_module_definition(
         "google.cloud.firestore_v1.collection",
         "newrelic.hooks.datastore_firestore",
         "instrument_google_cloud_firestore_v1_collection",
     )
 
-    _process_module_definition(
-        "google.cloud.firestore_v1.base_client",
-        "newrelic.hooks.datastore_firestore",
-        "instrument_google_cloud_firestore_v1_base_client",
-    )
     _process_module_definition(
         "ariadne.asgi",
         "newrelic.hooks.framework_ariadne",

newrelic/hooks/datastore_firestore.py

@@ -22,9 +22,10 @@
 _get_object_id = lambda obj, *args, **kwargs: obj.id
 
 
-def wrap_generator_method(module, class_name, method_name):
+def wrap_generator_method(module, class_name, method_name, target=_get_object_id):
     def _wrapper(wrapped, instance, args, kwargs):
-        trace = DatastoreTrace(product="Firestore", target=instance.id, operation=method_name)
+        target_ = target(instance) if callable(target) else target
+        trace = DatastoreTrace(product="Firestore", target=target_, operation=method_name)
         wrapped = generator_wrapper(wrapped, trace)
         return wrapped(*args, **kwargs)
 
@@ -41,6 +42,14 @@ def instrument_google_cloud_firestore_v1_base_client(module):
     )
 
 
+def instrument_google_cloud_firestore_v1_client(module):
+    if hasattr(module, "Client"):
+        class_ = module.Client
+        for method in ("collections", "get_all"):
+            if hasattr(class_, method):
+                wrap_generator_method(module, "Client", method, target=None)
+
+
 def instrument_google_cloud_firestore_v1_collection(module):
     if hasattr(module, "CollectionReference"):
         class_ = module.CollectionReference

tests/datastore_firestore/conftest.py

@@ -18,6 +18,8 @@
 
 from google.cloud.firestore import Client
 
+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.fixtures import collector_agent_registration_fixture, collector_available_fixture  # noqa: F401; pylint: disable=W0611
 
@@ -57,5 +59,20 @@ def collection(client):
 @pytest.fixture(scope="function", autouse=True)
 def reset_firestore(client):
     for coll in client.collections():
-        for document in coll.list_documents():
-            document.delete()
+        client.recursive_delete(coll)
+
+
+@pytest.fixture(scope="session")
+def assert_trace_for_generator():
+    def _assert_trace_for_generator(generator_func, *args, **kwargs):
+        txn = current_trace()
+        assert not isinstance(txn, DatastoreTrace)
+
+        # Check for generator trace on collections
+        _trace_check = []
+        for _ in generator_func(*args, **kwargs):
+            _trace_check.append(isinstance(current_trace(), DatastoreTrace))
+        assert _trace_check and all(_trace_check)  # All checks are True, and at least 1 is present.
+        assert current_trace() is txn  # Generator trace has exited.
+
+    return _assert_trace_for_generator

tests/datastore_firestore/test_client.py

@@ -0,0 +1,76 @@
+# Copyright 2010 New Relic, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import pytest
+
+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.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,
+)
+
+
+@pytest.fixture()
+def existing_document(collection, reset_firestore):
+    # reset_firestore must be run before, not after this fixture
+    doc = collection.document("document")
+    doc.set({"x": 1})
+    return doc
+
+
+def _exercise_firestore(client, collection, existing_document):
+    assert len(list(client.collections())) == 1
+    doc = list(client.get_all([existing_document]))[0]
+    assert doc.to_dict()["x"] == 1
+
+
+def test_firestore_client(client, collection, existing_document):
+    _test_scoped_metrics = [
+        ("Datastore/operation/Firestore/collections", 1),
+        ("Datastore/operation/Firestore/get_all", 1),
+    ]
+
+    _test_rollup_metrics = [
+        ("Datastore/all", 2),
+        ("Datastore/allOther", 2),
+    ]
+
+    @validate_transaction_metrics(
+        "test_firestore_client",
+        scoped_metrics=_test_scoped_metrics,
+        rollup_metrics=_test_rollup_metrics,
+        background_task=True,
+    )
+    @background_task(name="test_firestore_client")
+    def _test():
+        _exercise_firestore(client, collection, existing_document)
+
+    _test()
+
+
+@background_task()
+def test_firestore_client_generators(client, collection, assert_trace_for_generator):
+    doc = collection.document("test")
+    doc.set({})
+
+    assert_trace_for_generator(client.collections)
+    assert_trace_for_generator(client.get_all, [doc])
+
+
+@validate_database_duration()
+@background_task()
+def test_firestore_client_db_duration(client, collection, existing_document):
+    _exercise_firestore(client, collection, existing_document)

tests/datastore_firestore/test_collections.py

@@ -64,25 +64,14 @@ def _test():
 
 
 @background_task()
-def test_firestore_collections_generators(collection):
+def test_firestore_collections_generators(collection, assert_trace_for_generator):
     txn = current_trace()
     collection.add({})
     collection.add({})
     assert len(list(collection.list_documents())) == 2
 
-    # Check for generator trace on stream
-    _trace_check = []
-    for _ in collection.stream():
-        _trace_check.append(isinstance(current_trace(), DatastoreTrace))
-    assert _trace_check and all(_trace_check)
-    assert current_trace() is txn
-
-    # Check for generator trace on list_documents
-    _trace_check = []
-    for _ in collection.list_documents():
-        _trace_check.append(isinstance(current_trace(), DatastoreTrace))
-    assert _trace_check and all(_trace_check)
-    assert current_trace() is txn
+    assert_trace_for_generator(collection.stream)
+    assert_trace_for_generator(collection.list_documents)
 
 
 @validate_database_duration()

tests/datastore_firestore/test_documents.py

@@ -73,7 +73,7 @@ def _test():
 
 
 @background_task()
-def test_firestore_documents_generators(collection):
+def test_firestore_documents_generators(collection, assert_trace_for_generator):
     txn = current_trace()
 
     subcollection_doc = collection.document("SubCollections")
@@ -82,12 +82,7 @@ def test_firestore_documents_generators(collection):
     subcollection_doc.collection("collection2").add({})
     assert len(list(subcollection_doc.collections())) == 2
 
-    # Check for generator trace on collections
-    _trace_check = []
-    for _ in subcollection_doc.collections():
-        _trace_check.append(isinstance(current_trace(), DatastoreTrace))
-    assert _trace_check and all(_trace_check)
-    assert current_trace() is txn
+    assert_trace_for_generator(subcollection_doc.collections)
 
 
 @validate_database_duration()