Firestore Sync Client Instrumentation

newrelic/config.py

@@ -2269,23 +2269,52 @@ 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",
+        "google.cloud.firestore_v1.query",
         "newrelic.hooks.datastore_firestore",
-        "instrument_google_cloud_firestore_v1_base_client",
+        "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(
+        "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(
+        "google.cloud.firestore_v1.transaction",
+        "newrelic.hooks.datastore_firestore",
+        "instrument_google_cloud_firestore_v1_transaction",
+    )
+
     _process_module_definition(
         "ariadne.asgi",
         "newrelic.hooks.framework_ariadne",

newrelic/hooks/datastore_firestore.py

@@ -12,22 +12,24 @@
 # 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.datastore_trace import DatastoreTrace, 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
+from newrelic.common.object_wrapper import wrap_function_wrapper
 
 
-_get_object_id = lambda obj, *args, **kwargs: obj.id
+_get_object_id = lambda obj, *args, **kwargs: getattr(obj, "id", None)
+_get_parent_id = lambda obj, *args, **kwargs: getattr(getattr(obj, "_parent", None), "id", None)
+_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):
+def wrap_generator_method(module, class_name, method_name, target):
     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)
-    
+
     class_ = getattr(module, class_name)
     if class_ is not None:
         if hasattr(class_, method_name):
@@ -41,18 +43,30 @@ 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
         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"):
             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):
@@ -61,9 +75,82 @@ 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",):
             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)
+
+    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_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_collection_ref_id,
+                    operation=method,
+                )
+
+        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
+                )
+
+
+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
+                )

tests/datastore_firestore/conftest.py

@@ -15,47 +15,64 @@
 import uuid
 
 import pytest
-
 from google.cloud.firestore import Client
-
 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
+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")
+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.
 
-@pytest.fixture(scope="function", autouse=True)
-def reset_firestore(client):
-    for coll in client.collections():
-        for document in coll.list_documents():
-            document.delete()
+    return _assert_trace_for_generator

tests/datastore_firestore/test_batching.py

@@ -0,0 +1,103 @@
+# 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_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 =====
+
+
+@pytest.fixture()
+def exercise_write_batch(client, collection):
+    def _exercise_write_batch():
+        docs = [collection.document(str(x)) for x in range(1, 4)]
+        batch = client.batch()
+        for doc in docs:
+            batch.set(doc, {})
+
+        batch.commit()
+    return _exercise_write_batch
+
+
+def test_firestore_write_batch(exercise_write_batch):
+    _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()
+
+    _test()
+
+
+# ===== BulkWriteBatch =====
+
+
+@pytest.fixture()
+def exercise_bulk_write_batch(client, collection):
+    def _exercise_bulk_write_batch():
+        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()
+    return _exercise_bulk_write_batch
+
+
+def test_firestore_bulk_write_batch(exercise_bulk_write_batch):
+    _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()
+
+    _test()