conolidate botocore sessions into utils

Author: liustve

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

@@ -4,6 +4,7 @@
 import os
 from importlib.metadata import PackageNotFoundError, version
 from logging import Logger, getLogger
+from typing import Optional
 
 from packaging.requirements import Requirement
 
@@ -37,30 +38,18 @@ def is_agent_observability_enabled() -> bool:
     return os.environ.get(AGENT_OBSERVABILITY_ENABLED, "false").lower() == "true"
 
 
-def get_aws_region() -> str:
-    """Get AWS region using botocore session.
-
-    botocore automatically checks in the following priority order:
-    1. AWS_REGION environment variable
-    2. AWS_DEFAULT_REGION environment variable
-    3. AWS CLI config file (~/.aws/config)
-    4. EC2 instance metadata service
-
-    Returns:
-        The AWS region if found, None otherwise.
-    """
-    if is_installed("botocore"):
-        try:
-            from botocore import session  # pylint: disable=import-outside-toplevel
-
-            botocore_session = session.Session()
-            if botocore_session.region_name:
-                return botocore_session.region_name
-        except (ImportError, AttributeError):
-            # botocore failed to determine region
-            pass
-
-    _logger.warning(
-        "AWS region not found. Please set AWS_REGION environment variable or configure AWS CLI with 'aws configure'."
-    )
+IS_BOTOCORE_INSTALLED: bool = is_installed("botocore")
+
+
+def get_aws_session():
+    if IS_BOTOCORE_INSTALLED:
+        # pylint: disable=import-outside-toplevel
+        from botocore.session import Session
+
+        return Session()
     return None
+
+
+def get_aws_region() -> Optional[str]:
+    botocore_session = get_aws_session()
+    return botocore_session.get_config_variable("region") if botocore_session else None

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
 from amazon.opentelemetry.distro._aws_resource_attribute_configurator import get_service_attribute
-from amazon.opentelemetry.distro._utils import is_agent_observability_enabled, is_installed
+from amazon.opentelemetry.distro._utils import IS_BOTOCORE_INSTALLED, 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,11 +23,6 @@
     AwsMetricAttributesSpanExporterBuilder,
 )
 from amazon.opentelemetry.distro.aws_span_metrics_processor_builder import AwsSpanMetricsProcessorBuilder
-from amazon.opentelemetry.distro.exporter.otlp.aws.logs.aws_batch_log_record_processor import (
-    AwsCloudWatchOtlpBatchLogRecordProcessor,
-)
-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
@@ -214,8 +209,7 @@ 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(provider, log_exporter)
 
     event_logger_provider = EventLoggerProvider(logger_provider=provider)
     set_event_logger_provider(event_logger_provider)
