dstandish commented on a change in pull request #19665:
URL: https://github.com/apache/airflow/pull/19665#discussion_r764378031
##########
File path: tests/providers/amazon/aws/operators/test_redshift.py
##########
@@ -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):
Review comment:
we no longer subclass unittest.TestCase (no subclass required)
##########
File path: airflow/providers/amazon/aws/operators/redshift.py
##########
@@ -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)
Review comment:
```suggestion
cluster_state =
redshift_hook.cluster_status(cluster_identifier=self.cluster_identifier)
if cluster_state == 'paused':
self.log.info("Starting Redshift cluster %s",
self.cluster_identifier)
redshift_hook.get_conn().resume_cluster(ClusterIdentifier=self.cluster_identifier)
```
shoulld probably only say "starting" if it's actually gonna try to do it
##########
File path: airflow/providers/amazon/aws/operators/redshift.py
##########
@@ -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)
Review comment:
```suggestion
cluster_state =
redshift_hook.cluster_status(cluster_identifier=self.cluster_identifier)
if cluster_state == 'available':
self.log.info("Pausing Redshift cluster %s",
self.cluster_identifier)
redshift_hook.get_conn().pause_cluster(ClusterIdentifier=self.cluster_identifier)
```
##########
File path: tests/providers/amazon/aws/operators/test_redshift.py
##########
@@ -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):
Review comment:
you can (and should) actually add some reasonable tests for this
operator.
use standard mocking techniques
for the resume case you can at least verify that when the returned cluster
state is `paused`, that `resume_cluster` is called
similar for the pause case
##########
File path: tests/providers/amazon/aws/operators/test_redshift.py
##########
@@ -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')
Review comment:
are we using this? if not should chop it
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
