Trace Async Wrapper Argument (#886)

* Add async_wrapper to datastore_trace api

* Add async wrapper argument to all trace APIs

* Add testing for automatic and manual asyncwrappers

newrelic/api/database_trace.py

@@ -16,7 +16,7 @@
 import logging
 
 from newrelic.api.time_trace import TimeTrace, current_trace
-from newrelic.common.async_wrapper import async_wrapper
+from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper
 from newrelic.common.object_wrapper import FunctionWrapper, wrap_object
 from newrelic.core.database_node import DatabaseNode
 from newrelic.core.stack_trace import current_stack
@@ -244,9 +244,9 @@ def create_node(self):
         )
 
 
-def DatabaseTraceWrapper(wrapped, sql, dbapi2_module=None):
+def DatabaseTraceWrapper(wrapped, sql, dbapi2_module=None, async_wrapper=None):
     def _nr_database_trace_wrapper_(wrapped, instance, args, kwargs):
-        wrapper = async_wrapper(wrapped)
+        wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped)
         if not wrapper:
             parent = current_trace()
             if not parent:
@@ -273,9 +273,9 @@ def _nr_database_trace_wrapper_(wrapped, instance, args, kwargs):
     return FunctionWrapper(wrapped, _nr_database_trace_wrapper_)
 
 
-def database_trace(sql, dbapi2_module=None):
-    return functools.partial(DatabaseTraceWrapper, sql=sql, dbapi2_module=dbapi2_module)
+def database_trace(sql, dbapi2_module=None, async_wrapper=None):
+    return functools.partial(DatabaseTraceWrapper, sql=sql, dbapi2_module=dbapi2_module, async_wrapper=async_wrapper)
 
 
-def wrap_database_trace(module, object_path, sql, dbapi2_module=None):
-    wrap_object(module, object_path, DatabaseTraceWrapper, (sql, dbapi2_module))
+def wrap_database_trace(module, object_path, sql, dbapi2_module=None, async_wrapper=None):
+    wrap_object(module, object_path, DatabaseTraceWrapper, (sql, dbapi2_module, async_wrapper))

newrelic/api/datastore_trace.py

@@ -15,7 +15,7 @@
 import functools
 
 from newrelic.api.time_trace import TimeTrace, current_trace
-from newrelic.common.async_wrapper import async_wrapper
+from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper
 from newrelic.common.object_wrapper import FunctionWrapper, wrap_object
 from newrelic.core.datastore_node import DatastoreNode
 
@@ -135,7 +135,7 @@ def create_node(self):
         )
 
 
-def DatastoreTraceWrapper(wrapped, product, target, operation, host=None, port_path_or_id=None, database_name=None):
+def DatastoreTraceWrapper(wrapped, product, target, operation, host=None, port_path_or_id=None, database_name=None, async_wrapper=None):
     """Wraps a method to time datastore queries.
 
     :param wrapped: The function to apply the trace to.
@@ -158,6 +158,8 @@ def DatastoreTraceWrapper(wrapped, product, target, operation, host=None, port_p
     :param database_name: The name of database where the current query is being
                           executed.
     :type database_name: str
+    :param async_wrapper: An async trace wrapper from newrelic.common.async_wrapper.
+    :type async_wrapper: callable or None
     :rtype: :class:`newrelic.common.object_wrapper.FunctionWrapper`
 
     This is typically used to wrap datastore queries such as calls to Redis or
@@ -173,7 +175,7 @@ def DatastoreTraceWrapper(wrapped, product, target, operation, host=None, port_p
     """
 
     def _nr_datastore_trace_wrapper_(wrapped, instance, args, kwargs):
-        wrapper = async_wrapper(wrapped)
+        wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped)
         if not wrapper:
             parent = current_trace()
             if not parent:
@@ -242,7 +244,7 @@ def _nr_datastore_trace_wrapper_(wrapped, instance, args, kwargs):
     return FunctionWrapper(wrapped, _nr_datastore_trace_wrapper_)
 
 
