Instrument Firestore batching

Author: TimPansino

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

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

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