feat: add middleware

pyproject.toml

@@ -109,7 +109,13 @@ module = [
   "momento.internal.aio._scs_data_client",
   "momento.internal.aio._scs_grpc_manager",
   "momento.internal.aio._utilities",
+  "momento.internal.synchronous._utilities",
   "momento.responses.control.signing_key.*",
+  "momento.internal.aio._middleware_interceptor",
+  "momento.internal.synchronous._middleware_interceptor",
+  "momento.config.middleware.models",
+  "momento.config.middleware.aio.middleware_metadata",
+  "momento.config.middleware.synchronous.middleware_metadata",
 ]
 disallow_any_expr = false
 

src/momento/config/configuration.py

@@ -3,14 +3,16 @@
 from abc import ABC, abstractmethod
 from datetime import timedelta
 from pathlib import Path
+from typing import List, Optional
 
+import momento.config.middleware.aio
 from momento.retry import RetryStrategy
 
+from .middleware import Middleware
 from .transport.transport_strategy import TransportStrategy
 
 
 class ConfigurationBase(ABC):
-    # TODO: Middlewares
     @abstractmethod
     def get_retry_strategy(self) -> RetryStrategy:
         pass
@@ -35,20 +37,39 @@ def with_client_timeout(self, client_timeout: timedelta) -> Configuration:
     def with_root_certificates_pem(self, root_certificate_path: Path) -> Configuration:
         pass
 
+    @abstractmethod
+    def with_middlewares(self, middlewares: List[Middleware]) -> Configuration:
+        pass
+
+    @abstractmethod
+    def add_middleware(self, middleware: Middleware) -> Configuration:
+        pass
+
+    @abstractmethod
+    def get_middlewares(self) -> List[Middleware]:
+        pass
+
 
 class Configuration(ConfigurationBase):
     """Configuration options for Momento Simple Cache Client."""
 
-    def __init__(self, transport_strategy: TransportStrategy, retry_strategy: RetryStrategy):
+    def __init__(
+        self,
+        transport_strategy: TransportStrategy,
+        retry_strategy: RetryStrategy,
+        middlewares: Optional[List[Middleware]] = None,
+    ):
         """Instantiate a Configuration.
 
         Args:
             transport_strategy (TransportStrategy): Configuration options for networking with
             the Momento service.
             retry_strategy (RetryStrategy): the strategy to use when determining whether to retry a grpc call.
+            middlewares: Middleware that can intercept Momento calls. May be aio or synchronous.
         """
         self._transport_strategy = transport_strategy
         self._retry_strategy = retry_strategy
+        self._middlewares: List[Middleware] = list(middlewares or [])
 
     def get_retry_strategy(self) -> RetryStrategy:
         """Access the retry strategy.
@@ -67,7 +88,7 @@ def with_retry_strategy(self, retry_strategy: RetryStrategy) -> Configuration:
         Returns:
             Configuration: the new Configuration with the specified RetryStrategy.
         """
-        return Configuration(self._transport_strategy, retry_strategy)
+        return Configuration(self._transport_strategy, retry_strategy, self._middlewares)
 
     def get_transport_strategy(self) -> TransportStrategy:
         """Access the transport strategy.
@@ -86,7 +107,7 @@ def with_transport_strategy(self, transport_strategy: TransportStrategy) -> Conf
         Returns:
             Configuration: the new Configuration with the specified TransportStrategy.
         """
-        return Configuration(transport_strategy, self._retry_strategy)
+        return Configuration(transport_strategy, self._retry_strategy, self._middlewares)
 
     def with_client_timeout(self, client_timeout: timedelta) -> Configuration:
         """Copies the Configuration and sets the new client-side timeout in the copy's TransportStrategy.
@@ -97,7 +118,11 @@ def with_client_timeout(self, client_timeout: timedelta) -> Configuration:
         Return:
             Configuration: the new Configuration.
         """
-        return Configuration(self._transport_strategy.with_client_timeout(client_timeout), self._retry_strategy)
+        return Configuration(
+            self._transport_strategy.with_client_timeout(client_timeout),
+            self._retry_strategy,
+            self._middlewares,
+        )
 
     def with_root_certificates_pem(self, root_certificates_pem_path: Path) -> Configuration:
         """Copies the Configuration and sets the new root certificates in the copy's TransportStrategy.
@@ -106,10 +131,60 @@ def with_root_certificates_pem(self, root_certificates_pem_path: Path) -> Config
             root_certificates_pem_path (Path): the new root certificates.
 
         Returns:
-            ConfigurationBase: the new Configuration.
+            Configuration: the new Configuration.
         """
         grpc_configuration = self._transport_strategy.get_grpc_configuration().with_root_certificates_pem(
             root_certificates_pem_path
         )
         transport_strategy = self._transport_strategy.with_grpc_configuration(grpc_configuration)
         return self.with_transport_strategy(transport_strategy)
