From c49a1cf0b079c53f61192de589efa32044712b58 Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Wed, 26 Jul 2023 15:54:13 -0700 Subject: [PATCH] Add async client instrumentation --- newrelic/config.py | 5 ++ newrelic/hooks/datastore_firestore.py | 24 ++++++- tests/datastore_firestore/conftest.py | 40 ++++++++++- .../datastore_firestore/test_async_client.py | 68 +++++++++++++++++++ 4 files changed, 131 insertions(+), 6 deletions(-) create mode 100644 tests/datastore_firestore/test_async_client.py diff --git a/newrelic/config.py b/newrelic/config.py index 7083cc872b..ca21a97339 100644 --- a/newrelic/config.py +++ b/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", diff --git a/newrelic/hooks/datastore_firestore.py b/newrelic/hooks/datastore_firestore.py index b591176a98..e534d98ecc 100644 --- a/newrelic/hooks/datastore_firestore.py +++ b/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 diff --git a/tests/datastore_firestore/conftest.py b/tests/datastore_firestore/conftest.py index 0bf899c718..767c5c474a 100644 --- a/tests/datastore_firestore/conftest.py +++ b/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 \ No newline at end of file + 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 diff --git a/tests/datastore_firestore/test_async_client.py b/tests/datastore_firestore/test_async_client.py new file mode 100644 index 0000000000..6496e6a357 --- /dev/null +++ b/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])