Add gRPC metadata option (#16)

Signed-off-by: Deepanshu Agarwal <deepanshu.agarwal1984@gmail.com>
Co-authored-by: Chris Gillum <cgillum@gmail.com>

Author: DeepanshuA

CHANGELOG.md

@@ -1,4 +1,21 @@
-## v0.1.0
+## Unreleased
+
+### New
+
+- Add gRPC metadata option ([#16](https://github.com/microsoft/durabletask-python/pull/16)) - contributed by [@DeepanshuA](https://github.com/DeepanshuA)
+
+### Changes
+
+- Removed Python 3.7 support due to EOL ([#14](https://github.com/microsoft/durabletask-python/pull/14)) - contributed by [@berndverst](https://github.com/berndverst)
+
+## v0.1.0b
+
+### New
+
+- Continue-as-new ([#9](https://github.com/microsoft/durabletask-python/pull/9))
+- Support for Python 3.7+ ([#10](https://github.com/microsoft/durabletask-python/pull/10)) - contributed by [@DeepanshuA](https://github.com/DeepanshuA)
+
+## v0.1.0a
 
 Initial release, which includes the following features:
 

durabletask/client.py

@@ -6,7 +6,7 @@
 from dataclasses import dataclass
 from datetime import datetime
 from enum import Enum
-from typing import Any, TypeVar, Union
+from typing import Any, List, Tuple, TypeVar, Union
 
 import grpc
 from google.protobuf import wrappers_pb2
@@ -93,9 +93,10 @@ class TaskHubGrpcClient:
 
     def __init__(self, *,
                  host_address: Union[str, None] = None,
-                 log_handler=None,
+                 metadata: Union[List[Tuple[str, str]], None] = None,
+                 log_handler = None,
                  log_formatter: Union[logging.Formatter, None] = None):
-        channel = shared.get_grpc_channel(host_address)
+        channel = shared.get_grpc_channel(host_address, metadata)
         self._stub = stubs.TaskHubSidecarServiceStub(channel)
         self._logger = shared.get_logger("client", log_handler, log_formatter)
 

durabletask/internal/grpc_interceptor.py

@@ -0,0 +1,66 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from collections import namedtuple
+from typing import List, Tuple
+
+import grpc
+
+
+class _ClientCallDetails(
+        namedtuple(
+            '_ClientCallDetails',
+            ['method', 'timeout', 'metadata', 'credentials', 'wait_for_ready', 'compression']),
+        grpc.ClientCallDetails):
+    """This is an implementation of the ClientCallDetails interface needed for interceptors.
+    This class takes six named values and inherits the ClientCallDetails from grpc package.
+    This class encloses the values that describe a RPC to be invoked.
+    """
+    pass
+
+
+class DefaultClientInterceptorImpl (
+    grpc.UnaryUnaryClientInterceptor, grpc.UnaryStreamClientInterceptor,
+    grpc.StreamUnaryClientInterceptor, grpc.StreamStreamClientInterceptor):
+    """The class implements a UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor,
+    StreamUnaryClientInterceptor and StreamStreamClientInterceptor from grpc to add an 
+    interceptor to add additional headers to all calls as needed."""
+
+    def __init__(self, metadata: List[Tuple[str, str]]):
+        super().__init__()
+        self._metadata = metadata
+
+    def _intercept_call(
+                self, client_call_details: _ClientCallDetails) -> grpc.ClientCallDetails:
+        """Internal intercept_call implementation which adds metadata to grpc metadata in the RPC
+            call details."""
+        if self._metadata is None:
+            return client_call_details
+        
+        if client_call_details.metadata is not None:
+            metadata = list(client_call_details.metadata)
+        else:
+            metadata = []
+        
+        metadata.extend(self._metadata)
+        client_call_details = _ClientCallDetails(
+            client_call_details.method, client_call_details.timeout, metadata,
+            client_call_details.credentials, client_call_details.wait_for_ready, client_call_details.compression)
+
+        return client_call_details
+
+    def intercept_unary_unary(self, continuation, client_call_details, request):
+        new_client_call_details = self._intercept_call(client_call_details)
+        return continuation(new_client_call_details, request)
+
+    def intercept_unary_stream(self, continuation, client_call_details, request):
+        new_client_call_details = self._intercept_call(client_call_details)
+        return continuation(new_client_call_details, request)
+
+    def intercept_stream_unary(self, continuation, client_call_details, request):
+        new_client_call_details = self._intercept_call(client_call_details)
+        return continuation(new_client_call_details, request)
+
+    def intercept_stream_stream(self, continuation, client_call_details, request):
+        new_client_call_details = self._intercept_call(client_call_details)
+        return continuation(new_client_call_details, request)

durabletask/internal/shared.py

@@ -5,10 +5,12 @@
 import json
 import logging
 from types import SimpleNamespace
-from typing import Any, Dict, Union
+from typing import Any, Dict, List, Tuple, Union
 
 import grpc
 
+from durabletask.internal.grpc_interceptor import DefaultClientInterceptorImpl
+
 # Field name used to indicate that an object was automatically serialized
 # and should be deserialized as a SimpleNamespace
 AUTO_SERIALIZED = "__durabletask_autoobject__"
@@ -18,13 +20,15 @@ def get_default_host_address() -> str:
     return "localhost:4001"
 
 
-def get_grpc_channel(host_address: Union[str, None]) -> grpc.Channel:
+def get_grpc_channel(host_address: Union[str, None], metadata: Union[List[Tuple[str, str]], None]) -> grpc.Channel:
     if host_address is None:
         host_address = get_default_host_address()
     channel = grpc.insecure_channel(host_address)
+    if metadata is not None and len(metadata) > 0:
+        interceptors = [DefaultClientInterceptorImpl(metadata)]
+        channel = grpc.intercept_channel(channel, *interceptors)
     return channel
 
-
 def get_logger(
         name_suffix: str,
         log_handler: Union[logging.Handler, None] = None,

durabletask/worker.py

@@ -6,7 +6,7 @@
 from datetime import datetime, timedelta
 from threading import Event, Thread
 from types import GeneratorType
-from typing import Any, Dict, Generator, List, Sequence, TypeVar, Union
+from typing import Any, Dict, Generator, List, Sequence, Tuple, TypeVar, Union
 
 import grpc
 from google.protobuf import empty_pb2
@@ -85,10 +85,12 @@ class TaskHubGrpcWorker:
 
     def __init__(self, *,
                  host_address: Union[str, None] = None,
-                 log_handler=None,
+                 metadata: Union[List[Tuple[str, str]], None] = None,
+                 log_handler = None,
                  log_formatter: Union[logging.Formatter, None] = None):
         self._registry = _Registry()
         self._host_address = host_address if host_address else shared.get_default_host_address()
+        self._metadata = metadata
         self._logger = shared.get_logger("worker", log_handler, log_formatter)
         self._shutdown = Event()
         self._response_stream = None
@@ -114,7 +116,7 @@ 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)
+        channel = shared.get_grpc_channel(self._host_address, self._metadata)
         stub = stubs.TaskHubSidecarServiceStub(channel)
 
         if self._is_running: