diff --git a/instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/__init__.py b/instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/__init__.py index 24f0dc313b..2c5ae11b5a 100644 --- a/instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/__init__.py @@ -88,6 +88,7 @@ def _instrument(self, **kwargs): def _uninstrument(self, **kwargs): unwrap(BaseClient, "_make_api_call") + # pylint: disable=too-many-branches def _patched_api_call(self, original_func, instance, args, kwargs): # pylint: disable=protected-access @@ -105,35 +106,43 @@ def _patched_api_call(self, original_func, instance, args, kwargs): span.set_attribute("aws.region", instance.meta.region_name) span.set_attribute("aws.service", service_name) if "QueueUrl" in api_params: - span.set_attribute( - "aws.queue_url", - api_params["QueueUrl"] - ) + span.set_attribute("aws.queue_url", api_params["QueueUrl"]) if "TableName" in api_params: span.set_attribute( - "aws.table_name", - api_params["TableName"] + "aws.table_name", api_params["TableName"] ) try: result = original_func(*args, **kwargs) - except ClientError as e: - error = e + except ClientError as ex: + error = ex if error: result = error.response if span.is_recording(): if "ResponseMetadata" in result: - if "RequestId" in result["ResponseMetadata"]: + metadata = result["ResponseMetadata"] + req_id = None + if "RequestId" in metadata: + req_id = metadata["RequestId"] + elif "HTTPHeaders" in metadata: + headers = metadata["HTTPHeaders"] + if "x-amzn-RequestId" in headers: + req_id = headers["x-amzn-RequestId"] + elif "x-amz-request-id" in headers: + req_id = headers["x-amz-request-id"] + elif "x-amz-id-2" in headers: + req_id = headers["x-amz-id-2"] + + if req_id: span.set_attribute( - "aws.request_id", - result["ResponseMetadata"]["RequestId"], + "aws.request_id", req_id, ) - if "HTTPStatusCode" in result["ResponseMetadata"]: + + if "HTTPStatusCode" in metadata: span.set_attribute( - "http.status_code", - result["ResponseMetadata"]["HTTPStatusCode"], + "http.status_code", metadata["HTTPStatusCode"], ) if error: diff --git a/instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_instrumentation.py b/instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_instrumentation.py index afdf9827e2..dc7f945ea4 100644 --- a/instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_instrumentation.py +++ b/instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_instrumentation.py @@ -17,8 +17,8 @@ import botocore.session from botocore.exceptions import ParamValidationError from moto import ( # pylint: disable=import-error - mock_ec2, mock_dynamodb2, + mock_ec2, mock_kinesis, mock_kms, mock_lambda, @@ -203,22 +203,21 @@ def test_sqs_client(self): def test_sqs_send_message(self): sqs = self.session.create_client("sqs", region_name="us-east-1") - TEST_QUEUE_NAME = 'test_queue_name' + test_queue_name = "test_queue_name" - response = sqs.create_queue( - QueueName=TEST_QUEUE_NAME - ) + response = sqs.create_queue(QueueName=test_queue_name) sqs.send_message( - QueueUrl=response["QueueUrl"], - MessageBody="Test SQS 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}") + 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, @@ -230,7 +229,9 @@ def test_sqs_send_message(self): }, ) send_msg_attributes = dict(spans[1].attributes) - self.assertRegex(send_msg_attributes["aws.request_id"], r"[A-Z0-9]{52}") + 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, @@ -361,51 +362,31 @@ def test_sts_client(self): def test_dynamodb_client(self): ddb = self.session.create_client("dynamodb", region_name="us-west-2") - TEST_TABLE_NAME = "test_table_name" + test_table_name = "test_table_name" ddb.create_table( AttributeDefinitions=[ - { - 'AttributeName': 'id', - 'AttributeType': 'S', - }, - ], - KeySchema=[ - { - 'AttributeName': 'id', - 'KeyType': 'HASH', - }, + {"AttributeName": "id", "AttributeType": "S"}, ], + KeySchema=[{"AttributeName": "id", "KeyType": "HASH"}], ProvisionedThroughput={ - 'ReadCapacityUnits': 5, - 'WriteCapacityUnits': 5, + "ReadCapacityUnits": 5, + "WriteCapacityUnits": 5, }, - TableName=TEST_TABLE_NAME, + TableName=test_table_name, ) - ddb.put_item( - TableName=TEST_TABLE_NAME, - Item={ - 'id': { - 'S': 'test_key' - } - } - ) + ddb.put_item(TableName=test_table_name, Item={"id": {"S": "test_key"}}) - ddb.get_item( - TableName=TEST_TABLE_NAME, - Key={ - '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}") + 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, @@ -418,7 +399,9 @@ def test_dynamodb_client(self): }, ) put_item_attributes = dict(spans[1].attributes) - self.assertRegex(put_item_attributes["aws.request_id"], r"[A-Z0-9]{52}") + 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, @@ -431,7 +414,9 @@ def test_dynamodb_client(self): }, ) get_item_attributes = dict(spans[2].attributes) - self.assertRegex(get_item_attributes["aws.request_id"], r"[A-Z0-9]{52}") + 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,