Refactor Firestore Hooks (#879)

* Remove unnecessary instrumentation

* Simplify existing instrumentation

* Remove unnecessary settings lookups

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",

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)

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")

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")