Add two operators in AWS Providers: RedshiftResumeClusterOperator and RedshiftPauseClusterOperator

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

@@ -18,7 +18,7 @@
 from typing import Dict, Iterable, Optional, Union
 
 from airflow.models import BaseOperator
-from airflow.providers.amazon.aws.hooks.redshift import RedshiftSQLHook
+from airflow.providers.amazon.aws.hooks.redshift import RedshiftHook, RedshiftSQLHook
 
 
 class RedshiftSQLOperator(BaseOperator):
@@ -71,3 +71,85 @@ def execute(self, context: dict) -> None:
         self.log.info(f"Executing statement: {self.sql}")
         hook = self.get_hook()
         hook.run(self.sql, autocommit=self.autocommit, parameters=self.parameters)
+
+
+class RedshiftResumeClusterOperator(BaseOperator):
+    """
+    Resume a paused AWS Redshift Cluster
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:RedshiftResumeClusterOperator`
+
+    :param cluster_identifier: id of the AWS Redshift Cluster
+    :type cluster_identifier: str
+    :param aws_conn_id: aws connection to use
+    :type aws_conn_id: str
+    """
+
+    template_fields = ("cluster_identifier",)
+    ui_color = "#eeaa11"
+    ui_fgcolor = "#ffffff"
+
+    def __init__(
+        self,
+        *,
+        cluster_identifier: str,
+        aws_conn_id: str = "aws_default",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.cluster_identifier = cluster_identifier
+        self.aws_conn_id = aws_conn_id
+
+    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 = redshift_hook.cluster_status(cluster_identifier=self.cluster_identifier)
+        if cluster_state == 'paused':
+            redshift_hook.get_conn().resume_cluster(ClusterIdentifier=self.cluster_identifier)
+        else:
+            self.log.warning(
+                "Unable to resume cluster since cluster is currently in status: %s", cluster_state
+            )
+
+
+class RedshiftPauseClusterOperator(BaseOperator):
+    """
+    Pause an AWS Redshift Cluster if it has status `available`.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:RedshiftPauseClusterOperator`
+
+    :param cluster_identifier: id of the AWS Redshift Cluster
+    :type cluster_identifier: str
+    :param aws_conn_id: aws connection to use
+    :type aws_conn_id: str
+    """
+
+    template_fields = ("cluster_identifier",)
+    ui_color = "#eeaa11"
+    ui_fgcolor = "#ffffff"
+
+    def __init__(
+        self,
+        *,
+        cluster_identifier: str,
+        aws_conn_id: str = "aws_default",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.cluster_identifier = cluster_identifier
+        self.aws_conn_id = aws_conn_id
+
+    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 = redshift_hook.cluster_status(cluster_identifier=self.cluster_identifier)
+        if cluster_state == 'available':
+            redshift_hook.get_conn().pause_cluster(ClusterIdentifier=self.cluster_identifier)
+        else:
+            self.log.warning(
+                "Unable to pause cluster since cluster is currently in status: %s", cluster_state
+            )

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

@@ -48,7 +48,7 @@ def __init__(
         self.hook: Optional[RedshiftHook] = None
 
     def poke(self, context):
-        self.log.info('Poking for status : %s\nfor cluster %s', self.target_status, self.cluster_identifier)
+        self.log.info('Checking cluster %r for status %r', self.cluster_identifier, self.target_status)
         return self.get_hook().cluster_status(self.cluster_identifier) == self.target_status
 
     def get_hook(self) -> RedshiftHook:

docs/apache-airflow-providers-amazon/operators/redshift.rst

@@ -94,3 +94,25 @@ All together, here is our DAG:
     :language: python
     :start-after: [START redshift_operator_howto_guide]
     :end-before: [END redshift_operator_howto_guide]
+
+
+.. _howto/operator:RedshiftResumeClusterOperator:
+
+Resume a Redshift Cluster
+"""""""""""""""""""""""""""""""""""""""""""
+
+To resume a 'paused' AWS Redshift Cluster you can use
+:class:`RedshiftResumeClusterOperator <airflow.providers.amazon.aws.operators.redshift>`
+
+This Operator leverages the AWS CLI
+`resume-cluster <https://docs.aws.amazon.com/cli/latest/reference/redshift/resume-cluster.html>`__ API
+
+.. _howto/operator:RedshiftPauseClusterOperator:
+
+Pause a Redshift Cluster
+"""""""""""""""""""""""""""""""""""""""""""
+
+To pause an 'available' AWS Redshift Cluster you can use
+:class:`RedshiftPauseClusterOperator <airflow.providers.amazon.aws.operators.redshift>`
+This Operator leverages the AWS CLI
+`pause-cluster <https://docs.aws.amazon.com/cli/latest/reference/redshift/pause-cluster.html>`__ API

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

@@ -20,9 +20,14 @@
 from unittest import mock
 from unittest.mock import MagicMock
 
+import boto3
 from parameterized import parameterized
 
-from airflow.providers.amazon.aws.operators.redshift import RedshiftSQLOperator
+from airflow.providers.amazon.aws.operators.redshift import (
+    RedshiftPauseClusterOperator,
+    RedshiftResumeClusterOperator,
+    RedshiftSQLOperator,
+)
 
 
 class TestRedshiftSQLOperator(unittest.TestCase):
@@ -42,3 +47,69 @@ def test_redshift_operator(self, test_autocommit, test_parameters, mock_get_hook
             autocommit=test_autocommit,
             parameters=test_parameters,
         )
+
+
+class TestResumeClusterOperator(unittest.TestCase):
+    @staticmethod
+    def _create_clusters():
+        client = boto3.client('redshift', region_name='us-east-1')
+        client.create_cluster(
+            ClusterIdentifier='test_cluster_to_pause',
+            NodeType='dc1.large',
+            MasterUsername='admin',
+            MasterUserPassword='mock_password',
+        )
+        client.create_cluster(
+            ClusterIdentifier='test_cluster_to_resume',
+            NodeType='dc1.large',
+            MasterUsername='admin',
+            MasterUserPassword='mock_password',
+        )
+        if not client.describe_clusters()['Clusters']:
+            raise ValueError('AWS not properly mocked')
+
+    def test_init(self):
+        redshift_operator = RedshiftResumeClusterOperator(
+            task_id="task_test", cluster_identifier="test_cluster", aws_conn_id="aws_conn_test"
+        )
+        assert redshift_operator.task_id == "task_test"
+        assert redshift_operator.cluster_identifier == "test_cluster"
+        assert redshift_operator.aws_conn_id == "aws_conn_test"
+
+    def test_resume_cluster(self):
+        # TODO: Add test once moto library supports pause_cluster() or resume_cluster() boto api calls:
+        #  https://github.com/spulec/moto/issues/4591
+        pass
+
+
+class TestPauseClusterOperator(unittest.TestCase):
+    @staticmethod
+    def _create_clusters():
+        client = boto3.client('redshift', region_name='us-east-1')
+        client.create_cluster(
+            ClusterIdentifier='test_cluster_to_pause',
+            NodeType='dc1.large',
+            MasterUsername='admin',
+            MasterUserPassword='mock_password',
+        )
+        client.create_cluster(
+            ClusterIdentifier='test_cluster_to_resume',
+            NodeType='dc1.large',
+            MasterUsername='admin',
+            MasterUserPassword='mock_password',
+        )
+        if not client.describe_clusters()['Clusters']:
+            raise ValueError('AWS not properly mocked')
+
+    def test_init(self):
+        redshift_operator = RedshiftPauseClusterOperator(
+            task_id="task_test", cluster_identifier="test_cluster", aws_conn_id="aws_conn_test"
+        )
+        assert redshift_operator.task_id == "task_test"
+        assert redshift_operator.cluster_identifier == "test_cluster"
+        assert redshift_operator.aws_conn_id == "aws_conn_test"
+
+    def test_pause_cluster(self):
+        # TODO: Add test once moto library supports pause_cluster() or resume_cluster() boto api calls:
+        #  https://github.com/spulec/moto/issues/4591
+        pass