Add two operators in AWS Providers: RedshiftResumeClusterOperator and RedshiftPauseClusterOperator

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

@@ -16,14 +16,16 @@
 # specific language governing permissions and limitations
 # under the License.
 """Interact with AWS Redshift clusters."""
-
+import time
 from typing import Dict, List, Optional, Union
 
 try:
     from functools import cached_property
 except ImportError:
     from cached_property import cached_property
 
+from enum import Enum
+
 import redshift_connector
 from redshift_connector import Connection as RedshiftConnection
 from sqlalchemy import create_engine
@@ -33,6 +35,21 @@
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 
 
+class ClusterStates(Enum):
+    """Contains the possible State values of a Redshift Cluster."""
+
+    AVAILABLE = 'available'
+    CREATING = 'creating'
+    DELETING = 'deleting'
+    RESUMING = 'resuming'
+    MODIFYING = 'modifying'
+    PAUSED = 'paused'
+    REBOOTING = 'rebooting'
+    RENAMING = 'renaming'
+    RESIZING = 'resizing'
+    NONEXISTENT = 'nonexistent'
+
+
 class RedshiftHook(AwsBaseHook):
     """
     Interact with AWS Redshift, using the boto3 library
@@ -52,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
 
@@ -65,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

@@ -0,0 +1,58 @@
+#
+# 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 airflow.models import BaseOperator
+from airflow.providers.amazon.aws.hooks.redshift import ClusterStates, RedshiftHook
+
+
+class RedshiftPauseClusterOperator(BaseOperator):
+    """
+    Pause an AWS Redshift Cluster using boto3.
+
+    :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
+    :param check_interval: time in seconds that the job should wait in
+        between each instance state checks until operation is completed
+    :type check_interval: float
+    """
+
+    template_fields = ("cluster_identifier",)
+    ui_color = "#eeaa11"
+    ui_fgcolor = "#ffffff"
+
+    def __init__(
+        self,
+        *,
+        cluster_identifier: str,
+        aws_conn_id: str = "aws_default",
+        check_interval: float = 15,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.cluster_identifier = cluster_identifier
+        self.aws_conn_id = aws_conn_id
+        self.check_interval = check_interval
+
+    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 == ClusterStates.AVAILABLE:
+            redshift_hook.get_conn().pause_cluster(ClusterIdentifier=self.cluster_identifier)

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

@@ -0,0 +1,58 @@
+#
+# 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 airflow.models import BaseOperator
+from airflow.providers.amazon.aws.hooks.redshift import ClusterStates, RedshiftHook
+
+
+class RedshiftResumeClusterOperator(BaseOperator):
+    """
+    Resume an AWS Redshift Cluster using boto3.
+
+    :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
+    :param check_interval: time in seconds that the job should wait in
+        between each instance state checks until operation is completed
+    :type check_interval: float
+    """
+
+    template_fields = ("cluster_identifier",)
+    ui_color = "#eeaa11"
+    ui_fgcolor = "#ffffff"
+
+    def __init__(
+        self,
+        *,
+        cluster_identifier: str,
+        aws_conn_id: str = "aws_default",
+        check_interval: float = 15,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.cluster_identifier = cluster_identifier
+        self.aws_conn_id = aws_conn_id
+        self.check_interval = check_interval
+
+    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 == 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
 
 
@@ -28,7 +28,7 @@ class AwsRedshiftClusterSensor(BaseSensorOperator):
     :param cluster_identifier: The identifier for the cluster being pinged.
     :type cluster_identifier: str
     :param target_status: The cluster status desired.
-    :type target_status: str
+    :type target_status: ClusterStates
     """
 
     template_fields = ('cluster_identifier', 'target_status')
@@ -37,18 +37,21 @@ def __init__(
         self,
         *,
         cluster_identifier: str,
-        target_status: str = 'available',
+        target_status: ClusterStates = ClusterStates.AVAILABLE,
         aws_conn_id: str = 'aws_default',
         **kwargs,
     ):
         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
 
     def poke(self, context):