-def datastore_trace(product, target, operation, host=None, port_path_or_id=None, database_name=None):
+def datastore_trace(product, target, operation, host=None, port_path_or_id=None, database_name=None, async_wrapper=None):
     """Decorator allows datastore query to be timed.
 
     :param product: The name of the vendor.
@@ -263,6 +265,8 @@ def datastore_trace(product, target, operation, host=None, port_path_or_id=None,
     :param database_name: The name of database where the current query is being
                           executed.
     :type database_name: str
+    :param async_wrapper: An async trace wrapper from newrelic.common.async_wrapper.
+    :type async_wrapper: callable or None
 
     This is typically used to decorate datastore queries such as calls to Redis
     or ElasticSearch.
@@ -284,11 +288,12 @@ def datastore_trace(product, target, operation, host=None, port_path_or_id=None,
         host=host,
         port_path_or_id=port_path_or_id,
         database_name=database_name,
+        async_wrapper=async_wrapper,
     )
 
 
 def wrap_datastore_trace(
-    module, object_path, product, target, operation, host=None, port_path_or_id=None, database_name=None
+    module, object_path, product, target, operation, host=None, port_path_or_id=None, database_name=None, async_wrapper=None
 ):
     """Method applies custom timing to datastore query.
 
@@ -314,6 +319,8 @@ def wrap_datastore_trace(
     :param database_name: The name of database where the current query is being
                           executed.
     :type database_name: str
+    :param async_wrapper: An async trace wrapper from newrelic.common.async_wrapper.
+    :type async_wrapper: callable or None
 
     This is typically used to time database query method calls such as Redis
     GET.
@@ -327,5 +334,5 @@ def wrap_datastore_trace(
 
     """
     wrap_object(
-        module, object_path, DatastoreTraceWrapper, (product, target, operation, host, port_path_or_id, database_name)
+        module, object_path, DatastoreTraceWrapper, (product, target, operation, host, port_path_or_id, database_name, async_wrapper)
     )

newrelic/api/external_trace.py

@@ -16,7 +16,7 @@
 
 from newrelic.api.cat_header_mixin import CatHeaderMixin
 from newrelic.api.time_trace import TimeTrace, current_trace
-from newrelic.common.async_wrapper import async_wrapper
+from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper
 from newrelic.common.object_wrapper import FunctionWrapper, wrap_object
 from newrelic.core.external_node import ExternalNode
 
@@ -66,9 +66,9 @@ def create_node(self):
         )
 
 
-def ExternalTraceWrapper(wrapped, library, url, method=None):
+def ExternalTraceWrapper(wrapped, library, url, method=None, async_wrapper=None):
     def dynamic_wrapper(wrapped, instance, args, kwargs):
-        wrapper = async_wrapper(wrapped)
+        wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped)
         if not wrapper:
             parent = current_trace()
             if not parent:
@@ -103,7 +103,7 @@ def dynamic_wrapper(wrapped, instance, args, kwargs):
             return wrapped(*args, **kwargs)
 
     def literal_wrapper(wrapped, instance, args, kwargs):
-        wrapper = async_wrapper(wrapped)
+        wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped)
         if not wrapper:
             parent = current_trace()
             if not parent:
@@ -125,9 +125,9 @@ def literal_wrapper(wrapped, instance, args, kwargs):
     return FunctionWrapper(wrapped, literal_wrapper)
 
 
-def external_trace(library, url, method=None):
-    return functools.partial(ExternalTraceWrapper, library=library, url=url, method=method)
+def external_trace(library, url, method=None, async_wrapper=None):
+    return functools.partial(ExternalTraceWrapper, library=library, url=url, method=method, async_wrapper=async_wrapper)
 
 
-def wrap_external_trace(module, object_path, library, url, method=None):
-    wrap_object(module, object_path, ExternalTraceWrapper, (library, url, method))
+def wrap_external_trace(module, object_path, library, url, method=None, async_wrapper=None):
+    wrap_object(module, object_path, ExternalTraceWrapper, (library, url, method, async_wrapper))

newrelic/api/function_trace.py

@@ -15,7 +15,7 @@
 import functools
 
 from newrelic.api.time_trace import TimeTrace, current_trace
-from newrelic.common.async_wrapper import async_wrapper
+from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper
 from newrelic.common.object_names import callable_name
 from newrelic.common.object_wrapper import FunctionWrapper, wrap_object
 from newrelic.core.function_node import FunctionNode
@@ -89,9 +89,9 @@ def create_node(self):
         )
 
 
-def FunctionTraceWrapper(wrapped, name=None, group=None, label=None, params=None, terminal=False, rollup=None):
+def FunctionTraceWrapper(wrapped, name=None, group=None, label=None, params=None, terminal=False, rollup=None, async_wrapper=None):
     def dynamic_wrapper(wrapped, instance, args, kwargs):
-        wrapper = async_wrapper(wrapped)
+        wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped)
         if not wrapper:
             parent = current_trace()
             if not parent:
@@ -147,7 +147,7 @@ def dynamic_wrapper(wrapped, instance, args, kwargs):
             return wrapped(*args, **kwargs)
 
     def literal_wrapper(wrapped, instance, args, kwargs):
-        wrapper = async_wrapper(wrapped)
+        wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped)
         if not wrapper:
             parent = current_trace()
             if not parent:
@@ -171,13 +171,13 @@ def literal_wrapper(wrapped, instance, args, kwargs):
     return FunctionWrapper(wrapped, literal_wrapper)
 
 
-def function_trace(name=None, group=None, label=None, params=None, terminal=False, rollup=None):
+def function_trace(name=None, group=None, label=None, params=None, terminal=False, rollup=None, async_wrapper=None):
     return functools.partial(
-        FunctionTraceWrapper, name=name, group=group, label=label, params=params, terminal=terminal, rollup=rollup
+        FunctionTraceWrapper, name=name, group=group, label=label, params=params, terminal=terminal, rollup=rollup, async_wrapper=async_wrapper
     )
 
 
 def wrap_function_trace(
-    module, object_path, name=None, group=None, label=None, params=None, terminal=False, rollup=None
+    module, object_path, name=None, group=None, label=None, params=None, terminal=False, rollup=None, async_wrapper=None
 ):
-    return wrap_object(module, object_path, FunctionTraceWrapper, (name, group, label, params, terminal, rollup))
+    return wrap_object(module, object_path, FunctionTraceWrapper, (name, group, label, params, terminal, rollup, async_wrapper))

newrelic/api/graphql_trace.py

@@ -16,7 +16,7 @@
 
 from newrelic.api.time_trace import TimeTrace, current_trace
 from newrelic.api.transaction import current_transaction
-from newrelic.common.async_wrapper import async_wrapper
+from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper
 from newrelic.common.object_wrapper import FunctionWrapper, wrap_object
 from newrelic.core.graphql_node import GraphQLOperationNode, GraphQLResolverNode
 
@@ -109,9 +109,9 @@ def set_transaction_name(self, priority=None):
             transaction.set_transaction_name(name, "GraphQL", priority=priority)
 
 
-def GraphQLOperationTraceWrapper(wrapped):
+def GraphQLOperationTraceWrapper(wrapped, async_wrapper=None):
     def _nr_graphql_trace_wrapper_(wrapped, instance, args, kwargs):
-        wrapper = async_wrapper(wrapped)
+        wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped)
         if not wrapper:
             parent = current_trace()
             if not parent:
@@ -130,12 +130,12 @@ def _nr_graphql_trace_wrapper_(wrapped, instance, args, kwargs):
     return FunctionWrapper(wrapped, _nr_graphql_trace_wrapper_)
 
 
-def graphql_operation_trace():
-    return functools.partial(GraphQLOperationTraceWrapper)
+def graphql_operation_trace(async_wrapper=None):
+    return functools.partial(GraphQLOperationTraceWrapper, async_wrapper=async_wrapper)
 
 
-def wrap_graphql_operation_trace(module, object_path):
-    wrap_object(module, object_path, GraphQLOperationTraceWrapper)
+def wrap_graphql_operation_trace(module, object_path, async_wrapper=None):
+    wrap_object(module, object_path, GraphQLOperationTraceWrapper, (async_wrapper,))
 
 
 class GraphQLResolverTrace(TimeTrace):
@@ -193,9 +193,9 @@ def create_node(self):
         )
 
 
-def GraphQLResolverTraceWrapper(wrapped):
+def GraphQLResolverTraceWrapper(wrapped, async_wrapper=None):
     def _nr_graphql_trace_wrapper_(wrapped, instance, args, kwargs):
-        wrapper = async_wrapper(wrapped)
+        wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped)
         if not wrapper:
             parent = current_trace()
             if not parent:
@@ -214,9 +214,9 @@ def _nr_graphql_trace_wrapper_(wrapped, instance, args, kwargs):
     return FunctionWrapper(wrapped, _nr_graphql_trace_wrapper_)
 
 
-def graphql_resolver_trace():
-    return functools.partial(GraphQLResolverTraceWrapper)
+def graphql_resolver_trace(async_wrapper=None):
+    return functools.partial(GraphQLResolverTraceWrapper, async_wrapper=async_wrapper)
 
 
-def wrap_graphql_resolver_trace(module, object_path):
-    wrap_object(module, object_path, GraphQLResolverTraceWrapper)
+def wrap_graphql_resolver_trace(module, object_path, async_wrapper=None):
+    wrap_object(module, object_path, GraphQLResolverTraceWrapper, (async_wrapper,))

newrelic/api/memcache_trace.py

@@ -15,7 +15,7 @@
 import functools
 
 from newrelic.api.time_trace import TimeTrace, current_trace
-from newrelic.common.async_wrapper import async_wrapper
+from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper
 from newrelic.common.object_wrapper import FunctionWrapper, wrap_object
 from newrelic.core.memcache_node import MemcacheNode
 
@@ -51,9 +51,9 @@ def create_node(self):
         )
 
 
-def MemcacheTraceWrapper(wrapped, command):
+def MemcacheTraceWrapper(wrapped, command, async_wrapper=None):
     def _nr_wrapper_memcache_trace_(wrapped, instance, args, kwargs):
-        wrapper = async_wrapper(wrapped)
+        wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped)
         if not wrapper:
             parent = current_trace()
             if not parent:
@@ -80,9 +80,9 @@ def _nr_wrapper_memcache_trace_(wrapped, instance, args, kwargs):
     return FunctionWrapper(wrapped, _nr_wrapper_memcache_trace_)
 
 
-def memcache_trace(command):
-    return functools.partial(MemcacheTraceWrapper, command=command)
+def memcache_trace(command, async_wrapper=None):
+    return functools.partial(MemcacheTraceWrapper, command=command, async_wrapper=async_wrapper)
 
 
-def wrap_memcache_trace(module, object_path, command):
-    wrap_object(module, object_path, MemcacheTraceWrapper, (command,))
+def wrap_memcache_trace(module, object_path, command, async_wrapper=None):
+    wrap_object(module, object_path, MemcacheTraceWrapper, (command, async_wrapper))

newrelic/api/message_trace.py

@@ -16,7 +16,7 @@
 
 from newrelic.api.cat_header_mixin import CatHeaderMixin
 from newrelic.api.time_trace import TimeTrace, current_trace
-from newrelic.common.async_wrapper import async_wrapper
+from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper
 from newrelic.common.object_wrapper import FunctionWrapper, wrap_object
 from newrelic.core.message_node import MessageNode
 
@@ -91,9 +91,9 @@ def create_node(self):
         )
 
 
-def MessageTraceWrapper(wrapped, library, operation, destination_type, destination_name, params={}, terminal=True):
+def MessageTraceWrapper(wrapped, library, operation, destination_type, destination_name, params={}, terminal=True, async_wrapper=None):
     def _nr_message_trace_wrapper_(wrapped, instance, args, kwargs):
-        wrapper = async_wrapper(wrapped)
+        wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped)
         if not wrapper:
             parent = current_trace()
             if not parent:
@@ -144,7 +144,7 @@ def _nr_message_trace_wrapper_(wrapped, instance, args, kwargs):
     return FunctionWrapper(wrapped, _nr_message_trace_wrapper_)
 
 
-def message_trace(library, operation, destination_type, destination_name, params={}, terminal=True):
+def message_trace(library, operation, destination_type, destination_name, params={}, terminal=True, async_wrapper=None):
     return functools.partial(
         MessageTraceWrapper,
         library=library,
@@ -153,10 +153,11 @@ def message_trace(library, operation, destination_type, destination_name, params
         destination_name=destination_name,
         params=params,
         terminal=terminal,
+        async_wrapper=async_wrapper,
     )
 
 
-def wrap_message_trace(module, object_path, library, operation, destination_type, destination_name, params={}, terminal=True):
+def wrap_message_trace(module, object_path, library, operation, destination_type, destination_name, params={}, terminal=True, async_wrapper=None):
     wrap_object(
-        module, object_path, MessageTraceWrapper, (library, operation, destination_type, destination_name, params, terminal)
+        module, object_path, MessageTraceWrapper, (library, operation, destination_type, destination_name, params, terminal, async_wrapper)
     )