From 9883c2b0580ce3c5e1927a4ff83fc6f72f687a7a Mon Sep 17 00:00:00 2001
From: Lalleh Rafeei <84813886+lrafeei@users.noreply.github.com>
Date: Wed, 28 Jun 2023 16:09:09 -0700
Subject: [PATCH] Replace drop_transaction logic by using transaction context
 manager (#832)

* Replace drop_transaction call

* [Mega-Linter] Apply linters fixes

* Empty commit to start tests

* Change logic in BG Wrappers

---------

Co-authored-by: lrafeei <lrafeei@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
---
 newrelic/api/background_task.py     | 38 ++++++-----------------------
 newrelic/api/message_transaction.py | 26 +++-----------------
 2 files changed, 10 insertions(+), 54 deletions(-)

diff --git a/newrelic/api/background_task.py b/newrelic/api/background_task.py
index a4a9e8e6a6..4cdcd8a0d4 100644
--- a/newrelic/api/background_task.py
+++ b/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))
diff --git a/newrelic/api/message_transaction.py b/newrelic/api/message_transaction.py
index 291a3897e6..54a71f6eff 100644
--- a/newrelic/api/message_transaction.py
+++ b/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)