Add additional aws sdk and http attributes

Author: NathanielRN

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/__init__.py

@@ -54,7 +54,7 @@
 import logging
 
 from botocore.client import BaseClient
-from botocore.exceptions import ClientError
+from botocore.exceptions import ClientError, ParamValidationError
 from wrapt import ObjectProxy, wrap_function_wrapper
 
 from opentelemetry.instrumentation.botocore.version import __version__
@@ -67,7 +67,7 @@
 
 
 class BotocoreInstrumentor(BaseInstrumentor):
-    """A instrumentor for Botocore.
+    """An instrumentor for Botocore.
 
     See `BaseInstrumentor`
     """
@@ -100,14 +100,6 @@ def _patched_api_call(self, original_func, instance, args, kwargs):
             error = None
             result = None
 
-            try:
-                result = original_func(*args, **kwargs)
-            except ClientError as e:
-                error = e
-
-            if error:
-                result = error.response
-
             if span.is_recording():
                 span.set_attribute("aws.operation", operation_name)
                 span.set_attribute("aws.region", instance.meta.region_name)
@@ -122,6 +114,16 @@ def _patched_api_call(self, original_func, instance, args, kwargs):
                         "aws.table_name",
                         api_params["TableName"]
                     )
+
+            try:
+                result = original_func(*args, **kwargs)
+            except ClientError as e:
+                error = e
+
+            if error:
+                result = error.response
+
+            if span.is_recording():
                 if "ResponseMetadata" in result:
                     if "RequestId" in result["ResponseMetadata"]:
                         span.set_attribute(

instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_instrumentation.py

@@ -18,6 +18,7 @@
 from botocore.exceptions import ParamValidationError
 from moto import (  # pylint: disable=import-error
     mock_ec2,
+    mock_dynamodb2,
     mock_kinesis,
     mock_kms,
     mock_lambda,
@@ -61,6 +62,7 @@ def test_traced_client(self):
             {
                 "aws.operation": "DescribeInstances",
                 "aws.region": "us-west-2",
+                "aws.request_id": "fdcdcab1-ae5c-489e-9c33-4637c5dda355",
                 "aws.service": "ec2",
                 "http.status_code": 200,
             },
@@ -103,7 +105,6 @@ def test_s3_client(self):
         assert spans
         span = spans[0]
         self.assertEqual(len(spans), 2)
-        self.assertEqual(span.attributes["aws.operation"], "ListBuckets")
         self.assertEqual(
             dict(span.attributes),
             {
@@ -138,21 +139,24 @@ def test_s3_put(self):
         location = {"LocationConstraint": "us-west-2"}
         s3.create_bucket(Bucket="mybucket", CreateBucketConfiguration=location)
         s3.put_object(**params)
+        s3.get_object(Bucket="mybucket", Key="foo")
 
         spans = self.memory_exporter.get_finished_spans()
         assert spans
-        self.assertEqual(len(spans), 2)
+        self.assertEqual(len(spans), 3)
+        create_bucket_attributes = dict(spans[0].attributes)
         self.assertEqual(
-            spans[0].attributes,
+            create_bucket_attributes,
             {
                 "aws.operation": "CreateBucket",
                 "aws.region": "us-west-2",
                 "aws.service": "s3",
                 "http.status_code": 200,
             },
         )
+        put_object_attributes = dict(spans[1].attributes)
         self.assertEqual(
-            spans[1].attributes,
+            put_object_attributes,
             {
                 "aws.operation": "PutObject",
                 "aws.region": "us-west-2",
@@ -161,6 +165,16 @@ def test_s3_put(self):
             },
         )
         self.assertTrue("params.Body" not in spans[1].attributes.keys())
+        get_object_attributes = dict(spans[2].attributes)
+        self.assertEqual(
+            get_object_attributes,
+            {
+                "aws.operation": "GetObject",
+                "aws.region": "us-west-2",
+                "aws.service": "s3",
+                "http.status_code": 200,
+            },
+        )
 
     @mock_sqs
     def test_sqs_client(self):
@@ -172,8 +186,11 @@ def test_sqs_client(self):
         assert spans
         span = spans[0]
         self.assertEqual(len(spans), 1)
+        actual = dict(span.attributes)
+        self.assertRegex(actual["aws.request_id"], r"[A-Z0-9]{52}")
+        del actual["aws.request_id"]
         self.assertEqual(
-            dict(span.attributes),
+            actual,
             {
                 "aws.operation": "ListQueues",
                 "aws.region": "us-east-1",
@@ -182,6 +199,50 @@ def test_sqs_client(self):
             },
         )
 
+    @mock_sqs
+    def test_sqs_send_message(self):
+        sqs = self.session.create_client("sqs", region_name="us-east-1")
+
+        TEST_QUEUE_NAME = 'test_queue_name'
+
+        response = sqs.create_queue(
+            QueueName=TEST_QUEUE_NAME
+        )
+
+        sqs.send_message(
+            QueueUrl=response["QueueUrl"],
+            MessageBody="Test SQS MESSAGE!"
+        )
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        self.assertEqual(len(spans), 2)
+        create_queue_attributes = dict(spans[0].attributes)
+        self.assertRegex(create_queue_attributes["aws.request_id"], r"[A-Z0-9]{52}")
+        del create_queue_attributes["aws.request_id"]
+        self.assertEqual(
+            create_queue_attributes,
+            {
+                "aws.operation": "CreateQueue",
+                "aws.region": "us-east-1",
+                "aws.service": "sqs",
+                "http.status_code": 200,
+            },
+        )
+        send_msg_attributes = dict(spans[1].attributes)
+        self.assertRegex(send_msg_attributes["aws.request_id"], r"[A-Z0-9]{52}")
+        del send_msg_attributes["aws.request_id"]
+        self.assertEqual(
+            send_msg_attributes,
+            {
+                "aws.operation": "SendMessage",
+                "aws.queue_url": response["QueueUrl"],
+                "aws.region": "us-east-1",
+                "aws.service": "sqs",
+                "http.status_code": 200,
+            },
+        )
+
     @mock_kinesis
     def test_kinesis_client(self):
         kinesis = self.session.create_client(
@@ -287,10 +348,98 @@ def test_sts_client(self):
             {
                 "aws.operation": "GetCallerIdentity",
                 "aws.region": "us-east-1",
+                "aws.request_id": "c6104cbe-af31-11e0-8154-cbc7ccf896c7",
                 "aws.service": "sts",
                 "http.status_code": 200,
             },
         )
 
         # checking for protection on sts against security leak
         self.assertTrue("params" not in span.attributes.keys())
+
+    @mock_dynamodb2
+    def test_dynamodb_client(self):
+        ddb = self.session.create_client("dynamodb", region_name="us-west-2")
+
+        TEST_TABLE_NAME = "test_table_name"
+
+        ddb.create_table(
+            AttributeDefinitions=[
+                {
+                    'AttributeName': 'id',
+                    'AttributeType': 'S',
+                },
+            ],
+            KeySchema=[
+                {
+                    'AttributeName': 'id',
+                    'KeyType': 'HASH',
+                },
+            ],
+            ProvisionedThroughput={
+                'ReadCapacityUnits': 5,
+                'WriteCapacityUnits': 5,
+            },
+            TableName=TEST_TABLE_NAME,
+        )
+
+        ddb.put_item(
+            TableName=TEST_TABLE_NAME,
+            Item={
+                'id': {
+                    'S': 'test_key'
+                }
+            }
+        )
+
+        ddb.get_item(
+            TableName=TEST_TABLE_NAME,
+            Key={
+                'id': {
+                    'S': 'test_key'
+                }
+            }
+        )
+
+        spans = self.memory_exporter.get_finished_spans()
+        assert spans
+        self.assertEqual(len(spans), 3)
+        create_table_attributes = dict(spans[0].attributes)
+        self.assertRegex(create_table_attributes["aws.request_id"], r"[A-Z0-9]{52}")
+        del create_table_attributes["aws.request_id"]
+        self.assertEqual(
+            create_table_attributes,
+            {
+                "aws.operation": "CreateTable",
+                "aws.region": "us-west-2",
+                "aws.service": "dynamodb",
+                "aws.table_name": "test_table_name",
+                "http.status_code": 200,
+            },
+        )
+        put_item_attributes = dict(spans[1].attributes)
+        self.assertRegex(put_item_attributes["aws.request_id"], r"[A-Z0-9]{52}")
+        del put_item_attributes["aws.request_id"]
+        self.assertEqual(
+            put_item_attributes,
+            {
+                "aws.operation": "PutItem",
+                "aws.region": "us-west-2",
+                "aws.service": "dynamodb",
+                "aws.table_name": "test_table_name",
+                "http.status_code": 200,
+            },
+        )
+        get_item_attributes = dict(spans[2].attributes)
+        self.assertRegex(get_item_attributes["aws.request_id"], r"[A-Z0-9]{52}")
+        del get_item_attributes["aws.request_id"]
+        self.assertEqual(
+            get_item_attributes,
+            {
+                "aws.operation": "GetItem",
+                "aws.region": "us-west-2",
+                "aws.service": "dynamodb",
+                "aws.table_name": "test_table_name",
+                "http.status_code": 200,
+            },
+        )