Repository: apache/airflow

Issue:
Bring in more resolution to hivestats



Discussion:

Resolving patch:
diff --git a/airflow/providers/amazon/aws/operators/redshift_cluster.py b/airflow/providers/amazon/aws/operators/redshift_cluster.py
--- a/airflow/providers/amazon/aws/operators/redshift_cluster.py
+++ b/airflow/providers/amazon/aws/operators/redshift_cluster.py
@@ -25,6 +25,7 @@
 from airflow.providers.amazon.aws.hooks.redshift_cluster import RedshiftHook
 from airflow.providers.amazon.aws.triggers.redshift_cluster import (
     RedshiftClusterTrigger,
+    RedshiftCreateClusterSnapshotTrigger,
     RedshiftCreateClusterTrigger,
     RedshiftPauseClusterTrigger,
 )
@@ -306,6 +307,7 @@ class RedshiftCreateClusterSnapshotOperator(BaseOperator):
     :param max_attempt: The maximum number of attempts to be made to check the state
     :param aws_conn_id: The Airflow connection used for AWS credentials.
         The default connection id is ``aws_default``
+    :param deferrable: If True, the operator will run as a deferrable operator.
     """
 
     template_fields: Sequence[str] = (
@@ -324,6 +326,7 @@ def __init__(
         poll_interval: int = 15,
         max_attempt: int = 20,
         aws_conn_id: str = "aws_default",
+        deferrable: bool = False,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -334,6 +337,8 @@ def __init__(
         self.wait_for_completion = wait_for_completion
         self.poll_interval = poll_interval
         self.max_attempt = max_attempt
+        self.deferrable = deferrable
+        self.aws_conn_id = aws_conn_id
         self.redshift_hook = RedshiftHook(aws_conn_id=aws_conn_id)
 
     def execute(self, context: Context) -> Any:
@@ -350,6 +355,19 @@ def execute(self, context: Context) -> Any:
             retention_period=self.retention_period,
             tags=self.tags,
         )
+        if self.deferrable:
+            self.defer(
+                trigger=RedshiftCreateClusterSnapshotTrigger(
+                    cluster_identifier=self.cluster_identifier,
+                    poll_interval=self.poll_interval,
+                    max_attempts=self.max_attempt,
+                    aws_conn_id=self.aws_conn_id,
+                ),
+                method_name="execute_complete",
+                # timeout is set to ensure that if a trigger dies, the timeout does not restart
+                # 60 seconds is added to allow the trigger to exit gracefully (i.e. yield TriggerEvent)
+                timeout=timedelta(seconds=self.max_attempt * self.poll_interval + 60),
+            )
 
         if self.wait_for_completion:
             self.redshift_hook.get_conn().get_waiter("snapshot_available").wait(
@@ -360,6 +378,13 @@ def execute(self, context: Context) -> Any:
                 },
             )
 
+    def execute_complete(self, context, event=None):
+        if event["status"] != "success":
+            raise AirflowException(f"Error creating snapshot: {event}")
+        else:
+            self.log.info("Cluster snapshot created.")
+        return
+
 
 class RedshiftDeleteClusterSnapshotOperator(BaseOperator):
     """
diff --git a/airflow/providers/amazon/aws/triggers/redshift_cluster.py b/airflow/providers/amazon/aws/triggers/redshift_cluster.py
--- a/airflow/providers/amazon/aws/triggers/redshift_cluster.py
+++ b/airflow/providers/amazon/aws/triggers/redshift_cluster.py
@@ -210,3 +210,78 @@ async def run(self):
             )
         else:
             yield TriggerEvent({"status": "success", "message": "Cluster paused"})
