From 5e4aa6ac2892c0b6e22a389ee564f4e6ed92d0f7 Mon Sep 17 00:00:00 2001 From: ivanauth Date: Fri, 15 May 2026 13:26:45 -0400 Subject: [PATCH] feat: expose close() on Authzed clients to release gRPC channel Adds a public close() method to the v1 and materialize/v0 Client, SyncClient, and AsyncClient classes. SyncClient.close() shuts the underlying grpc.Channel synchronously; AsyncClient.close(grace=None) is a coroutine that awaits grpc.aio.Channel.close(grace). The auto-detecting Client returns whichever the active channel produces, so callers can use the same shape in both modes. Lets applications deterministically tear down connections during graceful shutdown, credential rotation, or recovery from stale connections, rather than relying on garbage collection. --- src/authzed/api/materialize/v0/__init__.py | 30 +++++++++++ src/authzed/api/materialize/v0/__init__.pyi | 5 +- src/authzed/api/v1/__init__.py | 30 +++++++++++ src/authzed/api/v1/__init__.pyi | 5 +- tests/close_test.py | 59 +++++++++++++++++++++ 5 files changed, 127 insertions(+), 2 deletions(-) create mode 100644 tests/close_test.py diff --git a/src/authzed/api/materialize/v0/__init__.py b/src/authzed/api/materialize/v0/__init__.py index 526aea2..a6c7527 100644 --- a/src/authzed/api/materialize/v0/__init__.py +++ b/src/authzed/api/materialize/v0/__init__.py @@ -44,6 +44,7 @@ def __init__(self, target, credentials, options=None, compression=None): self.init_stubs(channel) def init_stubs(self, channel): + self._channel = channel WatchPermissionsServiceStub.__init__(self, channel) WatchPermissionSetsServiceStub.__init__(self, channel) @@ -56,6 +57,21 @@ def create_channel(self, target, credentials, options=None, compression=None): return channelfn(target, credentials, options, compression) + def close(self): + """ + Close the underlying gRPC channel. + + For async channels (``grpc.aio.Channel``), this returns a coroutine that + must be awaited; the caller is expected to ``await client.close()``. + For sync channels (``grpc.Channel``), the channel is closed + synchronously and ``None`` is returned. + + Closing the channel cancels in-flight RPCs and prevents new RPCs from + being issued through this client. Calling ``close`` more than once is + safe. + """ + return self._channel.close() + class AsyncClient(Client): """ @@ -66,6 +82,16 @@ def __init__(self, target, credentials, options=None, compression=None): channel = grpc.aio.secure_channel(target, credentials, options, compression) self.init_stubs(channel) + async def close(self, grace=None): + """ + Close the underlying async gRPC channel. + + ``grace`` is forwarded to ``grpc.aio.Channel.close``; when set, the + channel waits up to ``grace`` seconds for pending RPCs to finish + before cancelling them. + """ + await self._channel.close(grace) + class SyncClient(Client): """ @@ -76,6 +102,10 @@ def __init__(self, target, credentials, options=None, compression=None): channel = grpc.secure_channel(target, credentials, options, compression) self.init_stubs(channel) + def close(self): + """Close the underlying sync gRPC channel.""" + self._channel.close() + class TokenAuthorization(ClientInterceptor): def __init__(self, token: str): diff --git a/src/authzed/api/materialize/v0/__init__.pyi b/src/authzed/api/materialize/v0/__init__.pyi index 0250839..6ad164b 100644 --- a/src/authzed/api/materialize/v0/__init__.pyi +++ b/src/authzed/api/materialize/v0/__init__.pyi @@ -1,4 +1,4 @@ -from typing import Any, Optional, Sequence, Tuple +from typing import Any, Awaitable, Optional, Sequence, Tuple, Union import grpc @@ -42,6 +42,7 @@ class Client(WatchPermissionsServiceStub, WatchPermissionSetsServiceStub): options: Optional[Sequence[Tuple[str, Any]]] = None, compression: Optional[grpc.Compression] = None, ) -> None: ... + def close(self) -> Union[None, Awaitable[None]]: ... class SyncClient(WatchPermissionsServiceStub, WatchPermissionSetsServiceStub): def __init__( @@ -51,6 +52,7 @@ class SyncClient(WatchPermissionsServiceStub, WatchPermissionSetsServiceStub): options: Optional[Sequence[Tuple[str, Any]]] = None, compression: Optional[grpc.Compression] = None, ) -> None: ... + def close(self) -> None: ... class AsyncClient(WatchPermissionsServiceStub, WatchPermissionSetsServiceStub): def __init__( @@ -60,6 +62,7 @@ class AsyncClient(WatchPermissionsServiceStub, WatchPermissionSetsServiceStub): options: Optional[Sequence[Tuple[str, Any]]] = None, compression: Optional[grpc.Compression] = None, ) -> None: ... + async def close(self, grace: Optional[float] = None) -> None: ... class InsecureClient(Client): def __init__( diff --git a/src/authzed/api/v1/__init__.py b/src/authzed/api/v1/__init__.py index 6d13570..7652ab5 100644 --- a/src/authzed/api/v1/__init__.py +++ b/src/authzed/api/v1/__init__.py @@ -96,6 +96,7 @@ def __init__(self, target, credentials, options=None, compression=None): self.init_stubs(channel) def init_stubs(self, channel): + self._channel = channel SchemaServiceStub.__init__(self, channel) PermissionsServiceStub.__init__(self, channel) ExperimentalServiceStub.__init__(self, channel) @@ -110,6 +111,21 @@ def create_channel(self, target, credentials, options=None, compression=None): return channelfn(target, credentials, options, compression) + def close(self): + """ + Close the underlying gRPC channel. + + For async channels (``grpc.aio.Channel``), this returns a coroutine that + must be awaited; the caller is expected to ``await client.close()``. + For sync channels (``grpc.Channel``), the channel is closed + synchronously and ``None`` is returned. + + Closing the channel cancels in-flight RPCs and prevents new RPCs from + being issued through this client. Calling ``close`` more than once is + safe. + """ + return self._channel.close() + class AsyncClient(Client): """ @@ -120,6 +136,16 @@ def __init__(self, target, credentials, options=None, compression=None): channel = grpc.aio.secure_channel(target, credentials, options, compression) self.init_stubs(channel) + async def close(self, grace=None): + """ + Close the underlying async gRPC channel. + + ``grace`` is forwarded to ``grpc.aio.Channel.close``; when set, the + channel waits up to ``grace`` seconds for pending RPCs to finish + before cancelling them. + """ + await self._channel.close(grace) + class SyncClient(Client): """ @@ -130,6 +156,10 @@ def __init__(self, target, credentials, options=None, compression=None): channel = grpc.secure_channel(target, credentials, options, compression) self.init_stubs(channel) + def close(self): + """Close the underlying sync gRPC channel.""" + self._channel.close() + class TokenAuthorization(ClientInterceptor): def __init__(self, token: str): diff --git a/src/authzed/api/v1/__init__.pyi b/src/authzed/api/v1/__init__.pyi index c2214b3..f8cdbd3 100644 --- a/src/authzed/api/v1/__init__.pyi +++ b/src/authzed/api/v1/__init__.pyi @@ -1,4 +1,4 @@ -from typing import Any, Optional, Sequence, Tuple +from typing import Any, Awaitable, Optional, Sequence, Tuple, Union import grpc @@ -81,6 +81,7 @@ class Client(SchemaServiceStub, PermissionsServiceStub, ExperimentalServiceStub, options: Optional[Sequence[Tuple[str, Any]]] = None, compression: Optional[grpc.Compression] = None, ) -> None: ... + def close(self) -> Union[None, Awaitable[None]]: ... class SyncClient( SchemaServiceStub, PermissionsServiceStub, ExperimentalServiceStub, WatchServiceStub @@ -92,6 +93,7 @@ class SyncClient( options: Optional[Sequence[Tuple[str, Any]]] = None, compression: Optional[grpc.Compression] = None, ) -> None: ... + def close(self) -> None: ... class AsyncClient( SchemaServiceAsyncStub, @@ -106,6 +108,7 @@ class AsyncClient( options: Optional[Sequence[Tuple[str, Any]]] = None, compression: Optional[grpc.Compression] = None, ) -> None: ... + async def close(self, grace: Optional[float] = None) -> None: ... class InsecureClient(Client): def __init__( diff --git a/tests/close_test.py b/tests/close_test.py new file mode 100644 index 0000000..61c6b17 --- /dev/null +++ b/tests/close_test.py @@ -0,0 +1,59 @@ +import asyncio + +import grpc +import grpc.aio +import pytest + +from authzed.api.v1 import AsyncClient, Client, InsecureClient, SyncClient +from grpcutil import insecure_bearer_token_credentials + + +def test_sync_client_close_closes_channel(token): + client = SyncClient("localhost:50051", insecure_bearer_token_credentials(token)) + assert isinstance(client._channel, grpc.Channel) + + # close should be idempotent and not raise + client.close() + client.close() + + +def test_insecure_client_close_closes_channel(token): + client = InsecureClient("localhost:50051", token) + # InsecureClient uses an intercepted channel, but it still exposes close() + client.close() + client.close() + + +async def test_async_client_close_closes_channel(token): + client = AsyncClient("localhost:50051", insecure_bearer_token_credentials(token)) + assert isinstance(client._channel, grpc.aio.Channel) + + # close should be awaitable and idempotent + await client.close() + await client.close() + + +async def test_async_client_close_accepts_grace(token): + client = AsyncClient("localhost:50051", insecure_bearer_token_credentials(token)) + await client.close(grace=0) + + +def test_autodetect_client_close_when_sync(token): + # Outside of an event loop, Client builds a sync channel; close() should + # return None (not a coroutine). + with pytest.raises(RuntimeError): + asyncio.get_running_loop() + client = Client("localhost:50051", insecure_bearer_token_credentials(token)) + assert isinstance(client._channel, grpc.Channel) + result = client.close() + assert result is None + + +async def test_autodetect_client_close_when_async(token): + # Inside an event loop, Client builds an async channel; close() returns a + # coroutine that must be awaited. + client = Client("localhost:50051", insecure_bearer_token_credentials(token)) + assert isinstance(client._channel, grpc.aio.Channel) + awaitable = client.close() + assert asyncio.iscoroutine(awaitable) + await awaitable