From e07ffc3efb351769614c67425f7352dc4217e6be Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Mon, 24 Jul 2023 15:44:10 -0700 Subject: [PATCH 01/31] Remove unnecessary instrumentation --- newrelic/config.py | 12 ------------ newrelic/hooks/datastore_firestore.py | 28 --------------------------- 2 files changed, 40 deletions(-) diff --git a/newrelic/config.py b/newrelic/config.py index 51aa2a44b2..3e258a19cb 100644 --- a/newrelic/config.py +++ b/newrelic/config.py @@ -2269,18 +2269,6 @@ def _process_module_builtin_defaults(): "instrument_graphql_validate", ) - _process_module_definition( - "google.cloud.firestore_v1.base_document", - "newrelic.hooks.datastore_firestore", - "instrument_google_cloud_firestore_v1_base_document", - ) - - _process_module_definition( - "google.cloud.firestore_v1.base_collection", - "newrelic.hooks.datastore_firestore", - "instrument_google_cloud_firestore_v1_base_collection", - ) - _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 672438c606..efb2292f6e 100644 --- a/newrelic/hooks/datastore_firestore.py +++ b/newrelic/hooks/datastore_firestore.py @@ -60,20 +60,6 @@ def instrument_google_cloud_firestore_v1_base_client(module): ) -def instrument_google_cloud_firestore_v1_base_collection(module): - for name in _firestore_collection_commands: - if hasattr(module.BaseCollectionReference, name) and not getattr(getattr(module.BaseCollectionReference, name), "_nr_wrapped", False): - wrap_datastore_trace( - module, "BaseCollectionReference.%s" % name, product="Firestore", target=_get_object_id, operation=name - ) - getattr(module.BaseCollectionReference, name)._nr_wrapped = True - - for name in _firestore_collection_generator_commands: - if hasattr(module.BaseCollectionReference, name) and not getattr(getattr(module.BaseCollectionReference, name), "_nr_wrapped", False): - wrap_generator_method(module, "BaseCollectionReference", name) - getattr(module.BaseCollectionReference, name)._nr_wrapped = True - - def instrument_google_cloud_firestore_v1_collection(module): for name in _firestore_collection_commands: if hasattr(module.CollectionReference, name) and not getattr(getattr(module.CollectionReference, name), "_nr_wrapped", False): @@ -88,20 +74,6 @@ def instrument_google_cloud_firestore_v1_collection(module): getattr(module.CollectionReference, name)._nr_wrapped = True -def instrument_google_cloud_firestore_v1_base_document(module): - for name in _firestore_document_commands: - if hasattr(module.BaseDocumentReference, name) and not getattr(getattr(module.BaseDocumentReference, name), "_nr_wrapped", False): - wrap_datastore_trace( - module, "BaseDocumentReference.%s" % name, product="Firestore", target=_get_object_id, operation=name - ) - getattr(module.BaseDocumentReference, name)._nr_wrapped = True - - for name in _firestore_document_generator_commands: - if hasattr(module.BaseDocumentReference, name) and not getattr(getattr(module.BaseDocumentReference, name), "_nr_wrapped", False): - wrap_generator_method(module, "BaseDocumentReference", name) - getattr(module.BaseDocumentReference, name)._nr_wrapped = True - - def instrument_google_cloud_firestore_v1_document(module): for name in _firestore_document_commands: if hasattr(module.DocumentReference, name) and not getattr(getattr(module.DocumentReference, name), "_nr_wrapped", False): From ba7850a06a48005612e59b44c1a509d28f99f86d Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Mon, 24 Jul 2023 15:44:54 -0700 Subject: [PATCH 02/31] Simplify existing instrumentation --- newrelic/hooks/datastore_firestore.py | 59 +++++++++------------------ 1 file changed, 20 insertions(+), 39 deletions(-) diff --git a/newrelic/hooks/datastore_firestore.py b/newrelic/hooks/datastore_firestore.py index efb2292f6e..b59df5d0dc 100644 --- a/newrelic/hooks/datastore_firestore.py +++ b/newrelic/hooks/datastore_firestore.py @@ -18,25 +18,6 @@ from newrelic.common.async_wrapper import generator_wrapper from newrelic.api.datastore_trace import DatastoreTrace -_firestore_document_commands = ( - "create", - "delete", - "get", - "set", - "update", -) -_firestore_document_generator_commands = ( - "collections", -) - -_firestore_collection_commands = ( - "add", - "get", -) -_firestore_collection_generator_commands = ( - "stream", - "list_documents" -) _get_object_id = lambda obj, *args, **kwargs: obj.id @@ -61,28 +42,28 @@ def instrument_google_cloud_firestore_v1_base_client(module): def instrument_google_cloud_firestore_v1_collection(module): - for name in _firestore_collection_commands: - if hasattr(module.CollectionReference, name) and not getattr(getattr(module.CollectionReference, name), "_nr_wrapped", False): - wrap_datastore_trace( - module, "CollectionReference.%s" % name, product="Firestore", target=_get_object_id, operation=name - ) - getattr(module.CollectionReference, name)._nr_wrapped = True + if hasattr(module, "CollectionReference"): + class_ = module.CollectionReference + for method in ("add", "get"): + if hasattr(class_, method): + wrap_datastore_trace( + module, "CollectionReference.%s" % method, product="Firestore", target=_get_object_id, operation=method + ) - for name in _firestore_collection_generator_commands: - if hasattr(module.CollectionReference, name) and not getattr(getattr(module.CollectionReference, name), "_nr_wrapped", False): - wrap_generator_method(module, "CollectionReference", name) - getattr(module.CollectionReference, name)._nr_wrapped = True + for method in ("stream", "list_documents"): + if hasattr(class_, method): + wrap_generator_method(module, "CollectionReference", method) def instrument_google_cloud_firestore_v1_document(module): - for name in _firestore_document_commands: - if hasattr(module.DocumentReference, name) and not getattr(getattr(module.DocumentReference, name), "_nr_wrapped", False): - wrap_datastore_trace( - module, "DocumentReference.%s" % name, product="Firestore", target=_get_object_id, operation=name - ) - getattr(module.DocumentReference, name)._nr_wrapped = True + if hasattr(module, "DocumentReference"): + class_ = module.DocumentReference + 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 + ) - for name in _firestore_document_generator_commands: - if hasattr(module.DocumentReference, name) and not getattr(getattr(module.DocumentReference, name), "_nr_wrapped", False): - wrap_generator_method(module, "DocumentReference", name) - getattr(module.DocumentReference, name)._nr_wrapped = True + for method in ("collections",): + if hasattr(class_, method): + wrap_generator_method(module, "DocumentReference", method) From 19f5a48326b6aa51c1deb7e3acc2e5e6ba6ef749 Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Mon, 24 Jul 2023 15:55:52 -0700 Subject: [PATCH 03/31] Remove unnecessary settings lookups --- tests/datastore_firestore/test_collections.py | 5 ----- tests/datastore_firestore/test_documents.py | 5 ----- 2 files changed, 10 deletions(-) diff --git a/tests/datastore_firestore/test_collections.py b/tests/datastore_firestore/test_collections.py index c416ec4e28..30f26eb8f9 100644 --- a/tests/datastore_firestore/test_collections.py +++ b/tests/datastore_firestore/test_collections.py @@ -14,17 +14,12 @@ 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, ) -DB_SETTINGS = firestore_settings()[0] -FIRESTORE_HOST = DB_SETTINGS["host"] -FIRESTORE_PORT = DB_SETTINGS["port"] - def _exercise_firestore(collection): collection.document("DoesNotExist") diff --git a/tests/datastore_firestore/test_documents.py b/tests/datastore_firestore/test_documents.py index 9850c2c58d..be47e820fd 100644 --- a/tests/datastore_firestore/test_documents.py +++ b/tests/datastore_firestore/test_documents.py @@ -14,17 +14,12 @@ 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, ) -DB_SETTINGS = firestore_settings()[0] -FIRESTORE_HOST = DB_SETTINGS["host"] -FIRESTORE_PORT = DB_SETTINGS["port"] - def _exercise_firestore(collection): italy_doc = collection.document("Italy") From b9eaa5b35144be48243e6315b8c64ad599d6a4de Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Tue, 25 Jul 2023 13:33:42 -0700 Subject: [PATCH 04/31] 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() From 44598cc271e4a8d5d2962284894a9547372efdbe Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Tue, 25 Jul 2023 15:15:13 -0700 Subject: [PATCH 05/31] Add query and aggregation query instrumentation --- newrelic/config.py | 10 +++ newrelic/hooks/datastore_firestore.py | 53 ++++++++++++-- .../test_aggregation_query.py | 71 +++++++++++++++++++ tests/datastore_firestore/test_collections.py | 1 - tests/datastore_firestore/test_documents.py | 2 - tests/datastore_firestore/test_query.py | 71 +++++++++++++++++++ 6 files changed, 201 insertions(+), 7 deletions(-) create mode 100644 tests/datastore_firestore/test_aggregation_query.py create mode 100644 tests/datastore_firestore/test_query.py diff --git a/newrelic/config.py b/newrelic/config.py index 3e129f6abc..80fb9949eb 100644 --- a/newrelic/config.py +++ b/newrelic/config.py @@ -2289,6 +2289,16 @@ def _process_module_builtin_defaults(): "newrelic.hooks.datastore_firestore", "instrument_google_cloud_firestore_v1_collection", ) + _process_module_definition( + "google.cloud.firestore_v1.query", + "newrelic.hooks.datastore_firestore", + "instrument_google_cloud_firestore_v1_query", + ) + _process_module_definition( + "google.cloud.firestore_v1.aggregation", + "newrelic.hooks.datastore_firestore", + "instrument_google_cloud_firestore_v1_aggregation", + ) _process_module_definition( "ariadne.asgi", diff --git a/newrelic/hooks/datastore_firestore.py b/newrelic/hooks/datastore_firestore.py index c64f36b2a2..c21f0bb368 100644 --- a/newrelic/hooks/datastore_firestore.py +++ b/newrelic/hooks/datastore_firestore.py @@ -19,10 +19,27 @@ from newrelic.api.datastore_trace import DatastoreTrace -_get_object_id = lambda obj, *args, **kwargs: obj.id +def _get_object_id(obj, *args, **kwargs): + try: + return obj.id + except Exception: + return None -def wrap_generator_method(module, class_name, method_name, target=_get_object_id): +def _get_parent_id(obj, *args, **kwargs): + try: + return obj._parent.id + except Exception: + return None + +def _get_nested_query_parent_id(obj, *args, **kwargs): + try: + return obj._nested_query._parent.id + except Exception: + return None + + +def wrap_generator_method(module, class_name, method_name, target): def _wrapper(wrapped, instance, args, kwargs): target_ = target(instance) if callable(target) else target trace = DatastoreTrace(product="Firestore", target=target_, operation=method_name) @@ -61,7 +78,7 @@ def instrument_google_cloud_firestore_v1_collection(module): for method in ("stream", "list_documents"): if hasattr(class_, method): - wrap_generator_method(module, "CollectionReference", method) + wrap_generator_method(module, "CollectionReference", method, target=_get_object_id) def instrument_google_cloud_firestore_v1_document(module): @@ -75,4 +92,32 @@ def instrument_google_cloud_firestore_v1_document(module): for method in ("collections",): if hasattr(class_, method): - wrap_generator_method(module, "DocumentReference", method) + wrap_generator_method(module, "DocumentReference", method, target=_get_object_id) + + +def instrument_google_cloud_firestore_v1_query(module): + if hasattr(module, "Query"): + class_ = module.Query + for method in ("get",): + if hasattr(class_, method): + wrap_datastore_trace( + module, "Query.%s" % method, product="Firestore", target=_get_parent_id, operation=method + ) + + for method in ("stream",): + if hasattr(class_, method): + wrap_generator_method(module, "Query", method, target=_get_parent_id) + + +def instrument_google_cloud_firestore_v1_aggregation(module): + if hasattr(module, "AggregationQuery"): + class_ = module.AggregationQuery + for method in ("get",): + if hasattr(class_, method): + wrap_datastore_trace( + module, "AggregationQuery.%s" % method, product="Firestore", target=_get_nested_query_parent_id, operation=method + ) + + for method in ("stream",): + if hasattr(class_, method): + wrap_generator_method(module, "AggregationQuery", method, target=_get_nested_query_parent_id) diff --git a/tests/datastore_firestore/test_aggregation_query.py b/tests/datastore_firestore/test_aggregation_query.py new file mode 100644 index 0000000000..9d276507f0 --- /dev/null +++ b/tests/datastore_firestore/test_aggregation_query.py @@ -0,0 +1,71 @@ +# 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(autouse=True) +def sample_data(collection, reset_firestore): + # reset_firestore must be run before, not after this fixture + for x in range(1, 6): + collection.add({"x": x}) + + +def _exercise_firestore(collection): + aggregation_query = collection.select("x").where("x", "<=", 3).count() + assert aggregation_query.get()[0][0].value == 3 + assert list(aggregation_query.stream())[0][0].value == 3 + + +def test_firestore_aggregation_query(collection): + _test_scoped_metrics = [ + ("Datastore/statement/Firestore/%s/stream" % collection.id, 1), + ("Datastore/statement/Firestore/%s/get" % collection.id, 1), + ] + + _test_rollup_metrics = [ + ("Datastore/operation/Firestore/get", 1), + ("Datastore/operation/Firestore/stream", 1), + ("Datastore/all", 2), + ("Datastore/allOther", 2), + ] + @validate_transaction_metrics( + "test_firestore_aggregation_query", + scoped_metrics=_test_scoped_metrics, + rollup_metrics=_test_rollup_metrics, + background_task=True, + ) + @background_task(name="test_firestore_aggregation_query") + def _test(): + _exercise_firestore(collection) + + _test() + + +@background_task() +def test_firestore_aggregation_query_generators(collection, assert_trace_for_generator): + aggregation_query = collection.select("x").where("x", "<=", 3).count() + assert_trace_for_generator(aggregation_query.stream) + + +@validate_database_duration() +@background_task() +def test_firestore_aggregation_query_db_duration(collection): + _exercise_firestore(collection) diff --git a/tests/datastore_firestore/test_collections.py b/tests/datastore_firestore/test_collections.py index 8a83d1679d..7faa30cc90 100644 --- a/tests/datastore_firestore/test_collections.py +++ b/tests/datastore_firestore/test_collections.py @@ -65,7 +65,6 @@ def _test(): @background_task() def test_firestore_collections_generators(collection, assert_trace_for_generator): - txn = current_trace() collection.add({}) collection.add({}) assert len(list(collection.list_documents())) == 2 diff --git a/tests/datastore_firestore/test_documents.py b/tests/datastore_firestore/test_documents.py index b2b971abf7..3e731371bd 100644 --- a/tests/datastore_firestore/test_documents.py +++ b/tests/datastore_firestore/test_documents.py @@ -74,8 +74,6 @@ def _test(): @background_task() def test_firestore_documents_generators(collection, assert_trace_for_generator): - txn = current_trace() - subcollection_doc = collection.document("SubCollections") subcollection_doc.set({}) subcollection_doc.collection("collection1").add({}) diff --git a/tests/datastore_firestore/test_query.py b/tests/datastore_firestore/test_query.py new file mode 100644 index 0000000000..e5de8bada6 --- /dev/null +++ b/tests/datastore_firestore/test_query.py @@ -0,0 +1,71 @@ +# 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(autouse=True) +def sample_data(collection, reset_firestore): + # reset_firestore must be run before, not after this fixture + for x in range(1, 6): + collection.add({"x": x}) + + +def _exercise_firestore(collection): + query = collection.select("x").limit(10).order_by("x").where("x", "<=", 3) + assert len(query.get()) == 3 + assert len(list(query.stream())) == 3 + + +def test_firestore_query(collection): + _test_scoped_metrics = [ + ("Datastore/statement/Firestore/%s/stream" % collection.id, 1), + ("Datastore/statement/Firestore/%s/get" % collection.id, 1), + ] + + _test_rollup_metrics = [ + ("Datastore/operation/Firestore/get", 1), + ("Datastore/operation/Firestore/stream", 1), + ("Datastore/all", 2), + ("Datastore/allOther", 2), + ] + @validate_transaction_metrics( + "test_firestore_query", + scoped_metrics=_test_scoped_metrics, + rollup_metrics=_test_rollup_metrics, + background_task=True, + ) + @background_task(name="test_firestore_query") + def _test(): + _exercise_firestore(collection) + + _test() + + +@background_task() +def test_firestore_query_generators(collection, assert_trace_for_generator): + query = collection.select("x").where("x", "<=", 3) + assert_trace_for_generator(query.stream) + + +@validate_database_duration() +@background_task() +def test_firestore_query_db_duration(collection): + _exercise_firestore(collection) From a0c78a22dbd4ac43d7ef0cb444614683ce76142b Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Tue, 25 Jul 2023 15:18:51 -0700 Subject: [PATCH 06/31] Fix deprecation warning --- tests/datastore_firestore/test_aggregation_query.py | 4 ++-- tests/datastore_firestore/test_query.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/datastore_firestore/test_aggregation_query.py b/tests/datastore_firestore/test_aggregation_query.py index 9d276507f0..1f3ade19fd 100644 --- a/tests/datastore_firestore/test_aggregation_query.py +++ b/tests/datastore_firestore/test_aggregation_query.py @@ -29,7 +29,7 @@ def sample_data(collection, reset_firestore): def _exercise_firestore(collection): - aggregation_query = collection.select("x").where("x", "<=", 3).count() + aggregation_query = collection.select("x").where(field_path="x", op_string="<=", value=3).count() assert aggregation_query.get()[0][0].value == 3 assert list(aggregation_query.stream())[0][0].value == 3 @@ -61,7 +61,7 @@ def _test(): @background_task() def test_firestore_aggregation_query_generators(collection, assert_trace_for_generator): - aggregation_query = collection.select("x").where("x", "<=", 3).count() + aggregation_query = collection.select("x").where(field_path="x", op_string="<=", value=3).count() assert_trace_for_generator(aggregation_query.stream) diff --git a/tests/datastore_firestore/test_query.py b/tests/datastore_firestore/test_query.py index e5de8bada6..3f5cab10a8 100644 --- a/tests/datastore_firestore/test_query.py +++ b/tests/datastore_firestore/test_query.py @@ -29,7 +29,7 @@ def sample_data(collection, reset_firestore): def _exercise_firestore(collection): - query = collection.select("x").limit(10).order_by("x").where("x", "<=", 3) + query = collection.select("x").limit(10).order_by("x").where(field_path="x", op_string="<=", value=3) assert len(query.get()) == 3 assert len(list(query.stream())) == 3 @@ -61,7 +61,7 @@ def _test(): @background_task() def test_firestore_query_generators(collection, assert_trace_for_generator): - query = collection.select("x").where("x", "<=", 3) + query = collection.select("x").where(field_path="x", op_string="<=", value=3) assert_trace_for_generator(query.stream) From b4e87005d6b15c777563dc9ba1885612b384c61e Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Tue, 25 Jul 2023 15:23:03 -0700 Subject: [PATCH 07/31] Simplify collection lookup --- newrelic/hooks/datastore_firestore.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/newrelic/hooks/datastore_firestore.py b/newrelic/hooks/datastore_firestore.py index c21f0bb368..d2008bce29 100644 --- a/newrelic/hooks/datastore_firestore.py +++ b/newrelic/hooks/datastore_firestore.py @@ -32,9 +32,10 @@ def _get_parent_id(obj, *args, **kwargs): except Exception: return None -def _get_nested_query_parent_id(obj, *args, **kwargs): + +def _get_collection_ref_id(obj, *args, **kwargs): try: - return obj._nested_query._parent.id + return obj._collection_ref.id except Exception: return None @@ -115,9 +116,9 @@ 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_nested_query_parent_id, operation=method + module, "AggregationQuery.%s" % method, product="Firestore", target=_get_collection_ref_id, operation=method ) for method in ("stream",): if hasattr(class_, method): - wrap_generator_method(module, "AggregationQuery", method, target=_get_nested_query_parent_id) + wrap_generator_method(module, "AggregationQuery", method, target=_get_collection_ref_id) From 6214f0bc5926b0b76acdf4bad612cc2710eeb3c7 Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Tue, 25 Jul 2023 15:30:23 -0700 Subject: [PATCH 08/31] Combine query test files --- .../test_aggregation_query.py | 71 ------------------- tests/datastore_firestore/test_query.py | 51 ++++++++++++- 2 files changed, 48 insertions(+), 74 deletions(-) delete mode 100644 tests/datastore_firestore/test_aggregation_query.py diff --git a/tests/datastore_firestore/test_aggregation_query.py b/tests/datastore_firestore/test_aggregation_query.py deleted file mode 100644 index 1f3ade19fd..0000000000 --- a/tests/datastore_firestore/test_aggregation_query.py +++ /dev/null @@ -1,71 +0,0 @@ -# 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(autouse=True) -def sample_data(collection, reset_firestore): - # reset_firestore must be run before, not after this fixture - for x in range(1, 6): - collection.add({"x": x}) - - -def _exercise_firestore(collection): - aggregation_query = collection.select("x").where(field_path="x", op_string="<=", value=3).count() - assert aggregation_query.get()[0][0].value == 3 - assert list(aggregation_query.stream())[0][0].value == 3 - - -def test_firestore_aggregation_query(collection): - _test_scoped_metrics = [ - ("Datastore/statement/Firestore/%s/stream" % collection.id, 1), - ("Datastore/statement/Firestore/%s/get" % collection.id, 1), - ] - - _test_rollup_metrics = [ - ("Datastore/operation/Firestore/get", 1), - ("Datastore/operation/Firestore/stream", 1), - ("Datastore/all", 2), - ("Datastore/allOther", 2), - ] - @validate_transaction_metrics( - "test_firestore_aggregation_query", - scoped_metrics=_test_scoped_metrics, - rollup_metrics=_test_rollup_metrics, - background_task=True, - ) - @background_task(name="test_firestore_aggregation_query") - def _test(): - _exercise_firestore(collection) - - _test() - - -@background_task() -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_firestore(collection) diff --git a/tests/datastore_firestore/test_query.py b/tests/datastore_firestore/test_query.py index 3f5cab10a8..52e1f47fec 100644 --- a/tests/datastore_firestore/test_query.py +++ b/tests/datastore_firestore/test_query.py @@ -27,8 +27,9 @@ def sample_data(collection, reset_firestore): for x in range(1, 6): collection.add({"x": x}) +# ===== Query ===== -def _exercise_firestore(collection): +def _exercise_firestore_query(collection): query = collection.select("x").limit(10).order_by("x").where(field_path="x", op_string="<=", value=3) assert len(query.get()) == 3 assert len(list(query.stream())) == 3 @@ -54,7 +55,7 @@ def test_firestore_query(collection): ) @background_task(name="test_firestore_query") def _test(): - _exercise_firestore(collection) + _exercise_firestore_query(collection) _test() @@ -68,4 +69,48 @@ def test_firestore_query_generators(collection, assert_trace_for_generator): @validate_database_duration() @background_task() def test_firestore_query_db_duration(collection): - _exercise_firestore(collection) + _exercise_firestore_query(collection) + +# ===== AggregationQuery ===== + +def _exercise_firestore_aggregation(collection): + aggregation_query = collection.select("x").where(field_path="x", op_string="<=", value=3).count() + assert aggregation_query.get()[0][0].value == 3 + assert list(aggregation_query.stream())[0][0].value == 3 + + +def test_firestore_aggregation_query(collection): + _test_scoped_metrics = [ + ("Datastore/statement/Firestore/%s/stream" % collection.id, 1), + ("Datastore/statement/Firestore/%s/get" % collection.id, 1), + ] + + _test_rollup_metrics = [ + ("Datastore/operation/Firestore/get", 1), + ("Datastore/operation/Firestore/stream", 1), + ("Datastore/all", 2), + ("Datastore/allOther", 2), + ] + @validate_transaction_metrics( + "test_firestore_aggregation_query", + scoped_metrics=_test_scoped_metrics, + rollup_metrics=_test_rollup_metrics, + background_task=True, + ) + @background_task(name="test_firestore_aggregation_query") + def _test(): + _exercise_firestore_aggregation(collection) + + _test() + + +@background_task() +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_firestore_aggregation(collection) From d17b62f720c98216fe5e80df13234ab84ccd9924 Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Tue, 25 Jul 2023 15:31:48 -0700 Subject: [PATCH 09/31] Rename methods for clarity --- tests/datastore_firestore/test_client.py | 6 +++--- tests/datastore_firestore/test_collections.py | 6 +++--- tests/datastore_firestore/test_documents.py | 6 +++--- tests/datastore_firestore/test_query.py | 12 ++++++------ 4 files changed, 15 insertions(+), 15 deletions(-) diff --git a/tests/datastore_firestore/test_client.py b/tests/datastore_firestore/test_client.py index 32a67b271b..4b69650f64 100644 --- a/tests/datastore_firestore/test_client.py +++ b/tests/datastore_firestore/test_client.py @@ -31,7 +31,7 @@ def existing_document(collection, reset_firestore): return doc -def _exercise_firestore(client, collection, existing_document): +def _exercise_client(client, collection, existing_document): assert len(list(client.collections())) == 1 doc = list(client.get_all([existing_document]))[0] assert doc.to_dict()["x"] == 1 @@ -56,7 +56,7 @@ def test_firestore_client(client, collection, existing_document): ) @background_task(name="test_firestore_client") def _test(): - _exercise_firestore(client, collection, existing_document) + _exercise_client(client, collection, existing_document) _test() @@ -73,4 +73,4 @@ def test_firestore_client_generators(client, collection, assert_trace_for_genera @validate_database_duration() @background_task() def test_firestore_client_db_duration(client, collection, existing_document): - _exercise_firestore(client, collection, existing_document) + _exercise_client(client, collection, existing_document) diff --git a/tests/datastore_firestore/test_collections.py b/tests/datastore_firestore/test_collections.py index 7faa30cc90..c7e7a346f2 100644 --- a/tests/datastore_firestore/test_collections.py +++ b/tests/datastore_firestore/test_collections.py @@ -21,7 +21,7 @@ ) -def _exercise_firestore(collection): +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") @@ -58,7 +58,7 @@ def test_firestore_collections(collection): ) @background_task(name="test_firestore_collections") def _test(): - _exercise_firestore(collection) + _exercise_collections(collection) _test() @@ -76,4 +76,4 @@ def test_firestore_collections_generators(collection, assert_trace_for_generator @validate_database_duration() @background_task() def test_firestore_collections_db_duration(collection): - _exercise_firestore(collection) + _exercise_collections(collection) diff --git a/tests/datastore_firestore/test_documents.py b/tests/datastore_firestore/test_documents.py index 3e731371bd..61d82d0f02 100644 --- a/tests/datastore_firestore/test_documents.py +++ b/tests/datastore_firestore/test_documents.py @@ -21,7 +21,7 @@ ) -def _exercise_firestore(collection): +def _exercise_documents(collection): italy_doc = collection.document("Italy") italy_doc.set({"capital": "Rome", "currency": "Euro", "language": "Italian"}) italy_doc.get() @@ -67,7 +67,7 @@ def test_firestore_documents(collection): ) @background_task(name="test_firestore_documents") def _test(): - _exercise_firestore(collection) + _exercise_documents(collection) _test() @@ -86,4 +86,4 @@ def test_firestore_documents_generators(collection, assert_trace_for_generator): @validate_database_duration() @background_task() def test_firestore_documents_db_duration(collection): - _exercise_firestore(collection) + _exercise_documents(collection) diff --git a/tests/datastore_firestore/test_query.py b/tests/datastore_firestore/test_query.py index 52e1f47fec..cf4909446e 100644 --- a/tests/datastore_firestore/test_query.py +++ b/tests/datastore_firestore/test_query.py @@ -29,7 +29,7 @@ def sample_data(collection, reset_firestore): # ===== Query ===== -def _exercise_firestore_query(collection): +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 assert len(list(query.stream())) == 3 @@ -55,7 +55,7 @@ def test_firestore_query(collection): ) @background_task(name="test_firestore_query") def _test(): - _exercise_firestore_query(collection) + _exercise_query(collection) _test() @@ -69,11 +69,11 @@ def test_firestore_query_generators(collection, assert_trace_for_generator): @validate_database_duration() @background_task() def test_firestore_query_db_duration(collection): - _exercise_firestore_query(collection) + _exercise_query(collection) # ===== AggregationQuery ===== -def _exercise_firestore_aggregation(collection): +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 assert list(aggregation_query.stream())[0][0].value == 3 @@ -99,7 +99,7 @@ def test_firestore_aggregation_query(collection): ) @background_task(name="test_firestore_aggregation_query") def _test(): - _exercise_firestore_aggregation(collection) + _exercise_aggregation_query(collection) _test() @@ -113,4 +113,4 @@ def test_firestore_aggregation_query_generators(collection, assert_trace_for_gen @validate_database_duration() @background_task() def test_firestore_aggregation_query_db_duration(collection): - _exercise_firestore_aggregation(collection) + _exercise_aggregation_query(collection) From 2ce45c85ebf0a6951884c675d6cad77486988b7b Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Tue, 25 Jul 2023 15:55:50 -0700 Subject: [PATCH 10/31] Instrument Firestore batching --- newrelic/config.py | 10 +++ newrelic/hooks/datastore_firestore.py | 19 +++++ tests/datastore_firestore/test_batching.py | 87 ++++++++++++++++++++++ 3 files changed, 116 insertions(+) create mode 100644 tests/datastore_firestore/test_batching.py diff --git a/newrelic/config.py b/newrelic/config.py index 80fb9949eb..9b15737fa9 100644 --- a/newrelic/config.py +++ b/newrelic/config.py @@ -2299,6 +2299,16 @@ def _process_module_builtin_defaults(): "newrelic.hooks.datastore_firestore", "instrument_google_cloud_firestore_v1_aggregation", ) + _process_module_definition( + "google.cloud.firestore_v1.batch", + "newrelic.hooks.datastore_firestore", + "instrument_google_cloud_firestore_v1_batch", + ) + _process_module_definition( + "google.cloud.firestore_v1.bulk_batch", + "newrelic.hooks.datastore_firestore", + "instrument_google_cloud_firestore_v1_bulk_batch", + ) _process_module_definition( "ariadne.asgi", diff --git a/newrelic/hooks/datastore_firestore.py b/newrelic/hooks/datastore_firestore.py index d2008bce29..9c3997bb3e 100644 --- a/newrelic/hooks/datastore_firestore.py +++ b/newrelic/hooks/datastore_firestore.py @@ -122,3 +122,22 @@ def instrument_google_cloud_firestore_v1_aggregation(module): for method in ("stream",): if hasattr(class_, method): wrap_generator_method(module, "AggregationQuery", method, target=_get_collection_ref_id) + + +def instrument_google_cloud_firestore_v1_batch(module): + if hasattr(module, "WriteBatch"): + class_ = module.WriteBatch + for method in ("commit",): + if hasattr(class_, method): + wrap_datastore_trace( + module, "WriteBatch.%s" % method, product="Firestore", target=None, operation=method + ) + +def instrument_google_cloud_firestore_v1_bulk_batch(module): + if hasattr(module, "BulkWriteBatch"): + class_ = module.BulkWriteBatch + for method in ("commit",): + if hasattr(class_, method): + wrap_datastore_trace( + module, "BulkWriteBatch.%s" % method, product="Firestore", target=None, operation=method + ) diff --git a/tests/datastore_firestore/test_batching.py b/tests/datastore_firestore/test_batching.py new file mode 100644 index 0000000000..9c2d492fd4 --- /dev/null +++ b/tests/datastore_firestore/test_batching.py @@ -0,0 +1,87 @@ +# 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. + +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, +) + +# ===== WriteBatch ===== + +def _exercise_write_batch(client, collection): + docs = [collection.document(str(x)) for x in range(1, 4)] + batch = client.batch() + for doc in docs: + batch.set(doc, {}) + + batch.commit() + + +def test_firestore_write_batch(client, collection): + _test_scoped_metrics = [ + ("Datastore/operation/Firestore/commit", 1), + ] + + _test_rollup_metrics = [ + ("Datastore/all", 1), + ("Datastore/allOther", 1), + ] + @validate_database_duration() + @validate_transaction_metrics( + "test_firestore_write_batch", + scoped_metrics=_test_scoped_metrics, + rollup_metrics=_test_rollup_metrics, + background_task=True, + ) + @background_task(name="test_firestore_write_batch") + def _test(): + _exercise_write_batch(client, collection) + + _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: + batch.set(doc, {}) + + batch.commit() + + +def test_firestore_bulk_write_batch(client, collection): + _test_scoped_metrics = [ + ("Datastore/operation/Firestore/commit", 1), + ] + + _test_rollup_metrics = [ + ("Datastore/all", 1), + ("Datastore/allOther", 1), + ] + @validate_database_duration() + @validate_transaction_metrics( + "test_firestore_bulk_write_batch", + scoped_metrics=_test_scoped_metrics, + rollup_metrics=_test_rollup_metrics, + background_task=True, + ) + @background_task(name="test_firestore_bulk_write_batch") + def _test(): + _exercise_bulk_write_batch(client, collection) + + _test() From ef06df5dca7d6e6f0f7e96700544514b99e9c132 Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Wed, 26 Jul 2023 12:01:25 -0700 Subject: [PATCH 11/31] Add transaction instrumentation --- newrelic/config.py | 5 + newrelic/hooks/datastore_firestore.py | 12 ++ tests/datastore_firestore/test_transaction.py | 122 ++++++++++++++++++ 3 files changed, 139 insertions(+) create mode 100644 tests/datastore_firestore/test_transaction.py diff --git a/newrelic/config.py b/newrelic/config.py index 9b15737fa9..7083cc872b 100644 --- a/newrelic/config.py +++ b/newrelic/config.py @@ -2309,6 +2309,11 @@ def _process_module_builtin_defaults(): "newrelic.hooks.datastore_firestore", "instrument_google_cloud_firestore_v1_bulk_batch", ) + _process_module_definition( + "google.cloud.firestore_v1.transaction", + "newrelic.hooks.datastore_firestore", + "instrument_google_cloud_firestore_v1_transaction", + ) _process_module_definition( "ariadne.asgi", diff --git a/newrelic/hooks/datastore_firestore.py b/newrelic/hooks/datastore_firestore.py index 9c3997bb3e..b591176a98 100644 --- a/newrelic/hooks/datastore_firestore.py +++ b/newrelic/hooks/datastore_firestore.py @@ -133,6 +133,7 @@ def instrument_google_cloud_firestore_v1_batch(module): module, "WriteBatch.%s" % method, product="Firestore", target=None, operation=method ) + def instrument_google_cloud_firestore_v1_bulk_batch(module): if hasattr(module, "BulkWriteBatch"): class_ = module.BulkWriteBatch @@ -141,3 +142,14 @@ def instrument_google_cloud_firestore_v1_bulk_batch(module): wrap_datastore_trace( module, "BulkWriteBatch.%s" % method, product="Firestore", target=None, operation=method ) + + +def instrument_google_cloud_firestore_v1_transaction(module): + if hasattr(module, "Transaction"): + class_ = module.Transaction + for method in ("_commit", "_rollback"): + if hasattr(class_, method): + operation = method[1:] # Trim leading underscore + wrap_datastore_trace( + module, "Transaction.%s" % method, product="Firestore", target=None, operation=operation + ) diff --git a/tests/datastore_firestore/test_transaction.py b/tests/datastore_firestore/test_transaction.py new file mode 100644 index 0000000000..3e462b3244 --- /dev/null +++ b/tests/datastore_firestore/test_transaction.py @@ -0,0 +1,122 @@ +# 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(autouse=True) +def sample_data(collection, reset_firestore): + # reset_firestore must be run before, not after this fixture + for x in range(1, 4): + collection.add({"x": x}, "doc%d" % x) + + +def _exercise_transaction_commit(client, collection): + from google.cloud.firestore_v1.transaction import transactional + + @transactional + def _exercise(transaction): + # get a DocumentReference + list(transaction.get(collection.document("doc1"))) + + # get a Query + query = collection.select("x").where(field_path="x", op_string=">", value=2) + assert len(list(transaction.get(query))) == 1 + + # get_all on a list of DocumentReferences + all_docs = transaction.get_all([collection.document("doc%d" % x) for x in range(1, 4)]) + assert len(list(all_docs)) == 3 + + # set and delete methods + transaction.set(collection.document("doc2"), {"x": 0}) + transaction.delete(collection.document("doc3")) + + _exercise(client.transaction()) + assert len(list(collection.list_documents())) == 2 + + +def _exercise_transaction_rollback(client, collection): + from google.cloud.firestore_v1.transaction import transactional + + @transactional + def _exercise(transaction): + # set and delete methods + transaction.set(collection.document("doc2"), {"x": 99}) + transaction.delete(collection.document("doc1")) + raise RuntimeError() + + with pytest.raises(RuntimeError): + _exercise(client.transaction()) + assert len(list(collection.list_documents())) == 3 + + +def test_firestore_transaction_commit(client, collection): + _test_scoped_metrics = [ + ("Datastore/operation/Firestore/commit", 1), + ("Datastore/operation/Firestore/get_all", 2), + ("Datastore/statement/Firestore/%s/stream" % collection.id, 1), + ("Datastore/statement/Firestore/%s/list_documents" % collection.id, 1), + ] + + _test_rollup_metrics = [ + ("Datastore/operation/Firestore/stream", 1), + ("Datastore/operation/Firestore/list_documents", 1), + ("Datastore/all", 5), + ("Datastore/allOther", 5), + ] + @validate_database_duration() + @validate_transaction_metrics( + "test_firestore_transaction", + scoped_metrics=_test_scoped_metrics, + rollup_metrics=_test_rollup_metrics, + background_task=True, + ) + @background_task(name="test_firestore_transaction") + def _test(): + _exercise_transaction_commit(client, collection) + + _test() + + +def test_firestore_transaction_rollback(client, collection): + _test_scoped_metrics = [ + ("Datastore/operation/Firestore/rollback", 1), + ("Datastore/statement/Firestore/%s/list_documents" % collection.id, 1), + ] + + _test_rollup_metrics = [ + ("Datastore/operation/Firestore/list_documents", 1), + ("Datastore/all", 2), + ("Datastore/allOther", 2), + ] + @validate_database_duration() + @validate_transaction_metrics( + "test_firestore_transaction", + scoped_metrics=_test_scoped_metrics, + rollup_metrics=_test_rollup_metrics, + background_task=True, + ) + @background_task(name="test_firestore_transaction") + def _test(): + _exercise_transaction_rollback(client, collection) + + _test() From b9a91e574a8e183249549f223bb4090226467f80 Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Wed, 26 Jul 2023 12:21:25 -0700 Subject: [PATCH 12/31] Consumer iterators on <=Py38 --- tests/datastore_firestore/test_client.py | 4 ++-- tests/datastore_firestore/test_collections.py | 6 +++--- tests/datastore_firestore/test_documents.py | 4 ++-- tests/datastore_firestore/test_query.py | 4 ++-- tests/datastore_firestore/test_transaction.py | 10 +++++----- 5 files changed, 14 insertions(+), 14 deletions(-) diff --git a/tests/datastore_firestore/test_client.py b/tests/datastore_firestore/test_client.py index 4b69650f64..c4bd872a08 100644 --- a/tests/datastore_firestore/test_client.py +++ b/tests/datastore_firestore/test_client.py @@ -32,8 +32,8 @@ def existing_document(collection, reset_firestore): def _exercise_client(client, collection, existing_document): - assert len(list(client.collections())) == 1 - doc = list(client.get_all([existing_document]))[0] + assert len([_ for _ in client.collections()]) == 1 + doc = [_ for _ in client.get_all([existing_document])][0] assert doc.to_dict()["x"] == 1 diff --git a/tests/datastore_firestore/test_collections.py b/tests/datastore_firestore/test_collections.py index c7e7a346f2..ae2b48f04f 100644 --- a/tests/datastore_firestore/test_collections.py +++ b/tests/datastore_firestore/test_collections.py @@ -28,9 +28,9 @@ def _exercise_collections(collection): documents_get = collection.get() assert len(documents_get) == 2 - documents_stream = list(collection.stream()) + documents_stream = [_ for _ in collection.stream()] assert len(documents_stream) == 2 - documents_list = list(collection.list_documents()) + documents_list = [_ for _ in collection.list_documents()] assert len(documents_list) == 2 @@ -67,7 +67,7 @@ def _test(): def test_firestore_collections_generators(collection, assert_trace_for_generator): collection.add({}) collection.add({}) - assert len(list(collection.list_documents())) == 2 + assert len([_ for _ in collection.list_documents()]) == 2 assert_trace_for_generator(collection.stream) assert_trace_for_generator(collection.list_documents) diff --git a/tests/datastore_firestore/test_documents.py b/tests/datastore_firestore/test_documents.py index 61d82d0f02..67ad42dd64 100644 --- a/tests/datastore_firestore/test_documents.py +++ b/tests/datastore_firestore/test_documents.py @@ -27,7 +27,7 @@ def _exercise_documents(collection): italy_doc.get() italian_cities = italy_doc.collection("cities") italian_cities.add({"capital": "Rome"}) - retrieved_coll = list(italy_doc.collections()) + retrieved_coll = [_ for _ in italy_doc.collections()] assert len(retrieved_coll) == 1 usa_doc = collection.document("USA") @@ -78,7 +78,7 @@ def test_firestore_documents_generators(collection, assert_trace_for_generator): subcollection_doc.set({}) subcollection_doc.collection("collection1").add({}) subcollection_doc.collection("collection2").add({}) - assert len(list(subcollection_doc.collections())) == 2 + assert len([_ for _ in subcollection_doc.collections()]) == 2 assert_trace_for_generator(subcollection_doc.collections) diff --git a/tests/datastore_firestore/test_query.py b/tests/datastore_firestore/test_query.py index cf4909446e..128042a99e 100644 --- a/tests/datastore_firestore/test_query.py +++ b/tests/datastore_firestore/test_query.py @@ -32,7 +32,7 @@ def sample_data(collection, reset_firestore): 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 - assert len(list(query.stream())) == 3 + assert len([_ for _ in query.stream()]) == 3 def test_firestore_query(collection): @@ -76,7 +76,7 @@ def test_firestore_query_db_duration(collection): 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 - assert list(aggregation_query.stream())[0][0].value == 3 + assert [_ for _ in aggregation_query.stream()][0][0].value == 3 def test_firestore_aggregation_query(collection): diff --git a/tests/datastore_firestore/test_transaction.py b/tests/datastore_firestore/test_transaction.py index 3e462b3244..679821dfc4 100644 --- a/tests/datastore_firestore/test_transaction.py +++ b/tests/datastore_firestore/test_transaction.py @@ -36,22 +36,22 @@ def _exercise_transaction_commit(client, collection): @transactional def _exercise(transaction): # get a DocumentReference - list(transaction.get(collection.document("doc1"))) + [_ for _ in transaction.get(collection.document("doc1"))] # get a Query query = collection.select("x").where(field_path="x", op_string=">", value=2) - assert len(list(transaction.get(query))) == 1 + assert len([_ for _ in transaction.get(query)]) == 1 # get_all on a list of DocumentReferences all_docs = transaction.get_all([collection.document("doc%d" % x) for x in range(1, 4)]) - assert len(list(all_docs)) == 3 + assert len([_ for _ in all_docs]) == 3 # set and delete methods transaction.set(collection.document("doc2"), {"x": 0}) transaction.delete(collection.document("doc3")) _exercise(client.transaction()) - assert len(list(collection.list_documents())) == 2 + assert len([_ for _ in collection.list_documents()]) == 2 def _exercise_transaction_rollback(client, collection): @@ -66,7 +66,7 @@ def _exercise(transaction): with pytest.raises(RuntimeError): _exercise(client.transaction()) - assert len(list(collection.list_documents())) == 3 + assert len([_ for _ in collection.list_documents()]) == 3 def test_firestore_transaction_commit(client, collection): From fbe40eaf4eb9da57fd6cb881328087fedc0dc2d9 Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Wed, 26 Jul 2023 14:22:53 -0700 Subject: [PATCH 13/31] Add async generator wrapper --- newrelic/common/async_wrapper.py | 51 +- newrelic/common/coroutine.py | 8 + .../_test_async_generator_trace.py | 456 ++++++++++++++++++ .../test_async_generator_trace.py | 19 + 4 files changed, 530 insertions(+), 4 deletions(-) create mode 100644 tests/agent_features/_test_async_generator_trace.py create mode 100644 tests/agent_features/test_async_generator_trace.py diff --git a/newrelic/common/async_wrapper.py b/newrelic/common/async_wrapper.py index c5f95308da..29d7809d5e 100644 --- a/newrelic/common/async_wrapper.py +++ b/newrelic/common/async_wrapper.py @@ -18,18 +18,23 @@ is_coroutine_callable, is_asyncio_coroutine, is_generator_function, + is_async_generator_function, ) +try: + import asyncio +except ImportError: + asyncio = None + def evaluate_wrapper(wrapper_string, wrapped, trace): values = {'wrapper': None, 'wrapped': wrapped, - 'trace': trace, 'functools': functools} + 'trace': trace, 'functools': functools, "asyncio": asyncio} exec(wrapper_string, values) return values['wrapper'] def coroutine_wrapper(wrapped, trace): - WRAPPER = textwrap.dedent(""" @functools.wraps(wrapped) async def wrapper(*args, **kwargs): @@ -45,8 +50,6 @@ async def wrapper(*args, **kwargs): def awaitable_generator_wrapper(wrapped, trace): WRAPPER = textwrap.dedent(""" - import asyncio - @functools.wraps(wrapped) @asyncio.coroutine def wrapper(*args, **kwargs): @@ -81,9 +84,49 @@ def wrapper(*args, **kwargs): return wrapper +def async_generator_wrapper(wrapped, trace): + WRAPPER = textwrap.dedent(""" + @functools.wraps(wrapped) + async def wrapper(*args, **kwargs): + g = wrapped(*args, **kwargs) + value = None + with trace: + while True: + try: + g.asend(value).send(None) + except StopAsyncIteration as e: + # The underlying async generator has finished, return propagates a new StopAsyncIteration + return + except StopIteration as e: + # The call to async_generator_asend.send() should raise a StopIteration containing the yielded value + yielded = e.value + + try: + value = yield yielded + except BaseException as e: + # An exception was thrown with .athrow(), propagate to the original async generator. + # Return value logic must be identical to .asend() + try: + g.athrow(type(e), e).send(None) + except StopAsyncIteration as e: + # The underlying async generator has finished, return propagates a new StopAsyncIteration + return + except StopIteration as e: + # The call to async_generator_athrow.send() should raise a StopIteration containing a yielded value + value = yield e.value + """) + + try: + return evaluate_wrapper(WRAPPER, wrapped, trace) + except: + return wrapped + + def async_wrapper(wrapped): if is_coroutine_callable(wrapped): return coroutine_wrapper + elif is_async_generator_function(wrapped): + return async_generator_wrapper elif is_generator_function(wrapped): if is_asyncio_coroutine(wrapped): return awaitable_generator_wrapper diff --git a/newrelic/common/coroutine.py b/newrelic/common/coroutine.py index cf4c91f85c..33a4922f56 100644 --- a/newrelic/common/coroutine.py +++ b/newrelic/common/coroutine.py @@ -43,3 +43,11 @@ def _iscoroutinefunction_tornado(fn): def is_coroutine_callable(wrapped): return is_coroutine_function(wrapped) or is_coroutine_function(getattr(wrapped, "__call__", None)) + + +if hasattr(inspect, 'isasyncgenfunction'): + def is_async_generator_function(wrapped): + return inspect.isasyncgenfunction(wrapped) +else: + def is_async_generator_function(wrapped): + return False diff --git a/tests/agent_features/_test_async_generator_trace.py b/tests/agent_features/_test_async_generator_trace.py new file mode 100644 index 0000000000..11efef1112 --- /dev/null +++ b/tests/agent_features/_test_async_generator_trace.py @@ -0,0 +1,456 @@ +# 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 functools +import gc +import sys +import time + +import pytest +from testing_support.fixtures import capture_transaction_metrics, validate_tt_parenting +from testing_support.validators.validate_transaction_errors import ( + validate_transaction_errors, +) +from testing_support.validators.validate_transaction_metrics import ( + validate_transaction_metrics, +) + +from newrelic.api.background_task import background_task +from newrelic.api.database_trace import database_trace +from newrelic.api.datastore_trace import datastore_trace +from newrelic.api.external_trace import external_trace +from newrelic.api.function_trace import function_trace +from newrelic.api.memcache_trace import memcache_trace +from newrelic.api.message_trace import message_trace + +is_pypy = hasattr(sys, "pypy_version_info") +asyncio = pytest.importorskip("asyncio") + + +@pytest.mark.parametrize( + "trace,metric", + [ + (functools.partial(function_trace, name="simple_gen"), "Function/simple_gen"), + (functools.partial(external_trace, library="lib", url="http://foo.com"), "External/foo.com/lib/"), + (functools.partial(database_trace, "select * from foo"), "Datastore/statement/None/foo/select"), + (functools.partial(datastore_trace, "lib", "foo", "bar"), "Datastore/statement/lib/foo/bar"), + (functools.partial(message_trace, "lib", "op", "typ", "name"), "MessageBroker/lib/typ/op/Named/name"), + (functools.partial(memcache_trace, "cmd"), "Memcache/cmd"), + ], +) +def test_async_generator_timing(event_loop, trace, metric): + @trace() + async def simple_gen(): + time.sleep(0.1) + yield + time.sleep(0.1) + + metrics = [] + full_metrics = {} + + @capture_transaction_metrics(metrics, full_metrics) + @validate_transaction_metrics( + "test_async_generator_timing", background_task=True, scoped_metrics=[(metric, 1)], rollup_metrics=[(metric, 1)] + ) + @background_task(name="test_async_generator_timing") + def _test_async_generator_timing(): + async def _test(): + async for _ in simple_gen(): + pass + + event_loop.run_until_complete(_test()) + _test_async_generator_timing() + + # Check that coroutines time the total call time (including pauses) + metric_key = (metric, "") + assert full_metrics[metric_key].total_call_time >= 0.2 + + +class MyException(Exception): + pass + + +@validate_transaction_metrics( + "test_async_generator_error", + background_task=True, + scoped_metrics=[("Function/agen", 1)], + rollup_metrics=[("Function/agen", 1)], +) +@validate_transaction_errors(errors=["_test_async_generator_trace:MyException"]) +def test_async_generator_error(event_loop): + @function_trace(name="agen") + async def agen(): + yield + + @background_task(name="test_async_generator_error") + async def _test(): + gen = agen() + await gen.asend(None) + await gen.athrow(MyException) + + with pytest.raises(MyException): + event_loop.run_until_complete(_test()) + + +@validate_transaction_metrics( + "test_async_generator_caught_exception", + background_task=True, + scoped_metrics=[("Function/agen", 1)], + rollup_metrics=[("Function/agen", 1)], +) +@validate_transaction_errors(errors=[]) +def test_async_generator_caught_exception(event_loop): + @function_trace(name="agen") + async def agen(): + for _ in range(2): + time.sleep(0.1) + try: + yield + except ValueError: + pass + + metrics = [] + full_metrics = {} + + @capture_transaction_metrics(metrics, full_metrics) + @background_task(name="test_async_generator_caught_exception") + def _test_async_generator_caught_exception(): + async def _test(): + gen = agen() + # kickstart the generator (the try/except logic is inside the + # generator) + await anext(gen) + await gen.athrow(ValueError) + + # consume the generator + async for _ in gen: + pass + + # The ValueError should not be reraised + event_loop.run_until_complete(_test()) + _test_async_generator_caught_exception() + + assert full_metrics[("Function/agen", "")].total_call_time >= 0.2 + + +@validate_transaction_metrics( + "test_async_generator_handles_terminal_nodes", + background_task=True, + scoped_metrics=[("Function/parent", 1), ("Function/agen", None)], + rollup_metrics=[("Function/parent", 1), ("Function/agen", None)], +) +def test_async_generator_handles_terminal_nodes(event_loop): + # sometimes coroutines can be called underneath terminal nodes + # In this case, the trace shouldn't actually be created and we also + # shouldn't get any errors + + @function_trace(name="agen") + async def agen(): + yield + time.sleep(0.1) + + @function_trace(name="parent", terminal=True) + async def parent(): + # parent calls child + async for _ in agen(): + pass + + metrics = [] + full_metrics = {} + + @capture_transaction_metrics(metrics, full_metrics) + @background_task(name="test_async_generator_handles_terminal_nodes") + def _test_async_generator_handles_terminal_nodes(): + async def _test(): + await parent() + + event_loop.run_until_complete(_test()) + _test_async_generator_handles_terminal_nodes() + + metric_key = ("Function/parent", "") + assert full_metrics[metric_key].total_exclusive_call_time >= 0.1 + + +@validate_transaction_metrics( + "test_async_generator_close_ends_trace", + background_task=True, + scoped_metrics=[("Function/agen", 1)], + rollup_metrics=[("Function/agen", 1)], +) +def test_async_generator_close_ends_trace(event_loop): + @function_trace(name="agen") + async def agen(): + yield + + @background_task(name="test_async_generator_close_ends_trace") + async def _test(): + gen = agen() + + # kickstart the coroutine + await anext(gen) + + # trace should be ended/recorded by close + await gen.aclose() + + # We may call gen.close as many times as we want + await gen.aclose() + + event_loop.run_until_complete(_test()) + +@validate_tt_parenting( + ( + "TransactionNode", + [ + ( + "FunctionNode", + [ + ("FunctionNode", []), + ], + ), + ], + ) +) +@validate_transaction_metrics( + "test_async_generator_parents", + background_task=True, + scoped_metrics=[("Function/child", 1), ("Function/parent", 1)], + rollup_metrics=[("Function/child", 1), ("Function/parent", 1)], +) +def test_async_generator_parents(event_loop): + @function_trace(name="child") + async def child(): + yield + time.sleep(0.1) + yield + + @function_trace(name="parent") + async def parent(): + time.sleep(0.1) + yield + async for _ in child(): + pass + + metrics = [] + full_metrics = {} + + @capture_transaction_metrics(metrics, full_metrics) + @background_task(name="test_async_generator_parents") + def _test_async_generator_parents(): + async def _test(): + async for _ in parent(): + pass + + event_loop.run_until_complete(_test()) + _test_async_generator_parents() + + # Check that the child time is subtracted from the parent time (parenting + # relationship is correctly established) + key = ("Function/parent", "") + assert full_metrics[key].total_exclusive_call_time < 0.2 + + +@validate_transaction_metrics( + "test_athrow_yields_a_value", + background_task=True, + scoped_metrics=[("Function/agen", 1)], + rollup_metrics=[("Function/agen", 1)], +) +def test_athrow_yields_a_value(event_loop): + @function_trace(name="agen") + async def agen(): + for _ in range(2): + try: + yield + except MyException: + yield "foobar" + + @background_task(name="test_athrow_yields_a_value") + async def _test(): + gen = agen() + + # kickstart the coroutine + await anext(gen) + + assert await gen.athrow(MyException) == "foobar" + + # finish consumption of the coroutine if necessary + async for _ in gen: + pass + + event_loop.run_until_complete(_test()) + + +@validate_transaction_metrics( + "test_athrow_does_not_yield_a_value", + background_task=True, + scoped_metrics=[("Function/agen", 1)], + rollup_metrics=[("Function/agen", 1)], +) +def test_athrow_does_not_yield_a_value(event_loop): + @function_trace(name="agen") + async def agen(): + for _ in range(2): + try: + yield + except MyException: + return + + @background_task(name="test_athrow_does_not_yield_a_value") + async def _test(): + gen = agen() + + # kickstart the coroutine + await anext(gen) + + # async generator will raise StopAsyncIteration + with pytest.raises(StopAsyncIteration): + await gen.athrow(MyException) + + + event_loop.run_until_complete(_test()) + + +@pytest.mark.parametrize( + "trace", + [ + function_trace(name="simple_gen"), + external_trace(library="lib", url="http://foo.com"), + database_trace("select * from foo"), + datastore_trace("lib", "foo", "bar"), + message_trace("lib", "op", "typ", "name"), + memcache_trace("cmd"), + ], +) +def test_async_generator_functions_outside_of_transaction(event_loop, trace): + @trace + async def agen(): + for _ in range(2): + yield "foo" + + async def _test(): + assert [_ async for _ in agen()] == ["foo", "foo"] + + event_loop.run_until_complete(_test()) + + +@validate_transaction_metrics( + "test_catching_generator_exit_causes_runtime_error", + background_task=True, + scoped_metrics=[("Function/agen", 1)], + rollup_metrics=[("Function/agen", 1)], +) +def test_catching_generator_exit_causes_runtime_error(event_loop): + @function_trace(name="agen") + async def agen(): + try: + yield + except GeneratorExit: + yield + + @background_task(name="test_catching_generator_exit_causes_runtime_error") + async def _test(): + gen = agen() + + # kickstart the coroutine (we're inside the try now) + await gen.asend(None) + + # Generators cannot catch generator exit exceptions (which are injected by + # close). This will result in a runtime error. + with pytest.raises(RuntimeError): + await gen.aclose() + + if is_pypy: + gen = None + gc.collect() + + event_loop.run_until_complete(_test()) + + +@validate_transaction_metrics( + "test_async_generator_time_excludes_creation_time", + background_task=True, + scoped_metrics=[("Function/agen", 1)], + rollup_metrics=[("Function/agen", 1)], +) +def test_async_generator_time_excludes_creation_time(event_loop): + @function_trace(name="agen") + async def agen(): + yield + + metrics = [] + full_metrics = {} + + @capture_transaction_metrics(metrics, full_metrics) + @background_task(name="test_async_generator_time_excludes_creation_time") + def _test_async_generator_time_excludes_creation_time(): + async def _test(): + gen = agen() + time.sleep(0.1) + async for _ in gen: + pass + + event_loop.run_until_complete(_test()) + _test_async_generator_time_excludes_creation_time() + + # check that the trace does not include the time between creation and + # consumption + assert full_metrics[("Function/agen", "")].total_call_time < 0.1 + + +@pytest.mark.parametrize("nr_transaction", [True, False]) +def test_incomplete_async_generator(event_loop, nr_transaction): + @function_trace(name="agen") + async def agen(): + for _ in range(5): + yield + + def _test_incomplete_async_generator(): + async def _test(): + c = agen() + + async for _ in c: + break + + if is_pypy: + # pypy is not guaranteed to delete the coroutine when it goes out + # of scope. This code "helps" pypy along. The test above is really + # just to verify that incomplete coroutines will "eventually" be + # cleaned up. In pypy, unfortunately that means it may not be + # reported all the time. A customer would be expected to call gc + # directly; however, they already have to handle this case since + # incomplete generators are well documented as having problems with + # pypy's gc. + + # See: + # http://doc.pypy.org/en/latest/cpython_differences.html#differences-related-to-garbage-collection-strategies + # https://bitbucket.org/pypy/pypy/issues/736 + del c + import gc + + gc.collect() + + if nr_transaction: + _test = background_task(name="test_incomplete_coroutine")(_test) + + event_loop.run_until_complete(_test()) + + if nr_transaction: + _test_incomplete_async_generator = validate_transaction_metrics( + "test_incomplete_coroutine", + background_task=True, + scoped_metrics=[("Function/agen", 1)], + rollup_metrics=[("Function/agen", 1)], + )(_test_incomplete_async_generator) + + _test_incomplete_async_generator() diff --git a/tests/agent_features/test_async_generator_trace.py b/tests/agent_features/test_async_generator_trace.py new file mode 100644 index 0000000000..208cf1588a --- /dev/null +++ b/tests/agent_features/test_async_generator_trace.py @@ -0,0 +1,19 @@ +# 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 sys + +# Async Generators were introduced in Python 3.6, but some APIs weren't completely stable until Python 3.7. +if sys.version_info >= (3, 7): + from _test_async_generator_trace import * # NOQA From 5693dd2f3ca2c23bc170b1e2cd9ea87862d9d80f Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Wed, 26 Jul 2023 14:27:01 -0700 Subject: [PATCH 14/31] Allow better parallelization in firestore tests --- tests/datastore_firestore/test_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/datastore_firestore/test_client.py b/tests/datastore_firestore/test_client.py index c4bd872a08..e7c731f4a1 100644 --- a/tests/datastore_firestore/test_client.py +++ b/tests/datastore_firestore/test_client.py @@ -32,7 +32,7 @@ def existing_document(collection, reset_firestore): def _exercise_client(client, collection, existing_document): - assert len([_ for _ in client.collections()]) == 1 + assert len([_ for _ in client.collections()]) >= 1 doc = [_ for _ in client.get_all([existing_document])][0] assert doc.to_dict()["x"] == 1 From c857358cc89d064fa7dddb5a6a0f2069496db708 Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Wed, 26 Jul 2023 15:53:21 -0700 Subject: [PATCH 15/31] Fix issue in async generator wrapper --- newrelic/common/async_wrapper.py | 4 +-- .../_test_async_generator_trace.py | 31 +++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/newrelic/common/async_wrapper.py b/newrelic/common/async_wrapper.py index 29d7809d5e..a2fd22f02c 100644 --- a/newrelic/common/async_wrapper.py +++ b/newrelic/common/async_wrapper.py @@ -93,7 +93,7 @@ async def wrapper(*args, **kwargs): with trace: while True: try: - g.asend(value).send(None) + yielded = await g.asend(value) except StopAsyncIteration as e: # The underlying async generator has finished, return propagates a new StopAsyncIteration return @@ -107,7 +107,7 @@ async def wrapper(*args, **kwargs): # An exception was thrown with .athrow(), propagate to the original async generator. # Return value logic must be identical to .asend() try: - g.athrow(type(e), e).send(None) + value = yield await g.athrow(type(e), e) except StopAsyncIteration as e: # The underlying async generator has finished, return propagates a new StopAsyncIteration return diff --git a/tests/agent_features/_test_async_generator_trace.py b/tests/agent_features/_test_async_generator_trace.py index 11efef1112..8e7703a0db 100644 --- a/tests/agent_features/_test_async_generator_trace.py +++ b/tests/agent_features/_test_async_generator_trace.py @@ -260,6 +260,37 @@ async def _test(): assert full_metrics[key].total_exclusive_call_time < 0.2 +@validate_transaction_metrics( + "test_asend_receives_a_value", + background_task=True, + scoped_metrics=[("Function/agen", 1)], + rollup_metrics=[("Function/agen", 1)], +) +def test_asend_receives_a_value(event_loop): + _received = [] + @function_trace(name="agen") + async def agen(): + value = yield + _received.append(value) + yield value + + @background_task(name="test_asend_receives_a_value") + async def _test(): + gen = agen() + + # kickstart the coroutine + await anext(gen) + + assert await gen.asend("foobar") == "foobar" + assert _received and _received[0] == "foobar" + + # finish consumption of the coroutine if necessary + async for _ in gen: + pass + + event_loop.run_until_complete(_test()) + + @validate_transaction_metrics( "test_athrow_yields_a_value", background_task=True, From c49a1cf0b079c53f61192de589efa32044712b58 Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Wed, 26 Jul 2023 15:54:13 -0700 Subject: [PATCH 16/31] 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]) From 7851baf92ece9d7aa85c0286b32aa8249d3b2191 Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Wed, 26 Jul 2023 15:58:12 -0700 Subject: [PATCH 17/31] Squashed commit of the following: commit 9d411e00e37476be4ce0c40c7e64e71c4a09cfc6 Author: Tim Pansino Date: Wed Jul 26 15:57:39 2023 -0700 Clean out unnecessary code commit cb550bad9bb9e15edfdcef5dd361022448e0348f Author: Tim Pansino Date: Wed Jul 26 14:27:01 2023 -0700 Allow better parallelization in firestore tests --- tests/datastore_firestore/test_client.py | 10 +--------- tests/datastore_firestore/test_collections.py | 9 +-------- tests/datastore_firestore/test_documents.py | 9 +-------- tests/datastore_firestore/test_query.py | 14 ++------------ tests/datastore_firestore/test_transaction.py | 3 --- 5 files changed, 5 insertions(+), 40 deletions(-) diff --git a/tests/datastore_firestore/test_client.py b/tests/datastore_firestore/test_client.py index e7c731f4a1..8409091ecb 100644 --- a/tests/datastore_firestore/test_client.py +++ b/tests/datastore_firestore/test_client.py @@ -13,9 +13,6 @@ # 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 ( @@ -48,6 +45,7 @@ def test_firestore_client(client, collection, existing_document): ("Datastore/allOther", 2), ] + @validate_database_duration() @validate_transaction_metrics( "test_firestore_client", scoped_metrics=_test_scoped_metrics, @@ -68,9 +66,3 @@ def test_firestore_client_generators(client, collection, assert_trace_for_genera 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_client(client, collection, existing_document) diff --git a/tests/datastore_firestore/test_collections.py b/tests/datastore_firestore/test_collections.py index ae2b48f04f..997675c9e3 100644 --- a/tests/datastore_firestore/test_collections.py +++ b/tests/datastore_firestore/test_collections.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from newrelic.api.time_trace import current_trace -from newrelic.api.datastore_trace import DatastoreTrace 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 ( @@ -50,6 +48,7 @@ def test_firestore_collections(collection): ("Datastore/all", 5), ("Datastore/allOther", 5), ] + @validate_database_duration() @validate_transaction_metrics( "test_firestore_collections", scoped_metrics=_test_scoped_metrics, @@ -71,9 +70,3 @@ def test_firestore_collections_generators(collection, assert_trace_for_generator assert_trace_for_generator(collection.stream) assert_trace_for_generator(collection.list_documents) - - -@validate_database_duration() -@background_task() -def test_firestore_collections_db_duration(collection): - _exercise_collections(collection) diff --git a/tests/datastore_firestore/test_documents.py b/tests/datastore_firestore/test_documents.py index 67ad42dd64..a648b3f92c 100644 --- a/tests/datastore_firestore/test_documents.py +++ b/tests/datastore_firestore/test_documents.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from newrelic.api.time_trace import current_trace -from newrelic.api.datastore_trace import DatastoreTrace 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 ( @@ -59,6 +57,7 @@ def test_firestore_documents(collection): ("Datastore/all", 7), ("Datastore/allOther", 7), ] + @validate_database_duration() @validate_transaction_metrics( "test_firestore_documents", scoped_metrics=_test_scoped_metrics, @@ -81,9 +80,3 @@ def test_firestore_documents_generators(collection, assert_trace_for_generator): assert len([_ for _ in subcollection_doc.collections()]) == 2 assert_trace_for_generator(subcollection_doc.collections) - - -@validate_database_duration() -@background_task() -def test_firestore_documents_db_duration(collection): - _exercise_documents(collection) diff --git a/tests/datastore_firestore/test_query.py b/tests/datastore_firestore/test_query.py index 128042a99e..26426695c3 100644 --- a/tests/datastore_firestore/test_query.py +++ b/tests/datastore_firestore/test_query.py @@ -47,6 +47,7 @@ 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, @@ -65,12 +66,6 @@ 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): @@ -91,6 +86,7 @@ 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, @@ -108,9 +104,3 @@ 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) diff --git a/tests/datastore_firestore/test_transaction.py b/tests/datastore_firestore/test_transaction.py index 679821dfc4..c38d19f6a5 100644 --- a/tests/datastore_firestore/test_transaction.py +++ b/tests/datastore_firestore/test_transaction.py @@ -13,9 +13,6 @@ # 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 ( From 3fb6a6cd32c3a7fcfa1874aeb68e2cf3c23ea85c Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Wed, 26 Jul 2023 16:11:17 -0700 Subject: [PATCH 18/31] Add async collection instrumentation --- newrelic/config.py | 5 ++ newrelic/hooks/datastore_firestore.py | 14 ++++ .../test_async_collections.py | 72 +++++++++++++++++++ 3 files changed, 91 insertions(+) create mode 100644 tests/datastore_firestore/test_async_collections.py diff --git a/newrelic/config.py b/newrelic/config.py index ca21a97339..d271157f15 100644 --- a/newrelic/config.py +++ b/newrelic/config.py @@ -2294,6 +2294,11 @@ def _process_module_builtin_defaults(): "newrelic.hooks.datastore_firestore", "instrument_google_cloud_firestore_v1_collection", ) + _process_module_definition( + "google.cloud.firestore_v1.async_collection", + "newrelic.hooks.datastore_firestore", + "instrument_google_cloud_firestore_v1_async_collection", + ) _process_module_definition( "google.cloud.firestore_v1.query", "newrelic.hooks.datastore_firestore", diff --git a/newrelic/hooks/datastore_firestore.py b/newrelic/hooks/datastore_firestore.py index e534d98ecc..099752021f 100644 --- a/newrelic/hooks/datastore_firestore.py +++ b/newrelic/hooks/datastore_firestore.py @@ -100,6 +100,20 @@ def instrument_google_cloud_firestore_v1_collection(module): wrap_generator_method(module, "CollectionReference", method, target=_get_object_id) +def instrument_google_cloud_firestore_v1_async_collection(module): + if hasattr(module, "AsyncCollectionReference"): + class_ = module.AsyncCollectionReference + for method in ("add", "get"): + if hasattr(class_, method): + wrap_datastore_trace( + module, "AsyncCollectionReference.%s" % method, product="Firestore", target=_get_object_id, operation=method + ) + + for method in ("stream", "list_documents"): + if hasattr(class_, method): + wrap_async_generator_method(module, "AsyncCollectionReference", method, target=_get_object_id) + + def instrument_google_cloud_firestore_v1_document(module): if hasattr(module, "DocumentReference"): class_ = module.DocumentReference diff --git a/tests/datastore_firestore/test_async_collections.py b/tests/datastore_firestore/test_async_collections.py new file mode 100644 index 0000000000..e5e3e2bae9 --- /dev/null +++ b/tests/datastore_firestore/test_async_collections.py @@ -0,0 +1,72 @@ +# 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. + +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, +) + + +async def _exercise_collections(async_collection): + async_collection.document("DoesNotExist") + await async_collection.add({"capital": "Rome", "currency": "Euro", "language": "Italian"}, "Italy") + await async_collection.add({"capital": "Mexico City", "currency": "Peso", "language": "Spanish"}, "Mexico") + + documents_get = await async_collection.get() + assert len(documents_get) == 2 + documents_stream = [_ async for _ in async_collection.stream()] + assert len(documents_stream) == 2 + documents_list = [_ async for _ in async_collection.list_documents()] + assert len(documents_list) == 2 + + +def test_firestore_async_collections(loop, async_collection): + _test_scoped_metrics = [ + ("Datastore/statement/Firestore/%s/stream" % async_collection.id, 1), + ("Datastore/statement/Firestore/%s/get" % async_collection.id, 1), + ("Datastore/statement/Firestore/%s/list_documents" % async_collection.id, 1), + ("Datastore/statement/Firestore/%s/add" % async_collection.id, 2), + ] + + _test_rollup_metrics = [ + ("Datastore/operation/Firestore/add", 2), + ("Datastore/operation/Firestore/get", 1), + ("Datastore/operation/Firestore/stream", 1), + ("Datastore/operation/Firestore/list_documents", 1), + ("Datastore/all", 5), + ("Datastore/allOther", 5), + ] + @validate_database_duration() + @validate_transaction_metrics( + "test_firestore_async_collections", + scoped_metrics=_test_scoped_metrics, + rollup_metrics=_test_rollup_metrics, + background_task=True, + ) + @background_task(name="test_firestore_async_collections") + def _test(): + loop.run_until_complete(_exercise_collections(async_collection)) + + _test() + + +@background_task() +def test_firestore_async_collections_generators(loop, collection, async_collection, assert_trace_for_async_generator): + collection.add({}) + collection.add({}) + assert len([_ for _ in collection.list_documents()]) == 2 + + assert_trace_for_async_generator(async_collection.stream) + assert_trace_for_async_generator(async_collection.list_documents) From aab244bcb45cc5cb6cb2be870a8182da95128582 Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Wed, 26 Jul 2023 16:20:58 -0700 Subject: [PATCH 19/31] Add async document instrumentation --- newrelic/config.py | 5 ++ newrelic/hooks/datastore_firestore.py | 14 ++++ .../test_async_documents.py | 84 +++++++++++++++++++ 3 files changed, 103 insertions(+) create mode 100644 tests/datastore_firestore/test_async_documents.py diff --git a/newrelic/config.py b/newrelic/config.py index d271157f15..604fe51061 100644 --- a/newrelic/config.py +++ b/newrelic/config.py @@ -2289,6 +2289,11 @@ def _process_module_builtin_defaults(): "newrelic.hooks.datastore_firestore", "instrument_google_cloud_firestore_v1_document", ) + _process_module_definition( + "google.cloud.firestore_v1.async_document", + "newrelic.hooks.datastore_firestore", + "instrument_google_cloud_firestore_v1_async_document", + ) _process_module_definition( "google.cloud.firestore_v1.collection", "newrelic.hooks.datastore_firestore", diff --git a/newrelic/hooks/datastore_firestore.py b/newrelic/hooks/datastore_firestore.py index 099752021f..c4d7aba022 100644 --- a/newrelic/hooks/datastore_firestore.py +++ b/newrelic/hooks/datastore_firestore.py @@ -128,6 +128,20 @@ def instrument_google_cloud_firestore_v1_document(module): wrap_generator_method(module, "DocumentReference", method, target=_get_object_id) +def instrument_google_cloud_firestore_v1_async_document(module): + if hasattr(module, "AsyncDocumentReference"): + class_ = module.AsyncDocumentReference + for method in ("create", "delete", "get", "set", "update"): + if hasattr(class_, method): + wrap_datastore_trace( + module, "AsyncDocumentReference.%s" % method, product="Firestore", target=_get_object_id, operation=method + ) + + for method in ("collections",): + if hasattr(class_, method): + wrap_async_generator_method(module, "AsyncDocumentReference", method, target=_get_object_id) + + def instrument_google_cloud_firestore_v1_query(module): if hasattr(module, "Query"): class_ = module.Query diff --git a/tests/datastore_firestore/test_async_documents.py b/tests/datastore_firestore/test_async_documents.py new file mode 100644 index 0000000000..2df6940add --- /dev/null +++ b/tests/datastore_firestore/test_async_documents.py @@ -0,0 +1,84 @@ +# 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. + +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, +) + + +async def _exercise_async_documents(async_collection): + italy_doc = async_collection.document("Italy") + await italy_doc.set({"capital": "Rome", "currency": "Euro", "language": "Italian"}) + await italy_doc.get() + italian_cities = italy_doc.collection("cities") + await italian_cities.add({"capital": "Rome"}) + retrieved_coll = [_ async for _ in italy_doc.collections()] + assert len(retrieved_coll) == 1 + + usa_doc = async_collection.document("USA") + await usa_doc.create({"capital": "Washington D.C.", "currency": "Dollar", "language": "English"}) + await usa_doc.update({"president": "Joe Biden"}) + + await async_collection.document("USA").delete() + + +def test_firestore_async_documents(loop, async_collection): + _test_scoped_metrics = [ + ("Datastore/statement/Firestore/Italy/set", 1), + ("Datastore/statement/Firestore/Italy/get", 1), + ("Datastore/statement/Firestore/Italy/collections", 1), + ("Datastore/statement/Firestore/cities/add", 1), + ("Datastore/statement/Firestore/USA/create", 1), + ("Datastore/statement/Firestore/USA/update", 1), + ("Datastore/statement/Firestore/USA/delete", 1), + ] + + _test_rollup_metrics = [ + ("Datastore/operation/Firestore/set", 1), + ("Datastore/operation/Firestore/get", 1), + ("Datastore/operation/Firestore/add", 1), + ("Datastore/operation/Firestore/collections", 1), + ("Datastore/operation/Firestore/create", 1), + ("Datastore/operation/Firestore/update", 1), + ("Datastore/operation/Firestore/delete", 1), + ("Datastore/all", 7), + ("Datastore/allOther", 7), + ] + @validate_database_duration() + @validate_transaction_metrics( + "test_firestore_async_documents", + scoped_metrics=_test_scoped_metrics, + rollup_metrics=_test_rollup_metrics, + background_task=True, + ) + @background_task(name="test_firestore_async_documents") + def _test(): + loop.run_until_complete(_exercise_async_documents(async_collection)) + + _test() + + +@background_task() +def test_firestore_async_documents_generators(collection, async_collection, assert_trace_for_async_generator): + subcollection_doc = collection.document("SubCollections") + subcollection_doc.set({}) + subcollection_doc.collection("collection1").add({}) + subcollection_doc.collection("collection2").add({}) + assert len([_ for _ in subcollection_doc.collections()]) == 2 + + async_subcollection = async_collection.document(subcollection_doc.id) + + assert_trace_for_async_generator(async_subcollection.collections) From c392e78fba4cde9334dc7e1b40a7a6531e9b672c Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Wed, 26 Jul 2023 16:36:03 -0700 Subject: [PATCH 20/31] Async Query instrumentation --- newrelic/config.py | 10 ++ newrelic/hooks/datastore_firestore.py | 28 +++++ tests/datastore_firestore/test_async_query.py | 106 ++++++++++++++++++ 3 files changed, 144 insertions(+) create mode 100644 tests/datastore_firestore/test_async_query.py diff --git a/newrelic/config.py b/newrelic/config.py index 604fe51061..1921bcc9bc 100644 --- a/newrelic/config.py +++ b/newrelic/config.py @@ -2309,11 +2309,21 @@ def _process_module_builtin_defaults(): "newrelic.hooks.datastore_firestore", "instrument_google_cloud_firestore_v1_query", ) + _process_module_definition( + "google.cloud.firestore_v1.async_query", + "newrelic.hooks.datastore_firestore", + "instrument_google_cloud_firestore_v1_async_query", + ) _process_module_definition( "google.cloud.firestore_v1.aggregation", "newrelic.hooks.datastore_firestore", "instrument_google_cloud_firestore_v1_aggregation", ) + _process_module_definition( + "google.cloud.firestore_v1.async_aggregation", + "newrelic.hooks.datastore_firestore", + "instrument_google_cloud_firestore_v1_async_aggregation", + ) _process_module_definition( "google.cloud.firestore_v1.batch", "newrelic.hooks.datastore_firestore", diff --git a/newrelic/hooks/datastore_firestore.py b/newrelic/hooks/datastore_firestore.py index c4d7aba022..b9993350a8 100644 --- a/newrelic/hooks/datastore_firestore.py +++ b/newrelic/hooks/datastore_firestore.py @@ -156,6 +156,20 @@ def instrument_google_cloud_firestore_v1_query(module): wrap_generator_method(module, "Query", method, target=_get_parent_id) +def instrument_google_cloud_firestore_v1_async_query(module): + if hasattr(module, "AsyncQuery"): + class_ = module.AsyncQuery + for method in ("get",): + if hasattr(class_, method): + wrap_datastore_trace( + module, "AsyncQuery.%s" % method, product="Firestore", target=_get_parent_id, operation=method + ) + + for method in ("stream",): + if hasattr(class_, method): + wrap_async_generator_method(module, "AsyncQuery", method, target=_get_parent_id) + + def instrument_google_cloud_firestore_v1_aggregation(module): if hasattr(module, "AggregationQuery"): class_ = module.AggregationQuery @@ -170,6 +184,20 @@ def instrument_google_cloud_firestore_v1_aggregation(module): wrap_generator_method(module, "AggregationQuery", method, target=_get_collection_ref_id) +def instrument_google_cloud_firestore_v1_async_aggregation(module): + if hasattr(module, "AsyncAggregationQuery"): + class_ = module.AsyncAggregationQuery + for method in ("get",): + if hasattr(class_, method): + wrap_datastore_trace( + module, "AsyncAggregationQuery.%s" % method, product="Firestore", target=_get_collection_ref_id, operation=method + ) + + for method in ("stream",): + if hasattr(class_, method): + wrap_async_generator_method(module, "AsyncAggregationQuery", method, target=_get_collection_ref_id) + + def instrument_google_cloud_firestore_v1_batch(module): if hasattr(module, "WriteBatch"): class_ = module.WriteBatch diff --git a/tests/datastore_firestore/test_async_query.py b/tests/datastore_firestore/test_async_query.py new file mode 100644 index 0000000000..02efa87c69 --- /dev/null +++ b/tests/datastore_firestore/test_async_query.py @@ -0,0 +1,106 @@ +# 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(autouse=True) +def sample_data(collection, reset_firestore): + # reset_firestore must be run before, not after this fixture + for x in range(1, 6): + collection.add({"x": x}) + +# ===== AsyncQuery ===== + +async def _exercise_async_query(async_collection): + async_query = async_collection.select("x").limit(10).order_by("x").where(field_path="x", op_string="<=", value=3) + assert len(await async_query.get()) == 3 + assert len([_ async for _ in async_query.stream()]) == 3 + + +def test_firestore_async_query(loop, async_collection): + _test_scoped_metrics = [ + ("Datastore/statement/Firestore/%s/stream" % async_collection.id, 1), + ("Datastore/statement/Firestore/%s/get" % async_collection.id, 1), + ] + + _test_rollup_metrics = [ + ("Datastore/operation/Firestore/get", 1), + ("Datastore/operation/Firestore/stream", 1), + ("Datastore/all", 2), + ("Datastore/allOther", 2), + ] + # @validate_database_duration() + @validate_transaction_metrics( + "test_firestore_async_query", + scoped_metrics=_test_scoped_metrics, + rollup_metrics=_test_rollup_metrics, + background_task=True, + ) + @background_task(name="test_firestore_async_query") + def _test(): + loop.run_until_complete(_exercise_async_query(async_collection)) + + _test() + + +@background_task() +def test_firestore_async_query_generators(async_collection, assert_trace_for_async_generator): + async_query = async_collection.select("x").where(field_path="x", op_string="<=", value=3) + assert_trace_for_async_generator(async_query.stream) + +# ===== AsyncAggregationQuery ===== + +async def _exercise_async_aggregation_query(async_collection): + async_aggregation_query = async_collection.select("x").where(field_path="x", op_string="<=", value=3).count() + assert (await async_aggregation_query.get())[0][0].value == 3 + assert [_ async for _ in async_aggregation_query.stream()][0][0].value == 3 + + +def test_firestore_async_aggregation_query(loop, async_collection): + _test_scoped_metrics = [ + ("Datastore/statement/Firestore/%s/stream" % async_collection.id, 1), + ("Datastore/statement/Firestore/%s/get" % async_collection.id, 1), + ] + + _test_rollup_metrics = [ + ("Datastore/operation/Firestore/get", 1), + ("Datastore/operation/Firestore/stream", 1), + ("Datastore/all", 2), + ("Datastore/allOther", 2), + ] + @validate_database_duration() + @validate_transaction_metrics( + "test_firestore_async_aggregation_query", + scoped_metrics=_test_scoped_metrics, + rollup_metrics=_test_rollup_metrics, + background_task=True, + ) + @background_task(name="test_firestore_async_aggregation_query") + def _test(): + loop.run_until_complete(_exercise_async_aggregation_query(async_collection)) + + _test() + + +@background_task() +def test_firestore_async_aggregation_query_generators(async_collection, assert_trace_for_async_generator): + async_aggregation_query = async_collection.select("x").where(field_path="x", op_string="<=", value=3).count() + assert_trace_for_async_generator(async_aggregation_query.stream) From 6b7fc79b2466bc729d07878193643f989f95bf04 Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Wed, 26 Jul 2023 16:56:04 -0700 Subject: [PATCH 21/31] Add async batch instrumentation --- newrelic/config.py | 5 ++ newrelic/hooks/datastore_firestore.py | 10 ++++ .../test_async_batching.py | 50 +++++++++++++++++++ 3 files changed, 65 insertions(+) create mode 100644 tests/datastore_firestore/test_async_batching.py diff --git a/newrelic/config.py b/newrelic/config.py index 1921bcc9bc..042f224611 100644 --- a/newrelic/config.py +++ b/newrelic/config.py @@ -2329,6 +2329,11 @@ def _process_module_builtin_defaults(): "newrelic.hooks.datastore_firestore", "instrument_google_cloud_firestore_v1_batch", ) + _process_module_definition( + "google.cloud.firestore_v1.async_batch", + "newrelic.hooks.datastore_firestore", + "instrument_google_cloud_firestore_v1_async_batch", + ) _process_module_definition( "google.cloud.firestore_v1.bulk_batch", "newrelic.hooks.datastore_firestore", diff --git a/newrelic/hooks/datastore_firestore.py b/newrelic/hooks/datastore_firestore.py index b9993350a8..5dc8f62869 100644 --- a/newrelic/hooks/datastore_firestore.py +++ b/newrelic/hooks/datastore_firestore.py @@ -208,6 +208,16 @@ def instrument_google_cloud_firestore_v1_batch(module): ) +def instrument_google_cloud_firestore_v1_async_batch(module): + if hasattr(module, "AsyncWriteBatch"): + class_ = module.AsyncWriteBatch + for method in ("commit",): + if hasattr(class_, method): + wrap_datastore_trace( + module, "AsyncWriteBatch.%s" % method, product="Firestore", target=None, operation=method + ) + + def instrument_google_cloud_firestore_v1_bulk_batch(module): if hasattr(module, "BulkWriteBatch"): class_ = module.BulkWriteBatch diff --git a/tests/datastore_firestore/test_async_batching.py b/tests/datastore_firestore/test_async_batching.py new file mode 100644 index 0000000000..0d51e99524 --- /dev/null +++ b/tests/datastore_firestore/test_async_batching.py @@ -0,0 +1,50 @@ +# 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. + +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, +) + +async def _exercise_async_write_batch(async_client, async_collection): + docs = [async_collection.document(str(x)) for x in range(1, 4)] + async_batch = async_client.batch() + for doc in docs: + async_batch.set(doc, {}) + + await async_batch.commit() + + +def test_firestore_async_write_batch(loop, async_client, async_collection): + _test_scoped_metrics = [ + ("Datastore/operation/Firestore/commit", 1), + ] + + _test_rollup_metrics = [ + ("Datastore/all", 1), + ("Datastore/allOther", 1), + ] + @validate_database_duration() + @validate_transaction_metrics( + "test_firestore_async_write_batch", + scoped_metrics=_test_scoped_metrics, + rollup_metrics=_test_rollup_metrics, + background_task=True, + ) + @background_task(name="test_firestore_async_write_batch") + def _test(): + loop.run_until_complete(_exercise_async_write_batch(async_client, async_collection)) + + _test() From e04ec6f7959001558951bb0b716bf7c2f9062380 Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Thu, 27 Jul 2023 11:55:44 -0700 Subject: [PATCH 22/31] Add instrumentation for AsyncTransaction --- newrelic/config.py | 5 + newrelic/hooks/datastore_firestore.py | 11 ++ .../test_async_transaction.py | 122 ++++++++++++++++++ 3 files changed, 138 insertions(+) create mode 100644 tests/datastore_firestore/test_async_transaction.py diff --git a/newrelic/config.py b/newrelic/config.py index 042f224611..03afd456e5 100644 --- a/newrelic/config.py +++ b/newrelic/config.py @@ -2344,6 +2344,11 @@ def _process_module_builtin_defaults(): "newrelic.hooks.datastore_firestore", "instrument_google_cloud_firestore_v1_transaction", ) + _process_module_definition( + "google.cloud.firestore_v1.async_transaction", + "newrelic.hooks.datastore_firestore", + "instrument_google_cloud_firestore_v1_async_transaction", + ) _process_module_definition( "ariadne.asgi", diff --git a/newrelic/hooks/datastore_firestore.py b/newrelic/hooks/datastore_firestore.py index 5dc8f62869..46b9eb0bc6 100644 --- a/newrelic/hooks/datastore_firestore.py +++ b/newrelic/hooks/datastore_firestore.py @@ -237,3 +237,14 @@ def instrument_google_cloud_firestore_v1_transaction(module): wrap_datastore_trace( module, "Transaction.%s" % method, product="Firestore", target=None, operation=operation ) + + +def instrument_google_cloud_firestore_v1_async_transaction(module): + if hasattr(module, "AsyncTransaction"): + class_ = module.AsyncTransaction + for method in ("_commit", "_rollback"): + if hasattr(class_, method): + operation = method[1:] # Trim leading underscore + wrap_datastore_trace( + module, "AsyncTransaction.%s" % method, product="Firestore", target=None, operation=operation + ) diff --git a/tests/datastore_firestore/test_async_transaction.py b/tests/datastore_firestore/test_async_transaction.py new file mode 100644 index 0000000000..4916669b89 --- /dev/null +++ b/tests/datastore_firestore/test_async_transaction.py @@ -0,0 +1,122 @@ +# 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(autouse=True) +def sample_data(collection, reset_firestore): + # reset_firestore must be run before, not after this fixture + for x in range(1, 4): + collection.add({"x": x}, "doc%d" % x) + + +async def _exercise_async_transaction_commit(async_client, async_collection): + from google.cloud.firestore_v1.async_transaction import async_transactional + + @async_transactional + async def _exercise(async_transaction): + # get a DocumentReference + with pytest.raises(TypeError): # get is currently broken. It attempts to await an async_generator instead of consuming it. + [_ async for _ in async_transaction.get(async_collection.document("doc1"))] + + # get a Query + with pytest.raises(TypeError): # get is currently broken. It attempts to await an async_generator instead of consuming it. + async_query = async_collection.select("x").where(field_path="x", op_string=">", value=2) + assert len([_ async for _ in async_transaction.get(async_query)]) == 1 + + # get_all on a list of DocumentReferences + with pytest.raises(TypeError): # get_all is currently broken. It attempts to await an async_generator instead of consuming it. + all_docs = async_transaction.get_all([async_collection.document("doc%d" % x) for x in range(1, 4)]) + assert len([_ async for _ in all_docs]) == 3 + + # set and delete methods + async_transaction.set(async_collection.document("doc2"), {"x": 0}) + async_transaction.delete(async_collection.document("doc3")) + + await _exercise(async_client.transaction()) + assert len([_ async for _ in async_collection.list_documents()]) == 2 + + +async def _exercise_async_transaction_rollback(async_client, async_collection): + from google.cloud.firestore_v1.async_transaction import async_transactional + + @async_transactional + async def _exercise(async_transaction): + # set and delete methods + async_transaction.set(async_collection.document("doc2"), {"x": 99}) + async_transaction.delete(async_collection.document("doc1")) + raise RuntimeError() + + with pytest.raises(RuntimeError): + await _exercise(async_client.transaction()) + assert len([_ async for _ in async_collection.list_documents()]) == 3 + + +def test_firestore_async_transaction_commit(loop, async_client, async_collection): + _test_scoped_metrics = [ + ("Datastore/operation/Firestore/commit", 1), + # ("Datastore/operation/Firestore/get_all", 2), + # ("Datastore/statement/Firestore/%s/stream" % async_collection.id, 1), + ("Datastore/statement/Firestore/%s/list_documents" % async_collection.id, 1), + ] + + _test_rollup_metrics = [ + # ("Datastore/operation/Firestore/stream", 1), + ("Datastore/operation/Firestore/list_documents", 1), + ("Datastore/all", 2), # Should be 5 if not for broken APIs + ("Datastore/allOther", 2), + ] + @validate_database_duration() + @validate_transaction_metrics( + "test_firestore_async_transaction", + scoped_metrics=_test_scoped_metrics, + rollup_metrics=_test_rollup_metrics, + background_task=True, + ) + @background_task(name="test_firestore_async_transaction") + def _test(): + loop.run_until_complete(_exercise_async_transaction_commit(async_client, async_collection)) + + _test() + + +def test_firestore_async_transaction_rollback(loop, async_client, async_collection): + _test_scoped_metrics = [ + ("Datastore/operation/Firestore/rollback", 1), + ("Datastore/statement/Firestore/%s/list_documents" % async_collection.id, 1), + ] + + _test_rollup_metrics = [ + ("Datastore/operation/Firestore/list_documents", 1), + ("Datastore/all", 2), + ("Datastore/allOther", 2), + ] + # @validate_database_duration() + @validate_transaction_metrics( + "test_firestore_async_transaction", + scoped_metrics=_test_scoped_metrics, + rollup_metrics=_test_rollup_metrics, + background_task=True, + ) + @background_task(name="test_firestore_async_transaction") + def _test(): + loop.run_until_complete(_exercise_async_transaction_rollback(async_client, async_collection)) + + _test() From 87fbe6203def1e583ff8fea58d4d6a0e70bfa606 Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Thu, 27 Jul 2023 13:00:37 -0700 Subject: [PATCH 23/31] Squashed commit of the following: commit c836f8f377f9391af86452e81b59f834330b18fb Author: TimPansino Date: Thu Jul 27 19:54:35 2023 +0000 [Mega-Linter] Apply linters fixes commit 02a55a11017fd27b45f06ab719a33917cf185aac Author: Tim Pansino Date: Thu Jul 27 12:46:46 2023 -0700 Add collection group instrumentation commit ab1f4ff5d2e88e6def42fa3c99c619f9673ce918 Author: Tim Pansino Date: Thu Jul 27 12:00:33 2023 -0700 Better parallelization safeguards commit fa5f39a2b037421cf017a062901c0ea1ec2b9723 Author: TimPansino Date: Wed Jul 26 22:59:11 2023 +0000 [Mega-Linter] Apply linters fixes commit 9d411e00e37476be4ce0c40c7e64e71c4a09cfc6 Author: Tim Pansino Date: Wed Jul 26 15:57:39 2023 -0700 Clean out unnecessary code commit cb550bad9bb9e15edfdcef5dd361022448e0348f Author: Tim Pansino Date: Wed Jul 26 14:27:01 2023 -0700 Allow better parallelization in firestore tests --- newrelic/hooks/datastore_firestore.py | 26 +++++++++-- tests/datastore_firestore/conftest.py | 45 ++++++++++--------- tests/datastore_firestore/test_batching.py | 14 ++++-- tests/datastore_firestore/test_client.py | 28 ++++++------ tests/datastore_firestore/test_collections.py | 11 +++-- tests/datastore_firestore/test_documents.py | 9 ++-- tests/datastore_firestore/test_query.py | 32 ++++++++++--- tests/datastore_firestore/test_transaction.py | 13 +++--- 8 files changed, 115 insertions(+), 63 deletions(-) diff --git a/newrelic/hooks/datastore_firestore.py b/newrelic/hooks/datastore_firestore.py index 46b9eb0bc6..ea92b4e5ec 100644 --- a/newrelic/hooks/datastore_firestore.py +++ b/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",): diff --git a/tests/datastore_firestore/conftest.py b/tests/datastore_firestore/conftest.py index 767c5c474a..95f2d6cc54 100644 --- a/tests/datastore_firestore/conftest.py +++ b/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): diff --git a/tests/datastore_firestore/test_batching.py b/tests/datastore_firestore/test_batching.py index 9c2d492fd4..c0f52530b2 100644 --- a/tests/datastore_firestore/test_batching.py +++ b/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", diff --git a/tests/datastore_firestore/test_client.py b/tests/datastore_firestore/test_client.py index 8409091ecb..d9d46ea206 100644 --- a/tests/datastore_firestore/test_client.py +++ b/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]) diff --git a/tests/datastore_firestore/test_collections.py b/tests/datastore_firestore/test_collections.py index 997675c9e3..49a97810af 100644 --- a/tests/datastore_firestore/test_collections.py +++ b/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) diff --git a/tests/datastore_firestore/test_documents.py b/tests/datastore_firestore/test_documents.py index a648b3f92c..7873cf024d 100644 --- a/tests/datastore_firestore/test_documents.py +++ b/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) diff --git a/tests/datastore_firestore/test_query.py b/tests/datastore_firestore/test_query.py index 26426695c3..7e8759bb6e 100644 --- a/tests/datastore_firestore/test_query.py +++ b/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) diff --git a/tests/datastore_firestore/test_transaction.py b/tests/datastore_firestore/test_transaction.py index c38d19f6a5..121fac67ab 100644 --- a/tests/datastore_firestore/test_transaction.py +++ b/tests/datastore_firestore/test_transaction.py @@ -12,17 +12,18 @@ # 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(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, 4): collection.add({"x": x}, "doc%d" % x) @@ -80,6 +81,7 @@ def test_firestore_transaction_commit(client, collection): ("Datastore/all", 5), ("Datastore/allOther", 5), ] + @validate_database_duration() @validate_transaction_metrics( "test_firestore_transaction", @@ -105,6 +107,7 @@ def test_firestore_transaction_rollback(client, collection): ("Datastore/all", 2), ("Datastore/allOther", 2), ] + @validate_database_duration() @validate_transaction_metrics( "test_firestore_transaction", From b6bc9a47f28f2da29a8c5bf77dd37273c83a3757 Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Thu, 27 Jul 2023 13:01:27 -0700 Subject: [PATCH 24/31] Remove reset_firestore --- tests/datastore_firestore/test_async_client.py | 3 +-- tests/datastore_firestore/test_async_query.py | 3 +-- tests/datastore_firestore/test_async_transaction.py | 3 +-- 3 files changed, 3 insertions(+), 6 deletions(-) diff --git a/tests/datastore_firestore/test_async_client.py b/tests/datastore_firestore/test_async_client.py index 6496e6a357..c8e7b1dccf 100644 --- a/tests/datastore_firestore/test_async_client.py +++ b/tests/datastore_firestore/test_async_client.py @@ -21,8 +21,7 @@ @pytest.fixture() -def existing_document(collection, reset_firestore): - # reset_firestore must be run before, not after this fixture +def existing_document(collection): doc = collection.document("document") doc.set({"x": 1}) return doc diff --git a/tests/datastore_firestore/test_async_query.py b/tests/datastore_firestore/test_async_query.py index 02efa87c69..3319e73a53 100644 --- a/tests/datastore_firestore/test_async_query.py +++ b/tests/datastore_firestore/test_async_query.py @@ -22,8 +22,7 @@ @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}) diff --git a/tests/datastore_firestore/test_async_transaction.py b/tests/datastore_firestore/test_async_transaction.py index 4916669b89..fd9130a0ee 100644 --- a/tests/datastore_firestore/test_async_transaction.py +++ b/tests/datastore_firestore/test_async_transaction.py @@ -21,8 +21,7 @@ @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, 4): collection.add({"x": x}, "doc%d" % x) From 9266924d8ef965852dec415973d7d9699031f011 Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Thu, 27 Jul 2023 13:04:19 -0700 Subject: [PATCH 25/31] Re-merge of test_query --- tests/datastore_firestore/test_query.py | 87 ++++++++++++++++++++++--- 1 file changed, 79 insertions(+), 8 deletions(-) diff --git a/tests/datastore_firestore/test_query.py b/tests/datastore_firestore/test_query.py index 7e8759bb6e..144b582f49 100644 --- a/tests/datastore_firestore/test_query.py +++ b/tests/datastore_firestore/test_query.py @@ -54,6 +54,8 @@ 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, @@ -73,11 +75,6 @@ def test_firestore_query_generators(collection, assert_trace_for_generator): assert_trace_for_generator(query.stream) -@validate_database_duration() -@background_task() -def test_firestore_query_db_duration(collection): - _exercise_query(collection) - # ===== AggregationQuery ===== @@ -99,6 +96,8 @@ 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, @@ -118,7 +117,79 @@ def test_firestore_aggregation_query_generators(collection, assert_trace_for_gen assert_trace_for_generator(aggregation_query.stream) -@validate_database_duration() +# ===== CollectionGroup ===== + + +@pytest.fixture() +def patch_partition_queries(monkeypatch, client, collection, sample_data): + """ + Partitioning is not implemented in the Firestore emulator. + + Ordinarily this method would return a generator of Cursor objects. Each Cursor must point at a valid document path. + To test this, we can patch the RPC to return 1 Cursor which is pointed at any document available. + The get_partitions will take that and make 2 QueryPartition objects out of it, which should be enough to ensure + we can exercise the generator's tracing. + """ + from google.cloud.firestore_v1.types.document import Value + from google.cloud.firestore_v1.types.query import Cursor + + subcollection = collection.document("subcollection").collection("subcollection1") + documents = [d for d in subcollection.list_documents()] + + def mock_partition_query(*args, **kwargs): + yield Cursor(before=False, values=[Value(reference_value=documents[0].path)]) + + monkeypatch.setattr(client._firestore_api, "partition_query", mock_partition_query) + yield + + +def _exercise_collection_group(collection): + from google.cloud.firestore import CollectionGroup + + collection_group = CollectionGroup(collection) + assert len(collection_group.get()) + assert len([d for d in collection_group.stream()]) + + partitions = [p for p in collection_group.get_partitions(1)] + assert len(partitions) == 2 + documents = [] + while partitions: + documents.extend(partitions.pop().query().get()) + assert len(documents) == 6 + + +def test_firestore_collection_group(collection, patch_partition_queries): + _test_scoped_metrics = [ + ("Datastore/statement/Firestore/%s/get" % collection.id, 3), + ("Datastore/statement/Firestore/%s/stream" % collection.id, 1), + ("Datastore/statement/Firestore/%s/get_partitions" % collection.id, 1), + ] + + _test_rollup_metrics = [ + ("Datastore/operation/Firestore/get", 3), + ("Datastore/operation/Firestore/stream", 1), + ("Datastore/operation/Firestore/get_partitions", 1), + ("Datastore/all", 5), + ("Datastore/allOther", 5), + ] + + @validate_database_duration() + @validate_transaction_metrics( + "test_firestore_collection_group", + scoped_metrics=_test_scoped_metrics, + rollup_metrics=_test_rollup_metrics, + background_task=True, + ) + @background_task(name="test_firestore_collection_group") + def _test(): + _exercise_collection_group(collection) + + _test() + + @background_task() -def test_firestore_aggregation_query_db_duration(collection): - _exercise_aggregation_query(collection) +def test_firestore_collection_group_generators(collection, assert_trace_for_generator, patch_partition_queries): + from google.cloud.firestore import CollectionGroup + + collection_group = CollectionGroup(collection) + assert_trace_for_generator(collection_group.get_partitions, 1) From 5902515202f7c7985b787bd00bb66b7f89699e19 Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Thu, 27 Jul 2023 13:09:13 -0700 Subject: [PATCH 26/31] Use public API imports --- tests/datastore_firestore/test_async_transaction.py | 4 ++-- tests/datastore_firestore/test_transaction.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/datastore_firestore/test_async_transaction.py b/tests/datastore_firestore/test_async_transaction.py index fd9130a0ee..76667538b8 100644 --- a/tests/datastore_firestore/test_async_transaction.py +++ b/tests/datastore_firestore/test_async_transaction.py @@ -27,7 +27,7 @@ def sample_data(collection): async def _exercise_async_transaction_commit(async_client, async_collection): - from google.cloud.firestore_v1.async_transaction import async_transactional + from google.cloud.firestore import async_transactional @async_transactional async def _exercise(async_transaction): @@ -54,7 +54,7 @@ async def _exercise(async_transaction): async def _exercise_async_transaction_rollback(async_client, async_collection): - from google.cloud.firestore_v1.async_transaction import async_transactional + from google.cloud.firestore import async_transactional @async_transactional async def _exercise(async_transaction): diff --git a/tests/datastore_firestore/test_transaction.py b/tests/datastore_firestore/test_transaction.py index 121fac67ab..a32ad90754 100644 --- a/tests/datastore_firestore/test_transaction.py +++ b/tests/datastore_firestore/test_transaction.py @@ -29,7 +29,7 @@ def sample_data(collection): def _exercise_transaction_commit(client, collection): - from google.cloud.firestore_v1.transaction import transactional + from google.cloud.firestore import transactional @transactional def _exercise(transaction): @@ -53,7 +53,7 @@ def _exercise(transaction): def _exercise_transaction_rollback(client, collection): - from google.cloud.firestore_v1.transaction import transactional + from google.cloud.firestore import transactional @transactional def _exercise(transaction): From d3e473204bb2d840d6a73ec1b5de897e11e193ee Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Thu, 27 Jul 2023 13:20:37 -0700 Subject: [PATCH 27/31] Add async collection group instrumentation --- newrelic/hooks/datastore_firestore.py | 6 ++ tests/datastore_firestore/test_async_query.py | 81 +++++++++++++++++++ 2 files changed, 87 insertions(+) diff --git a/newrelic/hooks/datastore_firestore.py b/newrelic/hooks/datastore_firestore.py index ea92b4e5ec..9f28e7caf3 100644 --- a/newrelic/hooks/datastore_firestore.py +++ b/newrelic/hooks/datastore_firestore.py @@ -183,6 +183,12 @@ def instrument_google_cloud_firestore_v1_async_query(module): if hasattr(class_, method): wrap_async_generator_method(module, "AsyncQuery", method, target=_get_parent_id) + if hasattr(module, "AsyncCollectionGroup"): + class_ = module.AsyncCollectionGroup + for method in ("get_partitions",): + if hasattr(class_, method): + wrap_async_generator_method(module, "AsyncCollectionGroup", method, target=_get_parent_id) + def instrument_google_cloud_firestore_v1_aggregation(module): if hasattr(module, "AggregationQuery"): diff --git a/tests/datastore_firestore/test_async_query.py b/tests/datastore_firestore/test_async_query.py index 3319e73a53..d00449a122 100644 --- a/tests/datastore_firestore/test_async_query.py +++ b/tests/datastore_firestore/test_async_query.py @@ -26,6 +26,11 @@ 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({}) + + # ===== AsyncQuery ===== async def _exercise_async_query(async_collection): @@ -103,3 +108,79 @@ def _test(): def test_firestore_async_aggregation_query_generators(async_collection, assert_trace_for_async_generator): async_aggregation_query = async_collection.select("x").where(field_path="x", op_string="<=", value=3).count() assert_trace_for_async_generator(async_aggregation_query.stream) + + +# ===== CollectionGroup ===== + + +@pytest.fixture() +def patch_partition_queries(monkeypatch, async_client, collection, sample_data): + """ + Partitioning is not implemented in the Firestore emulator. + + Ordinarily this method would return a coroutine that returns an async_generator of Cursor objects. + Each Cursor must point at a valid document path. To test this, we can patch the RPC to return 1 Cursor + which is pointed at any document available. The get_partitions will take that and make 2 QueryPartition + objects out of it, which should be enough to ensure we can exercise the generator's tracing. + """ + from google.cloud.firestore_v1.types.document import Value + from google.cloud.firestore_v1.types.query import Cursor + + subcollection = collection.document("subcollection").collection("subcollection1") + documents = [d for d in subcollection.list_documents()] + + async def mock_partition_query(*args, **kwargs): + async def _mock_partition_query(): + yield Cursor(before=False, values=[Value(reference_value=documents[0].path)]) + return _mock_partition_query() + + monkeypatch.setattr(async_client._firestore_api, "partition_query", mock_partition_query) + yield + + +async def _exercise_async_collection_group(async_client, async_collection): + async_collection_group = async_client.collection_group(async_collection.id) + assert len(await async_collection_group.get()) + assert len([d async for d in async_collection_group.stream()]) + + partitions = [p async for p in async_collection_group.get_partitions(1)] + assert len(partitions) == 2 + documents = [] + while partitions: + documents.extend(await partitions.pop().query().get()) + assert len(documents) == 6 + + +def test_firestore_async_collection_group(loop, async_client, async_collection, patch_partition_queries): + _test_scoped_metrics = [ + ("Datastore/statement/Firestore/%s/get" % async_collection.id, 3), + ("Datastore/statement/Firestore/%s/stream" % async_collection.id, 1), + ("Datastore/statement/Firestore/%s/get_partitions" % async_collection.id, 1), + ] + + _test_rollup_metrics = [ + ("Datastore/operation/Firestore/get", 3), + ("Datastore/operation/Firestore/stream", 1), + ("Datastore/operation/Firestore/get_partitions", 1), + ("Datastore/all", 5), + ("Datastore/allOther", 5), + ] + + @validate_database_duration() + @validate_transaction_metrics( + "test_firestore_async_collection_group", + scoped_metrics=_test_scoped_metrics, + rollup_metrics=_test_rollup_metrics, + background_task=True, + ) + @background_task(name="test_firestore_async_collection_group") + def _test(): + loop.run_until_complete(_exercise_async_collection_group(async_client, async_collection)) + + _test() + + +@background_task() +def test_firestore_async_collection_group_generators(async_client, async_collection, assert_trace_for_async_generator, patch_partition_queries): + async_collection_group = async_client.collection_group(async_collection.id) + assert_trace_for_async_generator(async_collection_group.get_partitions, 1) From 7bf6f4978f058206c3cfb2b9c0efed963ca610ef Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Mon, 31 Jul 2023 14:34:26 -0700 Subject: [PATCH 28/31] Refactor exercise functions to fixtures --- .../test_async_batching.py | 22 +++-- .../datastore_firestore/test_async_client.py | 16 ++-- .../test_async_collections.py | 33 ++++--- .../test_async_documents.py | 33 ++++--- tests/datastore_firestore/test_async_query.py | 63 ++++++------ .../test_async_transaction.py | 95 ++++++++++--------- 6 files changed, 149 insertions(+), 113 deletions(-) diff --git a/tests/datastore_firestore/test_async_batching.py b/tests/datastore_firestore/test_async_batching.py index 0d51e99524..37b4e26e94 100644 --- a/tests/datastore_firestore/test_async_batching.py +++ b/tests/datastore_firestore/test_async_batching.py @@ -12,22 +12,28 @@ # 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, ) -async def _exercise_async_write_batch(async_client, async_collection): - docs = [async_collection.document(str(x)) for x in range(1, 4)] - async_batch = async_client.batch() - for doc in docs: - async_batch.set(doc, {}) - await async_batch.commit() +@pytest.fixture() +def exercise_async_write_batch(async_client, async_collection): + async def _exercise_async_write_batch(): + docs = [async_collection.document(str(x)) for x in range(1, 4)] + async_batch = async_client.batch() + for doc in docs: + async_batch.set(doc, {}) + + await async_batch.commit() + return _exercise_async_write_batch -def test_firestore_async_write_batch(loop, async_client, async_collection): +def test_firestore_async_write_batch(loop, exercise_async_write_batch): _test_scoped_metrics = [ ("Datastore/operation/Firestore/commit", 1), ] @@ -45,6 +51,6 @@ def test_firestore_async_write_batch(loop, async_client, async_collection): ) @background_task(name="test_firestore_async_write_batch") def _test(): - loop.run_until_complete(_exercise_async_write_batch(async_client, async_collection)) + loop.run_until_complete(exercise_async_write_batch()) _test() diff --git a/tests/datastore_firestore/test_async_client.py b/tests/datastore_firestore/test_async_client.py index c8e7b1dccf..2dc101c0b6 100644 --- a/tests/datastore_firestore/test_async_client.py +++ b/tests/datastore_firestore/test_async_client.py @@ -11,6 +11,7 @@ # 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 @@ -27,13 +28,16 @@ def existing_document(collection): 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 +@pytest.fixture() +def exercise_async_client(async_client, existing_document): + async def _exercise_async_client(): + 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 + return _exercise_async_client -def test_firestore_async_client(loop, async_client, existing_document): +def test_firestore_async_client(loop, exercise_async_client): _test_scoped_metrics = [ ("Datastore/operation/Firestore/collections", 1), ("Datastore/operation/Firestore/get_all", 1), @@ -53,7 +57,7 @@ def test_firestore_async_client(loop, async_client, existing_document): ) @background_task(name="test_firestore_async_client") def _test(): - loop.run_until_complete(_exercise_async_client(async_client, existing_document)) + loop.run_until_complete(exercise_async_client()) _test() diff --git a/tests/datastore_firestore/test_async_collections.py b/tests/datastore_firestore/test_async_collections.py index e5e3e2bae9..09472816f9 100644 --- a/tests/datastore_firestore/test_async_collections.py +++ b/tests/datastore_firestore/test_async_collections.py @@ -12,6 +12,8 @@ # 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 ( @@ -19,20 +21,23 @@ ) -async def _exercise_collections(async_collection): - async_collection.document("DoesNotExist") - await async_collection.add({"capital": "Rome", "currency": "Euro", "language": "Italian"}, "Italy") - await async_collection.add({"capital": "Mexico City", "currency": "Peso", "language": "Spanish"}, "Mexico") - - documents_get = await async_collection.get() - assert len(documents_get) == 2 - documents_stream = [_ async for _ in async_collection.stream()] - assert len(documents_stream) == 2 - documents_list = [_ async for _ in async_collection.list_documents()] - assert len(documents_list) == 2 +@pytest.fixture() +def exercise_collections(async_collection): + async def _exercise_collections(): + async_collection.document("DoesNotExist") + await async_collection.add({"capital": "Rome", "currency": "Euro", "language": "Italian"}, "Italy") + await async_collection.add({"capital": "Mexico City", "currency": "Peso", "language": "Spanish"}, "Mexico") + + documents_get = await async_collection.get() + assert len(documents_get) == 2 + documents_stream = [_ async for _ in async_collection.stream()] + assert len(documents_stream) == 2 + documents_list = [_ async for _ in async_collection.list_documents()] + assert len(documents_list) == 2 + return _exercise_collections -def test_firestore_async_collections(loop, async_collection): +def test_firestore_async_collections(loop, exercise_collections, async_collection): _test_scoped_metrics = [ ("Datastore/statement/Firestore/%s/stream" % async_collection.id, 1), ("Datastore/statement/Firestore/%s/get" % async_collection.id, 1), @@ -57,13 +62,13 @@ def test_firestore_async_collections(loop, async_collection): ) @background_task(name="test_firestore_async_collections") def _test(): - loop.run_until_complete(_exercise_collections(async_collection)) + loop.run_until_complete(exercise_collections()) _test() @background_task() -def test_firestore_async_collections_generators(loop, collection, async_collection, assert_trace_for_async_generator): +def test_firestore_async_collections_generators(collection, async_collection, assert_trace_for_async_generator): collection.add({}) collection.add({}) assert len([_ for _ in collection.list_documents()]) == 2 diff --git a/tests/datastore_firestore/test_async_documents.py b/tests/datastore_firestore/test_async_documents.py index 2df6940add..fd913d2ed6 100644 --- a/tests/datastore_firestore/test_async_documents.py +++ b/tests/datastore_firestore/test_async_documents.py @@ -12,6 +12,8 @@ # 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 ( @@ -19,23 +21,26 @@ ) -async def _exercise_async_documents(async_collection): - italy_doc = async_collection.document("Italy") - await italy_doc.set({"capital": "Rome", "currency": "Euro", "language": "Italian"}) - await italy_doc.get() - italian_cities = italy_doc.collection("cities") - await italian_cities.add({"capital": "Rome"}) - retrieved_coll = [_ async for _ in italy_doc.collections()] - assert len(retrieved_coll) == 1 +@pytest.fixture() +def exercise_async_documents(async_collection): + async def _exercise_async_documents(): + italy_doc = async_collection.document("Italy") + await italy_doc.set({"capital": "Rome", "currency": "Euro", "language": "Italian"}) + await italy_doc.get() + italian_cities = italy_doc.collection("cities") + await italian_cities.add({"capital": "Rome"}) + retrieved_coll = [_ async for _ in italy_doc.collections()] + assert len(retrieved_coll) == 1 - usa_doc = async_collection.document("USA") - await usa_doc.create({"capital": "Washington D.C.", "currency": "Dollar", "language": "English"}) - await usa_doc.update({"president": "Joe Biden"}) + usa_doc = async_collection.document("USA") + await usa_doc.create({"capital": "Washington D.C.", "currency": "Dollar", "language": "English"}) + await usa_doc.update({"president": "Joe Biden"}) - await async_collection.document("USA").delete() + await async_collection.document("USA").delete() + return _exercise_async_documents -def test_firestore_async_documents(loop, async_collection): +def test_firestore_async_documents(loop, exercise_async_documents): _test_scoped_metrics = [ ("Datastore/statement/Firestore/Italy/set", 1), ("Datastore/statement/Firestore/Italy/get", 1), @@ -66,7 +71,7 @@ def test_firestore_async_documents(loop, async_collection): ) @background_task(name="test_firestore_async_documents") def _test(): - loop.run_until_complete(_exercise_async_documents(async_collection)) + loop.run_until_complete(exercise_async_documents()) _test() diff --git a/tests/datastore_firestore/test_async_query.py b/tests/datastore_firestore/test_async_query.py index d00449a122..f7b2bad472 100644 --- a/tests/datastore_firestore/test_async_query.py +++ b/tests/datastore_firestore/test_async_query.py @@ -33,13 +33,16 @@ def sample_data(collection): # ===== AsyncQuery ===== -async def _exercise_async_query(async_collection): - async_query = async_collection.select("x").limit(10).order_by("x").where(field_path="x", op_string="<=", value=3) - assert len(await async_query.get()) == 3 - assert len([_ async for _ in async_query.stream()]) == 3 +@pytest.fixture() +def exercise_async_query(async_collection): + async def _exercise_async_query(): + async_query = async_collection.select("x").limit(10).order_by("x").where(field_path="x", op_string="<=", value=3) + assert len(await async_query.get()) == 3 + assert len([_ async for _ in async_query.stream()]) == 3 + return _exercise_async_query -def test_firestore_async_query(loop, async_collection): +def test_firestore_async_query(loop, exercise_async_query, async_collection): _test_scoped_metrics = [ ("Datastore/statement/Firestore/%s/stream" % async_collection.id, 1), ("Datastore/statement/Firestore/%s/get" % async_collection.id, 1), @@ -60,7 +63,7 @@ def test_firestore_async_query(loop, async_collection): ) @background_task(name="test_firestore_async_query") def _test(): - loop.run_until_complete(_exercise_async_query(async_collection)) + loop.run_until_complete(exercise_async_query()) _test() @@ -72,13 +75,16 @@ def test_firestore_async_query_generators(async_collection, assert_trace_for_asy # ===== AsyncAggregationQuery ===== -async def _exercise_async_aggregation_query(async_collection): - async_aggregation_query = async_collection.select("x").where(field_path="x", op_string="<=", value=3).count() - assert (await async_aggregation_query.get())[0][0].value == 3 - assert [_ async for _ in async_aggregation_query.stream()][0][0].value == 3 +@pytest.fixture() +def exercise_async_aggregation_query(async_collection): + async def _exercise_async_aggregation_query(): + async_aggregation_query = async_collection.select("x").where(field_path="x", op_string="<=", value=3).count() + assert (await async_aggregation_query.get())[0][0].value == 3 + assert [_ async for _ in async_aggregation_query.stream()][0][0].value == 3 + return _exercise_async_aggregation_query -def test_firestore_async_aggregation_query(loop, async_collection): +def test_firestore_async_aggregation_query(loop, exercise_async_aggregation_query, async_collection): _test_scoped_metrics = [ ("Datastore/statement/Firestore/%s/stream" % async_collection.id, 1), ("Datastore/statement/Firestore/%s/get" % async_collection.id, 1), @@ -99,7 +105,7 @@ def test_firestore_async_aggregation_query(loop, async_collection): ) @background_task(name="test_firestore_async_aggregation_query") def _test(): - loop.run_until_complete(_exercise_async_aggregation_query(async_collection)) + loop.run_until_complete(exercise_async_aggregation_query()) _test() @@ -138,20 +144,23 @@ async def _mock_partition_query(): yield -async def _exercise_async_collection_group(async_client, async_collection): - async_collection_group = async_client.collection_group(async_collection.id) - assert len(await async_collection_group.get()) - assert len([d async for d in async_collection_group.stream()]) - - partitions = [p async for p in async_collection_group.get_partitions(1)] - assert len(partitions) == 2 - documents = [] - while partitions: - documents.extend(await partitions.pop().query().get()) - assert len(documents) == 6 - - -def test_firestore_async_collection_group(loop, async_client, async_collection, patch_partition_queries): +@pytest.fixture() +def exercise_async_collection_group(async_client, async_collection): + async def _exercise_async_collection_group(): + async_collection_group = async_client.collection_group(async_collection.id) + assert len(await async_collection_group.get()) + assert len([d async for d in async_collection_group.stream()]) + + partitions = [p async for p in async_collection_group.get_partitions(1)] + assert len(partitions) == 2 + documents = [] + while partitions: + documents.extend(await partitions.pop().query().get()) + assert len(documents) == 6 + return _exercise_async_collection_group + + +def test_firestore_async_collection_group(loop, exercise_async_collection_group, async_collection, patch_partition_queries): _test_scoped_metrics = [ ("Datastore/statement/Firestore/%s/get" % async_collection.id, 3), ("Datastore/statement/Firestore/%s/stream" % async_collection.id, 1), @@ -175,7 +184,7 @@ def test_firestore_async_collection_group(loop, async_client, async_collection, ) @background_task(name="test_firestore_async_collection_group") def _test(): - loop.run_until_complete(_exercise_async_collection_group(async_client, async_collection)) + loop.run_until_complete(exercise_async_collection_group()) _test() diff --git a/tests/datastore_firestore/test_async_transaction.py b/tests/datastore_firestore/test_async_transaction.py index 76667538b8..3bf80c5478 100644 --- a/tests/datastore_firestore/test_async_transaction.py +++ b/tests/datastore_firestore/test_async_transaction.py @@ -11,6 +11,7 @@ # 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 @@ -26,49 +27,55 @@ def sample_data(collection): collection.add({"x": x}, "doc%d" % x) -async def _exercise_async_transaction_commit(async_client, async_collection): - from google.cloud.firestore import async_transactional - - @async_transactional - async def _exercise(async_transaction): - # get a DocumentReference - with pytest.raises(TypeError): # get is currently broken. It attempts to await an async_generator instead of consuming it. - [_ async for _ in async_transaction.get(async_collection.document("doc1"))] - - # get a Query - with pytest.raises(TypeError): # get is currently broken. It attempts to await an async_generator instead of consuming it. - async_query = async_collection.select("x").where(field_path="x", op_string=">", value=2) - assert len([_ async for _ in async_transaction.get(async_query)]) == 1 - - # get_all on a list of DocumentReferences - with pytest.raises(TypeError): # get_all is currently broken. It attempts to await an async_generator instead of consuming it. - all_docs = async_transaction.get_all([async_collection.document("doc%d" % x) for x in range(1, 4)]) - assert len([_ async for _ in all_docs]) == 3 - - # set and delete methods - async_transaction.set(async_collection.document("doc2"), {"x": 0}) - async_transaction.delete(async_collection.document("doc3")) - - await _exercise(async_client.transaction()) - assert len([_ async for _ in async_collection.list_documents()]) == 2 - - -async def _exercise_async_transaction_rollback(async_client, async_collection): - from google.cloud.firestore import async_transactional - - @async_transactional - async def _exercise(async_transaction): - # set and delete methods - async_transaction.set(async_collection.document("doc2"), {"x": 99}) - async_transaction.delete(async_collection.document("doc1")) - raise RuntimeError() - - with pytest.raises(RuntimeError): +@pytest.fixture() +def exercise_async_transaction_commit(async_client, async_collection): + async def _exercise_async_transaction_commit(): + from google.cloud.firestore import async_transactional + + @async_transactional + async def _exercise(async_transaction): + # get a DocumentReference + with pytest.raises(TypeError): # get is currently broken. It attempts to await an async_generator instead of consuming it. + [_ async for _ in async_transaction.get(async_collection.document("doc1"))] + + # get a Query + with pytest.raises(TypeError): # get is currently broken. It attempts to await an async_generator instead of consuming it. + async_query = async_collection.select("x").where(field_path="x", op_string=">", value=2) + assert len([_ async for _ in async_transaction.get(async_query)]) == 1 + + # get_all on a list of DocumentReferences + with pytest.raises(TypeError): # get_all is currently broken. It attempts to await an async_generator instead of consuming it. + all_docs = async_transaction.get_all([async_collection.document("doc%d" % x) for x in range(1, 4)]) + assert len([_ async for _ in all_docs]) == 3 + + # set and delete methods + async_transaction.set(async_collection.document("doc2"), {"x": 0}) + async_transaction.delete(async_collection.document("doc3")) + await _exercise(async_client.transaction()) - assert len([_ async for _ in async_collection.list_documents()]) == 3 + assert len([_ async for _ in async_collection.list_documents()]) == 2 + return _exercise_async_transaction_commit -def test_firestore_async_transaction_commit(loop, async_client, async_collection): +@pytest.fixture() +def exercise_async_transaction_rollback(async_client, async_collection): + async def _exercise_async_transaction_rollback(): + from google.cloud.firestore import async_transactional + + @async_transactional + async def _exercise(async_transaction): + # set and delete methods + async_transaction.set(async_collection.document("doc2"), {"x": 99}) + async_transaction.delete(async_collection.document("doc1")) + raise RuntimeError() + + with pytest.raises(RuntimeError): + await _exercise(async_client.transaction()) + assert len([_ async for _ in async_collection.list_documents()]) == 3 + return _exercise_async_transaction_rollback + + +def test_firestore_async_transaction_commit(loop, exercise_async_transaction_commit, async_collection): _test_scoped_metrics = [ ("Datastore/operation/Firestore/commit", 1), # ("Datastore/operation/Firestore/get_all", 2), @@ -91,12 +98,12 @@ def test_firestore_async_transaction_commit(loop, async_client, async_collection ) @background_task(name="test_firestore_async_transaction") def _test(): - loop.run_until_complete(_exercise_async_transaction_commit(async_client, async_collection)) + loop.run_until_complete(exercise_async_transaction_commit()) _test() -def test_firestore_async_transaction_rollback(loop, async_client, async_collection): +def test_firestore_async_transaction_rollback(loop, exercise_async_transaction_rollback, async_collection): _test_scoped_metrics = [ ("Datastore/operation/Firestore/rollback", 1), ("Datastore/statement/Firestore/%s/list_documents" % async_collection.id, 1), @@ -107,7 +114,7 @@ def test_firestore_async_transaction_rollback(loop, async_client, async_collecti ("Datastore/all", 2), ("Datastore/allOther", 2), ] - # @validate_database_duration() + @validate_database_duration() @validate_transaction_metrics( "test_firestore_async_transaction", scoped_metrics=_test_scoped_metrics, @@ -116,6 +123,6 @@ def test_firestore_async_transaction_rollback(loop, async_client, async_collecti ) @background_task(name="test_firestore_async_transaction") def _test(): - loop.run_until_complete(_exercise_async_transaction_rollback(async_client, async_collection)) + loop.run_until_complete(exercise_async_transaction_rollback()) _test() From ad2999ff50b6b17b5774f69ca8116ee901f47474 Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Wed, 2 Aug 2023 14:58:38 -0700 Subject: [PATCH 29/31] Squashed commit of the following: commit 09c5e11498b4c200057190e859f8151241c421f3 Author: Tim Pansino Date: Wed Aug 2 14:33:24 2023 -0700 Add testing for automatic and manual asyncwrappers commit fc3ef6bfb8cb2f9cd6c8ffdf7bfd953be41cc974 Author: Tim Pansino Date: Wed Aug 2 14:33:05 2023 -0700 Add async wrapper argument to all trace APIs commit 479f9e236e2212e0f9cdf51627996068027acd82 Merge: faf3cccea edd1f94b0 Author: Tim Pansino Date: Wed Aug 2 13:44:24 2023 -0700 Merge remote-tracking branch 'origin/develop-google-firestore-instrumentation' into feature-async-wrapper-argument commit edd1f94b0f601a2674da4e594b777bae0eed6643 Author: Timothy Pansino <11214426+TimPansino@users.noreply.github.com> Date: Wed Aug 2 13:40:51 2023 -0700 Async Generator Wrapper (#884) * Add async generator wrapper * Add no harm test * Remove anext calls * Add graphql traces to decorator testing * Remove pypy generator gc logic --------- Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> commit faf3ccceae127128aff81fc59d95dd3f49699a3c Author: Tim Pansino Date: Mon Jul 31 15:10:56 2023 -0700 Add async_wrapper to datastore_trace api --- newrelic/api/database_trace.py | 14 +-- newrelic/api/datastore_trace.py | 19 ++-- newrelic/api/external_trace.py | 16 +-- newrelic/api/function_trace.py | 16 +-- newrelic/api/graphql_trace.py | 26 ++--- newrelic/api/memcache_trace.py | 14 +-- newrelic/api/message_trace.py | 13 +-- .../_test_async_coroutine_trace.py | 5 + .../_test_async_generator_trace.py | 88 +++++++++------ .../test_async_wrapper_detection.py | 102 ++++++++++++++++++ tests/agent_features/test_coroutine_trace.py | 3 + 11 files changed, 230 insertions(+), 86 deletions(-) create mode 100644 tests/agent_features/test_async_wrapper_detection.py diff --git a/newrelic/api/database_trace.py b/newrelic/api/database_trace.py index 2bc4976887..1069be506f 100644 --- a/newrelic/api/database_trace.py +++ b/newrelic/api/database_trace.py @@ -16,7 +16,7 @@ import logging from newrelic.api.time_trace import TimeTrace, current_trace -from newrelic.common.async_wrapper import async_wrapper +from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper from newrelic.common.object_wrapper import FunctionWrapper, wrap_object from newrelic.core.database_node import DatabaseNode from newrelic.core.stack_trace import current_stack @@ -244,9 +244,9 @@ def create_node(self): ) -def DatabaseTraceWrapper(wrapped, sql, dbapi2_module=None): +def DatabaseTraceWrapper(wrapped, sql, dbapi2_module=None, async_wrapper=None): def _nr_database_trace_wrapper_(wrapped, instance, args, kwargs): - wrapper = async_wrapper(wrapped) + wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped) if not wrapper: parent = current_trace() if not parent: @@ -273,9 +273,9 @@ def _nr_database_trace_wrapper_(wrapped, instance, args, kwargs): return FunctionWrapper(wrapped, _nr_database_trace_wrapper_) -def database_trace(sql, dbapi2_module=None): - return functools.partial(DatabaseTraceWrapper, sql=sql, dbapi2_module=dbapi2_module) +def database_trace(sql, dbapi2_module=None, async_wrapper=None): + return functools.partial(DatabaseTraceWrapper, sql=sql, dbapi2_module=dbapi2_module, async_wrapper=async_wrapper) -def wrap_database_trace(module, object_path, sql, dbapi2_module=None): - wrap_object(module, object_path, DatabaseTraceWrapper, (sql, dbapi2_module)) +def wrap_database_trace(module, object_path, sql, dbapi2_module=None, async_wrapper=None): + wrap_object(module, object_path, DatabaseTraceWrapper, (sql, dbapi2_module, async_wrapper)) diff --git a/newrelic/api/datastore_trace.py b/newrelic/api/datastore_trace.py index 1561293f64..0401c79ea5 100644 --- a/newrelic/api/datastore_trace.py +++ b/newrelic/api/datastore_trace.py @@ -15,7 +15,7 @@ import functools from newrelic.api.time_trace import TimeTrace, current_trace -from newrelic.common.async_wrapper import async_wrapper +from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper from newrelic.common.object_wrapper import FunctionWrapper, wrap_object from newrelic.core.datastore_node import DatastoreNode @@ -135,7 +135,7 @@ def create_node(self): ) -def DatastoreTraceWrapper(wrapped, product, target, operation, host=None, port_path_or_id=None, database_name=None): +def DatastoreTraceWrapper(wrapped, product, target, operation, host=None, port_path_or_id=None, database_name=None, async_wrapper=None): """Wraps a method to time datastore queries. :param wrapped: The function to apply the trace to. @@ -158,6 +158,8 @@ def DatastoreTraceWrapper(wrapped, product, target, operation, host=None, port_p :param database_name: The name of database where the current query is being executed. :type database_name: str + :param async_wrapper: An async trace wrapper from newrelic.common.async_wrapper. + :type async_wrapper: callable or None :rtype: :class:`newrelic.common.object_wrapper.FunctionWrapper` This is typically used to wrap datastore queries such as calls to Redis or @@ -173,7 +175,7 @@ def DatastoreTraceWrapper(wrapped, product, target, operation, host=None, port_p """ def _nr_datastore_trace_wrapper_(wrapped, instance, args, kwargs): - wrapper = async_wrapper(wrapped) + wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped) if not wrapper: parent = current_trace() if not parent: @@ -242,7 +244,7 @@ def _nr_datastore_trace_wrapper_(wrapped, instance, args, kwargs): return FunctionWrapper(wrapped, _nr_datastore_trace_wrapper_) -def datastore_trace(product, target, operation, host=None, port_path_or_id=None, database_name=None): +def datastore_trace(product, target, operation, host=None, port_path_or_id=None, database_name=None, async_wrapper=None): """Decorator allows datastore query to be timed. :param product: The name of the vendor. @@ -263,6 +265,8 @@ def datastore_trace(product, target, operation, host=None, port_path_or_id=None, :param database_name: The name of database where the current query is being executed. :type database_name: str + :param async_wrapper: An async trace wrapper from newrelic.common.async_wrapper. + :type async_wrapper: callable or None This is typically used to decorate datastore queries such as calls to Redis or ElasticSearch. @@ -284,11 +288,12 @@ def datastore_trace(product, target, operation, host=None, port_path_or_id=None, host=host, port_path_or_id=port_path_or_id, database_name=database_name, + async_wrapper=async_wrapper, ) def wrap_datastore_trace( - module, object_path, product, target, operation, host=None, port_path_or_id=None, database_name=None + module, object_path, product, target, operation, host=None, port_path_or_id=None, database_name=None, async_wrapper=None ): """Method applies custom timing to datastore query. @@ -314,6 +319,8 @@ def wrap_datastore_trace( :param database_name: The name of database where the current query is being executed. :type database_name: str + :param async_wrapper: An async trace wrapper from newrelic.common.async_wrapper. + :type async_wrapper: callable or None This is typically used to time database query method calls such as Redis GET. @@ -327,5 +334,5 @@ def wrap_datastore_trace( """ wrap_object( - module, object_path, DatastoreTraceWrapper, (product, target, operation, host, port_path_or_id, database_name) + module, object_path, DatastoreTraceWrapper, (product, target, operation, host, port_path_or_id, database_name, async_wrapper) ) diff --git a/newrelic/api/external_trace.py b/newrelic/api/external_trace.py index c43c560c6c..2e147df450 100644 --- a/newrelic/api/external_trace.py +++ b/newrelic/api/external_trace.py @@ -16,7 +16,7 @@ from newrelic.api.cat_header_mixin import CatHeaderMixin from newrelic.api.time_trace import TimeTrace, current_trace -from newrelic.common.async_wrapper import async_wrapper +from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper from newrelic.common.object_wrapper import FunctionWrapper, wrap_object from newrelic.core.external_node import ExternalNode @@ -66,9 +66,9 @@ def create_node(self): ) -def ExternalTraceWrapper(wrapped, library, url, method=None): +def ExternalTraceWrapper(wrapped, library, url, method=None, async_wrapper=None): def dynamic_wrapper(wrapped, instance, args, kwargs): - wrapper = async_wrapper(wrapped) + wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped) if not wrapper: parent = current_trace() if not parent: @@ -103,7 +103,7 @@ def dynamic_wrapper(wrapped, instance, args, kwargs): return wrapped(*args, **kwargs) def literal_wrapper(wrapped, instance, args, kwargs): - wrapper = async_wrapper(wrapped) + wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped) if not wrapper: parent = current_trace() if not parent: @@ -125,9 +125,9 @@ def literal_wrapper(wrapped, instance, args, kwargs): return FunctionWrapper(wrapped, literal_wrapper) -def external_trace(library, url, method=None): - return functools.partial(ExternalTraceWrapper, library=library, url=url, method=method) +def external_trace(library, url, method=None, async_wrapper=None): + return functools.partial(ExternalTraceWrapper, library=library, url=url, method=method, async_wrapper=async_wrapper) -def wrap_external_trace(module, object_path, library, url, method=None): - wrap_object(module, object_path, ExternalTraceWrapper, (library, url, method)) +def wrap_external_trace(module, object_path, library, url, method=None, async_wrapper=None): + wrap_object(module, object_path, ExternalTraceWrapper, (library, url, method, async_wrapper)) diff --git a/newrelic/api/function_trace.py b/newrelic/api/function_trace.py index 474c1b2266..85d7617b68 100644 --- a/newrelic/api/function_trace.py +++ b/newrelic/api/function_trace.py @@ -15,7 +15,7 @@ import functools from newrelic.api.time_trace import TimeTrace, current_trace -from newrelic.common.async_wrapper import async_wrapper +from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper from newrelic.common.object_names import callable_name from newrelic.common.object_wrapper import FunctionWrapper, wrap_object from newrelic.core.function_node import FunctionNode @@ -89,9 +89,9 @@ def create_node(self): ) -def FunctionTraceWrapper(wrapped, name=None, group=None, label=None, params=None, terminal=False, rollup=None): +def FunctionTraceWrapper(wrapped, name=None, group=None, label=None, params=None, terminal=False, rollup=None, async_wrapper=None): def dynamic_wrapper(wrapped, instance, args, kwargs): - wrapper = async_wrapper(wrapped) + wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped) if not wrapper: parent = current_trace() if not parent: @@ -147,7 +147,7 @@ def dynamic_wrapper(wrapped, instance, args, kwargs): return wrapped(*args, **kwargs) def literal_wrapper(wrapped, instance, args, kwargs): - wrapper = async_wrapper(wrapped) + wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped) if not wrapper: parent = current_trace() if not parent: @@ -171,13 +171,13 @@ def literal_wrapper(wrapped, instance, args, kwargs): return FunctionWrapper(wrapped, literal_wrapper) -def function_trace(name=None, group=None, label=None, params=None, terminal=False, rollup=None): +def function_trace(name=None, group=None, label=None, params=None, terminal=False, rollup=None, async_wrapper=None): return functools.partial( - FunctionTraceWrapper, name=name, group=group, label=label, params=params, terminal=terminal, rollup=rollup + FunctionTraceWrapper, name=name, group=group, label=label, params=params, terminal=terminal, rollup=rollup, async_wrapper=async_wrapper ) def wrap_function_trace( - module, object_path, name=None, group=None, label=None, params=None, terminal=False, rollup=None + module, object_path, name=None, group=None, label=None, params=None, terminal=False, rollup=None, async_wrapper=None ): - return wrap_object(module, object_path, FunctionTraceWrapper, (name, group, label, params, terminal, rollup)) + return wrap_object(module, object_path, FunctionTraceWrapper, (name, group, label, params, terminal, rollup, async_wrapper)) diff --git a/newrelic/api/graphql_trace.py b/newrelic/api/graphql_trace.py index 7a2c9ec02f..3d6ae6b09f 100644 --- a/newrelic/api/graphql_trace.py +++ b/newrelic/api/graphql_trace.py @@ -16,7 +16,7 @@ from newrelic.api.time_trace import TimeTrace, current_trace from newrelic.api.transaction import current_transaction -from newrelic.common.async_wrapper import async_wrapper +from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper from newrelic.common.object_wrapper import FunctionWrapper, wrap_object from newrelic.core.graphql_node import GraphQLOperationNode, GraphQLResolverNode @@ -109,9 +109,9 @@ def set_transaction_name(self, priority=None): transaction.set_transaction_name(name, "GraphQL", priority=priority) -def GraphQLOperationTraceWrapper(wrapped): +def GraphQLOperationTraceWrapper(wrapped, async_wrapper=None): def _nr_graphql_trace_wrapper_(wrapped, instance, args, kwargs): - wrapper = async_wrapper(wrapped) + wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped) if not wrapper: parent = current_trace() if not parent: @@ -130,12 +130,12 @@ def _nr_graphql_trace_wrapper_(wrapped, instance, args, kwargs): return FunctionWrapper(wrapped, _nr_graphql_trace_wrapper_) -def graphql_operation_trace(): - return functools.partial(GraphQLOperationTraceWrapper) +def graphql_operation_trace(async_wrapper=None): + return functools.partial(GraphQLOperationTraceWrapper, async_wrapper=async_wrapper) -def wrap_graphql_operation_trace(module, object_path): - wrap_object(module, object_path, GraphQLOperationTraceWrapper) +def wrap_graphql_operation_trace(module, object_path, async_wrapper=None): + wrap_object(module, object_path, GraphQLOperationTraceWrapper, (async_wrapper,)) class GraphQLResolverTrace(TimeTrace): @@ -193,9 +193,9 @@ def create_node(self): ) -def GraphQLResolverTraceWrapper(wrapped): +def GraphQLResolverTraceWrapper(wrapped, async_wrapper=None): def _nr_graphql_trace_wrapper_(wrapped, instance, args, kwargs): - wrapper = async_wrapper(wrapped) + wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped) if not wrapper: parent = current_trace() if not parent: @@ -214,9 +214,9 @@ def _nr_graphql_trace_wrapper_(wrapped, instance, args, kwargs): return FunctionWrapper(wrapped, _nr_graphql_trace_wrapper_) -def graphql_resolver_trace(): - return functools.partial(GraphQLResolverTraceWrapper) +def graphql_resolver_trace(async_wrapper=None): + return functools.partial(GraphQLResolverTraceWrapper, async_wrapper=async_wrapper) -def wrap_graphql_resolver_trace(module, object_path): - wrap_object(module, object_path, GraphQLResolverTraceWrapper) +def wrap_graphql_resolver_trace(module, object_path, async_wrapper=None): + wrap_object(module, object_path, GraphQLResolverTraceWrapper, (async_wrapper,)) diff --git a/newrelic/api/memcache_trace.py b/newrelic/api/memcache_trace.py index 6657a9ce27..87f12f9fc7 100644 --- a/newrelic/api/memcache_trace.py +++ b/newrelic/api/memcache_trace.py @@ -15,7 +15,7 @@ import functools from newrelic.api.time_trace import TimeTrace, current_trace -from newrelic.common.async_wrapper import async_wrapper +from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper from newrelic.common.object_wrapper import FunctionWrapper, wrap_object from newrelic.core.memcache_node import MemcacheNode @@ -51,9 +51,9 @@ def create_node(self): ) -def MemcacheTraceWrapper(wrapped, command): +def MemcacheTraceWrapper(wrapped, command, async_wrapper=None): def _nr_wrapper_memcache_trace_(wrapped, instance, args, kwargs): - wrapper = async_wrapper(wrapped) + wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped) if not wrapper: parent = current_trace() if not parent: @@ -80,9 +80,9 @@ def _nr_wrapper_memcache_trace_(wrapped, instance, args, kwargs): return FunctionWrapper(wrapped, _nr_wrapper_memcache_trace_) -def memcache_trace(command): - return functools.partial(MemcacheTraceWrapper, command=command) +def memcache_trace(command, async_wrapper=None): + return functools.partial(MemcacheTraceWrapper, command=command, async_wrapper=async_wrapper) -def wrap_memcache_trace(module, object_path, command): - wrap_object(module, object_path, MemcacheTraceWrapper, (command,)) +def wrap_memcache_trace(module, object_path, command, async_wrapper=None): + wrap_object(module, object_path, MemcacheTraceWrapper, (command, async_wrapper)) diff --git a/newrelic/api/message_trace.py b/newrelic/api/message_trace.py index be819d7044..f564c41cb4 100644 --- a/newrelic/api/message_trace.py +++ b/newrelic/api/message_trace.py @@ -16,7 +16,7 @@ from newrelic.api.cat_header_mixin import CatHeaderMixin from newrelic.api.time_trace import TimeTrace, current_trace -from newrelic.common.async_wrapper import async_wrapper +from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper from newrelic.common.object_wrapper import FunctionWrapper, wrap_object from newrelic.core.message_node import MessageNode @@ -91,9 +91,9 @@ def create_node(self): ) -def MessageTraceWrapper(wrapped, library, operation, destination_type, destination_name, params={}, terminal=True): +def MessageTraceWrapper(wrapped, library, operation, destination_type, destination_name, params={}, terminal=True, async_wrapper=None): def _nr_message_trace_wrapper_(wrapped, instance, args, kwargs): - wrapper = async_wrapper(wrapped) + wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped) if not wrapper: parent = current_trace() if not parent: @@ -144,7 +144,7 @@ def _nr_message_trace_wrapper_(wrapped, instance, args, kwargs): return FunctionWrapper(wrapped, _nr_message_trace_wrapper_) -def message_trace(library, operation, destination_type, destination_name, params={}, terminal=True): +def message_trace(library, operation, destination_type, destination_name, params={}, terminal=True, async_wrapper=None): return functools.partial( MessageTraceWrapper, library=library, @@ -153,10 +153,11 @@ def message_trace(library, operation, destination_type, destination_name, params destination_name=destination_name, params=params, terminal=terminal, + async_wrapper=async_wrapper, ) -def wrap_message_trace(module, object_path, library, operation, destination_type, destination_name, params={}, terminal=True): +def wrap_message_trace(module, object_path, library, operation, destination_type, destination_name, params={}, terminal=True, async_wrapper=None): wrap_object( - module, object_path, MessageTraceWrapper, (library, operation, destination_type, destination_name, params, terminal) + module, object_path, MessageTraceWrapper, (library, operation, destination_type, destination_name, params, terminal, async_wrapper) ) diff --git a/tests/agent_features/_test_async_coroutine_trace.py b/tests/agent_features/_test_async_coroutine_trace.py index 51b81f5f64..1250b8c254 100644 --- a/tests/agent_features/_test_async_coroutine_trace.py +++ b/tests/agent_features/_test_async_coroutine_trace.py @@ -28,6 +28,7 @@ from newrelic.api.datastore_trace import datastore_trace from newrelic.api.external_trace import external_trace from newrelic.api.function_trace import function_trace +from newrelic.api.graphql_trace import graphql_operation_trace, graphql_resolver_trace from newrelic.api.memcache_trace import memcache_trace from newrelic.api.message_trace import message_trace @@ -41,6 +42,8 @@ (functools.partial(datastore_trace, "lib", "foo", "bar"), "Datastore/statement/lib/foo/bar"), (functools.partial(message_trace, "lib", "op", "typ", "name"), "MessageBroker/lib/typ/op/Named/name"), (functools.partial(memcache_trace, "cmd"), "Memcache/cmd"), + (functools.partial(graphql_operation_trace), "GraphQL/operation/GraphQL///"), + (functools.partial(graphql_resolver_trace), "GraphQL/resolve/GraphQL/"), ], ) def test_awaitable_timing(event_loop, trace, metric): @@ -79,6 +82,8 @@ def _test(): (functools.partial(datastore_trace, "lib", "foo", "bar"), "Datastore/statement/lib/foo/bar"), (functools.partial(message_trace, "lib", "op", "typ", "name"), "MessageBroker/lib/typ/op/Named/name"), (functools.partial(memcache_trace, "cmd"), "Memcache/cmd"), + (functools.partial(graphql_operation_trace), "GraphQL/operation/GraphQL///"), + (functools.partial(graphql_resolver_trace), "GraphQL/resolve/GraphQL/"), ], ) @pytest.mark.parametrize("yield_from", [True, False]) diff --git a/tests/agent_features/_test_async_generator_trace.py b/tests/agent_features/_test_async_generator_trace.py index 8e7703a0db..f6097bc434 100644 --- a/tests/agent_features/_test_async_generator_trace.py +++ b/tests/agent_features/_test_async_generator_trace.py @@ -13,7 +13,6 @@ # limitations under the License. import functools -import gc import sys import time @@ -31,10 +30,10 @@ from newrelic.api.datastore_trace import datastore_trace from newrelic.api.external_trace import external_trace from newrelic.api.function_trace import function_trace +from newrelic.api.graphql_trace import graphql_operation_trace, graphql_resolver_trace from newrelic.api.memcache_trace import memcache_trace from newrelic.api.message_trace import message_trace -is_pypy = hasattr(sys, "pypy_version_info") asyncio = pytest.importorskip("asyncio") @@ -47,6 +46,8 @@ (functools.partial(datastore_trace, "lib", "foo", "bar"), "Datastore/statement/lib/foo/bar"), (functools.partial(message_trace, "lib", "op", "typ", "name"), "MessageBroker/lib/typ/op/Named/name"), (functools.partial(memcache_trace, "cmd"), "Memcache/cmd"), + (functools.partial(graphql_operation_trace), "GraphQL/operation/GraphQL///"), + (functools.partial(graphql_resolver_trace), "GraphQL/resolve/GraphQL/"), ], ) def test_async_generator_timing(event_loop, trace, metric): @@ -130,7 +131,7 @@ async def _test(): gen = agen() # kickstart the generator (the try/except logic is inside the # generator) - await anext(gen) + await gen.asend(None) await gen.athrow(ValueError) # consume the generator @@ -198,7 +199,7 @@ async def _test(): gen = agen() # kickstart the coroutine - await anext(gen) + await gen.asend(None) # trace should be ended/recorded by close await gen.aclose() @@ -279,7 +280,7 @@ async def _test(): gen = agen() # kickstart the coroutine - await anext(gen) + await gen.asend(None) assert await gen.asend("foobar") == "foobar" assert _received and _received[0] == "foobar" @@ -311,7 +312,7 @@ async def _test(): gen = agen() # kickstart the coroutine - await anext(gen) + await gen.asend(None) assert await gen.athrow(MyException) == "foobar" @@ -342,7 +343,7 @@ async def _test(): gen = agen() # kickstart the coroutine - await anext(gen) + await gen.asend(None) # async generator will raise StopAsyncIteration with pytest.raises(StopAsyncIteration): @@ -401,10 +402,6 @@ async def _test(): with pytest.raises(RuntimeError): await gen.aclose() - if is_pypy: - gen = None - gc.collect() - event_loop.run_until_complete(_test()) @@ -439,6 +436,26 @@ async def _test(): assert full_metrics[("Function/agen", "")].total_call_time < 0.1 +@validate_transaction_metrics( + "test_complete_async_generator", + background_task=True, + scoped_metrics=[("Function/agen", 1)], + rollup_metrics=[("Function/agen", 1)], +) +@background_task(name="test_complete_async_generator") +def test_complete_async_generator(event_loop): + @function_trace(name="agen") + async def agen(): + for i in range(5): + yield i + + async def _test(): + gen = agen() + assert [x async for x in gen] == [x for x in range(5)] + + event_loop.run_until_complete(_test()) + + @pytest.mark.parametrize("nr_transaction", [True, False]) def test_incomplete_async_generator(event_loop, nr_transaction): @function_trace(name="agen") @@ -453,35 +470,44 @@ async def _test(): async for _ in c: break - if is_pypy: - # pypy is not guaranteed to delete the coroutine when it goes out - # of scope. This code "helps" pypy along. The test above is really - # just to verify that incomplete coroutines will "eventually" be - # cleaned up. In pypy, unfortunately that means it may not be - # reported all the time. A customer would be expected to call gc - # directly; however, they already have to handle this case since - # incomplete generators are well documented as having problems with - # pypy's gc. - - # See: - # http://doc.pypy.org/en/latest/cpython_differences.html#differences-related-to-garbage-collection-strategies - # https://bitbucket.org/pypy/pypy/issues/736 - del c - import gc - - gc.collect() - if nr_transaction: - _test = background_task(name="test_incomplete_coroutine")(_test) + _test = background_task(name="test_incomplete_async_generator")(_test) event_loop.run_until_complete(_test()) if nr_transaction: _test_incomplete_async_generator = validate_transaction_metrics( - "test_incomplete_coroutine", + "test_incomplete_async_generator", background_task=True, scoped_metrics=[("Function/agen", 1)], rollup_metrics=[("Function/agen", 1)], )(_test_incomplete_async_generator) _test_incomplete_async_generator() + + +def test_incomplete_async_generator_transaction_exited(event_loop): + @function_trace(name="agen") + async def agen(): + for _ in range(5): + yield + + @validate_transaction_metrics( + "test_incomplete_async_generator", + background_task=True, + scoped_metrics=[("Function/agen", 1)], + rollup_metrics=[("Function/agen", 1)], + ) + def _test_incomplete_async_generator(): + c = agen() + @background_task(name="test_incomplete_async_generator") + async def _test(): + async for _ in c: + break + + event_loop.run_until_complete(_test()) + + # Remove generator after transaction completes + del c + + _test_incomplete_async_generator() diff --git a/tests/agent_features/test_async_wrapper_detection.py b/tests/agent_features/test_async_wrapper_detection.py new file mode 100644 index 0000000000..bb1fd3f1e3 --- /dev/null +++ b/tests/agent_features/test_async_wrapper_detection.py @@ -0,0 +1,102 @@ +# 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 + +import functools +import time + +from newrelic.api.background_task import background_task +from newrelic.api.database_trace import database_trace +from newrelic.api.datastore_trace import datastore_trace +from newrelic.api.external_trace import external_trace +from newrelic.api.function_trace import function_trace +from newrelic.api.graphql_trace import graphql_operation_trace, graphql_resolver_trace +from newrelic.api.memcache_trace import memcache_trace +from newrelic.api.message_trace import message_trace + +from newrelic.common.async_wrapper import generator_wrapper + +from testing_support.fixtures import capture_transaction_metrics +from testing_support.validators.validate_transaction_metrics import ( + validate_transaction_metrics, +) + +trace_metric_cases = [ + (functools.partial(function_trace, name="simple_gen"), "Function/simple_gen"), + (functools.partial(external_trace, library="lib", url="http://foo.com"), "External/foo.com/lib/"), + (functools.partial(database_trace, "select * from foo"), "Datastore/statement/None/foo/select"), + (functools.partial(datastore_trace, "lib", "foo", "bar"), "Datastore/statement/lib/foo/bar"), + (functools.partial(message_trace, "lib", "op", "typ", "name"), "MessageBroker/lib/typ/op/Named/name"), + (functools.partial(memcache_trace, "cmd"), "Memcache/cmd"), + (functools.partial(graphql_operation_trace), "GraphQL/operation/GraphQL///"), + (functools.partial(graphql_resolver_trace), "GraphQL/resolve/GraphQL/"), +] + + +@pytest.mark.parametrize("trace,metric", trace_metric_cases) +def test_automatic_generator_trace_wrapper(trace, metric): + metrics = [] + full_metrics = {} + + @capture_transaction_metrics(metrics, full_metrics) + @validate_transaction_metrics( + "test_automatic_generator_trace_wrapper", background_task=True, scoped_metrics=[(metric, 1)], rollup_metrics=[(metric, 1)] + ) + @background_task(name="test_automatic_generator_trace_wrapper") + def _test(): + @trace() + def gen(): + time.sleep(0.1) + yield + time.sleep(0.1) + + for _ in gen(): + pass + + _test() + + # Check that generators time the total call time (including pauses) + metric_key = (metric, "") + assert full_metrics[metric_key].total_call_time >= 0.2 + + +@pytest.mark.parametrize("trace,metric", trace_metric_cases) +def test_manual_generator_trace_wrapper(trace, metric): + metrics = [] + full_metrics = {} + + @capture_transaction_metrics(metrics, full_metrics) + @validate_transaction_metrics( + "test_automatic_generator_trace_wrapper", background_task=True, scoped_metrics=[(metric, 1)], rollup_metrics=[(metric, 1)] + ) + @background_task(name="test_automatic_generator_trace_wrapper") + def _test(): + @trace(async_wrapper=generator_wrapper) + def wrapper_func(): + """Function that returns a generator object, obscuring the automatic introspection of async_wrapper()""" + def gen(): + time.sleep(0.1) + yield + time.sleep(0.1) + return gen() + + for _ in wrapper_func(): + pass + + _test() + + # Check that generators time the total call time (including pauses) + metric_key = (metric, "") + assert full_metrics[metric_key].total_call_time >= 0.2 diff --git a/tests/agent_features/test_coroutine_trace.py b/tests/agent_features/test_coroutine_trace.py index 36e365bc46..a30f1e70a8 100644 --- a/tests/agent_features/test_coroutine_trace.py +++ b/tests/agent_features/test_coroutine_trace.py @@ -31,6 +31,7 @@ from newrelic.api.datastore_trace import datastore_trace from newrelic.api.external_trace import external_trace from newrelic.api.function_trace import function_trace +from newrelic.api.graphql_trace import graphql_operation_trace, graphql_resolver_trace from newrelic.api.memcache_trace import memcache_trace from newrelic.api.message_trace import message_trace @@ -47,6 +48,8 @@ (functools.partial(datastore_trace, "lib", "foo", "bar"), "Datastore/statement/lib/foo/bar"), (functools.partial(message_trace, "lib", "op", "typ", "name"), "MessageBroker/lib/typ/op/Named/name"), (functools.partial(memcache_trace, "cmd"), "Memcache/cmd"), + (functools.partial(graphql_operation_trace), "GraphQL/operation/GraphQL///"), + (functools.partial(graphql_resolver_trace), "GraphQL/resolve/GraphQL/"), ], ) def test_coroutine_timing(trace, metric): From 844e556abfbca63573e51a2647141e07ce9e942f Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Wed, 2 Aug 2023 15:09:49 -0700 Subject: [PATCH 30/31] Remove custom wrapper code from firestore --- newrelic/hooks/datastore_firestore.py | 49 +++++++-------------------- 1 file changed, 12 insertions(+), 37 deletions(-) diff --git a/newrelic/hooks/datastore_firestore.py b/newrelic/hooks/datastore_firestore.py index 69fd8f1cc9..3b83716f9d 100644 --- a/newrelic/hooks/datastore_firestore.py +++ b/newrelic/hooks/datastore_firestore.py @@ -12,13 +12,9 @@ # 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, async_generator_wrapper -from newrelic.api.datastore_trace import DatastoreTrace _get_object_id = lambda obj, *args, **kwargs: getattr(obj, "id", None) @@ -26,27 +22,6 @@ _get_collection_ref_id = lambda obj, *args, **kwargs: getattr(getattr(obj, "_collection_ref", None), "id", None) -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 = async_wrapper(wrapped, trace) - return wrapped(*args, **kwargs) - - class_ = getattr(module, class_name) - if class_ is not None: - if hasattr(class_, method_name): - 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( @@ -59,7 +34,7 @@ def instrument_google_cloud_firestore_v1_client(module): class_ = module.Client for method in ("collections", "get_all"): if hasattr(class_, method): - wrap_generator_method(module, "Client", method, target=None) + wrap_datastore_trace(module, "Client.%s" % method, operation=method, product="Firestore", target=None, async_wrapper=generator_wrapper) def instrument_google_cloud_firestore_v1_async_client(module): @@ -67,7 +42,7 @@ def instrument_google_cloud_firestore_v1_async_client(module): class_ = module.AsyncClient for method in ("collections", "get_all"): if hasattr(class_, method): - wrap_async_generator_method(module, "AsyncClient", method, target=None) + wrap_datastore_trace(module, "AsyncClient.%s" % method, operation=method, product="Firestore", target=None, async_wrapper=async_generator_wrapper) def instrument_google_cloud_firestore_v1_collection(module): @@ -85,7 +60,7 @@ def instrument_google_cloud_firestore_v1_collection(module): for method in ("stream", "list_documents"): if hasattr(class_, method): - wrap_generator_method(module, "CollectionReference", method, target=_get_object_id) + wrap_datastore_trace(module, "CollectionReference.%s" % method, operation=method, product="Firestore", target=_get_object_id, async_wrapper=generator_wrapper) def instrument_google_cloud_firestore_v1_async_collection(module): @@ -99,7 +74,7 @@ def instrument_google_cloud_firestore_v1_async_collection(module): for method in ("stream", "list_documents"): if hasattr(class_, method): - wrap_async_generator_method(module, "AsyncCollectionReference", method, target=_get_object_id) + wrap_datastore_trace(module, "AsyncCollectionReference.%s" % method, operation=method, product="Firestore", target=_get_object_id, async_wrapper=async_generator_wrapper) def instrument_google_cloud_firestore_v1_document(module): @@ -117,7 +92,7 @@ def instrument_google_cloud_firestore_v1_document(module): for method in ("collections",): if hasattr(class_, method): - wrap_generator_method(module, "DocumentReference", method, target=_get_object_id) + wrap_datastore_trace(module, "DocumentReference.%s" % method, operation=method, product="Firestore", target=_get_object_id, async_wrapper=generator_wrapper) def instrument_google_cloud_firestore_v1_async_document(module): @@ -131,7 +106,7 @@ def instrument_google_cloud_firestore_v1_async_document(module): for method in ("collections",): if hasattr(class_, method): - wrap_async_generator_method(module, "AsyncDocumentReference", method, target=_get_object_id) + wrap_datastore_trace(module, "AsyncDocumentReference.%s" % method, operation=method, product="Firestore", target=_get_object_id, async_wrapper=async_generator_wrapper) def instrument_google_cloud_firestore_v1_query(module): @@ -145,13 +120,13 @@ def instrument_google_cloud_firestore_v1_query(module): for method in ("stream",): if hasattr(class_, method): - wrap_generator_method(module, "Query", method, target=_get_parent_id) + wrap_datastore_trace(module, "Query.%s" % method, operation=method, product="Firestore", target=_get_parent_id, async_wrapper=generator_wrapper) 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) + wrap_datastore_trace(module, "CollectionGroup.%s" % method, operation=method, product="Firestore", target=_get_parent_id, async_wrapper=generator_wrapper) def instrument_google_cloud_firestore_v1_async_query(module): @@ -165,13 +140,13 @@ def instrument_google_cloud_firestore_v1_async_query(module): for method in ("stream",): if hasattr(class_, method): - wrap_async_generator_method(module, "AsyncQuery", method, target=_get_parent_id) + wrap_datastore_trace(module, "AsyncQuery.%s" % method, operation=method, product="Firestore", target=_get_parent_id, async_wrapper=async_generator_wrapper) if hasattr(module, "AsyncCollectionGroup"): class_ = module.AsyncCollectionGroup for method in ("get_partitions",): if hasattr(class_, method): - wrap_async_generator_method(module, "AsyncCollectionGroup", method, target=_get_parent_id) + wrap_datastore_trace(module, "AsyncCollectionGroup.%s" % method, operation=method, product="Firestore", target=_get_parent_id, async_wrapper=async_generator_wrapper) def instrument_google_cloud_firestore_v1_aggregation(module): @@ -189,7 +164,7 @@ def instrument_google_cloud_firestore_v1_aggregation(module): for method in ("stream",): if hasattr(class_, method): - wrap_generator_method(module, "AggregationQuery", method, target=_get_collection_ref_id) + wrap_datastore_trace(module, "AggregationQuery.%s" % method, operation=method, product="Firestore", target=_get_collection_ref_id, async_wrapper=generator_wrapper) def instrument_google_cloud_firestore_v1_async_aggregation(module): @@ -203,7 +178,7 @@ def instrument_google_cloud_firestore_v1_async_aggregation(module): for method in ("stream",): if hasattr(class_, method): - wrap_async_generator_method(module, "AsyncAggregationQuery", method, target=_get_collection_ref_id) + wrap_datastore_trace(module, "AsyncAggregationQuery.%s" % method, operation=method, product="Firestore", target=_get_collection_ref_id, async_wrapper=async_generator_wrapper) def instrument_google_cloud_firestore_v1_batch(module): From dabfc4af5c586c05ff01bba740ded0cf83fe5dfa Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Thu, 3 Aug 2023 09:35:22 -0700 Subject: [PATCH 31/31] Undo wrapper edits --- newrelic/common/async_wrapper.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/newrelic/common/async_wrapper.py b/newrelic/common/async_wrapper.py index a2fd22f02c..575b30e308 100644 --- a/newrelic/common/async_wrapper.py +++ b/newrelic/common/async_wrapper.py @@ -21,15 +21,10 @@ is_async_generator_function, ) -try: - import asyncio -except ImportError: - asyncio = None - def evaluate_wrapper(wrapper_string, wrapped, trace): values = {'wrapper': None, 'wrapped': wrapped, - 'trace': trace, 'functools': functools, "asyncio": asyncio} + 'trace': trace, 'functools': functools} exec(wrapper_string, values) return values['wrapper'] @@ -50,6 +45,8 @@ async def wrapper(*args, **kwargs): def awaitable_generator_wrapper(wrapped, trace): WRAPPER = textwrap.dedent(""" + import asyncio + @functools.wraps(wrapped) @asyncio.coroutine def wrapper(*args, **kwargs):