97 lines
3.8 KiB
Python
97 lines
3.8 KiB
Python
import json
|
|
|
|
from azure.monitor.opentelemetry import configure_azure_monitor
|
|
from dynaconf import Dynaconf
|
|
from fastapi import FastAPI
|
|
from kn_utils.logging import logger
|
|
from opentelemetry import trace
|
|
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
|
|
from opentelemetry.instrumentation.aio_pika import AioPikaInstrumentor
|
|
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
|
|
from opentelemetry.instrumentation.pika import PikaInstrumentor
|
|
from opentelemetry.sdk.resources import Resource
|
|
from opentelemetry.sdk.trace import TracerProvider
|
|
from opentelemetry.sdk.trace.export import (
|
|
BatchSpanProcessor,
|
|
ConsoleSpanExporter,
|
|
SpanExporter,
|
|
SpanExportResult,
|
|
)
|
|
|
|
from pyinfra.config.loader import validate_settings
|
|
from pyinfra.config.validators import opentelemetry_validators
|
|
|
|
|
|
class JsonSpanExporter(SpanExporter):
|
|
def __init__(self):
|
|
self.traces = []
|
|
|
|
def export(self, spans):
|
|
for span in spans:
|
|
self.traces.append(json.loads(span.to_json()))
|
|
return SpanExportResult.SUCCESS
|
|
|
|
def shutdown(self):
|
|
pass
|
|
|
|
|
|
def setup_trace(settings: Dynaconf, service_name: str = None, exporter: SpanExporter = None):
|
|
tracing_type = settings.tracing.type
|
|
if tracing_type == "azure_monitor":
|
|
# Configure OpenTelemetry to use Azure Monitor with the
|
|
# APPLICATIONINSIGHTS_CONNECTION_STRING environment variable.
|
|
try:
|
|
configure_azure_monitor()
|
|
logger.info("Azure Monitor tracing enabled.")
|
|
except Exception as exception:
|
|
logger.warning(f"Azure Monitor tracing could not be enabled: {exception}")
|
|
elif tracing_type == "opentelemetry":
|
|
configure_opentelemtry_tracing(settings, service_name, exporter)
|
|
logger.info("OpenTelemetry tracing enabled.")
|
|
else:
|
|
logger.warning(f"Unknown tracing type: {tracing_type}. Tracing could not be enabled.")
|
|
|
|
|
|
def configure_opentelemtry_tracing(settings: Dynaconf, service_name: str = None, exporter: SpanExporter = None):
|
|
service_name = service_name or settings.tracing.opentelemetry.service_name
|
|
exporter = exporter or get_exporter(settings)
|
|
|
|
resource = Resource(attributes={"service.name": service_name})
|
|
provider = TracerProvider(resource=resource, shutdown_on_exit=True)
|
|
|
|
processor = BatchSpanProcessor(exporter)
|
|
provider.add_span_processor(processor)
|
|
|
|
# TODO: trace.set_tracer_provider produces a warning if trying to set the provider twice.
|
|
# "WARNING opentelemetry.trace:__init__.py:521 Overriding of current TracerProvider is not allowed"
|
|
# This doesn't seem to affect the functionality since we only want to use the tracer provided set in the beginning.
|
|
# We work around the log message by using the protected method with log=False.
|
|
trace._set_tracer_provider(provider, log=False)
|
|
|
|
|
|
def get_exporter(settings: Dynaconf):
|
|
validate_settings(settings, validators=opentelemetry_validators)
|
|
|
|
if settings.tracing.opentelemetry.exporter == "json":
|
|
return JsonSpanExporter()
|
|
elif settings.tracing.opentelemetry.exporter == "otlp":
|
|
return OTLPSpanExporter(endpoint=settings.tracing.opentelemetry.endpoint)
|
|
elif settings.tracing.opentelemetry.exporter == "console":
|
|
return ConsoleSpanExporter()
|
|
else:
|
|
raise ValueError(
|
|
f"Invalid OpenTelemetry exporter {settings.tracing.opentelemetry.exporter}. "
|
|
f"Valid values are 'json', 'otlp' and 'console'."
|
|
)
|
|
|
|
|
|
def instrument_pika(dynamic_queues: bool):
|
|
if dynamic_queues:
|
|
AioPikaInstrumentor().instrument()
|
|
else:
|
|
PikaInstrumentor().instrument()
|
|
|
|
|
|
def instrument_app(app: FastAPI, excluded_urls: str = "/health,/ready,/prometheus"):
|
|
FastAPIInstrumentor().instrument_app(app, excluded_urls=excluded_urls)
|