+
+
+class RedshiftCreateClusterSnapshotTrigger(BaseTrigger):
+    """
+    Trigger for RedshiftCreateClusterSnapshotOperator.
+    The trigger will asynchronously poll the boto3 API and wait for the
+    Redshift cluster snapshot to be in the `available` state.
+
+    :param cluster_identifier:  A unique identifier for the cluster.
+    :param poll_interval: The amount of time in seconds to wait between attempts.
+    :param max_attempts: The maximum number of attempts to be made.
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+    """
+
+    def __init__(
+        self,
+        cluster_identifier: str,
+        poll_interval: int,
+        max_attempts: int,
+        aws_conn_id: str,
+    ):
+        self.cluster_identifier = cluster_identifier
+        self.poll_interval = poll_interval
+        self.max_attempts = max_attempts
+        self.aws_conn_id = aws_conn_id
+
+    def serialize(self) -> tuple[str, dict[str, Any]]:
+        return (
+            "airflow.providers.amazon.aws.triggers.redshift_cluster.RedshiftCreateClusterSnapshotTrigger",
+            {
+                "cluster_identifier": self.cluster_identifier,
+                "poll_interval": str(self.poll_interval),
+                "max_attempts": str(self.max_attempts),
+                "aws_conn_id": self.aws_conn_id,
+            },
+        )
+
+    @cached_property
+    def hook(self) -> RedshiftHook:
+        return RedshiftHook(aws_conn_id=self.aws_conn_id)
+
+    async def run(self):
+        async with self.hook.async_conn as client:
+            attempt = 0
+            waiter = client.get_waiter("snapshot_available")
+            while attempt < int(self.max_attempts):
+                attempt = attempt + 1
+                try:
+                    await waiter.wait(
+                        ClusterIdentifier=self.cluster_identifier,
+                        WaiterConfig={
+                            "Delay": int(self.poll_interval),
+                            "MaxAttempts": 1,
+                        },
+                    )
+                    break
+                except WaiterError as error:
+                    if "terminal failure" in str(error):
+                        yield TriggerEvent(
+                            {"status": "failure", "message": f"Create Cluster Snapshot Failed: {error}"}
+                        )
+                        break
+                    self.log.info(
+                        "Status of cluster snapshot is %s", error.last_response["Snapshots"][0]["Status"]
+                    )
+                    await asyncio.sleep(int(self.poll_interval))
+        if attempt >= int(self.max_attempts):
+            yield TriggerEvent(
+                {
+                    "status": "failure",
+                    "message": "Create Cluster Snapshot Cluster Failed - max attempts reached.",
+                }
+            )
+        else:
+            yield TriggerEvent({"status": "success", "message": "Cluster Snapshot Created"})


Test patch:
diff --git a/tests/providers/amazon/aws/operators/test_redshift_cluster.py b/tests/providers/amazon/aws/operators/test_redshift_cluster.py
--- a/tests/providers/amazon/aws/operators/test_redshift_cluster.py
+++ b/tests/providers/amazon/aws/operators/test_redshift_cluster.py
@@ -33,6 +33,7 @@
 )
 from airflow.providers.amazon.aws.triggers.redshift_cluster import (
     RedshiftClusterTrigger,
+    RedshiftCreateClusterSnapshotTrigger,
     RedshiftPauseClusterTrigger,
 )
 
@@ -136,10 +137,10 @@ def test_create_cluster_deferrable(self, mock_get_conn):
 
 
 class TestRedshiftCreateClusterSnapshotOperator:
-    @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.cluster_status")
-    @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.get_conn")
+    @mock.patch.object(RedshiftHook, "cluster_status")
+    @mock.patch.object(RedshiftHook, "conn")
     def test_create_cluster_snapshot_is_called_when_cluster_is_available(
-        self, mock_get_conn, mock_cluster_status
+        self, mock_conn, mock_cluster_status
     ):
         mock_cluster_status.return_value = "available"
         create_snapshot = RedshiftCreateClusterSnapshotOperator(
@@ -155,7 +156,7 @@ def test_create_cluster_snapshot_is_called_when_cluster_is_available(
             ],
         )
         create_snapshot.execute(None)
