fixing consolidation changes

Author: liustve

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

@@ -12,7 +12,7 @@
 
 from amazon.opentelemetry.distro._aws_attribute_keys import AWS_LOCAL_SERVICE, AWS_SERVICE_TYPE
 from amazon.opentelemetry.distro._aws_resource_attribute_configurator import get_service_attribute
-from amazon.opentelemetry.distro._utils import IS_BOTOCORE_INSTALLED, get_aws_session, is_agent_observability_enabled
+from amazon.opentelemetry.distro._utils import get_aws_session, is_agent_observability_enabled
 from amazon.opentelemetry.distro.always_record_sampler import AlwaysRecordSampler
 from amazon.opentelemetry.distro.attribute_propagating_span_processor_builder import (
     AttributePropagatingSpanProcessorBuilder,
@@ -23,8 +23,6 @@
     AwsMetricAttributesSpanExporterBuilder,
 )
 from amazon.opentelemetry.distro.aws_span_metrics_processor_builder import AwsSpanMetricsProcessorBuilder
-from amazon.opentelemetry.distro.exporter.otlp.aws.logs.otlp_aws_logs_exporter import OTLPAwsLogExporter
-from amazon.opentelemetry.distro.exporter.otlp.aws.traces.otlp_aws_span_exporter import OTLPAwsSpanExporter
 from amazon.opentelemetry.distro.otlp_udp_exporter import OTLPUdpSpanExporter
 from amazon.opentelemetry.distro.sampler.aws_xray_remote_sampler import AwsXRayRemoteSampler
 from amazon.opentelemetry.distro.scope_based_exporter import ScopeBasedPeriodicExportingMetricReader
