Skip to content

Commit

Permalink
Merge branch 'develop-google-firestore-instrumentation' into feature-…
Browse files Browse the repository at this point in the history
…firstore-async-instrumentation
  • Loading branch information
mergify[bot] authored Aug 2, 2023
2 parents 4a8a3fe + 7596fb4 commit 29579fc
Show file tree
Hide file tree
Showing 3 changed files with 183 additions and 10 deletions.
90 changes: 81 additions & 9 deletions newrelic/api/datastore_trace.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,9 @@ def __enter__(self):
self.product = transaction._intern_string(self.product)
self.target = transaction._intern_string(self.target)
self.operation = transaction._intern_string(self.operation)
self.host = transaction._intern_string(self.host)
self.port_path_or_id = transaction._intern_string(self.port_path_or_id)
self.database_name = transaction._intern_string(self.database_name)

datastore_tracer_settings = transaction.settings.datastore_tracer
self.instance_reporting_enabled = datastore_tracer_settings.instance_reporting.enabled
Expand All @@ -92,7 +95,14 @@ def __repr__(self):
return "<%s object at 0x%x %s>" % (
self.__class__.__name__,
id(self),
dict(product=self.product, target=self.target, operation=self.operation),
dict(
product=self.product,
target=self.target,
operation=self.operation,
host=self.host,
port_path_or_id=self.port_path_or_id,
database_name=self.database_name,
),
)

def finalize_data(self, transaction, exc=None, value=None, tb=None):
Expand Down Expand Up @@ -125,7 +135,7 @@ def create_node(self):
)


