refactor: streamline tracing types

pyinfra/examples.py

@@ -1,7 +1,6 @@
 from dynaconf import Dynaconf
 from fastapi import FastAPI
 from kn_utils.logging import logger
-from azure.monitor.opentelemetry import configure_azure_monitor
 
 from pyinfra.config.loader import get_pyinfra_validators, validate_settings
 from pyinfra.queue.callback import Callback
@@ -42,15 +41,8 @@ def start_standard_queue_consumer(
         app = add_prometheus_endpoint(app)
         callback = make_prometheus_processing_time_decorator_from_settings(settings)(callback)
 
-    if settings.tracing.opentelemetry.enabled:
-        logger.info("OpenTelemetry tracing enabled.")
-        if settings.tracing.opentelemetry.azure_monitoring:
-            # Configure OpenTelemetry to use Azure Monitor with the 
-            # APPLICATIONINSIGHTS_CONNECTION_STRING environment variable.
-            logger.info("Azure Monitor tracing enabled.")
-            configure_azure_monitor()
-        else:
-            setup_trace(settings)
+    if settings.tracing.enabled:
+        setup_trace(settings)
             
         instrument_pika()
         instrument_app(app)

pyinfra/utils/opentelemetry.py

@@ -1,5 +1,6 @@
 import json
 
+from azure.monitor.opentelemetry import configure_azure_monitor
 from dynaconf import Dynaconf
 from fastapi import FastAPI
 from opentelemetry import trace
@@ -17,6 +18,7 @@ from opentelemetry.sdk.trace.export import (
 
 from pyinfra.config.loader import validate_settings
 from pyinfra.config.validators import opentelemetry_validators
+from kn_utils.logging import logger
 
 
 class JsonSpanExporter(SpanExporter):
@@ -33,6 +35,20 @@ class JsonSpanExporter(SpanExporter):
 
 
 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.
+        logger.info("Azure Monitor tracing enabled.")
+        configure_azure_monitor()
+    elif tracing_type == "opentelemetry":
+        logger.info("OpenTelemetry tracing enabled.")
+        configure_opentelemtry_tracing(settings, service_name, exporter)
+    else:
+        raise Exception(f"Unknown tracing type: {tracing_type}")
+
+
+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)
 

pyinfra/webserver/utils.py

@@ -14,7 +14,7 @@ from pyinfra.utils.opentelemetry import instrument_app, setup_trace
 def create_webserver_thread_from_settings(app: FastAPI, settings: Dynaconf) -> threading.Thread:
     validate_settings(settings, validators=webserver_validators)
 
-    if settings.tracing.opentelemetry.enabled:
+    if settings.tracing.enabled:
         return create_webserver_thread_with_tracing(app, settings)
 
     return create_webserver_thread(app=app, port=settings.webserver.port, host=settings.webserver.host)

tests/unit_test/opentelemetry_test.py

@@ -34,7 +34,7 @@ class TestOpenTelemetry:
 
     # def test_webserver_requests_are_traced(self, settings):
     #     settings.tracing.opentelemetry.exporter = "console"
-    #     settings.tracing.opentelemetry.enabled = True
+    #     settings.tracing.enabled = True
     #
     #     app = FastAPI()
     #