Add two operators in AWS Providers: RedshiftResumeClusterOperator and RedshiftPauseClusterOperator

airflow/providers/amazon/aws/hooks/redshift.py

@@ -47,6 +47,7 @@ class ClusterStates(Enum):
     REBOOTING = 'rebooting'
     RENAMING = 'renaming'
     RESIZING = 'resizing'
+    NONEXISTENT = 'nonexistent'
 
 
 class RedshiftHook(AwsBaseHook):
@@ -68,7 +69,7 @@ def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
 
     # TODO: Wrap create_cluster_snapshot
-    def cluster_status(self, cluster_identifier: str) -> str:
+    def cluster_status(self, cluster_identifier: str) -> ClusterStates:
         """
         Return status of a cluster
 
@@ -81,9 +82,9 @@ def cluster_status(self, cluster_identifier: str) -> str:
         """
         try:
             response = self.get_conn().describe_clusters(ClusterIdentifier=cluster_identifier)['Clusters']
-            return response[0]['ClusterStatus'] if response else None
+            return ClusterStates(response[0]['ClusterStatus']) if response else None
         except self.get_conn().exceptions.ClusterNotFoundFault:
-            return 'cluster_not_found'
+            return ClusterStates.NONEXISTENT
 
     def delete_cluster(
         self,

airflow/providers/amazon/aws/operators/redshift_pause_cluster.py

@@ -53,6 +53,6 @@ def __init__(
     def execute(self, context):
         redshift_hook = RedshiftHook(aws_conn_id=self.aws_conn_id)
         self.log.info("Pausing Redshift cluster %s", self.cluster_identifier)
-        cluster_state = ClusterStates(redshift_hook.cluster_status(cluster_identifier=self.cluster_identifier))
+        cluster_state = redshift_hook.cluster_status(cluster_identifier=self.cluster_identifier)
         if cluster_state == ClusterStates.AVAILABLE:
             redshift_hook.get_conn().pause_cluster(ClusterIdentifier=self.cluster_identifier)

airflow/providers/amazon/aws/operators/redshift_resume_cluster.py

@@ -53,6 +53,6 @@ def __init__(
     def execute(self, context):
         redshift_hook = RedshiftHook(aws_conn_id=self.aws_conn_id)
         self.log.info("Starting Redshift cluster %s", self.cluster_identifier)
-        cluster_state = ClusterStates(redshift_hook.cluster_status(cluster_identifier=self.cluster_identifier))
+        cluster_state = redshift_hook.cluster_status(cluster_identifier=self.cluster_identifier)
         if cluster_state == ClusterStates.PAUSED:
             redshift_hook.get_conn().resume_cluster(ClusterIdentifier=self.cluster_identifier)

airflow/providers/amazon/aws/sensors/redshift.py

@@ -17,7 +17,7 @@
 # under the License.
 from typing import Optional
 
-from airflow.providers.amazon.aws.hooks.redshift import RedshiftHook
+from airflow.providers.amazon.aws.hooks.redshift import ClusterStates, RedshiftHook
 from airflow.sensors.base import BaseSensorOperator
 
 
@@ -43,7 +43,10 @@ def __init__(
     ):
         super().__init__(**kwargs)
         self.cluster_identifier = cluster_identifier
-        self.target_status = target_status
+        self.target_status = (
+            target_status if isinstance(target_status, ClusterStates) else ClusterStates(str(target_status))
+        )
+
         self.aws_conn_id = aws_conn_id
         self.hook: Optional[RedshiftHook] = None
 

tests/providers/amazon/aws/hooks/test_redshift.py

@@ -26,7 +26,7 @@
 
 from airflow.models import Connection
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
-from airflow.providers.amazon.aws.hooks.redshift import RedshiftHook, RedshiftSQLHook
+from airflow.providers.amazon.aws.hooks.redshift import ClusterStates, RedshiftHook, RedshiftSQLHook
 
 try:
     from moto import mock_redshift
@@ -98,15 +98,15 @@ def test_cluster_status_returns_cluster_not_found(self):
         self._create_clusters()
         hook = RedshiftHook(aws_conn_id='aws_default')
         status = hook.cluster_status('test_cluster_not_here')
-        assert status == 'cluster_not_found'
+        assert status == ClusterStates.NONEXISTENT
 
     @unittest.skipIf(mock_redshift is None, 'mock_redshift package not present')
     @mock_redshift
     def test_cluster_status_returns_available_cluster(self):
         self._create_clusters()
         hook = RedshiftHook(aws_conn_id='aws_default')
         status = hook.cluster_status('test_cluster')
-        assert status == 'available'
+        assert status == ClusterStates.AVAILABLE
 
 
 class TestRedshiftSQLHookConn(unittest.TestCase):

tests/providers/amazon/aws/sensors/test_redshift.py

@@ -21,6 +21,7 @@
 
 import boto3
 
+from airflow.providers.amazon.aws.hooks.redshift import ClusterStates
 from airflow.providers.amazon.aws.sensors.redshift import AwsRedshiftClusterSensor
 
 try:
@@ -56,6 +57,20 @@ def test_poke(self):
         )
         assert op.poke(None)
 
+    @unittest.skipIf(mock_redshift is None, 'mock_redshift package not present')
+    @mock_redshift
+    def test_poke_with_cluster_state(self):
+        self._create_cluster()
+        op = AwsRedshiftClusterSensor(
+            task_id='test_cluster_sensor',
+            poke_interval=1,
+            timeout=5,
+            aws_conn_id='aws_default',
+            cluster_identifier='test_cluster',
+            target_status=ClusterStates.AVAILABLE,
+        )
+        assert op.poke(None)
+
     @unittest.skipIf(mock_redshift is None, 'mock_redshift package not present')
     @mock_redshift
     def test_poke_false(self):
@@ -81,7 +96,7 @@ def test_poke_cluster_not_found(self):
             timeout=5,
             aws_conn_id='aws_default',
             cluster_identifier='test_cluster_not_found',
-            target_status='cluster_not_found',
+            target_status=ClusterStates.NONEXISTENT,
         )
 
         assert op.poke(None)