From 29096035c3371b1ce59c1774b30d3188950f6d02 Mon Sep 17 00:00:00 2001
From: Timothy Pansino <11214426+TimPansino@users.noreply.github.com>
Date: Thu, 8 Aug 2024 14:17:44 -0700
Subject: [PATCH] Botocore SQS Code Cleanup (#1187)

* Rename SQS specific function for clarity

* Trim out unnecessary extra function

* Fix bug where params was not passed through SQS wrapper (unused currently)

* Remove unused imports

---------

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Hannah Stepanek <hstepanek@newrelic.com>
---
 newrelic/hooks/external_botocore.py | 39 ++++++-----------------------
 1 file changed, 8 insertions(+), 31 deletions(-)

diff --git a/newrelic/hooks/external_botocore.py b/newrelic/hooks/external_botocore.py
index 4f0b7a6a71..e4ade6be4a 100644
--- a/newrelic/hooks/external_botocore.py
+++ b/newrelic/hooks/external_botocore.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import functools
 import json
 import logging
 import re
@@ -31,7 +30,6 @@
 from newrelic.api.transaction import current_transaction
 from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper
 from newrelic.common.object_wrapper import (
-    FunctionWrapper,
     ObjectProxy,
     function_wrapper,
     wrap_function_wrapper,
@@ -59,7 +57,7 @@ def extract_sqs(*args, **kwargs):
     return queue_value.rsplit("/", 1)[-1]
 
 
-def extract_agent_attrs(*args, **kwargs):
+def extract_sqs_agent_attrs(*args, **kwargs):
     # Try to capture AWS SQS info as agent attributes. Log any exception to debug.
     agent_attrs = {}
     try:
@@ -848,29 +846,8 @@ def sqs_message_trace(
     async_wrapper=None,
     extract_agent_attrs=None,
 ):
-    return functools.partial(
-        SQSMessageTraceWrapper,
-        operation=operation,
-        destination_type=destination_type,
-        destination_name=destination_name,
-        params=params,
-        terminal=terminal,
-        async_wrapper=async_wrapper,
-        extract_agent_attrs=extract_agent_attrs,
-    )
-
-
-def SQSMessageTraceWrapper(
-    wrapped,
-    operation,
-    destination_type,
-    destination_name,
-    params={},
-    terminal=True,
-    async_wrapper=None,
-    extract_agent_attrs=None,
-):
-    def _nr_message_trace_wrapper_(wrapped, instance, args, kwargs):
+    @function_wrapper
+    def _nr_sqs_message_trace_wrapper_(wrapped, instance, args, kwargs):
         wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped)
         if not wrapper:
             parent = current_trace()
@@ -889,7 +866,7 @@ def _nr_message_trace_wrapper_(wrapped, instance, args, kwargs):
             _operation,
             _destination_type,
             _destination_name,
-            params={},
+            params=params,
             terminal=terminal,
             parent=parent,
             source=wrapped,
@@ -905,7 +882,7 @@ def _nr_message_trace_wrapper_(wrapped, instance, args, kwargs):
         with trace:
             return wrapped(*args, **kwargs)
 
-    return FunctionWrapper(wrapped, _nr_message_trace_wrapper_)
+    return _nr_sqs_message_trace_wrapper_
 
 
 CUSTOM_TRACE_POINTS = {
@@ -919,13 +896,13 @@ def _nr_message_trace_wrapper_(wrapped, instance, args, kwargs):
     ("dynamodb", "query"): datastore_trace("DynamoDB", extract("TableName"), "query"),
     ("dynamodb", "scan"): datastore_trace("DynamoDB", extract("TableName"), "scan"),
     ("sqs", "send_message"): sqs_message_trace(
-        "Produce", "Queue", extract_sqs, extract_agent_attrs=extract_agent_attrs
+        "Produce", "Queue", extract_sqs, extract_agent_attrs=extract_sqs_agent_attrs
     ),
     ("sqs", "send_message_batch"): sqs_message_trace(
-        "Produce", "Queue", extract_sqs, extract_agent_attrs=extract_agent_attrs
+        "Produce", "Queue", extract_sqs, extract_agent_attrs=extract_sqs_agent_attrs
     ),
     ("sqs", "receive_message"): sqs_message_trace(
-        "Consume", "Queue", extract_sqs, extract_agent_attrs=extract_agent_attrs
+        "Consume", "Queue", extract_sqs, extract_agent_attrs=extract_sqs_agent_attrs
     ),
     ("bedrock-runtime", "invoke_model"): wrap_bedrock_runtime_invoke_model(response_streaming=False),
     ("bedrock-runtime", "invoke_model_with_response_stream"): wrap_bedrock_runtime_invoke_model(