feat: Add auto-instrumentation support for Kinesis Stream Consumer

Author: yiyuan-he

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py

@@ -19,6 +19,7 @@
 AWS_SQS_QUEUE_URL: str = "aws.sqs.queue.url"
 AWS_SQS_QUEUE_NAME: str = "aws.sqs.queue.name"
 AWS_KINESIS_STREAM_NAME: str = "aws.kinesis.stream.name"
+AWS_KINESIS_STREAM_CONSUMERNAME: str = "aws.kinesis.stream.consumer_name"
 AWS_BEDROCK_DATA_SOURCE_ID: str = "aws.bedrock.data_source.id"
 AWS_BEDROCK_KNOWLEDGE_BASE_ID: str = "aws.bedrock.knowledge_base.id"
 AWS_BEDROCK_AGENT_ID: str = "aws.bedrock.agent.id"

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py

@@ -11,6 +11,7 @@
     AWS_BEDROCK_GUARDRAIL_ID,
     AWS_BEDROCK_KNOWLEDGE_BASE_ID,
     AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER,
+    AWS_KINESIS_STREAM_CONSUMERNAME,
     AWS_KINESIS_STREAM_NAME,
     AWS_LOCAL_OPERATION,
     AWS_LOCAL_SERVICE,
@@ -393,6 +394,9 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri
         elif is_key_present(span, AWS_KINESIS_STREAM_NAME):
             remote_resource_type = _NORMALIZED_KINESIS_SERVICE_NAME + "::Stream"
             remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_KINESIS_STREAM_NAME))
+        elif is_key_present(span, AWS_KINESIS_STREAM_CONSUMERNAME):
+            remote_resource_type = _NORMALIZED_KINESIS_SERVICE_NAME + "::StreamConsumer"
+            remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_KINESIS_STREAM_CONSUMERNAME))
         elif is_key_present(span, _AWS_BUCKET_NAME):
             remote_resource_type = _NORMALIZED_S3_SERVICE_NAME + "::Bucket"
             remote_resource_identifier = _escape_delimiters(span.attributes.get(_AWS_BUCKET_NAME))

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py

@@ -4,6 +4,7 @@
 import importlib
 
 from amazon.opentelemetry.distro._aws_attribute_keys import (
+    AWS_KINESIS_STREAM_CONSUMERNAME,
     AWS_KINESIS_STREAM_NAME,
     AWS_SECRETSMANAGER_SECRET_ARN,
     AWS_SNS_TOPIC_ARN,
@@ -204,3 +205,6 @@ def extract_attributes(self, attributes: _AttributeMapT):
         stream_name = self._call_context.params.get("StreamName")
         if stream_name:
             attributes[AWS_KINESIS_STREAM_NAME] = stream_name
+        consumer_name = self._call_context.params.get("ConsumerName")
+        if consumer_name:
+            attributes[AWS_KINESIS_STREAM_CONSUMERNAME] = consumer_name

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py

@@ -13,6 +13,7 @@
     AWS_BEDROCK_GUARDRAIL_ID,
     AWS_BEDROCK_KNOWLEDGE_BASE_ID,
     AWS_CONSUMER_PARENT_SPAN_KIND,
+    AWS_KINESIS_STREAM_CONSUMERNAME,
     AWS_KINESIS_STREAM_NAME,
     AWS_LOCAL_OPERATION,
     AWS_LOCAL_SERVICE,
@@ -1142,6 +1143,30 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
         )
         self._mock_attribute([AWS_STEPFUNCTIONS_ACTIVITY_ARN], [None])
 
+        # Validate behaviour of AWS_KINESIS_STREAM_CONSUMERNAME present, then remove it.
+        self._mock_attribute(
+            [AWS_KINESIS_STREAM_CONSUMERNAME],
+            ["aws_stream_consumername"],
+            keys,
+            values,
+        )
+        self._validate_remote_resource_attributes(
+            "AWS::Kinesis::StreamConsumer", 
+            "aws_stream_consumername"
+        )
+        self._mock_attribute([AWS_KINESIS_STREAM_CONSUMERNAME], [None])
+
+        # Validate behaviour with both AWS_KINESIS_STREAM_NAME and AWS_KINESIS_STREAM_CONSUMERNAME
+        # present, then remove it.
+        self._mock_attribute(
+            [AWS_KINESIS_STREAM_NAME, AWS_KINESIS_STREAM_CONSUMERNAME],
+            ["aws_stream_name", "aws_stream_consumername"],
+            keys,
+            values,
+        )
+        self._validate_remote_resource_attributes("AWS::Kinesis::Stream", "aws_stream_name")
+        self._mock_attribute([AWS_KINESIS_STREAM_NAME, AWS_KINESIS_STREAM_CONSUMERNAME], [None, None])
+
         self._mock_attribute([SpanAttributes.RPC_SYSTEM], [None])
 
     def test_client_db_span_with_remote_resource_attributes(self):

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py

@@ -17,6 +17,7 @@
 from opentelemetry.trace.span import Span
 
 _STREAM_NAME: str = "streamName"
+_CONSUMER_NAME: str = "consumerName"
 _BUCKET_NAME: str = "bucketName"
 _QUEUE_NAME: str = "queueName"
 _QUEUE_URL: str = "queueUrl"
@@ -175,6 +176,8 @@ def _test_patched_botocore_instrumentation(self):
         kinesis_attributes: Dict[str, str] = _do_extract_kinesis_attributes()
         self.assertTrue("aws.kinesis.stream.name" in kinesis_attributes)
         self.assertEqual(kinesis_attributes["aws.kinesis.stream.name"], _STREAM_NAME)
+        self.assertTrue("aws.kinesis.stream.consumer_name" in kinesis_attributes)
+        self.assertEqual(kinesis_attributes["aws.kinesis.stream.consumer_name"], _CONSUMER_NAME)
 
         # S3
         self.assertTrue("s3" in _KNOWN_EXTENSIONS)
@@ -356,7 +359,7 @@ def _reset_mocks(self):
 
 def _do_extract_kinesis_attributes() -> Dict[str, str]:
     service_name: str = "kinesis"
-    params: Dict[str, str] = {"StreamName": _STREAM_NAME}
+    params: Dict[str, str] = {"StreamName": _STREAM_NAME, "ConsumerName": _CONSUMER_NAME}
     return _do_extract_attributes(service_name, params)