snapshot

legacy/aws_elasticache.py

@@ -0,0 +1,438 @@
+import dataclasses
+import functools
+import itertools
+import logging
+import re
+import time
+import typing
+import uuid
+from typing import Optional
+
+import boto3
+import botocore.exceptions
+from botocore.exceptions import ClientError
+
+logger = logging.getLogger(__name__)
+
+
+def retry(timeout: int, backoff_seconds: int = 10, bubble_errors: typing.List[typing.Type[Exception]] = None):
+    """
+    Tries calling a function {timeout} seconds until it succeeds or gives up and throw an error
+    :param timeout: Timeout in seconds
+    :param backoff_seconds: Wait time between retries
+    :param bubble_errors: List of exceptions to bubble up
+    :return: Wrapped function
+    """
+    bubble_errors = tuple(bubble_errors or [])
+
+    def wrapped(func: typing.Callable):
+        @functools.wraps(func)
+        def inner(*args, **kwargs):
+            waited = 0
+            while True:
+                try:
+                    logger.debug(f"Calling {func}")
+                    return func(*args, **kwargs)
+                except KeyboardInterrupt:
+                    raise
+                except tuple(bubble_errors):
+                    raise
+                except Exception as e:
+                    doze = min(timeout - waited, backoff_seconds)
+                    times = (timeout - waited) // backoff_seconds
+                    logger.debug(f"{func} failed, will try {times} times in {doze} sec")
+                    time.sleep(doze)
+                    waited += backoff_seconds
+                    if waited >= timeout:
+                        raise TimeoutError(f"Couldn't get a successful result from {func} in {timeout} seconds") from e
+
+        return inner
+
+    return wrapped
+
+
+@dataclasses.dataclass
+class RedisInstance:
+    cluster_id: str
+    host: str
+    port: int
+    arn: str
+    user_group_ids: typing.List[str] = dataclasses.field(default_factory=list)
+
+
+def paginate_by_marker(
+    func: callable,
+    list_field: str,
+) -> typing.Iterable:
+    """
+    Simplifies paginating over AWS results with a simpler interface than Paginators.
+    :param func: A function that accepts a parameter named `Marker`
+    :param list_field: Result field to iterate on
+    """
+    marker = ""
+    while True:
+        result = func(Marker=marker)
+
+        items = result.get(list_field, [])
+        yield from items
+
+        if "Marker" not in result:
+            break
+
+        marker = result["Marker"]
+
+
+@dataclasses.dataclass
+class RedisUser:
+    username: str
+    password: str
+    redis_key_prefix: str
+
+
+class CreateUserResult(typing.TypedDict):
+    user_group_id: str
+    users: typing.List[RedisUser]
+    status: str
+
+
+class RedisService:
+    def __init__(self, region: str):
+        self.elasticache = boto3.client("elasticache", region_name=region)
+
+    def create_redis_acl_user(
+        self,
+        replication_group_id: str,
+        users: typing.List[RedisUser],
+    ) -> None:
+        """
+        Creates a Redis ACL user. It will throw an error if the user group attached to the replication group is not ready.
+        This prevents us from quickly creating users. In that case, pass in a list of users instead.
+
+        It takes about 2 minutes until the user is ready.
+        """
+        repl_groups = self.elasticache.describe_replication_groups(ReplicationGroupId=replication_group_id)
+        user_group_id: str = repl_groups['ReplicationGroups'][0]['UserGroupIds'][0]
+        group = self.elasticache.describe_user_groups(UserGroupId=user_group_id)['UserGroups'][0]
+
+        """
+        Limits imposed by AWS:
+        
+        - User groups per replication group = 1
+        - Users per user group = 100
+        - Number of users = 1000
+        - Number of user groups = 100
+        https://docs.aws.amazon.com/AmazonElastiCache/latest/red-ug/Clusters.RBAC.html#Users-groups-to-RGs
+        
+        In short, we can't place more than 100 users on a single Redis instance.   
+        """
+
+        total_users = len(group['UserIds'])
+        if total_users >= 100:
+            raise Exception('Replication group reached the limit of 100 users')
+
+        self._update_or_create_user_group(user_group_id, users)
+
+    def delete_redis_acl_user(self, username: str) -> None:
+        logger.info(f'Deleting user {username}')
+        self.elasticache.delete_user(UserId=username)
+
+        # Even though we just deleted the user, it takes a couple of seconds until AWS updates user group status.
+        # So, we have to wait a bit before we can actually wait and check that the changes have propagated.
+        # Overall, it should take about 2-3 minutes to delete a user.
+
+    def create_redis_acl(
+        self,
+        *,
+        replication_group_id: str,
+        description: str = None,
+        node_type: str,
+        cache_subnet_group_name: str,
+        security_group_id: str,
+        tags: typing.Dict[str, str] = None,
+        users: typing.List[RedisUser],
+    ) -> dict:
+        user_group_id = f"{replication_group_id}-ug"
+        # we have to wait (~1m) until the user group is ready before we can create the ACL
+        self._update_or_create_user_group(user_group_id, users, wait=True)
+
+        cache_param_group = 'redis-acl-with-100-db'
+        self._ensure_cache_parameter_group(cache_param_group)
+
+        logger.info(f"Creating Redis replication group {replication_group_id=}")
+        default_replication_group_kwargs = dict(
+            ReplicasPerNodeGroup=0,
+            Engine="redis",
+            EngineVersion="6.x",
+            TransitEncryptionEnabled=True,  # must be True to use Redis ACL
+            MultiAZEnabled=False,
+            AutomaticFailoverEnabled=False,
+        )
+        repl_result = self.elasticache.create_replication_group(
+            **default_replication_group_kwargs,
+            ReplicationGroupId=replication_group_id,
+            ReplicationGroupDescription=description or f'{replication_group_id} replication group',
+            CacheNodeType=node_type,
+            CacheSubnetGroupName=cache_subnet_group_name,
+            CacheParameterGroupName=cache_param_group,
+            UserGroupIds=[
+                user_group_id,
+            ],
+            SecurityGroupIds=[
+                security_group_id,
+            ],
+            Tags=[{"Key": k, "Value": v} for k, v in (tags or {}).items()],
+        )
+
+        # replication group takes ~8m minutes to be ready
+
+        return {
+            "replication_group_id": repl_result["ReplicationGroup"]["ReplicationGroupId"],
+            'arn': repl_result["ReplicationGroup"]["ARN"],
+            'status': repl_result["ReplicationGroup"]["Status"],
+            'users': users,
+            '_response': repl_result,
+        }
+
+    def _update_or_create_user_group(
+        self, user_group_id: str, users: typing.List[RedisUser], wait: bool = False
+    ) -> None:
+        """
+        Creates or updates a user group with the given users. Waits until the user group is active.
+
+        :param user_group_id:
+        :param users:
+        """
+        # we have to add the `default` user for backwards compatibility
+        existing_users = self.elasticache.describe_users(
+            Filters=[{"Name": "user-id", "Values": [u.username for u in users]}],
+        )["Users"]
+        existing_user_ids = [it["UserId"] for it in existing_users]
+
+        created_user_ids = ["default"]
+        for user in users:
+            if user.username in existing_user_ids:
+                logger.info(f"User {user.username} is already exists")
+                created_user_ids.append(user.username)
+                continue
+
+            redis_acl = f"on +@all -@dangerous ~{user.redis_key_prefix}*"
+            logger.info(f"Creating tenant {user.username=}")
+            user_result = self.elasticache.create_user(
+                UserId=user.username,
+                UserName=user.username,
+                Passwords=[user.password],
+                AccessString=redis_acl,
+                Engine="redis",
+            )
+            created_user_ids.append(user_result["UserId"])
+
+        try:
+            group = self.elasticache.describe_user_groups(UserGroupId=user_group_id)["UserGroups"][0]
+            member_user_ids: typing.List[str] = group["UserIds"]
+            users_to_add = set(created_user_ids) - set(member_user_ids)
+
+            if users_to_add:
+                self.elasticache.modify_user_group(
+                    UserGroupId=user_group_id,
+                    UserIdsToAdd=list(users_to_add),
+                )
+                # this will take some time (~45s) until the changes propagate
+        except self.elasticache.exceptions.UserGroupNotFoundFault:
+            logger.info('User group does not exist, creating')
+            _ = self.elasticache.create_user_group(
+                UserGroupId=user_group_id,
+                Engine="redis",
+                UserIds=created_user_ids,
+            )
+            # user group creation takes around 60s
+            if wait:
+                self._wait_user_group(user_group_id)
+
+    @retry(timeout=60, backoff_seconds=5)
+    def _wait_user_group(self, user_group_id: str) -> None:
+        logger.info(f'Checking status of user group {user_group_id=}')
+        g = self.elasticache.describe_user_groups(UserGroupId=user_group_id)["UserGroups"][0]
+        assert g["Status"] == "active"
+
+    def is_redis_acl_ready(self, replication_group_id: str) -> bool:
+        repl = self.elasticache.describe_replication_groups(ReplicationGroupId=replication_group_id)[
+            'ReplicationGroups'
+        ][0]
+        is_redis_available = repl['Status'] == 'available'
+
+        user_group_id = repl['UserGroupIds'][0]
+        group = self.elasticache.describe_user_groups(UserGroupId=user_group_id)['UserGroups'][0]
+        is_group_available = group['Status'] == 'active'
+
+        return is_group_available and is_redis_available
+
+    def get_redis_acl(self, replication_group_id: str) -> dict:
+        redis = self.elasticache.describe_replication_groups(ReplicationGroupId=replication_group_id)[
+            'ReplicationGroups'
+        ][0]
+        return dict(
+            replication_group_id=redis['ReplicationGroupId'],
+            arn=redis['ARN'],
+            host=redis["NodeGroups"][0]["PrimaryEndpoint"]["Address"],
+            port=redis["NodeGroups"][0]["PrimaryEndpoint"]["Port"],
+        )
+
+    def delete_redis_acl(self, replication_group_id: str) -> None:
+        self.elasticache.delete_replication_group(ReplicationGroupId=replication_group_id)
+
+    def _ensure_cache_parameter_group(self, parameter_group_name: str) -> None:
+        try:
+            _ = self.elasticache.describe_cache_parameter_groups(CacheParameterGroupName=parameter_group_name)[
+                'CacheParameterGroups'
+            ][0]
+            return
+        except self.elasticache.exceptions.CacheParameterGroupNotFoundFault:
+            pass
+
+        logger.info(f"Creating cache parameter group {parameter_group_name=}")
+        _ = self.elasticache.create_cache_parameter_group(
+            CacheParameterGroupName=parameter_group_name,
+            CacheParameterGroupFamily="redis6.x",
+            Description="Redis parameter group for the Redis ACL",
+        )
+        _ = self.elasticache.modify_cache_parameter_group(
+            CacheParameterGroupName=parameter_group_name,
+            ParameterNameValues=[
+                {
+                    "ParameterName": "databases",
+                    "ParameterValue": "100",
+                },
+            ],
+        )
+
+
+def akinon_create_redis(
+    *,
+    k8s_cluster_name: str,
+    region: str,
+    node_type: str,
+    replicas: int,
+    owner_arn: str,
+    app_name: str,
+    role: str,
+):
+    subnet_group_name = f"{k8s_cluster_name}-redis-subg"
+    security_group_name = f"{k8s_cluster_name}-redis-sg"
+
+    # owner arn is formatted like:
+    # arn:aws:iam::412344683105:user/myusername
+    username = owner_arn.split("/")[1]
+    tags = {
+        "CostCenter": f"{username}-{app_name}-{role}-redis",
+    }
+
+    cache_cluster_id = uuid.uuid4().hex
+    r = RedisService(region)
+    return r.create_redis(
+        cluster_id=cache_cluster_id,
+        node_type=node_type,
+        replicas=replicas,
+        security_group_name=security_group_name,
+        subnet_group_name=subnet_group_name,
+        tags=tags,
+    )
+
+
+def akinon_create_redis_with_acl(
+    k8s_cluster_name: str,
+    owner_arn: str,
+    node_type: str,
+    app_name: str,
+    role: str,
+    users: typing.List[RedisUser],
+) -> RedisInstance:
+    cache_cluster_id = uuid.uuid4().hex
+    subnet_group_name = f"{k8s_cluster_name}-redis-subg"
+    security_group_name = f"{k8s_cluster_name}-redis-sg"
+
+    # owner arn is formatted as:
+    # arn:aws:iam::412344683105:user/myusername
+    username = owner_arn.split("/")[1]
+    tags = {
+        "CostCenter": f"{username}-{app_name}-{role}-redis",
+    }
+
+    return r.create_redis_acl(
+        cluster_id=cache_cluster_id,
+        node_type=node_type,
+        subnet_group_name=subnet_group_name,
+        security_group_name=security_group_name,
+        tags=tags,
+        users=users,
+    )
+
+
+def chunk(it, size):
+    it = iter(it)
+    sentinel = ()
+    return iter(lambda: tuple(itertools.islice(it, size)), sentinel)
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO, format=f"%(asctime)s: {logging.BASIC_FORMAT}")
+    r = RedisService(region="eu-central-1")
+    # res = r.create_redis_acl(
+    #     replication_group_id="zerodev3",
+    #     description="zero redis dev",
+    #     node_type="cache.t3.micro",
+    #     cache_subnet_group_name="dev",
+    #     security_group_id="sg-c11a77ac",
+    #     tags={"CostCenter": "zero123"},
+    #     users=[RedisUser(username="u1", password="zeropassword1234", redis_key_prefix="zero:")],
+    # )
+    # print(res)
+
+    r.create_redis_acl_user(
+        replication_group_id='zerodev3',
+        users=[
+            RedisUser(
+                username='u2',
+                password='testtesttesttest123',
+                redis_key_prefix='test:',
+            )
+        ],
+    )
+    # r.delete_redis_acl_user('test6')
+    while True:
+        if r.is_redis_acl_ready('zerodev3'):
+            logger.info('ready')
+            print(r.get_redis_acl('zerodev3'))
+            break
+        logger.info('not ready')
+        time.sleep(10)
+    exit()
+
+    res = r.create_redis_acl(
+        replication_group_id="zerodev",
+        description="zero redis dev",
+        node_type="cache.t3.micro",
+        cache_subnet_group_name="dev",
+        security_group_id="sg-c11a77ac",
+        tags={"CostCenter": "zero123"},
+        users=[RedisUser(username="zerouser", password="zeropassword1234", redis_key_prefix="zero:")],
+    )
+    print(res)
+    #
+    # redis = r.get_redis_acl(cluster_arn="arn:aws:elasticache:eu-central-1:400344683105:replicationgroup:zero")
+    # r.add_redis_user(
+    #     cluster_id=redis.cluster_id,
+    #     users=[RedisUser(username="zero27", password="zero27password123", redis_key_prefix="zero27:")],
+    # )
+    # print(r.add_redis_user(cluster_id="zero1", users=[RedisUser("zero26", "zero26password123", "zero26:")]))
+    #
+    # print(
+    #     r.create_redis_with_acl(
+    #         cluster_id="zero",
+    #         cluster_description="zero redis",
+    #         node_type="cache.t3.micro",
+    #         subnet_group_name="testing-subnet",
+    #         security_group_name="default",
+    #         tags={"CostCenter": "zero123"},
+    #         users=[RedisUser(username="zerouser", password="zeropassword1234", redis_key_prefix="zero:")],
+    #     )
+    # )