amq-adapter-python/amqp/service/tracing.py

"""
OpenTelemetry configuration for tracing and metrics.

This module provides functionality to initialize and configure OpenTelemetry
for both tracing (Jaeger via OTLP) and metrics (Prometheus) in the AMQ adapter.
"""

import logging
import os
from dataclasses import dataclass
from typing import Any, Dict, Optional

from opentelemetry import metrics, trace
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.exporter.prometheus import PrometheusMetricReader
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
from opentelemetry.sdk.resources import SERVICE_NAME, SERVICE_VERSION, Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
from opentelemetry.trace import Tracer
from prometheus_client import start_http_server

logger = logging.getLogger(__name__)


@dataclass
class TracingConfig:
    """Configuration for OpenTelemetry tracing."""

    service_name: str
    service_version: str
    otlp_endpoint: str = "http://localhost:4317"  # Default Jaeger OTLP endpoint
    insecure: bool = True
    debug: bool = False
    additional_attributes: Optional[Dict[str, Any]] = None


@dataclass
class MetricsConfig:
    """Configuration for OpenTelemetry metrics."""

    prometheus_port: int = 9464
    prometheus_host: str = "localhost"
    export_interval_millis: int = 30000
    additional_attributes: Optional[Dict[str, Any]] = None


def create_resource(config: TracingConfig) -> Resource:
    """Create an OpenTelemetry resource with service information."""
    attributes = {
        SERVICE_NAME: config.service_name,
        SERVICE_VERSION: config.service_version,
    }

    if config.additional_attributes:
        attributes.update(config.additional_attributes)

    return Resource.create(attributes)


def setup_otlp_exporter(config: TracingConfig) -> OTLPSpanExporter:
    """Configure and create an OTLP exporter for Jaeger."""
    return OTLPSpanExporter(endpoint=config.otlp_endpoint, insecure=config.insecure)


def setup_trace_provider(config: TracingConfig, resource: Resource) -> TracerProvider:
    """Set up the TracerProvider with OTLP exporter."""
    provider = TracerProvider(resource=resource)

    # Add OTLP exporter for Jaeger
    otlp_exporter = setup_otlp_exporter(config)
    provider.add_span_processor(BatchSpanProcessor(otlp_exporter))

    # Add console exporter in debug mode
    if config.debug:
        provider.add_span_processor(BatchSpanProcessor(ConsoleSpanExporter()))

    return provider


def setup_metrics_provider(config: MetricsConfig, resource: Resource) -> MeterProvider:
    """Set up the MeterProvider with Prometheus export."""
    try:
        # Start Prometheus HTTP server
        start_http_server(port=config.prometheus_port, addr=config.prometheus_host)
        logger.info(
            f"Started Prometheus metrics server on {config.prometheus_host}:{config.prometheus_port}"
        )

        # Create Prometheus reader
        prometheus_reader = PrometheusMetricReader()

        # Create periodic reader for custom metrics
        periodic_reader = PeriodicExportingMetricReader(
            prometheus_reader, export_interval_millis=config.export_interval_millis
        )

        # Create and configure meter provider
        return MeterProvider(resource=resource, metric_readers=[prometheus_reader, periodic_reader])
    except Exception as e:
        logger.error(f"Failed to set up metrics provider: {str(e)}")
        # Return a no-op meter provider
        return MeterProvider(resource=resource)


def initialize_telemetry(
    tracing_config: Optional[TracingConfig] = None,
    metrics_config: Optional[MetricsConfig] = None,
) -> Tracer:
    """
    Initialize OpenTelemetry with tracing and metrics.

    Args:
        tracing_config: Configuration for tracing
        metrics_config: Configuration for metrics

    Returns:
        Tracer: Configured OpenTelemetry tracer

    Example:
        ```python
        config = TracingConfig(
            service_name="amq-adapter",
            service_version="1.0.0",
            otlp_endpoint="http://jaeger:4317",
            debug=True
        )

        metrics_config = MetricsConfig(
            prometheus_port=9464,
            prometheus_host="0.0.0.0"
        )

        tracer = initialize_telemetry(config, metrics_config)
        ```
    """
    # Use default configs if none provided
    if tracing_config is None:
        tracing_config = TracingConfig(
            service_name=os.getenv("SERVICE_NAME", "amq-adapter"),
            service_version=os.getenv("SERVICE_VERSION", "unknown"),
            otlp_endpoint=os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4317"),
            insecure=os.getenv("OTEL_INSECURE_MODE", "true").lower() == "true",
            debug=os.getenv("OTEL_DEBUG", "false").lower() == "true",
        )

    if metrics_config is None:
        metrics_config = MetricsConfig(
            prometheus_port=int(os.getenv("PROMETHEUS_PORT", "9464")),
            prometheus_host=os.getenv("PROMETHEUS_HOST", "localhost"),
            export_interval_millis=int(os.getenv("METRICS_EXPORT_INTERVAL", "30000")),
        )

    try:
        # Create resource
        resource = create_resource(tracing_config)

        # Set up tracing
        provider = setup_trace_provider(tracing_config, resource)
        trace.set_tracer_provider(provider)
        logger.info(f"Initialized tracing for service {tracing_config.service_name}")

        # Set up metrics
        metrics_provider = setup_metrics_provider(metrics_config, resource)
        metrics.set_meter_provider(metrics_provider)
        logger.info("Initialized metrics export")

        # Return configured tracer
        return trace.get_tracer(
            tracing_config.service_name,
            tracing_config.service_version,
        )

    except Exception as e:
        logger.error(f"Failed to initialize telemetry: {str(e)}")
        # Return a no-op tracer in case of failure
        return trace.get_tracer("no-op")


def create_span(tracer: Tracer, name: str, attributes: Optional[Dict[str, Any]] = None):
    """
    Create a new span with the given name and attributes.

    Args:
        tracer: The OpenTelemetry tracer
        name: Name of the span
        attributes: Optional attributes to add to the span

    Returns:
        A context manager that creates and manages a span
    """
    return tracer.start_as_current_span(
        name,
        attributes=attributes or {},
    )