implement specific instrumentation for AWS SQS (#1123)

* implement specific instrumentation for AWS SQS

closes #1026

* added docs

Author: beniwohli

docs/supported-technologies.asciidoc

@@ -543,6 +543,12 @@ Additionally, some services collect more specific data
 
  * Topic name
 
+[float]
+[[automatic-instrumentation-sqs]]
+===== SQS
+
+ * Queue name
+
 [float]
 [[automatic-instrumentation-template-engines]]
 === Template Engines

elasticapm/instrumentation/packages/botocore.py

@@ -30,9 +30,14 @@
 
 from collections import namedtuple
 
+from elasticapm.conf import constants
 from elasticapm.instrumentation.packages.base import AbstractInstrumentedModule
 from elasticapm.traces import capture_span
 from elasticapm.utils.compat import urlparse
+from elasticapm.utils.logging import get_logger
+
+logger = get_logger("elasticapm.instrument")
+
 
 HandlerInfo = namedtuple("HandlerInfo", ("signature", "span_type", "span_subtype", "span_action", "context"))
 
@@ -81,7 +86,9 @@ def call(self, module, method, wrapped, instance, args, kwargs):
             span_subtype=handler_info.span_subtype,
             span_action=handler_info.span_action,
             extra=handler_info.context,
-        ):
+        ) as span:
+            if service in span_modifiers:
+                span_modifiers[service](span, args, kwargs)
             return wrapped(*args, **kwargs)
 
 
@@ -141,8 +148,46 @@ def handle_sns(operation_name, service, instance, args, kwargs, context):
     return HandlerInfo(signature, span_type, span_subtype, span_action, context)
 
 
-def handle_sqs(operation_name, service, instance, args, kwargs, destination):
-    pass
+def handle_sqs(operation_name, service, instance, args, kwargs, context):
+    if operation_name not in ("SendMessage", "SendMessageBatch", "ReceiveMessage"):
+        # only "publish" is handled specifically, other endpoints get the default treatment
+        return False
+    span_type = "messaging"
+    span_subtype = "sqs"
+    span_action = "send" if operation_name in ("SendMessage", "SendMessageBatch") else "receive"
+    topic_name = ""
+    batch = "_BATCH" if operation_name == "SendMessageBatch" else ""
+    signature_type = "RECEIVE from" if span_action == "receive" else f"SEND{batch} to"
+
+    if len(args) > 1:
+        topic_name = args[1]["QueueUrl"].rsplit("/", maxsplit=1)[-1]
+    signature = f"SQS {signature_type} {topic_name}".rstrip() if topic_name else f"SQS {signature_type}"
+    context["destination"]["service"] = {
+        "name": span_subtype,
+        "resource": f"{span_subtype}/{topic_name}" if topic_name else span_subtype,
+        "type": span_type,
+    }
+    return HandlerInfo(signature, span_type, span_subtype, span_action, context)
+
+
+def modify_span_sqs(span, args, kwargs):
+    trace_parent = span.transaction.trace_parent.copy_from(span_id=span.id)
+    attributes = {constants.TRACEPARENT_HEADER_NAME: {"DataType": "String", "StringValue": trace_parent.to_string()}}
+    if trace_parent.tracestate:
+        attributes[constants.TRACESTATE_HEADER_NAME] = {"DataType": "String", "StringValue": trace_parent.tracestate}
+    if len(args) > 1:
+        attributes_count = len(attributes)
+        if "MessageAttributes" in args[1]:
+            messages = [args[1]]
+        # elif "Entries" in args[1]:
+        #     messages = args[1]["Entries"]
+        else:
+            messages = []
+        for message in messages:
+            if len(message["MessageAttributes"]) + attributes_count <= 10:
+                message["MessageAttributes"].update(attributes)
+            else:
+                logger.info("Not adding disttracing headers to message due to attribute limit reached")
 
 
 def handle_default(operation_name, service, instance, args, kwargs, destination):
@@ -160,5 +205,10 @@ def handle_default(operation_name, service, instance, args, kwargs, destination)
     "S3": handle_s3,
     "DynamoDB": handle_dynamodb,
     "SNS": handle_sns,
+    "SQS": handle_sqs,
     "default": handle_default,
 }
+
+span_modifiers = {
+    "SQS": modify_span_sqs,
+}

elasticapm/traces.py

@@ -120,12 +120,16 @@ def set_success(self):
     def set_failure(self):
         self.outcome = "failure"
 
+    @staticmethod
+    def get_dist_tracing_id():
+        return "%016x" % random.getrandbits(64)
+
 
 class Transaction(BaseSpan):
     def __init__(
         self, tracer, transaction_type="custom", trace_parent=None, is_sampled=True, start=None, sample_rate=None
     ):
-        self.id = "%016x" % random.getrandbits(64)
+        self.id = self.get_dist_tracing_id()
         self.trace_parent = trace_parent
         if start:
             self.timestamp = self.start_time = start
@@ -402,7 +406,7 @@ def __init__(
         :param start: timestamp, mostly useful for testing
         """
         self.start_time = start or _time_func()
-        self.id = "%016x" % random.getrandbits(64)
+        self.id = self.get_dist_tracing_id()
         self.transaction = transaction
         self.name = name
         self.context = context if context is not None else {}

tests/instrumentation/botocore_tests.py

@@ -76,6 +76,17 @@ def dynamodb():
     db.delete_table(TableName="Movies")
 
 
+@pytest.fixture()
+def sqs_client_and_queue():
+    sqs = boto3.client("sqs", endpoint_url=LOCALSTACK_ENDPOINT)
+    response = sqs.create_queue(
+        QueueName="myqueue", Attributes={"DelaySeconds": "60", "MessageRetentionPeriod": "86400"}
+    )
+    queue_url = response["QueueUrl"]
+    yield sqs, queue_url
+    sqs.delete_queue(QueueUrl=queue_url)
+
+
 def test_botocore_instrumentation(instrument, elasticapm_client):
     elasticapm_client.begin_transaction("transaction.test")
     ec2 = boto3.client("ec2", endpoint_url=LOCALSTACK_ENDPOINT)
@@ -190,3 +201,80 @@ def test_sns(instrument, elasticapm_client):
     assert spans[2]["context"]["destination"]["service"]["name"] == "sns"
     assert spans[2]["context"]["destination"]["service"]["resource"] == "sns/mytopic"
     assert spans[2]["context"]["destination"]["service"]["type"] == "messaging"
+
+
+def test_sqs_send(instrument, elasticapm_client, sqs_client_and_queue):
+    sqs, queue_url = sqs_client_and_queue
+    elasticapm_client.begin_transaction("test")
+    sqs.send_message(
+        QueueUrl=queue_url,
+        MessageAttributes={
+            "Title": {"DataType": "String", "StringValue": "foo"},
+        },
+        MessageBody=("bar"),
+    )
+    elasticapm_client.end_transaction("test", "test")
+    span = elasticapm_client.events[constants.SPAN][0]
+    assert span["name"] == "SQS SEND to myqueue"
+    assert span["type"] == "messaging"
+    assert span["subtype"] == "sqs"
+    assert span["action"] == "send"
+    assert span["context"]["destination"]["cloud"]["region"] == "us-east-1"
+    assert span["context"]["destination"]["service"]["name"] == "sqs"
+    assert span["context"]["destination"]["service"]["resource"] == "sqs/myqueue"
+    assert span["context"]["destination"]["service"]["type"] == "messaging"
+
+
+def test_sqs_send_batch(instrument, elasticapm_client, sqs_client_and_queue):
+    sqs, queue_url = sqs_client_and_queue
+    elasticapm_client.begin_transaction("test")
+    response = sqs.send_message_batch(
+        QueueUrl=queue_url,
+        Entries=[
+            {
+                "Id": "foo",
+                "MessageBody": "foo",
+                "DelaySeconds": 123,
+                "MessageAttributes": {"string": {"StringValue": "foo", "DataType": "String"}},
+            },
+        ],
+    )
+    elasticapm_client.end_transaction("test", "test")
+    span = elasticapm_client.events[constants.SPAN][0]
+    assert span["name"] == "SQS SEND_BATCH to myqueue"
+    assert span["type"] == "messaging"
+    assert span["subtype"] == "sqs"
+    assert span["action"] == "send"
+    assert span["context"]["destination"]["cloud"]["region"] == "us-east-1"
+    assert span["context"]["destination"]["service"]["name"] == "sqs"
+    assert span["context"]["destination"]["service"]["resource"] == "sqs/myqueue"
+    assert span["context"]["destination"]["service"]["type"] == "messaging"
+
+
+def test_sqs_receive(instrument, elasticapm_client, sqs_client_and_queue):
+    sqs, queue_url = sqs_client_and_queue
+    sqs.send_message(
+        QueueUrl=queue_url,
+        MessageAttributes={
+            "Title": {"DataType": "String", "StringValue": "foo"},
+        },
+        MessageBody=("bar"),
+    )
+    elasticapm_client.begin_transaction("test")
+    response = sqs.receive_message(
+        QueueUrl=queue_url,
+        AttributeNames=["All"],
+        MessageAttributeNames=[
+            "All",
+        ],
+    )
+    elasticapm_client.end_transaction("test", "test")
+    span = elasticapm_client.events[constants.SPAN][0]
+    assert span["name"] == "SQS RECEIVE from myqueue"
+    assert span["type"] == "messaging"
+    assert span["subtype"] == "sqs"
+    assert span["action"] == "receive"
+    assert span["context"]["destination"]["cloud"]["region"] == "us-east-1"
+    assert span["context"]["destination"]["service"]["name"] == "sqs"
+    assert span["context"]["destination"]["service"]["resource"] == "sqs/myqueue"
+    assert span["context"]["destination"]["service"]["type"] == "messaging"