Replace drop_transaction logic by using transaction context manager

newrelic/api/background_task.py

@@ -13,19 +13,16 @@
 # limitations under the License.
 
 import functools
-import sys
 
 from newrelic.api.application import Application, application_instance
 from newrelic.api.transaction import Transaction, current_transaction
-from newrelic.common.async_proxy import async_proxy, TransactionContext
+from newrelic.common.async_proxy import TransactionContext, async_proxy
 from newrelic.common.object_names import callable_name
 from newrelic.common.object_wrapper import FunctionWrapper, wrap_object
 
 
 class BackgroundTask(Transaction):
-
     def __init__(self, application, name, group=None, source=None):
-
         # Initialise the common transaction base class.
 
         super(BackgroundTask, self).__init__(application, source=source)
@@ -53,7 +50,6 @@ def __init__(self, application, name, group=None, source=None):
 
 
 def BackgroundTaskWrapper(wrapped, application=None, name=None, group=None):
-
     def wrapper(wrapped, instance, args, kwargs):
         if callable(name):
             if instance is not None:
@@ -107,39 +103,19 @@ def create_transaction(transaction):
 
         manager = create_transaction(current_transaction(active_only=False))
 
+        # This means that a transaction already exists, so we want to return
         if not manager:
             return wrapped(*args, **kwargs)
-        success = True
-
-        try:
-            manager.__enter__()
-            try:
-                return wrapped(*args, **kwargs)
-            except:
-                success = False
-                if not manager.__exit__(*sys.exc_info()):
-                    raise
-        finally:
-            if success and manager._ref_count == 0:
-                manager._is_finalized = True
-                manager.__exit__(None, None, None)
-            else:
-                manager._request_handler_finalize = True
-                manager._server_adapter_finalize = True
 
-                old_transaction = current_transaction()
-                if old_transaction is not None:
-                    old_transaction.drop_transaction()
+        with manager:
+            return wrapped(*args, **kwargs)
 
     return FunctionWrapper(wrapped, wrapper)
 
 
 def background_task(application=None, name=None, group=None):
-    return functools.partial(BackgroundTaskWrapper,
-            application=application, name=name, group=group)
+    return functools.partial(BackgroundTaskWrapper, application=application, name=name, group=group)
 
 
-def wrap_background_task(module, object_path, application=None,
-        name=None, group=None):
-    wrap_object(module, object_path, BackgroundTaskWrapper,
-            (application, name, group))
+def wrap_background_task(module, object_path, application=None, name=None, group=None):
+    wrap_object(module, object_path, BackgroundTaskWrapper, (application, name, group))

newrelic/api/message_transaction.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import functools
-import sys
 
 from newrelic.api.application import Application, application_instance
 from newrelic.api.background_task import BackgroundTask
@@ -39,7 +38,6 @@ def __init__(
         transport_type="AMQP",
         source=None,
     ):
-
         name, group = self.get_transaction_name(library, destination_type, destination_name)
 
         super(MessageTransaction, self).__init__(application, name, group=group, source=source)
@@ -218,30 +216,12 @@ def create_transaction(transaction):
 
         manager = create_transaction(current_transaction(active_only=False))
 
+        # This means that transaction already exists and we want to return
         if not manager:
             return wrapped(*args, **kwargs)
 
-        success = True
-
-        try:
-            manager.__enter__()
-            try:
-                return wrapped(*args, **kwargs)
-            except:  # Catch all
-                success = False
-                if not manager.__exit__(*sys.exc_info()):
-                    raise
-        finally:
-            if success and manager._ref_count == 0:
-                manager._is_finalized = True
-                manager.__exit__(None, None, None)
-            else:
-                manager._request_handler_finalize = True
-                manager._server_adapter_finalize = True
-
-                old_transaction = current_transaction()
-                if old_transaction is not None:
-                    old_transaction.drop_transaction()
+        with manager:
+            return wrapped(*args, **kwargs)
 
     return FunctionWrapper(wrapped, wrapper)