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,20 @@
 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'
+
+
 class RedshiftHook(AwsBaseHook):
     """
     Interact with AWS Redshift, using the boto3 library

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 = ClusterStates(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 = ClusterStates(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/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/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

tests/providers/amazon/aws/operators/test_redshift_resume_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_resume_cluster import RedshiftResumeClusterOperator
+
+try:
+    from moto import mock_redshift
+except ImportError:
+    mock_redshift = None
+
+
+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",
+            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