Refactoring client and worker to pass around interceptors

Signed-off-by: Ryan Lettieri <[email protected]>

Author: RyanLettieri

durabletask/client.py

@@ -15,6 +15,8 @@
 import durabletask.internal.orchestrator_service_pb2 as pb
 import durabletask.internal.orchestrator_service_pb2_grpc as stubs
 import durabletask.internal.shared as shared
+from durabletask.internal.grpc_interceptor import DefaultClientInterceptorImpl
+
 from durabletask import task
 
 TInput = TypeVar('TInput')
@@ -96,8 +98,23 @@ def __init__(self, *,
                  metadata: Optional[list[tuple[str, str]]] = None,
                  log_handler: Optional[logging.Handler] = None,
                  log_formatter: Optional[logging.Formatter] = None,
-                 secure_channel: bool = False):
-        channel = shared.get_grpc_channel(host_address, metadata, secure_channel=secure_channel)
+                 secure_channel: bool = False,
+                 interceptors: Optional[list] = None):
+
+        # Determine the interceptors to use
+        if interceptors is not None:
+            self._interceptors = interceptors
+        elif metadata:
+            self._interceptors = [DefaultClientInterceptorImpl(metadata)]
+        else:
+            self._interceptors = None
+
+        channel = shared.get_grpc_channel(
+            host_address=host_address,
+            metadata=metadata,
+            secure_channel=secure_channel,
+            interceptors=self._interceptors
+        )
         self._stub = stubs.TaskHubSidecarServiceStub(channel)
         self._logger = shared.get_logger("client", log_handler, log_formatter)
 

durabletask/internal/shared.py

@@ -9,9 +9,6 @@
 
 import grpc
 
-from durabletask.internal.grpc_interceptor import DefaultClientInterceptorImpl
-from externalpackages.durabletaskscheduler.durabletask_grpc_interceptor import DTSDefaultClientInterceptorImpl
-
 # Field name used to indicate that an object was automatically serialized
 # and should be deserialized as a SimpleNamespace
 AUTO_SERIALIZED = "__durabletask_autoobject__"
@@ -26,8 +23,10 @@ def get_default_host_address() -> str:
 
 def get_grpc_channel(
         host_address: Optional[str],
-        metadata: Optional[list[tuple[str, str]]],
-        secure_channel: bool = False) -> grpc.Channel:
+        metadata: Optional[list[tuple[str, str]]] = None,
+        secure_channel: bool = False,
+        interceptors: Optional[list] = None) -> grpc.Channel:
+    
     if host_address is None:
         host_address = get_default_host_address()
 
@@ -45,19 +44,14 @@ def get_grpc_channel(
             host_address = host_address[len(protocol):]
             break
 
+    # Create the base channel
     if secure_channel:
         channel = grpc.secure_channel(host_address, grpc.ssl_channel_credentials())
     else:
         channel = grpc.insecure_channel(host_address)
 
-    if metadata is not None and len(metadata) > 0:
-        for key, _ in metadata:
-            # Check if we are using DTS as the backend and if so, construct the DTS specific interceptors
-            if key == "dts":
-                interceptors = [DTSDefaultClientInterceptorImpl(metadata)]
-                break
-            else:
-                interceptors = [DefaultClientInterceptorImpl(metadata)]
+    # Apply interceptors ONLY if they exist
+    if interceptors:
         channel = grpc.intercept_channel(channel, *interceptors)
     return channel
 

durabletask/worker.py

@@ -16,7 +16,9 @@
 import durabletask.internal.orchestrator_service_pb2 as pb
 import durabletask.internal.orchestrator_service_pb2_grpc as stubs
 import durabletask.internal.shared as shared
+
 from durabletask import task
+from durabletask.internal.grpc_interceptor import DefaultClientInterceptorImpl
 
 TInput = TypeVar('TInput')
 TOutput = TypeVar('TOutput')
@@ -88,15 +90,25 @@ def __init__(self, *,
                  metadata: Optional[list[tuple[str, str]]] = None,
                  log_handler=None,
                  log_formatter: Optional[logging.Formatter] = None,
-                 secure_channel: bool = False):
+                 secure_channel: bool = False,
+                 interceptors: Optional[list[grpc.ServerInterceptor]] = None):  # Add interceptors
         self._registry = _Registry()
         self._host_address = host_address if host_address else shared.get_default_host_address()
-        self._metadata = metadata
+        self._metadata = metadata or []  # Ensure metadata is never None
         self._logger = shared.get_logger("worker", log_handler, log_formatter)
         self._shutdown = Event()
         self._is_running = False
         self._secure_channel = secure_channel
 
+        # Determine the interceptors to use
+        if interceptors is not None:
+            self._interceptors = interceptors
+        elif self._metadata:
+            self._interceptors = [DefaultClientInterceptorImpl(self._metadata)]
+        else:
+            self._interceptors = None
+
+
     def __enter__(self):
         return self
 