def DatastoreTraceWrapper(wrapped, product, target, operation):
def DatastoreTraceWrapper(wrapped, product, target, operation, host=None, port_path_or_id=None, database_name=None):
"""Wraps a method to time datastore queries.
:param wrapped: The function to apply the trace to.
Expand All @@ -140,6 +150,14 @@ def DatastoreTraceWrapper(wrapped, product, target, operation):
or the name of any API function/method in the client
library.
:type operation: str or callable
:param host: The name of the server hosting the actual datastore.
:type host: str
:param port_path_or_id: The value passed in can represent either the port,
path, or id of the datastore being connected to.
:type port_path_or_id: str
:param database_name: The name of database where the current query is being
executed.
:type database_name: str
:rtype: :class:`newrelic.common.object_wrapper.FunctionWrapper`
This is typically used to wrap datastore queries such as calls to Redis or
Expand Down Expand Up @@ -187,7 +205,33 @@ def _nr_datastore_trace_wrapper_(wrapped, instance, args, kwargs):
else:
_operation = operation

trace = DatastoreTrace(_product, _target, _operation, parent=parent, source=wrapped)
if callable(host):
if instance is not None:
_host = host(instance, *args, **kwargs)
else:
_host = host(*args, **kwargs)
else:
_host = host

if callable(port_path_or_id):
if instance is not None:
_port_path_or_id = port_path_or_id(instance, *args, **kwargs)
else:
_port_path_or_id = port_path_or_id(*args, **kwargs)
else:
_port_path_or_id = port_path_or_id

if callable(database_name):
if instance is not None:
_database_name = database_name(instance, *args, **kwargs)
else:
_database_name = database_name(*args, **kwargs)
else:
_database_name = database_name

trace = DatastoreTrace(
_product, _target, _operation, _host, _port_path_or_id, _database_name, parent=parent, source=wrapped
)

if wrapper: # pylint: disable=W0125,W0126
return wrapper(wrapped, trace)(*args, **kwargs)
Expand All @@ -198,7 +242,7 @@ def _nr_datastore_trace_wrapper_(wrapped, instance, args, kwargs):
return FunctionWrapper(wrapped, _nr_datastore_trace_wrapper_)


def datastore_trace(product, target, operation):
def datastore_trace(product, target, operation, host=None, port_path_or_id=None, database_name=None):
"""Decorator allows datastore query to be timed.
:param product: The name of the vendor.
Expand All @@ -211,6 +255,14 @@ def datastore_trace(product, target, operation):
or the name of any API function/method in the client
library.
:type operation: str
:param host: The name of the server hosting the actual datastore.
:type host: str
:param port_path_or_id: The value passed in can represent either the port,
path, or id of the datastore being connected to.
:type port_path_or_id: str
:param database_name: The name of database where the current query is being
executed.
:type database_name: str
This is typically used to decorate datastore queries such as calls to Redis
or ElasticSearch.
Expand All @@ -224,10 +276,20 @@ def datastore_trace(product, target, operation):
... time.sleep(*args, **kwargs)
"""
return functools.partial(DatastoreTraceWrapper, product=product, target=target, operation=operation)


def wrap_datastore_trace(module, object_path, product, target, operation):
return functools.partial(
DatastoreTraceWrapper,
product=product,
target=target,
operation=operation,
host=host,
port_path_or_id=port_path_or_id,
database_name=database_name,
)


def wrap_datastore_trace(
module, object_path, product, target, operation, host=None, port_path_or_id=None, database_name=None
):
"""Method applies custom timing to datastore query.
:param module: Module containing the method to be instrumented.
Expand All @@ -244,6 +306,14 @@ def wrap_datastore_trace(module, object_path, product, target, operation):
or the name of any API function/method in the client
library.
:type operation: str
:param host: The name of the server hosting the actual datastore.
:type host: str
:param port_path_or_id: The value passed in can represent either the port,
path, or id of the datastore being connected to.
:type port_path_or_id: str
:param database_name: The name of database where the current query is being
executed.
:type database_name: str
This is typically used to time database query method calls such as Redis
GET.
Expand All @@ -256,4 +326,6 @@ def wrap_datastore_trace(module, object_path, product, target, operation):
... 'sleep')
"""
wrap_object(module, object_path, DatastoreTraceWrapper, (product, target, operation))
wrap_object(
module, object_path, DatastoreTraceWrapper, (product, target, operation, host, port_path_or_id, database_name)
)
89 changes: 89 additions & 0 deletions tests/agent_features/test_datastore_trace.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
# 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_datastore_trace_inputs import (
validate_datastore_trace_inputs,
)

from newrelic.api.background_task import background_task
from newrelic.api.datastore_trace import DatastoreTrace, DatastoreTraceWrapper


@validate_datastore_trace_inputs(
operation="test_operation",
target="test_target",
host="test_host",
port_path_or_id="test_port",
database_name="test_db_name",
)
@background_task()
def test_dt_trace_all_args():
with DatastoreTrace(
product="Agent Features",
target="test_target",
operation="test_operation",
host="test_host",
port_path_or_id="test_port",
database_name="test_db_name",
):
pass


@validate_datastore_trace_inputs(operation=None, target=None, host=None, port_path_or_id=None, database_name=None)
@background_task()
def test_dt_trace_empty():
with DatastoreTrace(product=None, target=None, operation=None):
pass


@background_task()
def test_dt_trace_callable_args():
def product_callable():
return "Agent Features"

def target_callable():
return "test_target"

def operation_callable():
return "test_operation"

def host_callable():
return "test_host"

def port_path_id_callable():
return "test_port"

def db_name_callable():
return "test_db_name"

@validate_datastore_trace_inputs(
operation="test_operation",
target="test_target",
host="test_host",
port_path_or_id="test_port",
database_name="test_db_name",
)
def _test():
pass

wrapped_fn = DatastoreTraceWrapper(
_test,
product=product_callable,
target=target_callable,
operation=operation_callable,
host=host_callable,
port_path_or_id=port_path_id_callable,
database_name=db_name_callable,
)
wrapped_fn()
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
"""


def validate_datastore_trace_inputs(operation=None, target=None):
def validate_datastore_trace_inputs(operation=None, target=None, host=None, port_path_or_id=None, database_name=None):
@transient_function_wrapper("newrelic.api.datastore_trace", "DatastoreTrace.__init__")
@catch_background_exceptions
def _validate_datastore_trace_inputs(wrapped, instance, args, kwargs):
Expand All @@ -44,6 +44,18 @@ def _bind_params(product, target, operation, host=None, port_path_or_id=None, da
assert captured_target == target, "%s didn't match expected %s" % (captured_target, target)
if operation is not None:
assert captured_operation == operation, "%s didn't match expected %s" % (captured_operation, operation)
if host is not None:
assert captured_host == host, "%s didn't match expected %s" % (captured_host, host)
if port_path_or_id is not None:
assert captured_port_path_or_id == port_path_or_id, "%s didn't match expected %s" % (
captured_port_path_or_id,
port_path_or_id,
)
if database_name is not None:
assert captured_database_name == database_name, "%s didn't match expected %s" % (
captured_database_name,
database_name,
)

return wrapped(*args, **kwargs)

Expand Down

0 comments on commit 29579fc

Please sign in to comment.