Merge branch 'main' into fix-add-missing-postgresql-tests

Author: TimPansino

newrelic/config.py

@@ -3742,7 +3742,12 @@ def _process_module_builtin_defaults():
         "valkey.commands.graph.commands", "newrelic.hooks.datastore_valkey", "instrument_valkey_commands_graph_commands"
     )
 
-    _process_module_definition("motor", "newrelic.hooks.datastore_motor", "patch_motor")
+    _process_module_definition(
+        "motor.motor_asyncio", "newrelic.hooks.datastore_motor", "instrument_motor_motor_asyncio"
+    )
+    _process_module_definition(
+        "motor.motor_tornado", "newrelic.hooks.datastore_motor", "instrument_motor_motor_tornado"
+    )
 
     _process_module_definition(
         "piston.resource",

newrelic/hooks/datastore_motor.py

@@ -12,35 +12,134 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from newrelic.api.datastore_trace import DatastoreTrace
+from newrelic.api.function_trace import wrap_function_trace
 from newrelic.common.object_wrapper import wrap_function_wrapper
 
-# This is NOT a fully-featured instrumentation for the motor library. Instead
-# this is a monkey-patch of the motor library to work around a bug that causes
-# the __name__ lookup on a MotorCollection object to fail. This bug was causing
-# customer's applications to fail when they used motor in Tornado applications.
+_motor_client_sync_methods = (
+    "aggregate_raw_batches",
+    "aggregate",
+    "find_raw_batches",
+    "find",
+    "list_indexes",
+    "list_search_indexes",
+    "watch",
+)
 
+_motor_client_async_methods = (
+    "bulk_write",
+    "count_documents",
+    "create_index",
+    "create_indexes",
+    "create_search_index",
+    "create_search_indexes",
+    "delete_many",
+    "delete_one",
+    "distinct",
+    "drop_index",
+    "drop_indexes",
+    "drop_search_index",
+    "drop",
+    "estimated_document_count",
+    "find_one_and_delete",
+    "find_one_and_replace",
+    "find_one_and_update",
+    "find_one",
+    "index_information",
+    "insert_many",
+    "insert_one",
+    "options",
+    "rename",
+    "replace_one",
+    "update_many",
+    "update_one",
+    "update_search_index",
+)
 
-def _nr_wrapper_Motor_getattr_(wrapped, instance, args, kwargs):
 
-    def _bind_params(name, *args, **kwargs):
-        return name
+def instance_info(collection):
+    try:
+        nodes = collection.database.client.nodes
+        if len(nodes) == 1:
+            return next(iter(nodes))
+    except Exception:
+        pass
 
-    name = _bind_params(*args, **kwargs)
+    # If there are 0 nodes we're not currently connected, return nothing.
+    # If there are 2+ nodes we're in a load balancing setup.
+    # Unfortunately we can't rely on a deeper method to determine the actual server we're connected to in all cases.
+    # We can't report more than 1 server for instance info, so we opt here to ignore reporting the host/port and
+    # leave it empty to avoid confusing customers by guessing and potentially reporting the wrong server.
+    return None, None
 
-    if name.startswith('__') or name.startswith('_nr_'):
-        raise AttributeError(f'{instance.__class__.__name__} class has no attribute {name}. To access use object[{name!r}].')
 
-    return wrapped(*args, **kwargs)
+def wrap_motor_method(module, class_name, method_name, is_async=False):
+    cls = getattr(module, class_name)
+    if not hasattr(cls, method_name):
+        return
 
+    # Define wrappers as closures to preserve method_name
+    def _wrap_motor_method_sync(wrapped, instance, args, kwargs):
+        target = getattr(instance, "name", None)
+        database_name = getattr(getattr(instance, "database", None), "name", None)
+        with DatastoreTrace(
+            product="MongoDB", target=target, operation=method_name, database_name=database_name
+        ) as trace:
+            response = wrapped(*args, **kwargs)
 
-def patch_motor(module):
-    if (hasattr(module, 'version_tuple') and
-            module.version_tuple >= (0, 6)):
-        return
+            # Gather instance info after response to ensure client is conncected
+            address = instance_info(instance)
+            trace.host = address[0]
+            trace.port_path_or_id = address[1]
+
+            return response
+
+    async def _wrap_motor_method_async(wrapped, instance, args, kwargs):
+        target = getattr(instance, "name", None)
+        database_name = getattr(getattr(instance, "database", None), "name", None)
+        with DatastoreTrace(
+            product="MongoDB", target=target, operation=method_name, database_name=database_name
+        ) as trace:
+            response = await wrapped(*args, **kwargs)
+
+            # Gather instance info after response to ensure client is conncected
+            address = instance_info(instance)
+            trace.host = address[0]
+            trace.port_path_or_id = address[1]
+
+            return response
+
+    wrapper = _wrap_motor_method_async if is_async else _wrap_motor_method_sync
+    wrap_function_wrapper(module, f"{class_name}.{method_name}", wrapper)
+
+
+def instrument_motor_motor_asyncio(module):
+    if hasattr(module, "AsyncIOMotorClient"):
+        rollup = ("Datastore/all", "Datastore/MongoDB/all")
+        # Name function explicitly as motor and pymongo have a history of overriding the
+        # __getattr__() method in a way that breaks introspection.
+        wrap_function_trace(
+            module, "AsyncIOMotorClient.__init__", name=f"{module.__name__}:AsyncIOMotorClient.__init__", rollup=rollup
+        )
+
+    if hasattr(module, "AsyncIOMotorCollection"):
+        for method_name in _motor_client_sync_methods:
+            wrap_motor_method(module, "AsyncIOMotorCollection", method_name, is_async=False)
+        for method_name in _motor_client_async_methods:
+            wrap_motor_method(module, "AsyncIOMotorCollection", method_name, is_async=True)
+
+
+def instrument_motor_motor_tornado(module):
+    if hasattr(module, "MotorClient"):
+        rollup = ("Datastore/all", "Datastore/MongoDB/all")
+        # Name function explicitly as motor and pymongo have a history of overriding the
+        # __getattr__() method in a way that breaks introspection.
+        wrap_function_trace(
+            module, "MotorClient.__init__", name=f"{module.__name__}:MotorClient.__init__", rollup=rollup
+        )
 
-    patched_classes = ['MotorClient', 'MotorReplicaSetClient', 'MotorDatabase',
-            'MotorCollection']
-    for patched_class in patched_classes:
-        if hasattr(module, patched_class):
-            wrap_function_wrapper(module, f"{patched_class}.__getattr__",
-                    _nr_wrapper_Motor_getattr_)
+    if hasattr(module, "MotorCollection"):
+        for method_name in _motor_client_sync_methods:
+            wrap_motor_method(module, "MotorCollection", method_name, is_async=False)
+        for method_name in _motor_client_async_methods:
+            wrap_motor_method(module, "MotorCollection", method_name, is_async=True)

tests/datastore_motor/conftest.py

@@ -0,0 +1,69 @@
+# 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.db_settings import mongodb_settings
+from testing_support.fixture.event_loop import event_loop as loop  # noqa
+from testing_support.fixtures import (  # noqa: F401; pylint: disable=W0611
+    collector_agent_registration_fixture,
+    collector_available_fixture,
+)
+
+_default_settings = {
+    "package_reporting.enabled": False,  # Turn off package reporting for testing as it causes slow downs.
+    "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,
+}
+
+collector_agent_registration = collector_agent_registration_fixture(
+    app_name="Python Agent Test (datastore_pymongo)",
+    default_settings=_default_settings,
+    linked_applications=["Python Agent Test (datastore)"],
+)
+
+DB_SETTINGS = mongodb_settings()[0]
+MONGODB_HOST = DB_SETTINGS["host"]
+MONGODB_PORT = DB_SETTINGS["port"]
+MONGODB_COLLECTION = DB_SETTINGS["collection"]
+
+
+@pytest.fixture(scope="session", params=["asyncio", "tornado"])
+def implementation(request):
+    return request.param
+
+
+@pytest.fixture(scope="session")
+def client(implementation):
+    from motor.motor_asyncio import AsyncIOMotorClient
+    from motor.motor_tornado import MotorClient as TornadoMotorClient
+
+    # Must be actually initialized in test function, so provide a callable that returns the client.
+    def _client():
+        if implementation == "asyncio":
+            return AsyncIOMotorClient(host=MONGODB_HOST, port=MONGODB_PORT)
+        else:
+            return TornadoMotorClient(host=MONGODB_HOST, port=MONGODB_PORT)
+
+    return _client
+
+
+@pytest.fixture(scope="session")
+def init_metric(implementation):
+    if implementation == "asyncio":
+        return ("Function/motor.motor_asyncio:AsyncIOMotorClient.__init__", 1)
+    else:
+        return ("Function/motor.motor_tornado:MotorClient.__init__", 1)