Add query and aggregation query instrumentation

Author: TimPansino

newrelic/config.py

@@ -2289,6 +2289,16 @@ def _process_module_builtin_defaults():
         "newrelic.hooks.datastore_firestore",
         "instrument_google_cloud_firestore_v1_collection",
     )
+    _process_module_definition(
+        "google.cloud.firestore_v1.query",
+        "newrelic.hooks.datastore_firestore",
+        "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(
         "ariadne.asgi",

newrelic/hooks/datastore_firestore.py

@@ -19,10 +19,27 @@
 from newrelic.api.datastore_trace import DatastoreTrace
 
 
-_get_object_id = lambda obj, *args, **kwargs: obj.id
+def _get_object_id(obj, *args, **kwargs):
+    try:
+        return obj.id
+    except Exception:
+        return None
 
 
-def wrap_generator_method(module, class_name, method_name, target=_get_object_id):
+def _get_parent_id(obj, *args, **kwargs):
+    try:
+        return obj._parent.id
+    except Exception:
+        return None
+
+def _get_nested_query_parent_id(obj, *args, **kwargs):
+    try:
+        return obj._nested_query._parent.id
+    except Exception:
+        return None
+
+
+def wrap_generator_method(module, class_name, method_name, target):
     def _wrapper(wrapped, instance, args, kwargs):
         target_ = target(instance) if callable(target) else target
         trace = DatastoreTrace(product="Firestore", target=target_, operation=method_name)
@@ -61,7 +78,7 @@ def instrument_google_cloud_firestore_v1_collection(module):
 
         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):
@@ -75,4 +92,32 @@ def instrument_google_cloud_firestore_v1_document(module):
 
         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)
+
+
+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_nested_query_parent_id, operation=method
+                )
+
+        for method in ("stream",):
+            if hasattr(class_, method):
+                wrap_generator_method(module, "AggregationQuery", method, target=_get_nested_query_parent_id)

tests/datastore_firestore/test_aggregation_query.py

@@ -0,0 +1,71 @@
+# 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_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,
+)
+
+
+@pytest.fixture(autouse=True)
+def sample_data(collection, reset_firestore):
+    # reset_firestore must be run before, not after this fixture
+    for x in range(1, 6):
+        collection.add({"x": x})
+
+
+def _exercise_firestore(collection):
+    aggregation_query = collection.select("x").where("x", "<=", 3).count()
+    assert aggregation_query.get()[0][0].value == 3
+    assert list(aggregation_query.stream())[0][0].value == 3
+
+
+def test_firestore_aggregation_query(collection):
+    _test_scoped_metrics = [
+        ("Datastore/statement/Firestore/%s/stream" % collection.id, 1),
+        ("Datastore/statement/Firestore/%s/get" % collection.id, 1),
+    ]
+
+    _test_rollup_metrics = [
+        ("Datastore/operation/Firestore/get", 1),
+        ("Datastore/operation/Firestore/stream", 1),
+        ("Datastore/all", 2),
+        ("Datastore/allOther", 2),
+    ]
+    @validate_transaction_metrics(
+        "test_firestore_aggregation_query",
+        scoped_metrics=_test_scoped_metrics,
+        rollup_metrics=_test_rollup_metrics,
+        background_task=True,
+    )
+    @background_task(name="test_firestore_aggregation_query")
+    def _test():
+        _exercise_firestore(collection)
+
+    _test()
+
+
+@background_task()
+def test_firestore_aggregation_query_generators(collection, assert_trace_for_generator):
+    aggregation_query = collection.select("x").where("x", "<=", 3).count()
+    assert_trace_for_generator(aggregation_query.stream)
+
+
+@validate_database_duration()
+@background_task()
+def test_firestore_aggregation_query_db_duration(collection):
+    _exercise_firestore(collection)

tests/datastore_firestore/test_collections.py

@@ -65,7 +65,6 @@ def _test():
 
 @background_task()
 def test_firestore_collections_generators(collection, assert_trace_for_generator):
-    txn = current_trace()
     collection.add({})
     collection.add({})
     assert len(list(collection.list_documents())) == 2

tests/datastore_firestore/test_documents.py

@@ -74,8 +74,6 @@ def _test():
 
 @background_task()
 def test_firestore_documents_generators(collection, assert_trace_for_generator):
-    txn = current_trace()
-
     subcollection_doc = collection.document("SubCollections")
     subcollection_doc.set({})
     subcollection_doc.collection("collection1").add({})

tests/datastore_firestore/test_query.py

@@ -0,0 +1,71 @@
+# 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_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,
+)
+
+
+@pytest.fixture(autouse=True)
+def sample_data(collection, reset_firestore):
+    # reset_firestore must be run before, not after this fixture
+    for x in range(1, 6):
+        collection.add({"x": x})
+
+
+def _exercise_firestore(collection):
+    query = collection.select("x").limit(10).order_by("x").where("x", "<=", 3)
+    assert len(query.get()) == 3
+    assert len(list(query.stream())) == 3
+
+
+def test_firestore_query(collection):
+    _test_scoped_metrics = [
+        ("Datastore/statement/Firestore/%s/stream" % collection.id, 1),
+        ("Datastore/statement/Firestore/%s/get" % collection.id, 1),
+    ]
+
+    _test_rollup_metrics = [
+        ("Datastore/operation/Firestore/get", 1),
+        ("Datastore/operation/Firestore/stream", 1),
+        ("Datastore/all", 2),
+        ("Datastore/allOther", 2),
+    ]
+    @validate_transaction_metrics(
+        "test_firestore_query",
+        scoped_metrics=_test_scoped_metrics,
+        rollup_metrics=_test_rollup_metrics,
+        background_task=True,
+    )
+    @background_task(name="test_firestore_query")
+    def _test():
+        _exercise_firestore(collection)
+
+    _test()
+
+
+@background_task()
+def test_firestore_query_generators(collection, assert_trace_for_generator):
+    query = collection.select("x").where("x", "<=", 3)
+    assert_trace_for_generator(query.stream)
+
+
+@validate_database_duration()
+@background_task()
+def test_firestore_query_db_duration(collection):
+    _exercise_firestore(collection)