diff --git a/newrelic/config.py b/newrelic/config.py index 51aa2a44b..3e258a19c 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 672438c60..b59df5d0d 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 @@ -60,57 +41,29 @@ 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): - 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 - - -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 + 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 - - 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 + 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 method in ("collections",): + if hasattr(class_, method): + wrap_generator_method(module, "DocumentReference", method) diff --git a/tests/datastore_firestore/test_collections.py b/tests/datastore_firestore/test_collections.py index c416ec4e2..30f26eb8f 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 9850c2c58..be47e820f 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")