VijayantSoni commented on a change in pull request #8701: URL: https://github.com/apache/airflow/pull/8701#discussion_r470951883
########## File path: airflow/providers/amazon/aws/hooks/elasticache_replication_group.py ########## @@ -0,0 +1,289 @@ +# +# 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 time import sleep + +from airflow.exceptions import AirflowException +from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook + + +class ElastiCacheReplicationGroupHook(AwsBaseHook): + """ + Interact with AWS ElastiCache + """ + + TERMINAL_STATES = frozenset({"available", "create-failed", "deleting"}) + + def __init__( + self, max_retries=10, exponential_back_off_factor=1, initial_poke_interval=60, *args, **kwargs + ): + """ + :param max_retries: Max retries for checking availability of and deleting replication group + :type max_retries: int + :param exponential_back_off_factor: Factor for deciding next sleep time + :type exponential_back_off_factor: float + :param initial_poke_interval: Initial sleep time in seconds + :type initial_poke_interval: float + """ + self.max_retries = max_retries + self.exponential_back_off_factor = exponential_back_off_factor + self.initial_poke_interval = initial_poke_interval + + super().__init__(client_type='elasticache', *args, **kwargs) + + def create_replication_group(self, config): + """ + Call ElastiCache API for creating a replication group + + :param config: Configuration for creating the replication group + :type config: dict + :return: Response from ElastiCache create replication group API + :rtype: dict + """ + return self.conn.create_replication_group(**config) + + def delete_replication_group(self, replication_group_id): + """ + Call ElastiCache API for deleting a replication group + + :param replication_group_id: ID of replication group to delete + :type replication_group_id: str + :return: Response from ElastiCache delete replication group API + :rtype: dict + """ + return self.conn.delete_replication_group(ReplicationGroupId=replication_group_id) + + def describe_replication_group(self, replication_group_id): + """ + Call ElastiCache API for describing a replication group + + :param replication_group_id: ID of replication group to describe + :type replication_group_id: str + :return: Response from ElastiCache describe replication group API + :rtype: dict + """ + return self.conn.describe_replication_groups(ReplicationGroupId=replication_group_id) + + def get_replication_group_status(self, replication_group_id): + """ + Get current status of replication group + + :param replication_group_id: ID of replication group to check for status + :type replication_group_id: str + :return: Current status of replication group + :rtype: str + """ + return self.describe_replication_group(replication_group_id)['ReplicationGroups'][0]['Status'] + + def is_replication_group_available(self, replication_group_id): + """ + Helper for checking if replication group is available or not + + :param replication_group_id: ID of replication group to check for availability + :type replication_group_id: str + :return: True if available else False + :rtype: bool + """ + return self.get_replication_group_status(replication_group_id) == 'available' + + def _has_reached_terminal_state(self, replication_group_id): + """ + Helper for checking if we should stop poking replication group for availability or not + + :param replication_group_id: ID of replication group to check for terminal state + :type replication_group_id: str + :return: Flag to check if availability-check should be stopped or not and current status + :rtype: (bool, str) + """ + status = self.get_replication_group_status(replication_group_id=replication_group_id) + + return status in self.TERMINAL_STATES, status Review comment: I think it makes sense. Since it was only getting used at one place, we can write 2 statements (as you suggested) and remove that method which was returning multiple values. Now that you pointed out, I think it was not correct because the method name was `_has_reached_terminal_state` but it was returning `status` also along with the bolen flag which I think was not correct. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected]