@@ -210,9 +208,9 @@ def _init_logging(
 
     for _, exporter_class in exporters.items():
         exporter_args = {}
-        log_exporter: LogExporter = _customize_logs_exporter(exporter_class(**exporter_args))
-        log_processor = _customize_log_record_processor(log_exporter)
-        provider.add_log_record_processor(log_processor)
+        _customize_log_record_processor(
+            logger_provider=provider, log_exporter=_customize_logs_exporter(exporter_class(**exporter_args))
+        )
 
     event_logger_provider = EventLoggerProvider(logger_provider=provider)
     set_event_logger_provider(event_logger_provider)
@@ -299,10 +297,12 @@ def _export_unsampled_span_for_agent_observability(trace_provider: TracerProvide
         return
 
     traces_endpoint = os.environ.get(OTEL_EXPORTER_OTLP_TRACES_ENDPOINT)
+    if traces_endpoint and _is_aws_otlp_endpoint(traces_endpoint):
+        endpoint = traces_endpoint.lower()
+        region = endpoint.split(".")[1]
 
-    span_exporter = _create_aws_exporter(endpoint=traces_endpoint)
-
-    trace_provider.add_span_processor(BatchUnsampledSpanProcessor(span_exporter=span_exporter))
+        span_exporter = _create_aws_otlp_exporter(endpoint=endpoint, service="xray", region=region)
+        trace_provider.add_span_processor(BatchUnsampledSpanProcessor(span_exporter=span_exporter))
 
 
 def _is_defer_to_workers_enabled():
@@ -397,11 +397,13 @@ def _customize_span_exporter(span_exporter: SpanExporter, resource: Resource) ->
             traces_endpoint = os.environ.get(AWS_XRAY_DAEMON_ADDRESS_CONFIG, "127.0.0.1:2000")
             span_exporter = OTLPUdpSpanExporter(endpoint=traces_endpoint)
 
-    if _is_aws_otlp_endpoint(traces_endpoint, "xray"):
+    if traces_endpoint and _is_aws_otlp_endpoint(traces_endpoint, "xray"):
         _logger.info("Detected using AWS OTLP Traces Endpoint.")
 
         if isinstance(span_exporter, OTLPSpanExporter):
-            return _create_aws_exporter(endpoint=traces_endpoint)
+            endpoint = traces_endpoint.lower()
+            region = endpoint.split(".")[1]
+            return _create_aws_otlp_exporter(endpoint=traces_endpoint, service="xray", region=region)
 
         else:
             _logger.warning(
@@ -415,24 +417,34 @@ def _customize_span_exporter(span_exporter: SpanExporter, resource: Resource) ->
     return AwsMetricAttributesSpanExporterBuilder(span_exporter, resource).build()
 
 
-def _customize_log_record_processor(log_exporter: LogExporter):
-    if isinstance(log_exporter, OTLPAwsLogExporter) and is_agent_observability_enabled():
-        return AwsCloudWatchOtlpBatchLogRecordProcessor(exporter=log_exporter)
+def _customize_log_record_processor(logger_provider: LoggerProvider, log_exporter: Optional[LogExporter]) -> None:
+    if not log_exporter:
+        return
 
-    return BatchLogRecordProcessor(exporter=log_exporter)
+    if is_agent_observability_enabled():
+        from amazon.opentelemetry.distro.exporter.otlp.aws.logs._aws_cw_otlp_batch_log_record_processor import (
+            AwsCloudWatchOtlpBatchLogRecordProcessor,
+        )
+
+        logger_provider.add_log_record_processor(AwsCloudWatchOtlpBatchLogRecordProcessor(exporter=log_exporter))
+    else:
+        logger_provider.add_log_record_processor(BatchLogRecordProcessor(exporter=log_exporter))
 
 
 def _customize_logs_exporter(log_exporter: LogExporter) -> LogExporter:
     logs_endpoint = os.environ.get(OTEL_EXPORTER_OTLP_LOGS_ENDPOINT)
 
-    if _is_aws_otlp_endpoint(logs_endpoint, "logs"):
+    if logs_endpoint and _is_aws_otlp_endpoint(logs_endpoint, "logs"):
+
         _logger.info("Detected using AWS OTLP Logs Endpoint.")
 
         if isinstance(log_exporter, OTLPLogExporter) and _validate_and_fetch_logs_header().is_valid:
+            endpoint = logs_endpoint.lower()
+            region = endpoint.split(".")[1]
             # Setting default compression mode to Gzip as this is the behavior in upstream's
             # collector otlp http exporter:
             # https://github.com/open-telemetry/opentelemetry-collector/tree/main/exporter/otlphttpexporter
-            return _create_aws_exporter(endpoint=logs_endpoint)
+            return _create_aws_otlp_exporter(endpoint=logs_endpoint, service="logs", region=region)
 
         _logger.warning(
             "Improper configuration see: please export/set "
@@ -594,11 +606,11 @@ def _is_lambda_environment():
 def _is_aws_otlp_endpoint(otlp_endpoint: Optional[str] = None, service: str = "xray") -> bool:
     """Is the given endpoint an AWS OTLP endpoint?"""
 
-    pattern = AWS_TRACES_OTLP_ENDPOINT_PATTERN if service == "xray" else AWS_LOGS_OTLP_ENDPOINT_PATTERN
-
     if not otlp_endpoint:
         return False
 
+    pattern = AWS_TRACES_OTLP_ENDPOINT_PATTERN if service == "xray" else AWS_LOGS_OTLP_ENDPOINT_PATTERN
+
     return bool(re.match(pattern, otlp_endpoint.lower()))
 
 
@@ -787,7 +799,7 @@ def create_emf_exporter():
         return None
 
 
-def _create_aws_exporter(endpoint: str):
+def _create_aws_otlp_exporter(endpoint: str, service: str, region: str):
     """Create and configure the AWS OTLP exporters."""
     try:
         session = get_aws_session()
@@ -800,12 +812,7 @@ def _create_aws_exporter(endpoint: str):
         from amazon.opentelemetry.distro.exporter.otlp.aws.logs.otlp_aws_logs_exporter import OTLPAwsLogExporter
         from amazon.opentelemetry.distro.exporter.otlp.aws.traces.otlp_aws_span_exporter import OTLPAwsSpanExporter
 
-        endpoint = endpoint.lower()
-        split = endpoint.split(".")
-        service = split[0]
-        region = split[1]
-
-        if "xray" in service:
+        if service == "xray":
             if is_agent_observability_enabled():
                 # Span exporter needs an instance of logger provider in ai agent
                 # observability case because we need to split input/output prompts
@@ -818,7 +825,7 @@ def _create_aws_exporter(endpoint: str):
 
             return OTLPAwsSpanExporter(session=session, endpoint=endpoint, aws_region=region)
 
-        if "logs" in service:
+        if service == "logs":
             return OTLPAwsLogExporter(session=session, aws_region=region)
 
     # pylint: disable=broad-exception-caught

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/exporter/aws/metrics/_cloudwatch_log_client.py

@@ -90,7 +90,7 @@ class CloudWatchLogClient:
     def __init__(
         self,
         log_group_name: str,
-        session: Session = Session(),
+        session: Session,
         log_stream_name: Optional[str] = None,
         aws_region: Optional[str] = None,
         **kwargs,

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/exporter/otlp/aws/logs/otlp_aws_logs_exporter.py

@@ -10,6 +10,7 @@
 from time import time
 from typing import Dict, Optional, Sequence
 
+from botocore.session import Session
 from requests import Response
 from requests.exceptions import ConnectionError as RequestsConnectionError
 from requests.structures import CaseInsensitiveDict
@@ -39,7 +40,7 @@ class OTLPAwsLogExporter(OTLPLogExporter):
     def __init__(
         self,
         aws_region: str,
-        session: Session = Session(),
+        session: Session,
         log_group: Optional[str] = None,
         log_stream: Optional[str] = None,
         endpoint: Optional[str] = None,

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/exporter/otlp/aws/traces/otlp_aws_span_exporter.py

@@ -31,7 +31,7 @@ class OTLPAwsSpanExporter(OTLPSpanExporter):
     def __init__(
         self,
         aws_region: str,
-        session: Session = Session(),
+        session: Session,
         endpoint: Optional[str] = None,
         certificate_file: Optional[str] = None,
         client_key_file: Optional[str] = None,

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/exporter/otlp/aws/logs/test_otlp_aws_logs_exporter.py

@@ -7,6 +7,7 @@
 import requests
 from requests.structures import CaseInsensitiveDict
 
+from amazon.opentelemetry.distro._utils import get_aws_session
 from amazon.opentelemetry.distro.exporter.otlp.aws.logs.otlp_aws_logs_exporter import _MAX_RETRYS, OTLPAwsLogExporter
 from opentelemetry._logs.severity import SeverityNumber
 from opentelemetry.sdk._logs import LogData, LogRecord
@@ -36,7 +37,7 @@ class TestOTLPAwsLogsExporter(TestCase):
 
     def setUp(self):
         self.logs = self.generate_test_log_data()
-        self.exporter = OTLPAwsLogExporter(aws_region="us-east-1", endpoint=self._ENDPOINT)
+        self.exporter = OTLPAwsLogExporter(session=get_aws_session(), aws_region="us-east-1", endpoint=self._ENDPOINT)
 
     @patch("requests.Session.post", return_value=good_response)
     def test_export_success(self, mock_request):

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/exporter/otlp/aws/traces/test_otlp_aws_span_exporter.py

@@ -4,6 +4,7 @@
 from unittest import TestCase
 from unittest.mock import MagicMock, patch
 
+from amazon.opentelemetry.distro._utils import get_aws_session
 from amazon.opentelemetry.distro.exporter.otlp.aws.traces.otlp_aws_span_exporter import OTLPAwsSpanExporter
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 from opentelemetry.sdk._logs import LoggerProvider
@@ -17,7 +18,9 @@ def test_init_with_logger_provider(self):
         mock_logger_provider = MagicMock(spec=LoggerProvider)
         endpoint = "https://xray.us-east-1.amazonaws.com/v1/traces"
 
-        exporter = OTLPAwsSpanExporter(aws_region="us-east-1", endpoint=endpoint, logger_provider=mock_logger_provider)
+        exporter = OTLPAwsSpanExporter(
+            session=get_aws_session(), aws_region="us-east-1", endpoint=endpoint, logger_provider=mock_logger_provider
+        )
 
         self.assertEqual(exporter._logger_provider, mock_logger_provider)
         self.assertEqual(exporter._aws_region, "us-east-1")
@@ -26,7 +29,7 @@ def test_init_without_logger_provider(self):
         # Test initialization without logger_provider (default behavior)
         endpoint = "https://xray.us-west-2.amazonaws.com/v1/traces"
 
-        exporter = OTLPAwsSpanExporter(aws_region="us-west-2", endpoint=endpoint)
+        exporter = OTLPAwsSpanExporter(session=get_aws_session(), aws_region="us-west-2", endpoint=endpoint)
 
         self.assertIsNone(exporter._logger_provider)
         self.assertEqual(exporter._aws_region, "us-west-2")
@@ -38,7 +41,7 @@ def test_ensure_llo_handler_when_disabled(self, mock_is_enabled):
         mock_is_enabled.return_value = False
         endpoint = "https://xray.us-east-1.amazonaws.com/v1/traces"
 
-        exporter = OTLPAwsSpanExporter(aws_region="us-east-1", endpoint=endpoint)
+        exporter = OTLPAwsSpanExporter(session=get_aws_session(), aws_region="us-east-1", endpoint=endpoint)
         result = exporter._ensure_llo_handler()
 
         self.assertFalse(result)
@@ -59,7 +62,7 @@ def test_ensure_llo_handler_lazy_initialization(
         mock_llo_handler_class.return_value = mock_llo_handler
 
         endpoint = "https://xray.us-east-1.amazonaws.com/v1/traces"
-        exporter = OTLPAwsSpanExporter(aws_region="us-east-1", endpoint=endpoint)
+        exporter = OTLPAwsSpanExporter(session=get_aws_session(), aws_region="us-east-1", endpoint=endpoint)
 
         # First call should initialize
         result = exporter._ensure_llo_handler()
@@ -87,7 +90,9 @@ def test_ensure_llo_handler_with_existing_logger_provider(self, mock_is_enabled,
         mock_logger_provider = MagicMock(spec=LoggerProvider)
 
         endpoint = "https://xray.us-east-1.amazonaws.com/v1/traces"
-        exporter = OTLPAwsSpanExporter(aws_region="us-east-1", endpoint=endpoint, logger_provider=mock_logger_provider)
+        exporter = OTLPAwsSpanExporter(
+            session=get_aws_session(), aws_region="us-east-1", endpoint=endpoint, logger_provider=mock_logger_provider
+        )
 
         with patch(
             "amazon.opentelemetry.distro.exporter.otlp.aws.traces.otlp_aws_span_exporter.LLOHandler"
@@ -110,7 +115,7 @@ def test_ensure_llo_handler_get_logger_provider_fails(self, mock_is_enabled, moc
         mock_get_logger_provider.side_effect = Exception("Failed to get logger provider")
 
         endpoint = "https://xray.us-east-1.amazonaws.com/v1/traces"
-        exporter = OTLPAwsSpanExporter(aws_region="us-east-1", endpoint=endpoint)
+        exporter = OTLPAwsSpanExporter(session=get_aws_session(), aws_region="us-east-1", endpoint=endpoint)
 
         result = exporter._ensure_llo_handler()
 
@@ -123,7 +128,7 @@ def test_export_with_llo_disabled(self, mock_is_enabled):
         mock_is_enabled.return_value = False
         endpoint = "https://xray.us-east-1.amazonaws.com/v1/traces"
 
-        exporter = OTLPAwsSpanExporter(aws_region="us-east-1", endpoint=endpoint)
+        exporter = OTLPAwsSpanExporter(session=get_aws_session(), aws_region="us-east-1", endpoint=endpoint)
 
         # Mock the parent class export method
         with patch.object(OTLPSpanExporter, "export") as mock_parent_export:
@@ -149,7 +154,7 @@ def test_export_with_llo_enabled(self, mock_llo_handler_class, mock_get_logger_p
         mock_llo_handler_class.return_value = mock_llo_handler
 
         endpoint = "https://xray.us-east-1.amazonaws.com/v1/traces"
-        exporter = OTLPAwsSpanExporter(aws_region="us-east-1", endpoint=endpoint)
+        exporter = OTLPAwsSpanExporter(session=get_aws_session(), aws_region="us-east-1", endpoint=endpoint)
 
         # Mock spans and processed spans
         original_spans = [MagicMock(spec=ReadableSpan), MagicMock(spec=ReadableSpan)]
@@ -182,7 +187,7 @@ def test_export_with_llo_processing_failure(
         mock_llo_handler.process_spans.side_effect = Exception("LLO processing failed")
 
         endpoint = "https://xray.us-east-1.amazonaws.com/v1/traces"
-        exporter = OTLPAwsSpanExporter(aws_region="us-east-1", endpoint=endpoint)
+        exporter = OTLPAwsSpanExporter(session=get_aws_session(), aws_region="us-east-1", endpoint=endpoint)
 
         spans = [MagicMock(spec=ReadableSpan), MagicMock(spec=ReadableSpan)]