#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations
from time import sleep
from airflow.exceptions import AirflowException
from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
[docs]class ElastiCacheReplicationGroupHook(AwsBaseHook):
"""
Interact with AWS ElastiCache
:param max_retries: Max retries for checking availability of and deleting replication group
If this is not supplied then this is defaulted to 10
:param exponential_back_off_factor: Multiplication factor for deciding next sleep time
If this is not supplied then this is defaulted to 1
:param initial_poke_interval: Initial sleep time in seconds
If this is not supplied then this is defaulted to 60 seconds
"""
[docs] TERMINAL_STATES = frozenset({"available", "create-failed", "deleting"})
def __init__(
self,
max_retries: int = 10,
exponential_back_off_factor: float = 1,
initial_poke_interval: float = 60,
*args,
**kwargs,
):
self.max_retries = max_retries
self.exponential_back_off_factor = exponential_back_off_factor
self.initial_poke_interval = initial_poke_interval
kwargs["client_type"] = "elasticache"
super().__init__(*args, **kwargs)
[docs] def create_replication_group(self, config: dict) -> dict:
"""
Call ElastiCache API for creating a replication group
:param config: Configuration for creating the replication group
:return: Response from ElastiCache create replication group API
"""
return self.conn.create_replication_group(**config)
[docs] def delete_replication_group(self, replication_group_id: str) -> dict:
"""
Call ElastiCache API for deleting a replication group
:param replication_group_id: ID of replication group to delete
:return: Response from ElastiCache delete replication group API
"""
return self.conn.delete_replication_group(ReplicationGroupId=replication_group_id)
[docs] def describe_replication_group(self, replication_group_id: str) -> dict:
"""
Call ElastiCache API for describing a replication group
:param replication_group_id: ID of replication group to describe
:return: Response from ElastiCache describe replication group API
"""
return self.conn.describe_replication_groups(ReplicationGroupId=replication_group_id)
[docs] def get_replication_group_status(self, replication_group_id: str) -> str:
"""
Get current status of replication group
:param replication_group_id: ID of replication group to check for status
:return: Current status of replication group
"""
return self.describe_replication_group(replication_group_id)["ReplicationGroups"][0]["Status"]
[docs] def is_replication_group_available(self, replication_group_id: str) -> bool:
"""
Helper for checking if replication group is available or not
:param replication_group_id: ID of replication group to check for availability
:return: True if available else False
"""
return self.get_replication_group_status(replication_group_id) == "available"
[docs] def wait_for_availability(
self,
replication_group_id: str,
initial_sleep_time: float | None = None,
exponential_back_off_factor: float | None = None,
max_retries: int | None = None,
) -> bool:
"""
Check if replication group is available or not by performing a describe over it
:param replication_group_id: ID of replication group to check for availability
:param initial_sleep_time: Initial sleep time in seconds
If this is not supplied then this is defaulted to class level value
:param exponential_back_off_factor: Multiplication factor for deciding next sleep time
If this is not supplied then this is defaulted to class level value
:param max_retries: Max retries for checking availability of replication group
If this is not supplied then this is defaulted to class level value
:return: True if replication is available else False
"""
sleep_time = initial_sleep_time or self.initial_poke_interval
exponential_back_off_factor = exponential_back_off_factor or self.exponential_back_off_factor
max_retries = max_retries or self.max_retries
num_tries = 0
status = "not-found"
stop_poking = False
while not stop_poking and num_tries <= max_retries:
status = self.get_replication_group_status(replication_group_id=replication_group_id)
stop_poking = status in self.TERMINAL_STATES
self.log.info(
"Current status of replication group with ID %s is %s", replication_group_id, status
)
if not stop_poking:
num_tries += 1
# No point in sleeping if all tries have exhausted
if num_tries > max_retries:
break
self.log.info("Poke retry %s. Sleep time %s seconds. Sleeping...", num_tries, sleep_time)
sleep(sleep_time)
sleep_time *= exponential_back_off_factor
if status != "available":
self.log.warning('Replication group is not available. Current status is "%s"', status)
return False
return True
[docs] def wait_for_deletion(
self,
replication_group_id: str,
initial_sleep_time: float | None = None,
exponential_back_off_factor: float | None = None,
max_retries: int | None = None,
):
"""
Helper for deleting a replication group ensuring it is either deleted or can't be deleted
:param replication_group_id: ID of replication to delete
:param initial_sleep_time: Initial sleep time in second
If this is not supplied then this is defaulted to class level value
:param exponential_back_off_factor: Multiplication factor for deciding next sleep time
If this is not supplied then this is defaulted to class level value
:param max_retries: Max retries for checking availability of replication group
If this is not supplied then this is defaulted to class level value
:return: Response from ElastiCache delete replication group API and flag to identify if deleted or not
"""
deleted = False
sleep_time = initial_sleep_time or self.initial_poke_interval
exponential_back_off_factor = exponential_back_off_factor or self.exponential_back_off_factor
max_retries = max_retries or self.max_retries
num_tries = 0
response = None
while not deleted and num_tries <= max_retries:
try:
status = self.get_replication_group_status(replication_group_id=replication_group_id)
self.log.info(
"Current status of replication group with ID %s is %s", replication_group_id, status
)
# Can only delete if status is `available`
# Status becomes `deleting` after this call so this will only run once
if status == "available":
self.log.info("Initiating delete and then wait for it to finish")
response = self.delete_replication_group(replication_group_id=replication_group_id)
except self.conn.exceptions.ReplicationGroupNotFoundFault:
self.log.info("Replication group with ID '%s' does not exist", replication_group_id)
deleted = True
# This should never occur as we only issue a delete request when status is `available`
# which is a valid status for deletion. Still handling for safety.
except self.conn.exceptions.InvalidReplicationGroupStateFault as exp:
# status Error Response
# creating - Cache cluster <cluster_id> is not in a valid state to be deleted.
# deleting - Replication group <replication_group_id> has status deleting which is not valid
# for deletion.
# modifying - Replication group <replication_group_id> has status deleting which is not valid
# for deletion.
message = exp.response["Error"]["Message"]
self.log.warning("Received error message from AWS ElastiCache API : %s", message)
if not deleted:
num_tries += 1
# No point in sleeping if all tries have exhausted
if num_tries > max_retries:
break
self.log.info("Poke retry %s. Sleep time %s seconds. Sleeping...", num_tries, sleep_time)
sleep(sleep_time)
sleep_time *= exponential_back_off_factor
return response, deleted
[docs] def ensure_delete_replication_group(
self,
replication_group_id: str,
initial_sleep_time: float | None = None,
exponential_back_off_factor: float | None = None,
max_retries: int | None = None,
) -> dict:
"""
Delete a replication group ensuring it is either deleted or can't be deleted
:param replication_group_id: ID of replication to delete
:param initial_sleep_time: Initial sleep time in second
If this is not supplied then this is defaulted to class level value
:param exponential_back_off_factor: Multiplication factor for deciding next sleep time
If this is not supplied then this is defaulted to class level value
:param max_retries: Max retries for checking availability of replication group
If this is not supplied then this is defaulted to class level value
:return: Response from ElastiCache delete replication group API
:raises AirflowException: If replication group is not deleted
"""
self.log.info("Deleting replication group with ID %s", replication_group_id)
response, deleted = self.wait_for_deletion(
replication_group_id=replication_group_id,
initial_sleep_time=initial_sleep_time,
exponential_back_off_factor=exponential_back_off_factor,
max_retries=max_retries,
)
if not deleted:
raise AirflowException(f'Replication group could not be deleted. Response "{response}"')
return response
