Add an option to disable OTLP tracing (#439)

Author: qianl15

dbos/_context.py

@@ -221,6 +221,8 @@ def get_current_span(self) -> Optional[Span]:
         return None
 
     def _start_span(self, attributes: TracedAttributes) -> None:
+        if dbos_tracer.disable_otlp:
+            return
         attributes["operationUUID"] = (
             self.workflow_id if len(self.workflow_id) > 0 else None
         )
@@ -246,6 +248,8 @@ def _start_span(self, attributes: TracedAttributes) -> None:
         cm.__enter__()
 
     def _end_span(self, exc_value: Optional[BaseException]) -> None:
+        if dbos_tracer.disable_otlp:
+            return
         context_span = self.context_spans.pop()
         if exc_value is None:
             context_span.span.set_status(Status(StatusCode.OK))

dbos/_dbos_config.py

@@ -33,6 +33,9 @@ class DBOSConfig(TypedDict, total=False):
         admin_port (int): Admin port
         run_admin_server (bool): Whether to run the DBOS admin server
         otlp_attributes (dict[str, str]): A set of custom attributes to apply OTLP-exported logs and traces
+        application_version (str): Application version
+        executor_id (str): Executor ID, used to identify the application instance in distributed environments
+        disable_otlp (bool): If True, disables OTLP tracing and logging. Defaults to False.
     """
 
     name: str
@@ -49,6 +52,7 @@ class DBOSConfig(TypedDict, total=False):
     otlp_attributes: Optional[dict[str, str]]
     application_version: Optional[str]
     executor_id: Optional[str]
+    disable_otlp: Optional[bool]
 
 
 class RuntimeConfig(TypedDict, total=False):
@@ -91,6 +95,7 @@ class TelemetryConfig(TypedDict, total=False):
     logs: Optional[LoggerConfig]
     OTLPExporter: Optional[OTLPExporterConfig]
     otlp_attributes: Optional[dict[str, str]]
+    disable_otlp: Optional[bool]
 
 
 class ConfigFile(TypedDict, total=False):
@@ -157,6 +162,7 @@ def translate_dbos_config_to_config_file(config: DBOSConfig) -> ConfigFile:
     telemetry: TelemetryConfig = {
         "OTLPExporter": {"tracesEndpoint": [], "logsEndpoint": []},
         "otlp_attributes": config.get("otlp_attributes", {}),
+        "disable_otlp": config.get("disable_otlp", False),
     }
     # For mypy
     assert telemetry["OTLPExporter"] is not None

dbos/_logger.py

@@ -77,7 +77,9 @@ def config_logger(config: "ConfigFile") -> None:
     otlp_logs_endpoints = (
         config.get("telemetry", {}).get("OTLPExporter", {}).get("logsEndpoint")  # type: ignore
     )
-    if otlp_logs_endpoints:
+    disable_otlp = config.get("telemetry", {}).get("disable_otlp", False)  # type: ignore
+
+    if not disable_otlp and otlp_logs_endpoints:
         log_provider = PatchedOTLPLoggerProvider(
             Resource.create(
                 attributes={

dbos/_tracer.py

@@ -24,10 +24,14 @@ class DBOSTracer:
     def __init__(self) -> None:
         self.app_id = os.environ.get("DBOS__APPID", None)
         self.provider: Optional[TracerProvider] = None
+        self.disable_otlp: bool = False
 
     def config(self, config: ConfigFile) -> None:
         self.otlp_attributes = config.get("telemetry", {}).get("otlp_attributes", {})  # type: ignore
-        if not isinstance(trace.get_tracer_provider(), TracerProvider):
+        self.disable_otlp = config.get("telemetry", {}).get("disable_otlp", False)  # type: ignore
+        if not self.disable_otlp and not isinstance(
+            trace.get_tracer_provider(), TracerProvider
+        ):
             resource = Resource(
                 attributes={
                     ResourceAttributes.SERVICE_NAME: config["name"],

tests/test_spans.py

@@ -254,3 +254,61 @@ def test_workflow_endpoint() -> str:
     assert spans[0].context is not None
     assert logs[0].log_record.span_id == spans[0].context.span_id
     assert logs[0].log_record.trace_id == spans[0].context.trace_id
+
+
+def test_disable_otlp_no_spans(config: DBOSConfig) -> None:
+    DBOS.destroy(destroy_registry=True)
+    config["otlp_attributes"] = {"foo": "bar"}
+    config["disable_otlp"] = True
+    DBOS(config=config)
+    DBOS.launch()
+
+    @DBOS.workflow()
+    def test_workflow() -> None:
+        test_step()
+        DBOS.logger.info("This is a test_workflow")
+
+    @DBOS.step()
+    def test_step() -> None:
+        DBOS.logger.info("This is a test_step")
+        return
+
+    exporter = InMemorySpanExporter()
+    span_processor = SimpleSpanProcessor(exporter)
+    provider = tracesdk.TracerProvider()
+    provider.add_span_processor(span_processor)
+    dbos_tracer.set_provider(provider)
+
+    # Set up in-memory log exporter
+    log_exporter = InMemoryLogExporter()  # type: ignore
+    log_processor = BatchLogRecordProcessor(log_exporter)
+    log_provider = LoggerProvider()
+    log_provider.add_log_record_processor(log_processor)
+    set_logger_provider(log_provider)
+    dbos_logger.addHandler(LoggingHandler(logger_provider=log_provider))
+
+    test_workflow()
+
+    log_processor.force_flush(timeout_millis=5000)
+    logs = log_exporter.get_finished_logs()
+    assert len(logs) == 2
+    for log in logs:
+        assert log.log_record.attributes is not None
+        assert (
+            log.log_record.attributes["applicationVersion"] == GlobalParams.app_version
+        )
+        assert log.log_record.attributes["executorID"] == GlobalParams.executor_id
+        assert log.log_record.attributes["foo"] == "bar"
+        # We disable OTLP, so no span_id or trace_id should be present
+        assert log.log_record.span_id is not None and log.log_record.span_id == 0
+        assert log.log_record.trace_id is not None and log.log_record.trace_id == 0
+        assert (
+            log.log_record.body == "This is a test_step"
+            or log.log_record.body == "This is a test_workflow"
+        )
+        assert log.log_record.attributes.get("traceId") is None
+
+    spans = exporter.get_finished_spans()
+
+    # No spans should be created since OTLP is disabled
+    assert len(spans) == 0