feat(redis) - refactoring connection and fixing mypy errors

Author: leandrodamascena

aws_lambda_powertools/utilities/connections/__init__.py

@@ -3,4 +3,4 @@
     RedisStandalone,
 )
 
-__all__ = (RedisStandalone, RedisCluster)
+__all__ = ["RedisStandalone", "RedisCluster"]

aws_lambda_powertools/utilities/connections/redis.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Optional
+from typing import Optional, Type, Union
 
 import redis
 
@@ -9,44 +9,29 @@
 logger = logging.getLogger(__name__)
 
 
-class RedisStandalone(BaseConnectionSync):
+class RedisConnection(BaseConnectionSync):
     def __init__(
         self,
+        client: Type[Union[redis.Redis, redis.RedisCluster]],
         host: Optional[str] = None,
         port: Optional[int] = None,
         username: Optional[str] = None,
         password: Optional[str] = None,
         db_index: Optional[int] = None,
         url: Optional[str] = None,
+        **extra_options,
     ) -> None:
-        """
-        Initialize the Redis standalone client
-        Parameters
-        ----------
-        host: str
-            Name of the host to connect to Redis instance/cluster
-        port: int
-            Number of the port to connect to Redis instance/cluster
-        username: str
-            Name of the username to connect to Redis instance/cluster in case of using ACL
-            See: https://redis.io/docs/management/security/acl/
-        password: str
-            Passwod to connect to Redis instance/cluster
-        db_index: int
-            Index of Redis database
-            See: https://redis.io/commands/select/
-        url: str
-            Redis client object configured from the given URL
-            See: https://redis.readthedocs.io/en/latest/connections.html#redis.Redis.from_url
-        """
+        self.extra_options: dict = {}
 
         self.url = url
         self.host = host
         self.port = port
         self.username = username
         self.password = password
         self.db_index = db_index
+        self.extra_options.update(**extra_options)
         self._connection = None
+        self._client = client
 
     def init_connection(self):
         """
@@ -55,36 +40,40 @@ def init_connection(self):
         if self._connection:
             return self._connection
 
-        logger.info(f"Trying to connect to Redis Host/Cluster: {self.host}")
+        logger.info(f"Trying to connect to Redis: {self.host}")
 
         try:
             if self.url:
                 logger.debug(f"Using URL format to connect to Redis: {self.host}")
-                self._connection = redis.Redis.from_url(url=self.url)
+                self._connection = self._client.from_url(url=self.url)
             else:
                 logger.debug(f"Using other parameters to connect to Redis: {self.host}")
-                self._connection = redis.Redis(
+                self._connection = self._client(
                     host=self.host,
                     port=self.port,
                     username=self.username,
                     password=self.password,
                     db=self.db_index,
                     decode_responses=True,
+                    **self.extra_options,
                 )
         except redis.exceptions.ConnectionError as exc:
-            logger.debug(f"Cannot connect in Redis Host: {self.host}")
-            raise RedisConnectionError("Could not to connect to Redis Standalone", exc) from exc
+            logger.debug(f"Cannot connect in Redis: {self.host}")
+            raise RedisConnectionError("Could not to connect to Redis", exc) from exc
 
         return self._connection
 
 
-class RedisCluster(BaseConnectionSync):
+class RedisStandalone(RedisConnection):
     def __init__(
         self,
         host: Optional[str] = None,
         port: Optional[int] = None,
-        read_from_replicas: Optional[bool] = False,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        db_index: Optional[int] = None,
         url: Optional[str] = None,
+        **extra_options,
     ) -> None:
         """
         Initialize the Redis standalone client
