From b9eaa5b35144be48243e6315b8c64ad599d6a4de Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Tue, 25 Jul 2023 13:33:42 -0700 Subject: [PATCH] Client instrumentation --- newrelic/config.py | 16 ++-- newrelic/hooks/datastore_firestore.py | 13 +++- tests/datastore_firestore/conftest.py | 21 ++++- tests/datastore_firestore/test_client.py | 76 +++++++++++++++++++ tests/datastore_firestore/test_collections.py | 17 +---- tests/datastore_firestore/test_documents.py | 9 +-- 6 files changed, 121 insertions(+), 31 deletions(-) create mode 100644 tests/datastore_firestore/test_client.py diff --git a/newrelic/config.py b/newrelic/config.py index 3e258a19cb..3e129f6abc 100644 --- a/newrelic/config.py +++ b/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", diff --git a/newrelic/hooks/datastore_firestore.py b/newrelic/hooks/datastore_firestore.py index b59df5d0dc..c64f36b2a2 100644 --- a/newrelic/hooks/datastore_firestore.py +++ b/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 diff --git a/tests/datastore_firestore/conftest.py b/tests/datastore_firestore/conftest.py index f4d76c3e41..0bf899c718 100644 --- a/tests/datastore_firestore/conftest.py +++ b/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 \ No newline at end of file diff --git a/tests/datastore_firestore/test_client.py b/tests/datastore_firestore/test_client.py new file mode 100644 index 0000000000..32a67b271b --- /dev/null +++ b/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) diff --git a/tests/datastore_firestore/test_collections.py b/tests/datastore_firestore/test_collections.py index 30f26eb8f9..8a83d1679d 100644 --- a/tests/datastore_firestore/test_collections.py +++ b/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() diff --git a/tests/datastore_firestore/test_documents.py b/tests/datastore_firestore/test_documents.py index be47e820fd..b2b971abf7 100644 --- a/tests/datastore_firestore/test_documents.py +++ b/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()