Make parse_grpc_uri() compatible with grpcio

Now `parse_grpc_uri()` takes the channel type as an extra argument and will create the appropriate type of channel based on this. Signed-off-by: Leandro Lucarella <[email protected]>
frequenz-floss · May 30, 2024 · cf57068 · cf57068
1 parent e36ffdb
commit cf57068
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 8 deletions.
diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md
@@ -6,11 +6,11 @@
 
 ## Upgrading
 
-<!-- Here goes notes on how to upgrade from previous versions, including deprecations and what they should be replaced with -->
+- `channel.parse_grpc_uri()` takes an extra argument, the channel type (which can be either `grpclib.client.Channel` or `grpcio.aio.Channel`).
 
 ## New Features
 
-<!-- Here goes the main new features and examples or instructions on how to use them -->
+- `channel.parse_grpc_uri()` can now be used with `grpcio` too.
 
 ## Bug Fixes
 

diff --git a/src/frequenz/client/base/channel.py b/src/frequenz/client/base/channel.py
@@ -3,10 +3,13 @@
 
 """Handling of gRPC channels."""
 
+from typing import TypeVar, assert_never
 from urllib.parse import parse_qs, urlparse
 
 from grpclib.client import Channel
 
+from . import _grpchacks
+
 
 def _to_bool(value: str) -> bool:
     value = value.lower()
@@ -17,7 +20,15 @@ def _to_bool(value: str) -> bool:
     raise ValueError(f"Invalid boolean value '{value}'")
 
 
-def parse_grpc_uri(uri: str, /, *, default_port: int = 9090) -> Channel:
+ChannelT = TypeVar(
+    "ChannelT", bound=_grpchacks.GrpclibChannel | _grpchacks.GrpcioChannel
+)
+"""A `grpclib` or `grpcio` channel type."""
+
+
+def parse_grpc_uri(
+    uri: str, channel_type: type[ChannelT], /, *, default_port: int = 9090
+) -> ChannelT:
     """Create a grpclib client channel from a URI.
 
     The URI must have the following format:
@@ -68,8 +79,19 @@ def parse_grpc_uri(uri: str, /, *, default_port: int = 9090) -> Channel:
             uri,
         )
 
-    return Channel(
-        host=parsed_uri.hostname,
-        port=parsed_uri.port or default_port,
-        ssl=ssl,
-    )
+    host = parsed_uri.hostname
+    port = parsed_uri.port or default_port
+    match channel_type:
+        case _grpchacks.GrpcioChannel():
+            target = f"{host}:{port}"
+            return (
+                _grpchacks.grpcio_secure_channel(
+                    target, _grpchacks.grpcio_ssl_channel_credentials()
+                )
+                if ssl
+                else _grpchacks.grpcio_insecure_channel(target)
+            )
+        case _grpchacks.GrpclibChannel():
+            return _grpchacks.GrpclibChannel(host=host, port=port, ssl=ssl)
+        case _:
+            assert False, "Unexpected channel type: {channel_type}"