+
+    def with_middlewares(self, middlewares: List[Middleware]) -> Configuration:
+        """Copies the Configuration and adds the new middlewares to the end of the list.
+
+        Args:
+            middlewares: the middleware list to be appended to the Configuration's existing middleware. These can be
+            aio or synchronous middleware.
+
+        Returns:
+            Configuration: the new Configuration.
+        """
+        new_middlewares = self._middlewares.copy() + middlewares
+        return Configuration(self._transport_strategy, self._retry_strategy, new_middlewares)
+
+    def add_middleware(self, middleware: Middleware) -> Configuration:
+        """Copies the Configuration and adds the new middleware to the end of the list.
+
+        Args:
+            middleware: the middleware to be appended to the Configuration's existing middleware. This can be aio or
+            synchronous middleware.
+
+        Returns:
+            Configuration: the new Configuration.
+        """
+        new_middlewares = self._middlewares.copy() + [middleware]
+        return Configuration(self._transport_strategy, self._retry_strategy, new_middlewares)
+
+    def get_middlewares(self) -> List[Middleware]:
+        """Access the middleware list.
+
+        Returns:
+            the configuration's list of middleware.
+        """
+        return self._middlewares.copy()
+
+    def get_aio_middlewares(self) -> List[momento.config.middleware.aio.Middleware]:
+        """Access the aio middleware from the middleware list.
+
+        Returns:
+            the configuration's list of aio middleware.
+        """
+        return [m for m in self._middlewares if isinstance(m, momento.config.middleware.aio.Middleware)]
+
+    def get_sync_middlewares(self) -> List[momento.config.middleware.synchronous.Middleware]:
+        """Access the synchronous middleware from the middleware list.
+
+        Returns:
+            the configuration's list of synchronous middleware.
+        """
+        return [m for m in self._middlewares if isinstance(m, momento.config.middleware.synchronous.Middleware)]

src/momento/config/middleware/__init__.py

@@ -0,0 +1,18 @@
+from typing import Union
+
+from momento.config.middleware.aio import Middleware as AsyncMiddleware
+from momento.config.middleware.models import (
+    MiddlewareMessage,
+    MiddlewareRequestHandlerContext,
+    MiddlewareStatus,
+)
+from momento.config.middleware.synchronous import Middleware as SyncMiddleware
+
+Middleware = Union[SyncMiddleware, AsyncMiddleware]
+
+__all__ = [
+    "Middleware",
+    "MiddlewareMessage",
+    "MiddlewareStatus",
+    "MiddlewareRequestHandlerContext",
+]

src/momento/config/middleware/aio/__init__.py

@@ -0,0 +1,4 @@
+from momento.config.middleware.aio.middleware import Middleware, MiddlewareRequestHandler
+from momento.config.middleware.aio.middleware_metadata import MiddlewareMetadata
+
+__all__ = ["Middleware", "MiddlewareMetadata", "MiddlewareRequestHandler"]

src/momento/config/middleware/aio/middleware.py

@@ -0,0 +1,36 @@
+import abc
+
+from momento.config.middleware.aio.middleware_metadata import MiddlewareMetadata
+from momento.config.middleware.models import MiddlewareMessage, MiddlewareRequestHandlerContext, MiddlewareStatus
+
+
+class MiddlewareRequestHandler(abc.ABC):
+    @abc.abstractmethod
+    async def on_request_metadata(self, metadata: MiddlewareMetadata) -> MiddlewareMetadata:
+        pass
+
+    @abc.abstractmethod
+    async def on_request_body(self, request: MiddlewareMessage) -> MiddlewareMessage:
+        pass
+
+    @abc.abstractmethod
+    async def on_response_metadata(self, metadata: MiddlewareMetadata) -> MiddlewareMetadata:
+        pass
+
+    @abc.abstractmethod
+    async def on_response_body(self, response: MiddlewareMessage) -> MiddlewareMessage:
+        pass
+
+    @abc.abstractmethod
+    async def on_response_status(self, status: MiddlewareStatus) -> MiddlewareStatus:
+        pass
+
+
+class Middleware(abc.ABC):
+    @abc.abstractmethod
+    async def on_new_request(self, context: MiddlewareRequestHandlerContext) -> MiddlewareRequestHandler:
+        pass
+
+    # noinspection PyMethodMayBeStatic
+    def close(self) -> None:
+        return None

