diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/otlp_aws_span_exporter.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/otlp_aws_span_exporter.py index 7c00b51a8..59f98edb9 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/otlp_aws_span_exporter.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/otlp_aws_span_exporter.py @@ -1,6 +1,7 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 import logging +import os from typing import Dict, Optional import requests @@ -10,6 +11,8 @@ from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter AWS_SERVICE = "xray" +AWS_CLOUDWATCH_LOG_GROUP_ENV = "AWS_CLOUDWATCH_LOG_GROUP" +AWS_CLOUDWATCH_LOG_STREAM_ENV = "AWS_CLOUDWATCH_LOG_STREAM" _logger = logging.getLogger(__name__) @@ -83,6 +86,18 @@ def _export(self, serialized_data: bytes): headers={"Content-Type": "application/x-protobuf"}, ) + # Add CloudWatch Log Group and Log Stream headers if configured + cloudwatch_log_group = os.environ.get(AWS_CLOUDWATCH_LOG_GROUP_ENV) + cloudwatch_log_stream = os.environ.get(AWS_CLOUDWATCH_LOG_STREAM_ENV) + + if cloudwatch_log_group: + request.headers["x-aws-log-group"] = cloudwatch_log_group + _logger.debug("Adding CloudWatch Log Group header: %s", cloudwatch_log_group) + + if cloudwatch_log_stream: + request.headers["x-aws-log-stream"] = cloudwatch_log_stream + _logger.debug("Adding CloudWatch Log Stream header: %s", cloudwatch_log_stream) + credentials = self.boto_session.get_credentials() if credentials is not None: diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_aws_span_exporter.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_aws_span_exporter.py index b0222bb7c..4ea886145 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_aws_span_exporter.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_aws_span_exporter.py @@ -26,6 +26,8 @@ AUTHORIZATION_HEADER = "Authorization" X_AMZ_DATE_HEADER = "X-Amz-Date" X_AMZ_SECURITY_TOKEN_HEADER = "X-Amz-Security-Token" +AWS_CLOUDWATCH_LOG_GROUP_ENV = "AWS_CLOUDWATCH_LOG_GROUP" +AWS_CLOUDWATCH_LOG_STREAM_ENV = "AWS_CLOUDWATCH_LOG_STREAM" class TestAwsSpanExporter(TestCase): @@ -171,6 +173,111 @@ def test_sigv4_exporter_export_adds_sigv4_authentication_if_valid_cw_endpoint( self.assertEqual(actual_headers[X_AMZ_DATE_HEADER], self.expected_auth_x_amz_date) self.assertEqual(actual_headers[X_AMZ_SECURITY_TOKEN_HEADER], self.expected_auth_security_token) + @patch("botocore.session.Session") + @patch("requests.Session") + @patch("botocore.auth.SigV4Auth.add_auth") + @patch.dict(os.environ, { + OTEL_EXPORTER_OTLP_TRACES_ENDPOINT: OTLP_XRAY_ENDPOINT, + AWS_CLOUDWATCH_LOG_GROUP_ENV: "my-log-group", + AWS_CLOUDWATCH_LOG_STREAM_ENV: "my-log-stream" + }, clear=True) + def test_sigv4_exporter_export_adds_cloudwatch_headers( + self, mock_sigv4_auth, requests_posts_mock, botocore_mock + ): + """Tests that CloudWatch Log Group and Log Stream headers are added when environment variables are set.""" + # Setting the exporter resopnse + mock_response = MagicMock() + mock_response.status_code = 200 + type(mock_response).ok = PropertyMock(return_value=True) + + # Setting the request session headers + mock_session = MagicMock() + mock_session.headers = {"User-Agent": USER_AGENT, "Content-Type": CONTENT_TYPE} + requests_posts_mock.return_value = mock_session + mock_session.post.return_value = mock_response + + mock_botocore_session = MagicMock() + botocore_mock.return_value = mock_botocore_session + mock_botocore_session.get_credentials.return_value = Credentials( + access_key="test_key", secret_key="test_secret", token="test_token" + ) + + # Mock for inspecting request headers before SigV4 signing + original_mock_add_auth = self.mock_add_auth + request_headers = {} + + def extended_mock_add_auth(request): + # Save headers for inspection before they're modified by SigV4 + nonlocal request_headers + request_headers = dict(request.headers) + # Call the original mock + original_mock_add_auth(request) + + mock_sigv4_auth.side_effect = extended_mock_add_auth + + # Initialize and call exporter + exporter = OTLPAwsSpanExporter(endpoint=OTLP_XRAY_ENDPOINT) + exporter.export(self.testing_spans) + + # Verify SigV4 auth was called + mock_sigv4_auth.assert_called_once() + + # Check that CloudWatch headers were added to the request before signing + self.assertEqual(request_headers.get("x-aws-log-group"), "my-log-group") + self.assertEqual(request_headers.get("x-aws-log-stream"), "my-log-stream") + + @patch("botocore.session.Session") + @patch("requests.Session") + @patch("botocore.auth.SigV4Auth.add_auth") + @patch.dict(os.environ, { + OTEL_EXPORTER_OTLP_TRACES_ENDPOINT: OTLP_XRAY_ENDPOINT, + AWS_CLOUDWATCH_LOG_GROUP_ENV: "my-log-group" + }, clear=True) + def test_sigv4_exporter_export_adds_only_log_group_header( + self, mock_sigv4_auth, requests_posts_mock, botocore_mock + ): + """Tests that only CloudWatch Log Group header is added when only that environment variable is set.""" + # Setting the exporter response + mock_response = MagicMock() + mock_response.status_code = 200 + type(mock_response).ok = PropertyMock(return_value=True) + + # Setting the request session headers + mock_session = MagicMock() + mock_session.headers = {"User-Agent": USER_AGENT, "Content-Type": CONTENT_TYPE} + requests_posts_mock.return_value = mock_session + mock_session.post.return_value = mock_response + + mock_botocore_session = MagicMock() + botocore_mock.return_value = mock_botocore_session + mock_botocore_session.get_credentials.return_value = Credentials( + access_key="test_key", secret_key="test_secret", token="test_token" + ) + + # Mock for inspecting request headers before SigV4 signing + original_mock_add_auth = self.mock_add_auth + request_headers = {} + + def extended_mock_add_auth(request): + # Save headers for inspection before they're modified by SigV4 + nonlocal request_headers + request_headers = dict(request.headers) + # Call the original mock + original_mock_add_auth(request) + + mock_sigv4_auth.side_effect = extended_mock_add_auth + + # Initialize and call exporter + exporter = OTLPAwsSpanExporter(endpoint=OTLP_XRAY_ENDPOINT) + exporter.export(self.testing_spans) + + # Verify SigV4 auth was called + mock_sigv4_auth.assert_called_once() + + # Check that only CloudWatch Log Group header was added + self.assertEqual(request_headers.get("x-aws-log-group"), "my-log-group") + self.assertNotIn("x-aws-log-stream", request_headers) + def validate_exporter_extends_http_span_exporter(self, exporter, endpoint): self.assertIsInstance(exporter, OTLPSpanExporter) self.assertEqual(exporter._endpoint, endpoint)