@@ -117,7 +129,12 @@ def add_activity(self, fn: task.Activity) -> str:
 
     def start(self):
         """Starts the worker on a background thread and begins listening for work items."""
-        channel = shared.get_grpc_channel(self._host_address, self._metadata, self._secure_channel)
+        
+        if self._metadata:
+            interceptors = [DefaultClientInterceptorImpl(self._metadata)]  
+        else:
+            interceptors =  None
+        channel = shared.get_grpc_channel(self._host_address, self._metadata, self._secure_channel, interceptors)
         stub = stubs.TaskHubSidecarServiceStub(channel)
 
         if self._is_running:

externalpackages/durabletaskscheduler/durabletask_scheduler_client.py

@@ -4,28 +4,40 @@
 from typing import Optional
 from durabletask.client import TaskHubGrpcClient
 from externalpackages.durabletaskscheduler.access_token_manager import AccessTokenManager
+from externalpackages.durabletaskscheduler.durabletask_grpc_interceptor import DTSDefaultClientInterceptorImpl
 
+# Client class used for Durable Task Scheduler (DTS)
 class DurableTaskSchedulerClient(TaskHubGrpcClient):
     def __init__(self,
                  host_address: str,
                  secure_channel: bool,
-                 metadata: Optional[list[tuple[str, str]]] = None,
+                 metadata: Optional[list[tuple[str, str]]] = [],
                  use_managed_identity: Optional[bool] = False,
                  client_id: Optional[str] = None,
                  taskhub: str = None,
                  **kwargs):
-        if metadata is None:
-            metadata = []  # Ensure metadata is initialized
-        self._metadata = metadata
-        self._use_managed_identity = use_managed_identity
-        self._client_id = client_id
-        self._metadata.append(("taskhub", taskhub))
+
+        # Ensure metadata is a list
+        metadata = metadata or []
+        self._metadata = metadata.copy()  # Use a copy to avoid modifying original
+
+        # Append DurableTask-specific metadata
+        self._metadata.append(("taskhub", taskhub or "default-taskhub"))
         self._metadata.append(("dts", "True"))
         self._metadata.append(("use_managed_identity", str(use_managed_identity)))
-        self._metadata.append(("client_id", str(client_id)))
+        self._metadata.append(("client_id", str(client_id or "None")))
+
         self._access_token_manager = AccessTokenManager(metadata=self._metadata)
         self.__update_metadata_with_token()
-        super().__init__(host_address=host_address, secure_channel=secure_channel, metadata=self._metadata, **kwargs)
+        interceptors = [DTSDefaultClientInterceptorImpl(self._metadata)]
+
+        super().__init__(
+            host_address=host_address,
+            secure_channel=secure_channel,
+            metadata=self._metadata,
+            interceptors=interceptors,  # Now explicitly passing interceptors
+            **kwargs
+        )
 
     def __update_metadata_with_token(self):
         """

externalpackages/durabletaskscheduler/durabletask_scheduler_worker.py

@@ -4,6 +4,7 @@
 from typing import Optional
 from durabletask.worker import TaskHubGrpcWorker
 from externalpackages.durabletaskscheduler.access_token_manager import AccessTokenManager
+from externalpackages.durabletaskscheduler.durabletask_grpc_interceptor import DTSDefaultClientInterceptorImpl
 
 # Worker class used for Durable Task Scheduler (DTS)
 class DurableTaskSchedulerWorker(TaskHubGrpcWorker):
@@ -15,18 +16,28 @@ def __init__(self,
                  client_id: Optional[str] = None,
                  taskhub: str = None,
                  **kwargs):
-        if metadata is None:
-            metadata = []  # Ensure metadata is initialized
-        self._metadata = metadata
-        self._use_managed_identity = use_managed_identity
-        self._client_id = client_id
-        self._metadata.append(("taskhub", taskhub))
+        
+        # Ensure metadata is a list
+        metadata = metadata or []
+        self._metadata = metadata.copy()  # Copy to prevent modifying input
+
+        # Append DurableTask-specific metadata
+        self._metadata.append(("taskhub", taskhub or "default-taskhub"))
         self._metadata.append(("dts", "True"))
         self._metadata.append(("use_managed_identity", str(use_managed_identity)))
-        self._metadata.append(("client_id", str(client_id)))
+        self._metadata.append(("client_id", str(client_id or "None")))
+
         self._access_token_manager = AccessTokenManager(metadata=self._metadata)
         self.__update_metadata_with_token()
-        super().__init__(host_address=host_address, secure_channel=secure_channel, metadata=self._metadata, **kwargs)
+        interceptors = [DTSDefaultClientInterceptorImpl(self._metadata)]
+
+        super().__init__(
+            host_address=host_address,
+            secure_channel=secure_channel,
+            metadata=self._metadata,
+            interceptors=interceptors,
+            **kwargs
+        )
 
     def __update_metadata_with_token(self):
         """