From 5328e8de035ab83fc4c15dbc277da013dea18108 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonathan=20K=C3=B6ssler?= <jonathan.koessler@knecon.com>
Date: Wed, 12 Jun 2024 10:41:52 +0200
Subject: [PATCH] refactor: streamline tracing types

---
 pyinfra/examples.py                   | 12 ++----------
 pyinfra/utils/opentelemetry.py        | 16 ++++++++++++++++
 pyinfra/webserver/utils.py            |  2 +-
 tests/unit_test/opentelemetry_test.py |  2 +-
 4 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/pyinfra/examples.py b/pyinfra/examples.py
index 7eede28..6b62f98 100644
--- a/pyinfra/examples.py
+++ b/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)
diff --git a/pyinfra/utils/opentelemetry.py b/pyinfra/utils/opentelemetry.py
index f23f6f5..7a05233 100644
--- a/pyinfra/utils/opentelemetry.py
+++ b/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)
 
diff --git a/pyinfra/webserver/utils.py b/pyinfra/webserver/utils.py
index bbf5ac5..8ca4e1d 100644
--- a/pyinfra/webserver/utils.py
+++ b/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)
diff --git a/tests/unit_test/opentelemetry_test.py b/tests/unit_test/opentelemetry_test.py
index 3a68ab9..7c55d45 100644
--- a/tests/unit_test/opentelemetry_test.py
+++ b/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()
     #