Fix broken tracing functionality

Author: davidigandan

src/workflows/services/common_service.py

@@ -197,17 +197,7 @@ def start_transport(self):
             otel_config = (
                 self.config._opentelemetry if self.config and hasattr(self.config, "opentelemetry") else None
             )
-            # debugging
-            with open("/scratch/logs.txt", 'w+') as file:
-                if otel_config:
-                    import json
-                    json.dump(otel_config, file, indent=4)
-                else:
-                    file.write("otel config was not truthy")
-            
-            if otel_config and "timeout" not in otel_config:
-                self.log.warning("Missing optional OTEL configuration field `timeout`. Will default to 10 seconds. ")
-                
+            if otel_config:
                 # Configure OTELTracing
                 resource = Resource.create(
                     {

src/workflows/transport/common_transport.py

@@ -10,7 +10,6 @@
 
 MessageCallback = Callable[[Mapping[str, Any], Any], None]
 
-
 class TemporarySubscription(NamedTuple):
     subscription_id: int
     queue_name: str

src/workflows/transport/middleware/__init__.py

@@ -233,6 +233,10 @@ def wrapped_callback(header, message):
 
 
 def wrap(f: Callable):
+    # debugging
+    if f.__name__ == "send":
+        print("we are wrapping send now")
+    
     @functools.wraps(f)
     def wrapper(self, *args, **kwargs):
         return functools.reduce(
@@ -243,4 +247,5 @@ def wrapper(self, *args, **kwargs):
             lambda *args, **kwargs: f(self, *args, **kwargs),
         )(*args, **kwargs)
 
-    return wrapper
+    print(wrapper.__wrapped__)
+    return wrapper

src/workflows/transport/middleware/otel_tracing.py

@@ -4,31 +4,228 @@
 from collections.abc import Callable
 
 from opentelemetry import trace
-from opentelemetry.propagate import extract
+from opentelemetry.propagate import extract, inject
+from opentelemetry.context import Context
 
 from workflows.transport.middleware import BaseTransportMiddleware
+from workflows.transport.common_transport import TemporarySubscription, MessageCallback
+import json
 
-
-class OTELTracingMiddleware(BaseTransportMiddleware):
+class OTELTracingMiddleware:
     def __init__(self, tracer: trace.Tracer, service_name: str):
         self.tracer = tracer
         self.service_name = service_name
 
-    def subscribe(self, call_next: Callable, channel, callback, **kwargs) -> int:
+    def send(self, call_next: Callable, destination: str, message: Any, **kwargs):
+        # Get current span context (may be None if this is the root span)
+        current_span = trace.get_current_span()
+        parent_context = trace.set_span_in_context(current_span) if current_span else None
+
+        with self.tracer.start_as_current_span(
+            "transport.send",
+            context=parent_context,
+        ) as span:
+            span.set_attribute("service_name", self.service_name)
+
+            span.set_attribute("message", json.dumps(message))
+            span.set_attribute("destination", destination)
+            print("parent_context is...",parent_context)
+
+
+            # Inject the current trace context into the message headers
+            headers = kwargs.get("headers", {})
+            if headers is None:
+                headers = {}
+            inject(headers)  # This modifies headers in-place
+            kwargs["headers"] = headers
+
+            return call_next(destination, message, **kwargs)
+
+    def subscribe(self, call_next: Callable, channel: str, callback: Callable, **kwargs) -> int:
         @functools.wraps(callback)
         def wrapped_callback(header, message):
             # Extract trace context from message headers
-            ctx = extract(header) if header else None
+            ctx = extract(header) if header else Context()
 
             # 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("message", json.dumps(message))
+                span.set_attribute("channel", channel)
+
+                # Call the original callback - this will process the message
+                # and potentially call send() which will pick up this context
+                return callback(header, message)
+
+        return call_next(channel, wrapped_callback, **kwargs)
+
+    def subscribe_broadcast(self, call_next: Callable, channel: str, callback: Callable, **kwargs) -> int:
+        @functools.wraps(callback)
+        def wrapped_callback(header, message):
+            # Extract trace context from message headers
+            ctx = extract(header) if header else Context()
+
+    #         # Start a new span with the extracted context
+            with self.tracer.start_as_current_span(
+                "transport.subscribe_broadcast",
+                context=ctx,
             ) as span:
                 span.set_attribute("service_name", self.service_name)
+
+                span.set_attribute("message", json.dumps(message))
                 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)
+
+    def subscribe_temporary(
+        self,
+        call_next: Callable,
+        channel_hint: str | None,
+        callback: MessageCallback,
+        **kwargs,
+    ) -> TemporarySubscription:
+        @functools.wraps(callback)
+        def wrapped_callback(header, message):
+            # Extract trace context from message headers
+            ctx = extract(header) if header else Context()
+
+            # Start a new span with the extracted context
+            with self.tracer.start_as_current_span(
+                "transport.subscribe_temporary",
+                context=ctx,
+            ) as span:
+                span.set_attribute("service_name", self.service_name)
+
+                span.set_attribute("message", json.dumps(message))
+                if channel_hint:
+                    span.set_attribute("channel_hint", channel_hint)
+
+                return callback(header, message)
+
+        return call_next(channel_hint, wrapped_callback, **kwargs)
+
+    def unsubscribe(
+        self,
+        call_next: Callable,
+        subscription: int,
+        drop_callback_reference=False,
+        **kwargs,
+    ):
+        # Get current span context
+        current_span = trace.get_current_span()
+        current_context = trace.set_span_in_context(current_span) if current_span else Context()
+
+        with self.tracer.start_as_current_span(
+            "transport.unsubscribe",
+            context=current_context,
+        ) as span:
+            span.set_attribute("service_name", self.service_name)
+            span.set_attribute("subscription_id", subscription)
+
+            call_next(
+                subscription, drop_callback_reference=drop_callback_reference, **kwargs
+            )
+
+    def ack(
+        self,
+        call_next: Callable,
+        message,
+        subscription_id: int | None = None,
+        **kwargs,
+    ):
+        # Get current span context
+        current_span = trace.get_current_span()
+        current_context = trace.set_span_in_context(current_span) if current_span else Context()
+
+        with self.tracer.start_as_current_span(
+            "transport.ack",
+            context=current_context,
+        ) as span:
+            span.set_attribute("service_name", self.service_name)
+            span.set_attribute("message", json.dumps(message))
+            if subscription_id:
+                span.set_attribute("subscription_id", subscription_id)
+
+            call_next(message, subscription_id=subscription_id, **kwargs)
+
+    def nack(
+        self,
+        call_next: Callable,
+        message,
+        subscription_id: int | None = None,
+        **kwargs,
+    ):
+        # Get current span context
+        current_span = trace.get_current_span()
+        current_context = trace.set_span_in_context(current_span) if current_span else Context()
+
+        with self.tracer.start_as_current_span(
+            "transport.nack",
+            context=current_context,
+        ) as span:
+            span.set_attribute("service_name", self.service_name)
+
+            span.set_attribute("message", json.dumps(message))
+            if subscription_id:
+                span.set_attribute("subscription_id", subscription_id)
+
+            call_next(message, subscription_id=subscription_id, **kwargs)
+
+    def transaction_begin(
+        self, call_next: Callable, subscription_id: int | None = None, **kwargs
+    ) -> int:
+        """Start a new transaction span"""
+        # Get current span context (may be None if this is the root span)
+        current_span = trace.get_current_span()
+        current_context = trace.set_span_in_context(current_span) if current_span else Context()
+
+        with self.tracer.start_as_current_span(
+            "transaction.begin",
+            context=current_context,
+        ) as span:
+            span.set_attribute("service_name", self.service_name)
+
+            if subscription_id:
+                span.set_attribute("subscription_id", subscription_id)
+
+            return call_next(subscription_id=subscription_id, **kwargs)
+
+    def transaction_abort(self, call_next: Callable, transaction_id: int | None = None, **kwargs):
+        """Abort a transaction span"""
+        # Get current span context
+        current_span = trace.get_current_span()
+        current_context = trace.set_span_in_context(current_span) if current_span else Context()
+
+        with self.tracer.start_as_current_span(
+            "transaction.abort",
+            context=current_context,
+        ) as span:
+            span.set_attribute("service_name", self.service_name)
+
+            if transaction_id:
+                span.set_attribute("transaction_id", transaction_id)
+
+            call_next(transaction_id=transaction_id, **kwargs)
+
+    def transaction_commit(self, call_next: Callable, transaction_id: int | None = None, **kwargs):
+        """Commit a transaction span"""
+        # Get current span context
+        current_span = trace.get_current_span()
+        current_context = trace.set_span_in_context(current_span) if current_span else Context()
+
+        with self.tracer.start_as_current_span(
+            "transaction.commit",
+            context=current_context,
+        ) as span:
+            span.set_attribute("service_name", self.service_name)
+            if transaction_id:
+                span.set_attribute("transaction_id", transaction_id)
+
+            call_next(transaction_id=transaction_id, **kwargs)
+