-        self.log.info('Poking for status : %s\nfor cluster %s', self.target_status, self.cluster_identifier)
+        self.log.info('Poking for status : %s\nfor cluster %s', self.target_status.value, self.cluster_identifier)
         return self.get_hook().cluster_status(self.cluster_identifier) == self.target_status
 
     def get_hook(self) -> RedshiftHook:

airflow/providers/amazon/provider.yaml

@@ -245,6 +245,8 @@ operators:
   - integration-name: Amazon Redshift
     python-modules:
       - airflow.providers.amazon.aws.operators.redshift
+      - airflow.providers.amazon.aws.operators.redshift_pause_cluster
+      - airflow.providers.amazon.aws.operators.redshift_resume_cluster
 
 sensors:
   - integration-name: Amazon Athena

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

@@ -15,17 +15,19 @@
     specific language governing permissions and limitations
     under the License.
 
-.. _howto/operator:RedshiftSQLOperator:
+Amazon Redshift Operators
+=================================================
 
-RedshiftSQLOperator
-===================
+`Amazon Redshift <https://aws.amazon.com/redshift/>`__ is a fully managed,
+petabyte-scale data warehouse service in the cloud. You can start with just a few hundred gigabytes
+of data and scale to a petabyte or more. This enables you to use your data to acquire new insights
+for your business and customers.
 
-.. contents::
-  :depth: 1
-  :local:
+Airflow provides operators to create and interact with the Redshift clusters and compute infrastructure.
 
-Overview
---------
+RedshiftSQLOperator
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+.. _howto/operator:RedshiftSQLOperator:
 
 Use the :class:`RedshiftSQLOperator <airflow.providers.amazon.aws.operators.redshift>` to execute
 statements against an Amazon Redshift cluster.
@@ -34,16 +36,6 @@ statements against an Amazon Redshift cluster.
 :class:`RedshiftSQLHook <airflow.providers.amazon.aws.hooks.redshift.RedshiftSQLHook>` to establish
 connections with Amazon Redshift.
 
-
-example_redshift.py
--------------------
-
-Purpose
-"""""""
-
-This is a basic example dag for using :class:`RedshiftSQLOperator <airflow.providers.amazon.aws.operators.redshift>`
-to execute statements against an Amazon Redshift cluster.
-
 Create a table
 """"""""""""""
 
@@ -86,11 +78,28 @@ parameters into SQL statements.
     :end-before: [END howto_operator_redshift_get_with_filter]
 
 The complete RedshiftSQLOperator DAG
-------------------------------------
+"""""""""""""""""""""""""""""""""""""""""""
 
 All together, here is our DAG:
 
 .. exampleinclude:: /../../airflow/providers/amazon/aws/example_dags/example_redshift.py
     :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 an existing AWS Redshift Cluster you can use
+:class:`~airflow.providers.amazon.aws.operators.redshift_resume_cluster.RedshiftResumeClusterOperator`.
+
+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 existing AWS Redshift Cluster you can use
+:class:`~airflow.providers.amazon.aws.operators.redshift_pause_cluster.RedshiftPauseClusterOperator`.
+
+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/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/operators/test_redshift_pause_cluster.py

@@ -0,0 +1,61 @@
+#
+# 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.
+#
+
+import unittest
+
+import boto3
+
+from airflow.providers.amazon.aws.operators.redshift_pause_cluster import RedshiftPauseClusterOperator
+
+try:
+    from moto import mock_redshift
+except ImportError:
+    mock_redshift = None
+
+
+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",
+            check_interval=3,
+        )
+        assert redshift_operator.task_id == "task_test"
+        assert redshift_operator.cluster_identifier == "test_cluster"
+        assert redshift_operator.aws_conn_id == "aws_conn_test"
+        assert redshift_operator.check_interval == 3