Add functionality to pass gRPC channel options when creating channel

pyzeebe/grpc_internals/channel_options.py

@@ -1,15 +1,37 @@
+""" gRPC Channel Options
+
+``grpc.keepalive_time_ms``
+--------------------------
+
+Time between keepalive pings. Following the official Zeebe Java/Go client, sending pings every 45 seconds.
+
+https://docs.camunda.io/docs/product-manuals/zeebe/deployment-guide/operations/setting-up-a-cluster/#keep-alive-intervals
+"""
+from typing import Dict, Any, Tuple
+
 GRPC_CHANNEL_OPTIONS = {
-    # https://docs.camunda.io/docs/product-manuals/zeebe/deployment-guide/operations/setting-up-a-cluster/#keep-alive-intervals
-    # "By default, the official Zeebe clients (Java and Go) send keep alive pings every 45 seconds."
     "grpc.keepalive_time_ms": 45_000
 }
 
 
-def get_channel_options():
+def get_channel_options(options: Dict[str, Any] = None) -> Tuple[Tuple[str, Any], ...]:
     """
-    Channel arguments are expected as a tuple of tuples:
-    https://grpc.github.io/grpc/python/glossary.html#term-channel_arguments
+    Convert options dict to tuple in expected format for creating the gRPC channel
+
+    Args:
+        options (Dict[str, Any]): A key/value representation of `gRPC channel arguments_`.
+            Default: None (will use library defaults)
+
+    Returns:
+        Tuple[Tuple[str, Any], ...]: Options for the gRPC channel
+
+    .. _gRPC channel arguments:
+        https://grpc.github.io/grpc/python/glossary.html#term-channel_arguments
     """
+    if options:
+        options = {**GRPC_CHANNEL_OPTIONS, **options}
+    else:
+        options = GRPC_CHANNEL_OPTIONS
     return tuple(
-        (k, v) for k, v in GRPC_CHANNEL_OPTIONS.items()
+        (k, v) for k, v in options.items()
     )

pyzeebe/grpc_internals/grpc_channel_utils.py

@@ -1,5 +1,5 @@
 import os
-from typing import Optional
+from typing import Optional, Dict, Any
 
 import grpc
 
@@ -21,11 +21,27 @@ def create_connection_uri(
 def create_channel(
     connection_uri: str,
     credentials: Optional[BaseCredentials] = None,
-    secure_connection: bool = False
+    secure_connection: bool = False,
+    options: Dict[str, Any] = None
 ) -> grpc.aio.Channel:
-    options = get_channel_options()
+    """
+    Create a gRPC channel.
+
+    Args:
+        connection_uri (str): The URI to Zeebe.
+        credentials (BaseCredentials): Credentials for accessing the channel. Default: None
+        secure_connection (bool): Create a secure channel (set to True when using credentials). Default: False
+        options (Dict[str, Any]): A key/value representation of `gRPC channel arguments_`. Default: None (will use library defaults)
+
+    Returns:
+        grpc.aio.Channel: A channel object set up to connect to Zeebe
+
+    .. _gRPC channel arguments:
+        https://grpc.github.io/grpc/python/glossary.html#term-channel_arguments
+    """
+    channel_options = get_channel_options(options)
     if credentials:
-        return grpc.aio.secure_channel(connection_uri, credentials.grpc_credentials, options=options)
+        return grpc.aio.secure_channel(connection_uri, credentials.grpc_credentials, options=channel_options)
     if secure_connection:
-        return grpc.aio.secure_channel(connection_uri, grpc.ssl_channel_credentials(), options=options)
-    return grpc.aio.insecure_channel(connection_uri, options=options)
+        return grpc.aio.secure_channel(connection_uri, grpc.ssl_channel_credentials(), options=channel_options)
+    return grpc.aio.insecure_channel(connection_uri, options=channel_options)

tests/unit/grpc_internals/grpc_channel_utils_test.py

@@ -26,6 +26,29 @@ def test_creates_secure_channel_when_secure_connection_is_enabled(self):
 
             channel_mock.assert_called_once()
 
+    def test_creates_secure_channel_with_default_options(self):
+        with patch("grpc.aio.insecure_channel") as channel_mock:
+            create_channel(str(uuid4()), options=None)
+
+            expected_options = (
+                ("grpc.keepalive_time_ms", 45000),
+            )
+            assert channel_mock.call_args[1]["options"] == expected_options
+
+    def test_creates_insecure_channel_with_options(self):
+        with patch("grpc.aio.insecure_channel") as channel_mock:
+            additional_options = {
+                "grpc.keepalive_timeout_ms": 120000,
+                "grpc.http2.min_time_between_pings_ms": 60000,
+            }
+            create_channel(str(uuid4()), options=additional_options)
+
+            expected_options = (
+                ("grpc.keepalive_time_ms", 45000),
+                ("grpc.keepalive_timeout_ms", 120000),
+                ("grpc.http2.min_time_between_pings_ms", 60000),
+            )
+            assert channel_mock.call_args[1]["options"] == expected_options
 
 class TestCreateConnectionUri:
     def test_uses_credentials_first(self):