src/momento/config/middleware/aio/middleware_metadata.py

@@ -0,0 +1,14 @@
+from typing import Optional
+
+from grpc.aio import Metadata
+
+
+class MiddlewareMetadata:
+    """Wrapper for gRPC metadata."""
+
+    def __init__(self, metadata: Optional[Metadata]):
+        self.grpc_metadata = metadata
+
+    def get_grpc_metadata(self) -> Optional[Metadata]:
+        """Get the underlying gRPC metadata."""
+        return self.grpc_metadata

src/momento/config/middleware/models.py

@@ -0,0 +1,46 @@
+from typing import Dict
+
+import grpc
+from google.protobuf.message import Message
+
+CONNECTION_ID_KEY = "connectionID"
+
+
+class MiddlewareMessage:
+    """Wrapper for a gRPC protobuf message."""
+
+    def __init__(self, message: Message):
+        self.grpc_message = message
+
+    def get_message_length(self) -> int:
+        """Get the length of the message in bytes."""
+        return len(self.grpc_message.SerializeToString())
+
+    def get_constructor_name(self) -> str:
+        """Get the class name of the message."""
+        return str(self.grpc_message.__class__.__name__)
+
+    def get_message(self) -> Message:
+        """Get the underlying gRPC message."""
+        return self.grpc_message
+
+
+class MiddlewareStatus:
+    """Wrapper for gRPC status."""
+
+    def __init__(self, status: grpc.StatusCode):
+        self.grpc_status = status
+
+    def get_code(self) -> grpc.StatusCode:
+        """Get the status code."""
+        return self.grpc_status
+
+
+class MiddlewareRequestHandlerContext:
+    """Context for middleware request handlers."""
+
+    def __init__(self, context: Dict[str, str]):
+        self.context = context
+
+    def get_context(self) -> Dict[str, str]:
+        return self.context

src/momento/config/middleware/synchronous/__init__.py

@@ -0,0 +1,4 @@
+from momento.config.middleware.synchronous.middleware import Middleware, MiddlewareRequestHandler
+from momento.config.middleware.synchronous.middleware_metadata import MiddlewareMetadata
+
+__all__ = ["Middleware", "MiddlewareMetadata", "MiddlewareRequestHandler"]

src/momento/config/middleware/synchronous/middleware.py

@@ -0,0 +1,36 @@
+import abc
+
+from momento.config.middleware.models import MiddlewareMessage, MiddlewareRequestHandlerContext, MiddlewareStatus
+from momento.config.middleware.synchronous.middleware_metadata import MiddlewareMetadata
+
+
+class MiddlewareRequestHandler(abc.ABC):
+    @abc.abstractmethod
+    def on_request_metadata(self, metadata: MiddlewareMetadata) -> MiddlewareMetadata:
+        pass
+
+    @abc.abstractmethod
+    def on_request_body(self, request: MiddlewareMessage) -> MiddlewareMessage:
+        pass
+
+    @abc.abstractmethod
+    def on_response_metadata(self, metadata: MiddlewareMetadata) -> MiddlewareMetadata:
+        pass
+
+    @abc.abstractmethod
+    def on_response_body(self, response: MiddlewareMessage) -> MiddlewareMessage:
+        pass
+
+    @abc.abstractmethod
+    def on_response_status(self, status: MiddlewareStatus) -> MiddlewareStatus:
+        pass
+
+
+class Middleware(abc.ABC):
+    @abc.abstractmethod
+    def on_new_request(self, context: MiddlewareRequestHandlerContext) -> MiddlewareRequestHandler:
+        pass
+
+    # noinspection PyMethodMayBeStatic
+    def close(self) -> None:
+        return None

src/momento/config/middleware/synchronous/middleware_metadata.py

@@ -0,0 +1,14 @@
+from typing import Optional
+
+from grpc._typing import MetadataType
+
+
+class MiddlewareMetadata:
+    """Wrapper for gRPC metadata."""
+
+    def __init__(self, metadata: Optional[MetadataType]):
+        self.grpc_metadata = metadata
+
+    def get_grpc_metadata(self) -> Optional[MetadataType]:
+        """Get the underlying gRPC metadata."""
+        return self.grpc_metadata