@@ -303,7 +297,7 @@ def _export_unsampled_span_for_agent_observability(trace_provider: TracerProvide
 
     traces_endpoint = os.environ.get(OTEL_EXPORTER_OTLP_TRACES_ENDPOINT)
 
-    span_exporter = OTLPAwsSpanExporter(endpoint=traces_endpoint, logger_provider=get_logger_provider())
+    span_exporter = _create_aws_exporter(endpoint=traces_endpoint)
 
     trace_provider.add_span_processor(BatchUnsampledSpanProcessor(span_exporter=span_exporter))
 
@@ -404,15 +398,7 @@ def _customize_span_exporter(span_exporter: SpanExporter, resource: Resource) ->
         _logger.info("Detected using AWS OTLP Traces Endpoint.")
 
         if isinstance(span_exporter, OTLPSpanExporter):
-            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
-                # from span attributes and send them to the logs pipeline per
-                # the new Gen AI semantic convention from OTel
-                # ref: https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-events/
-                span_exporter = OTLPAwsSpanExporter(endpoint=traces_endpoint, logger_provider=get_logger_provider())
-            else:
-                span_exporter = OTLPAwsSpanExporter(endpoint=traces_endpoint)
+            return _create_aws_exporter(endpoint=traces_endpoint)
 
         else:
             _logger.warning(
@@ -426,14 +412,20 @@ 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(provider: LoggerProvider, log_exporter: Optional[LogExporter]) -> None:
+    if log_exporter is None:
+        return
+    if is_agent_observability_enabled() and IS_BOTOCORE_INSTALLED:
+        from amazon.opentelemetry.distro.exporter.otlp.aws.logs.aws_batch_log_record_processor import (
+            AwsCloudWatchOtlpBatchLogRecordProcessor,
+        )
 
-    return BatchLogRecordProcessor(exporter=log_exporter)
+        provider.add_log_record_processor(AwsCloudWatchOtlpBatchLogRecordProcessor(exporter=log_exporter))
+    else:
+        provider.add_log_record_processor(BatchLogRecordProcessor(exporter=log_exporter))
 
 
-def _customize_logs_exporter(log_exporter: LogExporter) -> LogExporter:
+def _customize_logs_exporter(log_exporter: LogExporter) -> Optional[LogExporter]:
     logs_endpoint = os.environ.get(OTEL_EXPORTER_OTLP_LOGS_ENDPOINT)
 
     if _is_aws_otlp_endpoint(logs_endpoint, "logs"):
@@ -443,7 +435,7 @@ def _customize_logs_exporter(log_exporter: LogExporter) -> LogExporter:
             # 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 OTLPAwsLogExporter(endpoint=logs_endpoint)
+            return _create_aws_exporter(endpoint=logs_endpoint)
 
         _logger.warning(
             "Improper configuration see: please export/set "
@@ -762,8 +754,9 @@ def _check_emf_exporter_enabled() -> bool:
 def create_emf_exporter():
     """Create and configure the CloudWatch EMF exporter."""
     try:
+        session = get_aws_session()
         # Check if botocore is available before importing the EMF exporter
-        if not is_installed("botocore"):
+        if not session:
             _logger.warning("botocore is not installed. EMF exporter requires botocore")
             return None
 
@@ -778,6 +771,7 @@ def create_emf_exporter():
             return None
 
         return AwsCloudWatchEmfExporter(
+            session=session,
             namespace=log_header_setting.namespace,
             log_group_name=log_header_setting.log_group,
             log_stream_name=log_header_setting.log_stream,
@@ -786,3 +780,43 @@ def create_emf_exporter():
     except Exception as errors:
         _logger.error("Failed to create EMF exporter: %s", errors)
         return None
+
+
+def _create_aws_exporter(endpoint: str):
+    """Create and configure the AWS OTLP exporters."""
+    try:
+        session = get_aws_session()
+        # Check if botocore is available before importing the AWS exporter
+        if not session:
+            _logger.warning("SigV4 Auth requires botocore to be enabled")
+            return None
+
+        # pylint: disable=import-outside-toplevel
+        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 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
+                # from span attributes and send them to the logs pipeline per
+                # the new Gen AI semantic convention from OTel
+                # ref: https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-events/
+                return OTLPAwsSpanExporter(
+                    session=session, endpoint=endpoint, aws_region=region, logger_provider=get_logger_provider()
+                )
+
+            return OTLPAwsSpanExporter(session=session, endpoint=endpoint, aws_region=region)
+
+        if "logs" in service:
+            return OTLPAwsLogExporter(session=session, aws_region=region)
+
+    # pylint: disable=broad-exception-caught
+    except Exception as errors:
+        _logger.error("Failed to create AWS OTLP exporter: %s", errors)
+        return None

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

@@ -88,8 +88,9 @@ def _configure(self, **kwargs):
             # Set GenAI capture content default
             os.environ.setdefault(OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, "true")
 
-            # Set OTLP endpoints with AWS region if not already set
             region = get_aws_region()
+
+            # Set OTLP endpoints with AWS region if not already set
             if region:
                 os.environ.setdefault(
                     OTEL_EXPORTER_OTLP_TRACES_ENDPOINT, f"https://xray.{region}.amazonaws.com/v1/traces"

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

@@ -8,8 +8,8 @@
 import uuid
 from typing import Any, Dict, List, Optional
 
-import botocore.session
 from botocore.exceptions import ClientError
+from botocore.session import Session
 
 logger = logging.getLogger(__name__)
 
@@ -90,6 +90,7 @@ class CloudWatchLogClient:
     def __init__(
         self,
         log_group_name: str,
+        session: Session = Session(),
         log_stream_name: Optional[str] = None,
         aws_region: Optional[str] = None,
         **kwargs,
@@ -105,8 +106,6 @@ def __init__(
         """
         self.log_group_name = log_group_name
         self.log_stream_name = log_stream_name or self._generate_log_stream_name()
-
-        session = botocore.session.Session()
         self.logs_client = session.create_client("logs", region_name=aws_region, **kwargs)
 
         # Event batch to store logs before sending to CloudWatch

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/exporter/otlp/aws/common/aws_auth_session.py

@@ -4,8 +4,9 @@
 import logging
 
 import requests
-
-from amazon.opentelemetry.distro._utils import is_installed
+from botocore.auth import SigV4Auth
+from botocore.awsrequest import AWSRequest
+from botocore.session import Session
 
 _logger = logging.getLogger(__name__)
 
@@ -33,57 +34,36 @@ class AwsAuthSession(requests.Session):
         service (str): The AWS service name for signing (e.g., "logs" or "xray")
     """
 
-    def __init__(self, aws_region, service):
-
-        self._has_required_dependencies = False
-
-        # Requires botocore to be installed to sign the headers. However,
-        # some users might not need to use this authenticator. In order not conflict
-        # with existing behavior, we check for botocore before initializing this exporter.
-
-        if aws_region and service and is_installed("botocore"):
-            # pylint: disable=import-outside-toplevel
-            from botocore import auth, awsrequest, session
-
-            self._boto_auth = auth
-            self._boto_aws_request = awsrequest
-            self._boto_session = session.Session()
-
-            self._aws_region = aws_region
-            self._service = service
-            self._has_required_dependencies = True
-
-        else:
-            _logger.error(
-                "botocore is required to enable SigV4 Authentication. Please install it using `pip install botocore`",
-            )
+    def __init__(self, aws_region: str, service: str, session: Session = Session()):
+        self._aws_region: str = aws_region
+        self._service: str = service
+        self._session: Session = session
 
         super().__init__()
 
     def request(self, method, url, *args, data=None, headers=None, **kwargs):
-        if self._has_required_dependencies:
-
-            credentials = self._boto_session.get_credentials()
-
-            if credentials is not None:
-                signer = self._boto_auth.SigV4Auth(credentials, self._service, self._aws_region)
-
-                request = self._boto_aws_request.AWSRequest(
-                    method="POST",
-                    url=url,
-                    data=data,
-                    headers={"Content-Type": "application/x-protobuf"},
-                )
+        credentials = self._session.get_credentials()
+
+        if credentials:
+            signer = SigV4Auth(credentials, self._service, self._aws_region)
+            request = AWSRequest(
+                method="POST",
+                url=url,
+                data=data,
+                headers={"Content-Type": "application/x-protobuf"},
+            )
 
-                try:
-                    signer.add_auth(request)
+            try:
+                signer.add_auth(request)
 
-                    if headers is None:
-                        headers = {}
+                if headers is None:
+                    headers = {}
 
-                    headers.update(dict(request.headers))
+                headers.update(dict(request.headers))
 
-                except Exception as signing_error:  # pylint: disable=broad-except
-                    _logger.error("Failed to sign request: %s", signing_error)
+            except Exception as signing_error:  # pylint: disable=broad-except
+                _logger.error("Failed to sign request: %s", signing_error)
+        else:
+            _logger.error("Failed to load AWS Credentials: %s")
 
         return super().request(method=method, url=url, *args, data=data, headers=headers, **kwargs)

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

@@ -7,6 +7,7 @@
 from time import sleep
 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
@@ -28,17 +29,25 @@ class OTLPAwsLogExporter(OTLPLogExporter):
 
     def __init__(
         self,
+        aws_region: str,
+        session: Session = Session(),
+        log_group: Optional[str] = None,
+        log_stream: Optional[str] = None,
         endpoint: Optional[str] = None,
         certificate_file: Optional[str] = None,
         client_key_file: Optional[str] = None,
         client_certificate_file: Optional[str] = None,
         headers: Optional[Dict[str, str]] = None,
         timeout: Optional[int] = None,
     ):
-        self._aws_region = None
+        self._aws_region = aws_region
 
-        if endpoint:
-            self._aws_region = endpoint.split(".")[1]
+        if log_group and log_stream:
+            log_headers = {"x-aws-log-group": log_group, "x-aws-log-stream": log_stream}
+            if headers:
+                headers.update(log_headers)
+            else:
+                headers = log_headers
 
         OTLPLogExporter.__init__(
             self,
@@ -49,7 +58,7 @@ def __init__(
             headers,
             timeout,
             compression=Compression.Gzip,
-            session=AwsAuthSession(aws_region=self._aws_region, service="logs"),
+            session=AwsAuthSession(session=session, aws_region=self._aws_region, service="logs"),
         )
 
     def export(self, batch: Sequence[LogData]) -> LogExportResult: