Merge branch 'dev' of https://github.com/DiamondLightSource/python-workflows into dev

Author: davidigandan

src/workflows/recipe/__init__.py

@@ -3,9 +3,10 @@
 import functools
 import logging
 from collections.abc import Callable
-from opentelemetry import trace
 from typing import Any
 
+from opentelemetry import trace
+
 from workflows.recipe.recipe import Recipe
 from workflows.recipe.validate import validate_recipe
 from workflows.recipe.wrapper import RecipeWrapper
@@ -80,16 +81,18 @@ def unwrap_recipe(header, message):
                 environment = message.get("environment", {})
                 if isinstance(environment, dict):
                     recipe_id = environment.get("ID")
-            
+
             if recipe_id:
                 span.set_attribute("recipe_id", recipe_id)
-                span.add_event("recipe.id_extracted", attributes={"recipe_id": recipe_id})
+                span.add_event(
+                    "recipe.id_extracted", attributes={"recipe_id": recipe_id}
+                )
 
             # Extract span_id and trace_id for logging
             span_context = span.get_span_context()
             if span_context and span_context.is_valid:
-                span_id = format(span_context.span_id, '016x')
-                trace_id = format(span_context.trace_id, '032x')
+                span_id = format(span_context.span_id, "016x")
+                trace_id = format(span_context.trace_id, "032x")
 
                 log_extra = {
                     "span_id": span_id,

src/workflows/services/common_service.py

@@ -9,15 +9,15 @@
 import time
 from typing import Any
 
-import workflows
-import workflows.logging
-
 from opentelemetry import trace
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.resources import SERVICE_NAME, Resource
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
-from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+
+import workflows
+import workflows.logging
 from workflows.transport.middleware.otel_tracing import OTELTracingMiddleware
-from opentelemetry.sdk.resources import Resource, SERVICE_NAME
 
 
 class Status(enum.Enum):
@@ -216,7 +216,9 @@ def start_transport(self):
 
             # Add OTELTracingMiddleware to the transport layer
             tracer = trace.get_tracer(__name__)
-            otel_middleware = OTELTracingMiddleware(tracer, service_name=self._service_name)
+            otel_middleware = OTELTracingMiddleware(
+                tracer, service_name=self._service_name
+            )
             self._transport.add_middleware(otel_middleware)
 
 

src/workflows/transport/middleware/otel_tracing.py

@@ -1,8 +1,13 @@
+from __future__ import annotations
+
+import functools
+from collections.abc import Callable
+
 from opentelemetry import trace
+from opentelemetry.propagate import extract
+
 from workflows.transport.middleware import BaseTransportMiddleware
-from collections.abc import Callable
-import functools
-from opentelemetry.propagate import inject, extract
+
 
 class OTELTracingMiddleware(BaseTransportMiddleware):
     def __init__(self, tracer: trace.Tracer, service_name: str):
@@ -14,18 +19,16 @@ def subscribe(self, call_next: Callable, channel, callback, **kwargs) -> int:
         def wrapped_callback(header, message):
             # Extract trace context from message headers
             ctx = extract(header) if header else None
-            
+
             # Start a new span with the extracted context
             with self.tracer.start_as_current_span(
-                "transport.subscribe",
-                context=ctx
+                "transport.subscribe", context=ctx
             ) as span:
                 span.set_attribute("service_name", self.service_name)
                 span.set_attribute("channel", channel)
-        
-                
+
                 # Call the original callback
                 return callback(header, message)
-        
+
         # Call the next middleware with the wrapped callback
-        return call_next(channel, wrapped_callback, **kwargs)
+        return call_next(channel, wrapped_callback, **kwargs)