@@ -106,36 +95,40 @@ def __init__(
             Redis client object configured from the given URL
             See: https://redis.readthedocs.io/en/latest/connections.html#redis.Redis.from_url
         """
+        print(extra_options)
+        super().__init__(redis.Redis, host, port, username, password, db_index, url, **extra_options)
 
-        self.url = url
-        self.host = host
-        self.port = port
-        self.read_from_replicas = read_from_replicas
-        self._connection = None
 
-    def init_connection(self):
+class RedisCluster(RedisConnection):
+    def __init__(
+        self,
+        host: Optional[str] = None,
+        port: Optional[int] = None,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        db_index: Optional[int] = None,
+        url: Optional[str] = None,
+        **extra_options,
+    ) -> None:
         """
-        Connection is cached, so returning this
+        Initialize the Redis standalone client
+        Parameters
+        ----------
+        host: str
+            Name of the host to connect to Redis instance/cluster
+        port: int
+            Number of the port to connect to Redis instance/cluster
+        username: str
+            Name of the username to connect to Redis instance/cluster in case of using ACL
+            See: https://redis.io/docs/management/security/acl/
+        password: str
+            Passwod to connect to Redis instance/cluster
+        db_index: int
+            Index of Redis database
+            See: https://redis.io/commands/select/
+        url: str
+            Redis client object configured from the given URL
+            See: https://redis.readthedocs.io/en/latest/connections.html#redis.Redis.from_url
         """
-        if self._connection:
-            return self._connection
 
-        logger.info(f"Trying to connect to Redis Cluster: {self.host}")
-
-        try:
-            if self.url:
-                logger.debug(f"Using URL format to connect to Redis Cluster: {self.host}")
-                self._connection = redis.Redis.from_url(url=self.url)
-            else:
-                logger.debug(f"Using other parameters to connect to Redis Cluster: {self.host}")
-                self._connection = redis.cluster.RedisCluster(
-                    host=self.host,
-                    port=self.port,
-                    server_type=None,
-                    read_from_replicas=self.read_from_replicas,
-                )
-        except redis.exceptions.ConnectionError as exc:
-            logger.debug(f"Cannot connect in Redis Cluster: {self.host}")
-            raise RedisConnectionError("Could not to connect to Redis Cluster", exc) from exc
-
-        return self._connection
+        super().__init__(redis.cluster.RedisCluster, host, port, username, password, db_index, url, **extra_options)

aws_lambda_powertools/utilities/idempotency/persistence/base.py

@@ -37,7 +37,7 @@ class DataRecord:
 
     def __init__(
         self,
-        idempotency_key: Optional[str] = "",
+        idempotency_key: str = "",
         status: str = "",
         expiry_timestamp: Optional[int] = None,
         in_progress_expiry_timestamp: Optional[int] = None,

aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py

@@ -26,7 +26,7 @@ class DynamoDBPersistenceLayer(BasePersistenceLayer):
     def __init__(
         self,
         table_name: str,
-        key_attr: str = "id",
+        key_attr: Optional[str] = "id",
         static_pk_value: Optional[str] = None,
         sort_key_attr: Optional[str] = None,
         expiry_attr: str = "expiration",

aws_lambda_powertools/utilities/idempotency/persistence/redis.py

@@ -19,10 +19,10 @@ def __init__(
         connection,
         static_pk_value: Optional[str] = None,
         expiry_attr: str = "expiration",
-        in_progress_expiry_attr: str = "in_progress_expiration",
+        in_progress_expiry_attr="in_progress_expiration",
         status_attr: str = "status",
         data_attr: str = "data",
-        validation_key_attr: str = "validation",
+        validation_key_attr="validation",
     ):
         """
         Initialize the Redis Persistence Layer
@@ -54,12 +54,6 @@ def __init__(
         self.validation_key_attr = validation_key_attr
         super(RedisCachePersistenceLayer, self).__init__()
 
-    def _get_key(self, idempotency_key: str) -> dict:
-        # Need to review this after adding GETKEY logic
-        if self.sort_key_attr:
-            return {self.key_attr: self.static_pk_value, self.sort_key_attr: idempotency_key}
-        return {self.key_attr: idempotency_key}
-
     def _item_to_data_record(self, item: Dict[str, Any]) -> DataRecord:
         # Need to review this after adding GETKEY logic
         return DataRecord(
@@ -93,10 +87,10 @@ def _put_record(self, data_record: DataRecord) -> None:
         }
 
         if data_record.in_progress_expiry_timestamp is not None:
-            item["mapping"][self.in_progress_expiry_attr] = data_record.in_progress_expiry_timestamp
+            item.update({"mapping": {self.in_progress_expiry_attr: data_record.in_progress_expiry_timestamp}})
 
         if self.payload_validation_enabled:
-            item["mapping"][self.validation_key_attr] = data_record.payload_hash
+            item.update({"mapping": {self.validation_key_attr: data_record.payload_hash}})
 
         try:
             # |     LOCKED     |         RETRY if status = "INPROGRESS"                |     RETRY