feat: add support for CreateSigningKey, RevokeSigningKey, and ListSig…

…ningKeys APIs

.gitignore

@@ -14,3 +14,4 @@ src/*.egg-info
 
 build
 .idea
+.vscode

setup.cfg

@@ -11,7 +11,7 @@ package_dir =
 packages = find:
 python_requires = >=3.7
 install_requires =
-    momento-wire-types==0.8.1
+    momento-wire-types==0.12.0
     build
     setuptools
     pyjwt[crypto]

src/momento/_utilities/_data_validation.py

@@ -45,6 +45,11 @@ def _validate_ttl(ttl_seconds: int) -> None:
         raise errors.InvalidArgumentError("TTL Seconds must be a non-negative integer")
 
 
+def _validate_ttl_minutes(ttl_minutes: int) -> None:
+    if not isinstance(ttl_minutes, int) or ttl_minutes < 0:
+        raise errors.InvalidArgumentError("TTL Minutes must be a non-negative integer")
+
+
 def _validate_request_timeout(request_timeout_ms: Optional[int]) -> None:
     if request_timeout_ms is None:
         return

src/momento/aio/_scs_control_client.py

@@ -1,13 +1,23 @@
+from time import time
 from typing import Optional
 
 from momento_wire_types.controlclient_pb2 import _CreateCacheRequest
 from momento_wire_types.controlclient_pb2 import _DeleteCacheRequest
 from momento_wire_types.controlclient_pb2 import _ListCachesRequest
+from momento_wire_types.controlclient_pb2 import _CreateSigningKeyRequest
+from momento_wire_types.controlclient_pb2 import _RevokeSigningKeyRequest
+from momento_wire_types.controlclient_pb2 import _ListSigningKeysRequest
 
 from .._utilities._data_validation import _validate_cache_name
-from ..cache_operation_types import CreateCacheResponse
+from .._utilities._data_validation import _validate_ttl_minutes
+from ..cache_operation_types import (
+    CreateCacheResponse,
+    ListSigningKeysResponse,
+    RevokeSigningKeyResponse,
+)
 from ..cache_operation_types import DeleteCacheResponse
 from ..cache_operation_types import ListCachesResponse
+from ..cache_operation_types import CreateSigningKeyResponse
 
 from .. import _cache_service_errors_converter
 from .. import _momento_logger
@@ -70,5 +80,58 @@ async def list_caches(self, next_token: Optional[str] = None) -> ListCachesRespo
         except Exception as e:
             raise _cache_service_errors_converter.convert(e)
 
+    async def create_signing_key(
+        self, ttl_minutes: int, endpoint: str
+    ) -> CreateSigningKeyResponse:
+        _validate_ttl_minutes(ttl_minutes)
+        try:
+            _momento_logger.debug(
+                f"Creating signing key with ttl (in minutes): {ttl_minutes}"
+            )
+            create_signing_key_request = _CreateSigningKeyRequest()
+            create_signing_key_request.ttl_minutes = ttl_minutes
+            return CreateSigningKeyResponse(
+                await self._grpc_manager.async_stub().CreateSigningKey(
+                    create_signing_key_request, timeout=_DEADLINE_SECONDS
+                ),
+                endpoint,
+            )
+        except Exception as e:
+            _momento_logger.debug(f"Failed to create signing key with exception: {e}")
+            raise _cache_service_errors_converter.convert(e)
+
+    async def revoke_signing_key(self, key_id: str) -> RevokeSigningKeyResponse:
+        try:
+            _momento_logger.debug(f"Revoking signing key with key_id {key_id}")
+            request = _RevokeSigningKeyRequest()
+            request.key_id = key_id
+            return RevokeSigningKeyResponse(
+                await self._grpc_manager.async_stub().RevokeSigningKey(
+                    request, timeout=_DEADLINE_SECONDS
+                )
+            )
+        except Exception as e:
+            _momento_logger.debug(
+                f"Failed to revoke signing key with key_id {key_id} exception: {e}"
+            )
+            raise _cache_service_errors_converter.convert(e)
+
+    async def list_signing_keys(
+        self, endpoint: str, next_token: Optional[str] = None
+    ) -> ListSigningKeysResponse:
+        try:
+            list_signing_keys_request = _ListSigningKeysRequest()
+            list_signing_keys_request.next_token = (
+                next_token if next_token is not None else ""
+            )
+            return ListSigningKeysResponse(
+                await self._grpc_manager.async_stub().ListSigningKeys(
+                    list_signing_keys_request, timeout=_DEADLINE_SECONDS
+                ),
+                endpoint,
+            )
+        except Exception as e:
+            raise _cache_service_errors_converter.convert(e)
+
     async def close(self) -> None:
         await self._grpc_manager.close()

src/momento/aio/_scs_data_client.py

@@ -37,6 +37,10 @@ def __init__(
         self._grpc_manager = _scs_grpc_manager._DataGrpcManager(auth_token, endpoint)
         _validate_ttl(default_ttl_seconds)
         self._default_ttlSeconds = default_ttl_seconds
+        self._endpoint = endpoint
+
+    def get_endpoint(self) -> str:
+        return self._endpoint
 
     async def set(
         self,

src/momento/aio/simple_cache_client.py

@@ -30,6 +30,9 @@
     CreateCacheResponse,
     DeleteCacheResponse,
     ListCachesResponse,
+    CreateSigningKeyResponse,
+    RevokeSigningKeyResponse,
+    ListSigningKeysResponse,
     CacheSetResponse,
     CacheGetResponse,
     CacheMultiSetOperation,
@@ -121,6 +124,59 @@ async def list_caches(self, next_token: Optional[str] = None) -> ListCachesRespo
         """
         return await self._control_client.list_caches(next_token)
 
+    async def create_signing_key(self, ttl_minutes: int) -> CreateSigningKeyResponse:
+        """Creates a Momento signing key
+
+        Args:
+            ttl_minutes: The key's time-to-live in minutes
+
+        Returns:
+            CreateSigningKeyResponse
+
+        Raises:
+            InvalidArgumentError: If provided ttl minutes is negative.
+            BadRequestError: If the ttl provided is not accepted
+            AuthenticationError: If the provided Momento Auth Token is invalid.
+            ClientSdkError: For any SDK checks that fail.
+        """
+        return await self._control_client.create_signing_key(
+            ttl_minutes, self._data_client.get_endpoint()
+        )
+
+    async def revoke_signing_key(self, key_id: str) -> RevokeSigningKeyResponse:
+        """Revokes a Momento signing key, all tokens signed by which will be invalid
+
+        Args:
+            key_id: The id of the Momento signing key to revoke
+
+        Returns:
+            RevokeSigningKeyResponse
+
+        Raises:
+            AuthenticationError: If the provided Momento Auth Token is invalid.
+            ClientSdkError: For any SDK checks that fail.
+        """
+        return await self._control_client.revoke_signing_key(key_id)
+
+    async def list_signing_keys(
+        self, next_token: Optional[str] = None
+    ) -> ListSigningKeysResponse:
+        """Lists all Momento signing keys for the provided auth token.
+
+        Args:
+            next_token: Token to continue paginating through the list. It's used to handle large paginated lists.
+
+        Returns:
+            ListSigningKeysResponse
+
+        Raises:
+            AuthenticationError: If the provided Momento Auth Token is invalid.
+            ClientSdkError: For any SDK checks that fail.
+        """
+        return await self._control_client.list_signing_keys(
+            self._data_client.get_endpoint(), next_token
+        )
+
     async def multi_set(
         self,
         cache_name: str,

src/momento/cache_operation_types.py

@@ -1,11 +1,12 @@
+import json
+from datetime import datetime
 from enum import Enum
 from typing import Any, Optional, List, Union
 from dataclasses import dataclass
 
 from momento_wire_types import cacheclient_pb2 as cache_client_types
 from . import _cache_service_errors_converter as error_converter
 from . import _momento_logger
-from .errors import MomentoServiceError
 
 
 class CacheGetStatus(Enum):
@@ -209,3 +210,91 @@ def next_token(self) -> Optional[str]:
     def caches(self) -> List[CacheInfo]:
         """Returns all caches."""
         return self._caches
+
+
+class CreateSigningKeyResponse:
+    def __init__(self, grpc_create_signing_key_response: Any, endpoint: str):  # type: ignore[misc]
+        """Initializes CreateSigningKeyResponse to handle create signing key response.
+
+        Args:
+            grpc_create_signing_key_response: Protobuf based response returned by Scs.
+        """
+        self._key_id: str = json.loads(grpc_create_signing_key_response.key)["kid"]  # type: ignore[misc]
+        self._endpoint: str = endpoint
+        self._key: str = grpc_create_signing_key_response.key  # type: ignore[misc]
+        self._expires_at: datetime = datetime.fromtimestamp(
+            grpc_create_signing_key_response.expires_at  # type: ignore[misc]
+        )
+
+    def key_id(self) -> str:
+        """Returns the id of the signing key"""
+        return self._key_id
+
+    def endpoint(self) -> str:
+        """Returns the endpoint of the signing key"""
+        return self._endpoint
+
+    def key(self) -> str:
+        """Returns the JSON string of the key itself"""
+        return self._key
+
+    def expires_at(self) -> datetime:
+        """Returns the datetime representation of when the key expires"""
+        return self._expires_at
+
+
+class RevokeSigningKeyResponse:
+    def __init__(self, grpc_revoke_signing_key_response: Any):  # type: ignore[misc]
+        pass
+
+
+class SigningKey:
+    def __init__(self, grpc_listed_signing_key: Any, endpoint: str):  # type: ignore[misc]
+        """Initializes SigningKey to handle signing keys returned from list signing keys operation.
+
+        Args:
+            grpc_listed_signing_key: Protobuf based response returned by Scs.
+        """
+        self._key_id: str = grpc_listed_signing_key.key_id  # type: ignore[misc]
+        self._expires_at: datetime = datetime.fromtimestamp(
+            grpc_listed_signing_key.expires_at  # type: ignore[misc]
+        )
+        self._endpoint: str = endpoint
+
+    def key_id(self) -> str:
+        """Returns the id of the Momento signing key"""
+        return self._key_id
+
+    def expires_at(self) -> datetime:
+        """Returns the time the key expires"""
+        return self._expires_at
+
+    def endpoint(self) -> str:
+        """Returns the endpoint of the Momento signing key"""
+        return self._endpoint
+
+
+class ListSigningKeysResponse:
+    def __init__(self, grpc_list_signing_keys_response: Any, endpoint: str):  # type: ignore[misc]
+        """Initializes ListSigningKeysResponse to handle list signing keys response.
+
+        Args:
+            grpc_list_signing_keys_response: Protobuf based response returned by Scs.
+        """
+        self._next_token: Optional[str] = (
+            grpc_list_signing_keys_response.next_token  # type: ignore[misc]
+            if grpc_list_signing_keys_response.next_token != ""  # type: ignore[misc]
+            else None
+        )
+        self._signing_keys: List[SigningKey] = [  # type: ignore[misc]
+            SigningKey(signing_key, endpoint)  # type: ignore[misc]
+            for signing_key in grpc_list_signing_keys_response.signing_key  # type: ignore[misc]
+        ]
+
+    def next_token(self) -> Optional[str]:
+        """Returns next token."""
+        return self._next_token
+
+    def signing_keys(self) -> List[SigningKey]:
+        """Returns all signing keys."""
+        return self._signing_keys

src/momento/simple_cache_client.py

@@ -9,6 +9,7 @@
     CacheGetResponse,
     CacheSetResponse,
     CreateCacheResponse,
+    CreateSigningKeyResponse,
     DeleteCacheResponse,
     ListCachesResponse,
     CacheMultiGetResponse,
@@ -17,6 +18,8 @@
     CacheMultiSetResponse,
     CacheMultiSetFailureResponse,
     CacheMultiGetFailureResponse,
+    ListSigningKeysResponse,
+    RevokeSigningKeyResponse,
 )
 
 from ._utilities._data_validation import _validate_request_timeout
@@ -116,6 +119,58 @@ def list_caches(self, next_token: Optional[str] = None) -> ListCachesResponse:
         coroutine = self._momento_async_client.list_caches(next_token)
         return wait_for_coroutine(self._loop, coroutine)
 
+    def create_signing_key(self, ttl_minutes: int) -> CreateSigningKeyResponse:
+        """Creates a Momento signing key
+
+        Args:
+            ttl_minutes: The key's time-to-live in minutes
+
+        Returns:
+            CreateSigningKeyResponse
+
+        Raises:
+            InvalidArgumentError: If provided ttl minutes is negative.
+            BadRequestError: If the ttl provided is not accepted
+            AuthenticationError: If the provided Momento Auth Token is invalid.
+            ClientSdkError: For any SDK checks that fail.
+        """
+        coroutine = self._momento_async_client.create_signing_key(ttl_minutes)
+        return wait_for_coroutine(self._loop, coroutine)
+
+    def revoke_signing_key(self, key_id: str) -> RevokeSigningKeyResponse:
+        """Revokes a Momento signing key, all tokens signed by which will be invalid
+
+        Args:
+            key_id: The id of the Momento signing key to revoke
+
+        Returns:
+            RevokeSigningKeyResponse
+
+        Raises:
+            AuthenticationError: If the provided Momento Auth Token is invalid.
+            ClientSdkError: For any SDK checks that fail.
+        """
+        coroutine = self._momento_async_client.revoke_signing_key(key_id)
+        return wait_for_coroutine(self._loop, coroutine)
+
+    def list_signing_keys(
+        self, next_token: Optional[str] = None
+    ) -> ListSigningKeysResponse:
+        """Lists all Momento signing keys for the provided auth token.
+
+        Args:
+            next_token: Token to continue paginating through the list. It's used to handle large paginated lists.
+
+        Returns:
+            ListSigningKeysResponse
+
+        Raises:
+            AuthenticationError: If the provided Momento Auth Token is invalid.
+            ClientSdkError: For any SDK checks that fail.
+        """
+        coroutine = self._momento_async_client.list_signing_keys(next_token)
+        return wait_for_coroutine(self._loop, coroutine)
+
     def set(
         self,
         cache_name: str,

tests/test_momento_async.py

@@ -178,6 +178,22 @@ async def test_list_caches_with_next_token_works(self):
         # skip until pagination is actually implemented, see
         # https://github.com/momentohq/control-plane-service/issues/83
         self.skipTest("pagination not yet implemented")
+
+    # signing keys
+    async def test_create_list_revoke_signing_keys(self):
+        try:
+            create_resp = await self.client.create_signing_key(30)
+            list_resp = await self.client.list_signing_keys()
+            self.assertTrue(0 < len(list_resp.signing_keys()))
+            signing_key = list_resp.signing_keys()[0]
+            self.assertEqual(create_resp.endpoint(), signing_key.endpoint())
+            self.assertEqual(create_resp.key_id(), signing_key.key_id())
+        finally:
+            list_resp = await self.client.list_signing_keys()
+            for signing_key in list_resp.signing_keys():
+                await self.client.revoke_signing_key(signing_key.key_id())
+        list_resp = await self.client.list_signing_keys()
+        self.assertEqual(0, len(list_resp.signing_keys()))
 
     # setting and getting