Add instrumentation for Google Firestore documents and collections

newrelic/config.py

@@ -2269,6 +2269,35 @@ def _process_module_builtin_defaults():
         "instrument_graphql_validate",
     )
 
+    _process_module_definition(
+        "google.cloud.firestore_v1.base_document",
+        "newrelic.hooks.database_firestore",
+        "instrument_google_cloud_firestore_v1_base_document",
+    )
+
+    _process_module_definition(
+        "google.cloud.firestore_v1.base_collection",
+        "newrelic.hooks.database_firestore",
+        "instrument_google_cloud_firestore_v1_base_collection",
+    )
+
+    _process_module_definition(
+        "google.cloud.firestore_v1.document",
+        "newrelic.hooks.database_firestore",
+        "instrument_google_cloud_firestore_v1_document",
+    )
+
+    _process_module_definition(
+        "google.cloud.firestore_v1.collection",
+        "newrelic.hooks.database_firestore",
+        "instrument_google_cloud_firestore_v1_collection",
+    )
+
+    _process_module_definition(
+        "google.cloud.firestore_v1.base_client",
+        "newrelic.hooks.database_firestore",
+        "instrument_google_cloud_firestore_v1_base_client",
+    )
     _process_module_definition(
         "ariadne.asgi",
         "newrelic.hooks.framework_ariadne",

newrelic/hooks/datastore_firestore.py

@@ -0,0 +1,116 @@
+# 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 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.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
+
+
+def wrap_generator_method(module, class_name, method_name):
+    def _wrapper(wrapped, instance, args, kwargs):
+        trace = DatastoreTrace(product="Firestore", target=instance.id, operation=method_name)
+        wrapped = generator_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)
+
+
+def instrument_google_cloud_firestore_v1_base_client(module):
+    rollup = ("Datastore/all", "Datastore/Firestore/all")
+    wrap_function_trace(
+        module, "BaseClient.__init__", name="%s:BaseClient.__init__" % module.__name__, terminal=True, rollup=rollup
+    )
+
+
+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
+
+    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
+
+
+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

tests/datastore_firestore/conftest.py

@@ -0,0 +1,57 @@
+# 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 uuid
+
+import pytest
+
+import firebase_admin
+from firebase_admin import credentials
+from firebase_admin import firestore
+
+from testing_support.fixtures import collector_agent_registration_fixture, collector_available_fixture  # noqa: F401; pylint: disable=W0611
+
+
+_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
+}
+
+collector_agent_registration = collector_agent_registration_fixture(
+        app_name='Python Agent Test (datastore_firestore)',
+        default_settings=_default_settings,
+        linked_applications=['Python Agent Test (datastore)'])
+
+
+@pytest.fixture(scope="session")
+def client():
+    creds = credentials.ApplicationDefault()
+
+    firebase_admin.initialize_app(creds)
+    client = firestore.client()
+    return client
+
+
+@pytest.fixture(scope="function")
+def collection(client):
+    collection = client.collection("firestore_collection_" + str(uuid.uuid4()))
+
+    yield collection
+
+    for doc in collection.list_documents():
+        doc.delete()

tests/datastore_firestore/test_collections.py

@@ -0,0 +1,96 @@
+# 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 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")
+    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 = list(collection.stream())
+    assert len(documents_stream) == 2
+    documents_list = list(collection.list_documents())
+    assert len(documents_list) == 2
+
+
+def test_firestore_collections(collection):
+    _test_scoped_metrics = [
+        ("Datastore/statement/Firestore/%s/stream" % collection.id, 1),
+        ("Datastore/statement/Firestore/%s/get" % collection.id, 1),
+        ("Datastore/statement/Firestore/%s/list_documents" % collection.id, 1),
+        ("Datastore/statement/Firestore/%s/add" % 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_transaction_metrics(
+        "test_firestore_collections",
+        scoped_metrics=_test_scoped_metrics,
+        rollup_metrics=_test_rollup_metrics,
+        background_task=True,
+    )
+    @background_task(name="test_firestore_collections")
+    def _test():
+        _exercise_firestore(collection)
+
+    _test()
+
+
+@background_task()
+def test_firestore_collections_generators(collection):
+    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
+
+
+@validate_database_duration()
+@background_task()
+def test_firestore_collections_db_duration(collection):
+    _exercise_firestore(collection)