Add async client instrumentation

Author: TimPansino

newrelic/config.py

@@ -2279,6 +2279,11 @@ def _process_module_builtin_defaults():
         "newrelic.hooks.datastore_firestore",
         "instrument_google_cloud_firestore_v1_client",
     )
+    _process_module_definition(
+        "google.cloud.firestore_v1.async_client",
+        "newrelic.hooks.datastore_firestore",
+        "instrument_google_cloud_firestore_v1_async_client",
+    )
     _process_module_definition(
         "google.cloud.firestore_v1.document",
         "newrelic.hooks.datastore_firestore",

newrelic/hooks/datastore_firestore.py

@@ -12,10 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
+
 from newrelic.common.object_wrapper import wrap_function_wrapper
 from newrelic.api.datastore_trace import wrap_datastore_trace
 from newrelic.api.function_trace import wrap_function_trace
-from newrelic.common.async_wrapper import generator_wrapper
+from newrelic.common.async_wrapper import generator_wrapper, async_generator_wrapper
 from newrelic.api.datastore_trace import DatastoreTrace
 
 
@@ -40,11 +42,16 @@ def _get_collection_ref_id(obj, *args, **kwargs):
         return None
 
 
-def wrap_generator_method(module, class_name, method_name, target):
+def wrap_generator_method(module, class_name, method_name, target, is_async=False):
+    if is_async:
+        async_wrapper = async_generator_wrapper
+    else:
+        async_wrapper = generator_wrapper
+
     def _wrapper(wrapped, instance, args, kwargs):
         target_ = target(instance) if callable(target) else target
         trace = DatastoreTrace(product="Firestore", target=target_, operation=method_name)
-        wrapped = generator_wrapper(wrapped, trace)
+        wrapped = async_wrapper(wrapped, trace)
         return wrapped(*args, **kwargs)
 
     class_ = getattr(module, class_name)
@@ -53,6 +60,9 @@ def _wrapper(wrapped, instance, args, kwargs):
             wrap_function_wrapper(module, "%s.%s" % (class_name, method_name), _wrapper)
 
 
+wrap_async_generator_method = functools.partial(wrap_generator_method, is_async=True)
+
+
 def instrument_google_cloud_firestore_v1_base_client(module):
     rollup = ("Datastore/all", "Datastore/Firestore/all")
     wrap_function_trace(
@@ -68,6 +78,14 @@ def instrument_google_cloud_firestore_v1_client(module):
                 wrap_generator_method(module, "Client", method, target=None)
 
 
+def instrument_google_cloud_firestore_v1_async_client(module):
+    if hasattr(module, "AsyncClient"):
+        class_ = module.AsyncClient
+        for method in ("collections", "get_all"):
+            if hasattr(class_, method):
+                wrap_async_generator_method(module, "AsyncClient", method, target=None)
+
+
 def instrument_google_cloud_firestore_v1_collection(module):
     if hasattr(module, "CollectionReference"):
         class_ = module.CollectionReference

tests/datastore_firestore/conftest.py

@@ -16,14 +16,14 @@
 
 import pytest
 
-from google.cloud.firestore import Client
+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
 
-
 DB_SETTINGS = firestore_settings()[0]
 FIRESTORE_HOST = DB_SETTINGS["host"]
 FIRESTORE_PORT = DB_SETTINGS["port"]
@@ -56,6 +56,20 @@ def collection(client):
     yield client.collection("firestore_collection_" + str(uuid.uuid4()))
 
 
+@pytest.fixture(scope="session")
+def async_client(loop):
+    os.environ["FIRESTORE_EMULATOR_HOST"] = "%s:%d" % (FIRESTORE_HOST, FIRESTORE_PORT)
+    client = AsyncClient()
+    loop.run_until_complete(client.collection("healthcheck").document("healthcheck").set({}, retry=None, timeout=5))  # Ensure connection is available
+    return client
+
+
+@pytest.fixture(scope="function")
+def async_collection(async_client, collection):
+    # Use the same collection name as the collection fixture
+    yield async_client.collection(collection.id)
+
+
 @pytest.fixture(scope="function", autouse=True)
 def reset_firestore(client):
     for coll in client.collections():
@@ -75,4 +89,24 @@ def _assert_trace_for_generator(generator_func, *args, **kwargs):
         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
+    return _assert_trace_for_generator
+
+
+@pytest.fixture(scope="session")
+def assert_trace_for_async_generator(loop):
+    def _assert_trace_for_async_generator(generator_func, *args, **kwargs):
+        _trace_check = []
+        txn = current_trace()
+        assert not isinstance(txn, DatastoreTrace)
+        
+        async def coro():
+            # Check for generator trace on collections
+            async for _ in generator_func(*args, **kwargs):
+                _trace_check.append(isinstance(current_trace(), DatastoreTrace))
+
+        loop.run_until_complete(coro())
+
+        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_async_generator

tests/datastore_firestore/test_async_client.py

@@ -0,0 +1,68 @@
+# 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 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
+
+
+async def _exercise_async_client(async_client, existing_document):
+    assert len([_ async for _ in async_client.collections()]) >= 1
+    doc = [_ async for _ in async_client.get_all([existing_document])][0]
+    assert doc.to_dict()["x"] == 1
+
+
+def test_firestore_async_client(loop, async_client, 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_database_duration()
+    @validate_transaction_metrics(
+        "test_firestore_async_client",
+        scoped_metrics=_test_scoped_metrics,
+        rollup_metrics=_test_rollup_metrics,
+        background_task=True,
+    )
+    @background_task(name="test_firestore_async_client")
+    def _test():
+       loop.run_until_complete(_exercise_async_client(async_client, existing_document))
+
+    _test()
+
+
+@background_task()
+def test_firestore_async_client_generators(async_client, collection, assert_trace_for_async_generator):
+    doc = collection.document("test")
+    doc.set({})
+
+    assert_trace_for_async_generator(async_client.collections)
+    assert_trace_for_async_generator(async_client.get_all, [doc])