Custom Trace Attributes (#384)

Optionally define custom attributes that are added to every span.

Author: kraftp

dbos/_dbos_config.py

@@ -31,6 +31,7 @@ class DBOSConfig(TypedDict, total=False):
         otlp_logs_endpoints: List[str]: OTLP logs endpoints
         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
     """
 
     name: str
@@ -43,6 +44,7 @@ class DBOSConfig(TypedDict, total=False):
     otlp_logs_endpoints: Optional[List[str]]
     admin_port: Optional[int]
     run_admin_server: Optional[bool]
+    otlp_attributes: Optional[dict[str, str]]
 
 
 class RuntimeConfig(TypedDict, total=False):
@@ -84,6 +86,7 @@ class LoggerConfig(TypedDict, total=False):
 class TelemetryConfig(TypedDict, total=False):
     logs: Optional[LoggerConfig]
     OTLPExporter: Optional[OTLPExporterConfig]
+    otlp_attributes: Optional[dict[str, str]]
 
 
 class ConfigFile(TypedDict, total=False):
@@ -145,7 +148,8 @@ def translate_dbos_config_to_config_file(config: DBOSConfig) -> ConfigFile:
 
     # Telemetry config
     telemetry: TelemetryConfig = {
-        "OTLPExporter": {"tracesEndpoint": [], "logsEndpoint": []}
+        "OTLPExporter": {"tracesEndpoint": [], "logsEndpoint": []},
+        "otlp_attributes": config.get("otlp_attributes", {}),
     }
     # For mypy
     assert telemetry["OTLPExporter"] is not None

dbos/_logger.py

@@ -20,14 +20,17 @@
 
 
 class DBOSLogTransformer(logging.Filter):
-    def __init__(self) -> None:
+    def __init__(self, config: "ConfigFile") -> None:
         super().__init__()
         self.app_id = os.environ.get("DBOS__APPID", "")
+        self.otlp_attributes: dict[str, str] = config.get("telemetry", {}).get("otlp_attributes", {})  # type: ignore
 
     def filter(self, record: Any) -> bool:
         record.applicationID = self.app_id
         record.applicationVersion = GlobalParams.app_version
         record.executorID = GlobalParams.executor_id
+        for k, v in self.otlp_attributes.items():
+            setattr(record, k, v)
 
         # If available, decorate the log entry with Workflow ID and Trace ID
         from dbos._context import get_local_dbos_context
@@ -98,7 +101,7 @@ def config_logger(config: "ConfigFile") -> None:
 
     # Attach DBOS-specific attributes to all log entries.
     global _dbos_log_transformer
-    _dbos_log_transformer = DBOSLogTransformer()
+    _dbos_log_transformer = DBOSLogTransformer(config)
     dbos_logger.addFilter(_dbos_log_transformer)
 
 

dbos/_tracer.py

@@ -19,11 +19,14 @@
 
 class DBOSTracer:
 
+    otlp_attributes: dict[str, str] = {}
+
     def __init__(self) -> None:
         self.app_id = os.environ.get("DBOS__APPID", None)
         self.provider: Optional[TracerProvider] = None
 
     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):
             resource = Resource(
                 attributes={
@@ -63,6 +66,8 @@ def start_span(
         for k, v in attributes.items():
             if k != "name" and v is not None and isinstance(v, (str, bool, int, float)):
                 span.set_attribute(k, v)
+        for k, v in self.otlp_attributes.items():
+            span.set_attribute(k, v)
         return span
 
     def end_span(self, span: Span) -> None:

tests/test_spans.py

@@ -7,12 +7,16 @@
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 
-from dbos import DBOS
+from dbos import DBOS, DBOSConfig
 from dbos._tracer import dbos_tracer
 from dbos._utils import GlobalParams
 
 
-def test_spans(dbos: DBOS) -> None:
+def test_spans(config: DBOSConfig) -> None:
+    DBOS.destroy(destroy_registry=True)
+    config["otlp_attributes"] = {"foo": "bar"}
+    DBOS(config=config)
+    DBOS.launch()
 
     @DBOS.workflow()
     def test_workflow() -> None:
@@ -44,6 +48,7 @@ def test_step() -> None:
         assert span.attributes["applicationVersion"] == GlobalParams.app_version
         assert span.attributes["executorID"] == GlobalParams.executor_id
         assert span.context is not None
+        assert span.attributes["foo"] == "bar"
 
     assert spans[0].name == test_step.__name__
     assert spans[1].name == "a new span"