-        mock_get_conn.return_value.create_cluster_snapshot.assert_called_once_with(
+        mock_conn.create_cluster_snapshot.assert_called_once_with(
             ClusterIdentifier="test_cluster",
             SnapshotIdentifier="test_snapshot",
             ManualSnapshotRetentionPeriod=1,
@@ -167,9 +168,9 @@ def test_create_cluster_snapshot_is_called_when_cluster_is_available(
             ],
         )
 
-        mock_get_conn.return_value.get_waiter.assert_not_called()
+        mock_conn.get_waiter.assert_not_called()
 
-    @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.cluster_status")
+    @mock.patch.object(RedshiftHook, "cluster_status")
     def test_raise_exception_when_cluster_is_not_available(self, mock_cluster_status):
         mock_cluster_status.return_value = "paused"
         create_snapshot = RedshiftCreateClusterSnapshotOperator(
@@ -178,9 +179,9 @@ def test_raise_exception_when_cluster_is_not_available(self, mock_cluster_status
         with pytest.raises(AirflowException):
             create_snapshot.execute(None)
 
-    @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.cluster_status")
-    @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.get_conn")
-    def test_create_cluster_snapshot_with_wait(self, mock_get_conn, mock_cluster_status):
+    @mock.patch.object(RedshiftHook, "cluster_status")
+    @mock.patch.object(RedshiftHook, "conn")
+    def test_create_cluster_snapshot_with_wait(self, mock_conn, mock_cluster_status):
         mock_cluster_status.return_value = "available"
         create_snapshot = RedshiftCreateClusterSnapshotOperator(
             task_id="test_snapshot",
@@ -189,11 +190,28 @@ def test_create_cluster_snapshot_with_wait(self, mock_get_conn, mock_cluster_sta
             wait_for_completion=True,
         )
         create_snapshot.execute(None)
-        mock_get_conn.return_value.get_waiter.return_value.wait.assert_called_once_with(
+        mock_conn.get_waiter.return_value.wait.assert_called_once_with(
             ClusterIdentifier="test_cluster",
             WaiterConfig={"Delay": 15, "MaxAttempts": 20},
         )
 
+    @mock.patch.object(RedshiftHook, "cluster_status")
+    @mock.patch.object(RedshiftHook, "create_cluster_snapshot")
+    def test_create_cluster_snapshot_deferred(self, mock_create_cluster_snapshot, mock_cluster_status):
+        mock_cluster_status.return_value = "available"
+        mock_create_cluster_snapshot.return_value = True
+        create_snapshot = RedshiftCreateClusterSnapshotOperator(
+            task_id="test_snapshot",
+            cluster_identifier="test_cluster",
+            snapshot_identifier="test_snapshot",
+            deferrable=True,
+        )
+        with pytest.raises(TaskDeferred) as exc:
+            create_snapshot.execute(None)
+        assert isinstance(
+            exc.value.trigger, RedshiftCreateClusterSnapshotTrigger
+        ), "Trigger is not a RedshiftCreateClusterSnapshotTrigger"
+
 
 class TestRedshiftDeleteClusterSnapshotOperator:
     @mock.patch(
diff --git a/tests/providers/amazon/aws/triggers/test_redshift_cluster.py b/tests/providers/amazon/aws/triggers/test_redshift_cluster.py
--- a/tests/providers/amazon/aws/triggers/test_redshift_cluster.py
+++ b/tests/providers/amazon/aws/triggers/test_redshift_cluster.py
@@ -23,6 +23,7 @@
 
 from airflow.providers.amazon.aws.hooks.redshift_cluster import RedshiftHook
 from airflow.providers.amazon.aws.triggers.redshift_cluster import (
+    RedshiftCreateClusterSnapshotTrigger,
     RedshiftCreateClusterTrigger,
     RedshiftPauseClusterTrigger,
 )
@@ -217,3 +218,149 @@ async def test_redshift_pause_cluster_trigger_run_attempts_failed(
         assert response == TriggerEvent(
             {"status": "failure", "message": f"Pause Cluster Failed: {error_failed}"}
         )
+
+
+class TestRedshiftCreateClusterSnapshotTrigger:
+    def test_redshift_create_cluster_snapshot_trigger_serialize(self):
+        redshift_create_cluster_trigger = RedshiftCreateClusterSnapshotTrigger(
+            cluster_identifier=TEST_CLUSTER_IDENTIFIER,
+            poll_interval=TEST_POLL_INTERVAL,
+            max_attempts=TEST_MAX_ATTEMPT,
+            aws_conn_id=TEST_AWS_CONN_ID,
+        )
+        class_path, args = redshift_create_cluster_trigger.serialize()
+        assert (
+            class_path
+            == "airflow.providers.amazon.aws.triggers.redshift_cluster.RedshiftCreateClusterSnapshotTrigger"
+        )
+        assert args["cluster_identifier"] == TEST_CLUSTER_IDENTIFIER
+        assert args["poll_interval"] == str(TEST_POLL_INTERVAL)
+        assert args["max_attempts"] == str(TEST_MAX_ATTEMPT)
+        assert args["aws_conn_id"] == TEST_AWS_CONN_ID
+
+    @pytest.mark.asyncio
+    @async_mock.patch.object(RedshiftHook, "async_conn")
+    async def test_redshift_create_cluster_snapshot_trigger_run(self, mock_async_conn):
+        mock = async_mock.MagicMock()
+        mock_async_conn.__aenter__.return_value = mock
+        mock.get_waiter().wait = AsyncMock()
+
+        redshift_create_cluster_trigger = RedshiftCreateClusterSnapshotTrigger(
+            cluster_identifier=TEST_CLUSTER_IDENTIFIER,
+            poll_interval=TEST_POLL_INTERVAL,
+            max_attempts=TEST_MAX_ATTEMPT,
+            aws_conn_id=TEST_AWS_CONN_ID,
+        )
+
+        generator = redshift_create_cluster_trigger.run()
+        response = await generator.asend(None)
+
+        assert response == TriggerEvent({"status": "success", "message": "Cluster Snapshot Created"})
+
+    @pytest.mark.asyncio
+    @async_mock.patch("asyncio.sleep")
+    @async_mock.patch.object(RedshiftHook, "async_conn")
+    async def test_redshift_create_cluster_snapshot_trigger_run_multiple_attempts(
+        self, mock_async_conn, mock_sleep
+    ):
+        mock = async_mock.MagicMock()
+        mock_async_conn.__aenter__.return_value = mock
+
+        mock = async_mock.MagicMock()
+        mock_async_conn.__aenter__.return_value = mock
+        error = WaiterError(
+            name="test_name",
+            reason="test_reason",
+            last_response={"Snapshots": [{"Status": "available"}]},
+        )
+        mock.get_waiter().wait.side_effect = AsyncMock(side_effect=[error, error, True])
+        mock_sleep.return_value = True
+
+        redshift_create_cluster_snapshot_trigger = RedshiftCreateClusterSnapshotTrigger(
+            cluster_identifier=TEST_CLUSTER_IDENTIFIER,
+            poll_interval=TEST_POLL_INTERVAL,
+            max_attempts=TEST_MAX_ATTEMPT,
+            aws_conn_id=TEST_AWS_CONN_ID,
+        )
+
+        generator = redshift_create_cluster_snapshot_trigger.run()
+        response = await generator.asend(None)
+
+        assert mock.get_waiter().wait.call_count == 3
+        assert response == TriggerEvent({"status": "success", "message": "Cluster Snapshot Created"})
+
+    @pytest.mark.asyncio
+    @async_mock.patch("asyncio.sleep")
+    @async_mock.patch.object(RedshiftHook, "async_conn")
+    async def test_redshift_create_cluster_snapshot_trigger_run_attempts_exceeded(
+        self, mock_async_conn, mock_sleep
+    ):
+        mock = async_mock.MagicMock()
+        mock_async_conn.__aenter__.return_value = mock
+
+        mock = async_mock.MagicMock()
+        mock_async_conn.__aenter__.return_value = mock
+        error = WaiterError(
+            name="test_name",
+            reason="test_reason",
+            last_response={"Snapshots": [{"Status": "available"}]},
+        )
+        mock.get_waiter().wait.side_effect = AsyncMock(side_effect=[error, error, True])
+        mock_sleep.return_value = True
+
+        redshift_create_cluster_snapshot_trigger = RedshiftCreateClusterSnapshotTrigger(
+            cluster_identifier=TEST_CLUSTER_IDENTIFIER,
+            poll_interval=TEST_POLL_INTERVAL,
+            max_attempts=2,
+            aws_conn_id=TEST_AWS_CONN_ID,
+        )
+
+        generator = redshift_create_cluster_snapshot_trigger.run()
+        response = await generator.asend(None)
+
+        assert mock.get_waiter().wait.call_count == 2
+        assert response == TriggerEvent(
+            {"status": "failure", "message": "Create Cluster Snapshot Cluster Failed - max attempts reached."}
+        )
+
+    @pytest.mark.asyncio
+    @async_mock.patch("asyncio.sleep")
+    @async_mock.patch.object(RedshiftHook, "async_conn")
+    async def test_redshift_create_cluster_snapshot_trigger_run_attempts_failed(
+        self, mock_async_conn, mock_sleep
+    ):
+        mock = async_mock.MagicMock()
+        mock_async_conn.__aenter__.return_value = mock
+
+        mock = async_mock.MagicMock()
+        mock_async_conn.__aenter__.return_value = mock
+        error_available = WaiterError(
+            name="test_name",
+            reason="test_reason",
+            last_response={"Snapshots": [{"Status": "available"}]},
+        )
+
+        error_failed = WaiterError(
+            name="test_name",
+            reason="Waiter encountered a terminal failure state:",
+            last_response={"Snapshots": [{"Status": "available"}]},
+        )
+        mock.get_waiter().wait.side_effect = AsyncMock(
+            side_effect=[error_available, error_available, error_failed]
+        )
+        mock_sleep.return_value = True
+
+        redshift_create_cluster_snapshot_trigger = RedshiftCreateClusterSnapshotTrigger(
+            cluster_identifier=TEST_CLUSTER_IDENTIFIER,
+            poll_interval=TEST_POLL_INTERVAL,
+            max_attempts=TEST_MAX_ATTEMPT,
+            aws_conn_id=TEST_AWS_CONN_ID,
+        )
+
+        generator = redshift_create_cluster_snapshot_trigger.run()
+        response = await generator.asend(None)
+
+        assert mock.get_waiter().wait.call_count == 3
+        assert response == TriggerEvent(
+            {"status": "failure", "message": f"Create Cluster Snapshot Failed: {error_failed}"}
+        )