from abc import ABCMeta
from collections.abc import Awaitable, Callable
from typing import TypeVar
import grpc
from grpc.aio import ClientCallDetails, Metadata, UnaryUnaryCall, UnaryUnaryClientInterceptor
from .configuration import ClientConfiguration, SecretAccessToken
_RequestType = TypeVar("_RequestType")
_ResponseType = TypeVar("_ResponseType")
class RefreshInterceptor(UnaryUnaryClientInterceptor, metaclass=ABCMeta):
def __init__(self, client_configuration: ClientConfiguration | None = None):
if client_configuration is None:
client_configuration = ClientConfiguration.load_default()
self.configuration = client_configuration
async def _get_access_token(self) -> SecretAccessToken:
return await self.configuration.get_bearer_access_token_async()
async def intercept_unary_unary(
self,
continuation: Callable[
[ClientCallDetails, _RequestType], Awaitable[UnaryUnaryCall[_RequestType, _ResponseType]]
],
client_call_details: ClientCallDetails,
request: _RequestType,
) -> UnaryUnaryCall[_RequestType, _ResponseType]:
authorized_metadata = Metadata(*(client_call_details.metadata or ()))
authorized_metadata["authorization"] = f"Bearer {await self._get_access_token()}"
new_client_call_details = grpc.aio.ClientCallDetails(
client_call_details.method,
client_call_details.timeout,
authorized_metadata,
client_call_details.credentials,
client_call_details.wait_for_ready,
)
return await continuation(new_client_call_details, request)