From 15830f9af03cd302ba5bf4eed9e073fd38db2b7d Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Mon, 6 Jul 2026 19:56:33 +0000 Subject: [PATCH 1/2] Make RedshiftDeleteClusterOperator busy-retry window configurable The delete operator retries on InvalidClusterStateFault when the cluster is mid-transition (e.g. pausing/resuming). The retry budget was hardcoded to 10 attempts * 15s = 2.5 min, which expires long before a ~15 min pause settles into a deletable state, causing the task to fail and the cluster to leak. Promote the two hardcoded constants to __init__ params (busy_retry_attempts, busy_retry_interval) and raise the default to 60 * 15s = 15 min so the retry outlasts a pause. Both are overridable. poll_interval/max_attempts (completion waiter) are unchanged. Note: this synchronous retry loop still runs before the deferrable branch, so it blocks the worker in both sync and deferrable modes. Moving the busy retry into the async trigger for deferrable mode is a follow-up. Generated-by: Claude Code (Opus) --- .../amazon/aws/operators/redshift_cluster.py | 33 +++++++--- .../aws/operators/test_redshift_cluster.py | 65 ++++++++++++++++++- 2 files changed, 87 insertions(+), 11 deletions(-) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_cluster.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_cluster.py index ea0d1272db8c6..32d4ef00dff39 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_cluster.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_cluster.py @@ -837,6 +837,14 @@ class RedshiftDeleteClusterOperator(AwsBaseOperator[RedshiftHook]): :param poll_interval: Time (in seconds) to wait between two consecutive calls to check cluster state :param deferrable: Run operator in the deferrable mode. :param max_attempts: (Deferrable mode only) The maximum number of attempts to be made + :param busy_retry_attempts: Number of times to retry the delete when the cluster is busy and the + request cannot yet be accepted (an ``InvalidClusterStateFault`` is raised), e.g. while a + pause/resume is in progress. The default of 60 (combined with ``busy_retry_interval``) gives a + ~15 minute window, long enough to outlast a cluster pause. Note this retry is synchronous and + blocks the worker in both sync and deferrable modes (the async improvement for deferrable mode + is a follow-up). + :param busy_retry_interval: Time (in seconds) to wait between busy retries when the cluster is in a + transient state (``InvalidClusterStateFault``). See ``busy_retry_attempts``. """ template_fields: Sequence[str] = aws_template_fields( @@ -856,6 +864,8 @@ def __init__( poll_interval: int = 30, deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), max_attempts: int = 30, + busy_retry_attempts: int = 60, + busy_retry_interval: int = 15, **kwargs, ): super().__init__(**kwargs) @@ -864,16 +874,19 @@ def __init__( self.final_cluster_snapshot_identifier = final_cluster_snapshot_identifier self.wait_for_completion = wait_for_completion self.poll_interval = poll_interval - # These parameters are added to keep trying if there is a running operation in the cluster - # If there is a running operation in the cluster while trying to delete it, a InvalidClusterStateFault - # is thrown. In such case, retrying - self._attempts = 10 - self._attempt_interval = 15 + # These parameters are used to keep retrying if there is a running operation in the cluster. + # If there is a running operation in the cluster while trying to delete it (e.g. a pause/resume + # is in progress), an InvalidClusterStateFault is thrown. In such cases we retry until the + # cluster settles into a deletable state. The defaults give a ~15 minute window (60 * 15s), + # long enough to outlast a cluster pause. + self.busy_retry_attempts = busy_retry_attempts + self.busy_retry_interval = busy_retry_interval self.deferrable = deferrable self.max_attempts = max_attempts def execute(self, context: Context): - while self._attempts: + attempts_remaining = self.busy_retry_attempts + while attempts_remaining: try: self.hook.delete_cluster( cluster_identifier=self.cluster_identifier, @@ -882,18 +895,18 @@ def execute(self, context: Context): ) break except self.hook.conn.exceptions.InvalidClusterStateFault: - self._attempts -= 1 + attempts_remaining -= 1 - if self._attempts: + if attempts_remaining: current_state = self.hook.conn.describe_clusters( ClusterIdentifier=self.cluster_identifier )["Clusters"][0]["ClusterStatus"] self.log.error( "Cluster in %s state, unable to delete. %d attempts remaining.", current_state, - self._attempts, + attempts_remaining, ) - time.sleep(self._attempt_interval) + time.sleep(self.busy_retry_interval) else: raise diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_cluster.py b/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_cluster.py index db6466125a035..586d5ef55799e 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_cluster.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_cluster.py @@ -785,7 +785,70 @@ def test_delete_cluster_multiple_attempts_fail(self, _, mock_conn, mock_delete_c with pytest.raises(returned_exception): redshift_operator.execute(None) - assert mock_delete_cluster.call_count == 10 + assert mock_delete_cluster.call_count == 60 + + def test_busy_retry_defaults(self): + redshift_operator = RedshiftDeleteClusterOperator( + task_id="task_test", + cluster_identifier="test_cluster", + aws_conn_id="aws_conn_test", + ) + assert redshift_operator.busy_retry_attempts == 60 + assert redshift_operator.busy_retry_interval == 15 + + def test_busy_retry_custom_values(self): + redshift_operator = RedshiftDeleteClusterOperator( + task_id="task_test", + cluster_identifier="test_cluster", + aws_conn_id="aws_conn_test", + busy_retry_attempts=3, + busy_retry_interval=0, + ) + assert redshift_operator.busy_retry_attempts == 3 + assert redshift_operator.busy_retry_interval == 0 + + @mock.patch.object(RedshiftHook, "delete_cluster") + @mock.patch.object(RedshiftHook, "conn") + @mock.patch("time.sleep", return_value=None) + def test_delete_cluster_custom_busy_retry_attempts_fail(self, _, mock_conn, mock_delete_cluster): + exception = boto3.client("redshift").exceptions.InvalidClusterStateFault({}, "test") + returned_exception = type(exception) + mock_conn.exceptions.InvalidClusterStateFault = returned_exception + mock_delete_cluster.side_effect = exception + + redshift_operator = RedshiftDeleteClusterOperator( + task_id="task_test", + cluster_identifier="test_cluster", + aws_conn_id="aws_conn_test", + wait_for_completion=False, + busy_retry_attempts=3, + busy_retry_interval=0, + ) + with pytest.raises(returned_exception): + redshift_operator.execute(None) + + assert mock_delete_cluster.call_count == 3 + + @mock.patch.object(RedshiftHook, "delete_cluster") + @mock.patch.object(RedshiftHook, "conn") + @mock.patch("time.sleep", return_value=None) + def test_delete_cluster_succeeds_on_second_attempt(self, _, mock_conn, mock_delete_cluster): + exception = boto3.client("redshift").exceptions.InvalidClusterStateFault({}, "test") + returned_exception = type(exception) + mock_conn.exceptions.InvalidClusterStateFault = returned_exception + mock_delete_cluster.side_effect = [exception, True] + + redshift_operator = RedshiftDeleteClusterOperator( + task_id="task_test", + cluster_identifier="test_cluster", + aws_conn_id="aws_conn_test", + wait_for_completion=False, + busy_retry_attempts=3, + busy_retry_interval=0, + ) + redshift_operator.execute(None) + + assert mock_delete_cluster.call_count == 2 @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.cluster_status") @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.delete_cluster") From 737bb8796d1771cc5aba9c21118cf41ad8fcc0a0 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Mon, 6 Jul 2026 21:11:50 +0000 Subject: [PATCH 2/2] Make RedshiftDeleteClusterOperator busy-retry async in deferrable mode In deferrable mode the operator no longer runs the synchronous busy-retry loop that blocked the worker/triggerer for up to the ~15 minute busy-retry window. Instead it attempts the delete once and, if the cluster is mid-transition (InvalidClusterStateFault), defers to a new RedshiftClusterSettledTrigger that asynchronously polls until the cluster leaves every transitional lifecycle (creating/modifying/resizing/pausing/ resuming/deleting/rebooting/...). The _retry_delete_when_settled callback then re-issues the delete: on success it defers to the existing RedshiftDeleteClusterTrigger (cluster_deleted) -> execute_complete; on a race (still InvalidClusterStateFault) it re-defers to the settle-wait trigger, bounded by busy_retry_attempts. This mirrors the EksCreateCluster re-defer pattern. Sync mode (deferrable=False) keeps the existing synchronous busy-retry loop unchanged. A poll-until-settled trigger (rather than a single-target-status poller or a custom botocore "not-transitional" waiter) is used because a busy cluster settles into different terminal states depending on the in-flight operation (pausing -> paused, resizing -> available), and the delete is acceptable in any non-transitional lifecycle. Generated-by: Claude Code (Opus) --- docs/spelling_wordlist.txt | 1 + .../amazon/aws/operators/redshift_cluster.py | 119 ++++++++++++----- .../amazon/aws/triggers/redshift_cluster.py | 109 +++++++++++++++ .../aws/operators/test_redshift_cluster.py | 125 ++++++++++++++++-- .../aws/triggers/test_redshift_cluster.py | 120 +++++++++++++++++ 5 files changed, 433 insertions(+), 41 deletions(-) diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 6a9616c7c0dcd..2956a5d7ab99b 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -447,6 +447,7 @@ deferrable deidentify DeidentifyTemplate del +deletable delim deltalake denylist diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_cluster.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_cluster.py index 32d4ef00dff39..4537fcabe0121 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_cluster.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_cluster.py @@ -27,6 +27,7 @@ from airflow.providers.amazon.aws.hooks.redshift_cluster import RedshiftHook from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator from airflow.providers.amazon.aws.triggers.redshift_cluster import ( + RedshiftClusterSettledTrigger, RedshiftCreateClusterSnapshotTrigger, RedshiftCreateClusterTrigger, RedshiftDeleteClusterTrigger, @@ -837,14 +838,14 @@ class RedshiftDeleteClusterOperator(AwsBaseOperator[RedshiftHook]): :param poll_interval: Time (in seconds) to wait between two consecutive calls to check cluster state :param deferrable: Run operator in the deferrable mode. :param max_attempts: (Deferrable mode only) The maximum number of attempts to be made - :param busy_retry_attempts: Number of times to retry the delete when the cluster is busy and the - request cannot yet be accepted (an ``InvalidClusterStateFault`` is raised), e.g. while a - pause/resume is in progress. The default of 60 (combined with ``busy_retry_interval``) gives a - ~15 minute window, long enough to outlast a cluster pause. Note this retry is synchronous and - blocks the worker in both sync and deferrable modes (the async improvement for deferrable mode - is a follow-up). - :param busy_retry_interval: Time (in seconds) to wait between busy retries when the cluster is in a - transient state (``InvalidClusterStateFault``). See ``busy_retry_attempts``. + :param busy_retry_attempts: (Sync mode only) Number of times to retry the delete when the cluster is + busy and the request cannot yet be accepted (an ``InvalidClusterStateFault`` is raised), e.g. + while a pause/resume is in progress. The default of 60 (combined with ``busy_retry_interval``) + gives a ~15 minute window, long enough to outlast a cluster pause. This retry is synchronous and + blocks the worker; in deferrable mode the equivalent wait is performed asynchronously by the + triggerer instead (see ``deferrable``). + :param busy_retry_interval: (Sync mode only) Time (in seconds) to wait between busy retries when the + cluster is in a transient state (``InvalidClusterStateFault``). See ``busy_retry_attempts``. """ template_fields: Sequence[str] = aws_template_fields( @@ -885,6 +886,13 @@ def __init__( self.max_attempts = max_attempts def execute(self, context: Context): + if self.deferrable: + # In deferrable mode we must not block the worker with the synchronous busy-retry loop. + # Attempt the delete once; if the cluster is mid-transition (InvalidClusterStateFault), + # hand off to the triggerer to wait for it to settle and then re-issue the delete. + self._delete_or_defer_until_settled() + return + attempts_remaining = self.busy_retry_attempts while attempts_remaining: try: @@ -910,36 +918,83 @@ def execute(self, context: Context): else: raise - if self.deferrable: - cluster_state = self.hook.cluster_status(cluster_identifier=self.cluster_identifier) - if cluster_state == "cluster_not_found": - self.log.info("Cluster deleted successfully") - elif cluster_state in ("creating", "modifying"): - raise AirflowException( - f"Unable to delete cluster since cluster is currently in status: {cluster_state}" - ) - else: - self.defer( - timeout=timedelta(seconds=self.max_attempts * self.poll_interval + 60), - trigger=RedshiftDeleteClusterTrigger( - cluster_identifier=self.cluster_identifier, - waiter_delay=self.poll_interval, - waiter_max_attempts=self.max_attempts, - aws_conn_id=self.aws_conn_id, - region_name=self.region_name, - verify=self.verify, - botocore_config=self.botocore_config, - ), - method_name="execute_complete", - ) - - elif self.wait_for_completion: + if self.wait_for_completion: waiter = self.hook.conn.get_waiter("cluster_deleted") waiter.wait( ClusterIdentifier=self.cluster_identifier, WaiterConfig={"Delay": self.poll_interval, "MaxAttempts": self.max_attempts}, ) + def _delete_or_defer_until_settled(self) -> None: + """ + Issue the delete once (deferrable mode); defer to wait out a busy cluster if needed. + + If the delete is accepted, defer to :class:`RedshiftDeleteClusterTrigger` to wait for the + deletion to finish. If the cluster is mid-transition (``InvalidClusterStateFault``), defer to + :class:`RedshiftClusterSettledTrigger`, which fires once the cluster leaves every transitional + lifecycle; the ``_retry_delete_when_settled`` callback then re-issues the delete. + """ + try: + self.hook.delete_cluster( + cluster_identifier=self.cluster_identifier, + skip_final_cluster_snapshot=self.skip_final_cluster_snapshot, + final_cluster_snapshot_identifier=self.final_cluster_snapshot_identifier, + ) + except self.hook.conn.exceptions.InvalidClusterStateFault: + self.log.info( + "Cluster %s is busy; deferring until it settles into a deletable state.", + self.cluster_identifier, + ) + self.defer( + timeout=timedelta(seconds=self.busy_retry_attempts * self.busy_retry_interval + 60), + trigger=RedshiftClusterSettledTrigger( + cluster_identifier=self.cluster_identifier, + poke_interval=self.busy_retry_interval, + max_attempts=self.busy_retry_attempts, + aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, + ), + method_name="_retry_delete_when_settled", + ) + return + + self._defer_until_deleted() + + def _defer_until_deleted(self) -> None: + """Defer to the delete-completion waiter, short-circuiting if the cluster is already gone.""" + cluster_state = self.hook.cluster_status(cluster_identifier=self.cluster_identifier) + if cluster_state == "cluster_not_found": + self.log.info("Cluster deleted successfully") + return + self.defer( + timeout=timedelta(seconds=self.max_attempts * self.poll_interval + 60), + trigger=RedshiftDeleteClusterTrigger( + cluster_identifier=self.cluster_identifier, + waiter_delay=self.poll_interval, + waiter_max_attempts=self.max_attempts, + aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + botocore_config=self.botocore_config, + ), + method_name="execute_complete", + ) + + def _retry_delete_when_settled(self, context: Context, event: dict[str, Any] | None = None) -> None: + """ + Re-issue the delete once the cluster has settled, then defer until deletion completes. + + Callback for :class:`RedshiftClusterSettledTrigger`. If the delete is still rejected because of a + race (the cluster re-entered a transitional state), defer to the settle-wait trigger again. + """ + validated_event = validate_execute_complete_event(event) + if validated_event["status"] != "success": + raise AirflowException(f"Error waiting for cluster to become deletable: {validated_event}") + + self._delete_or_defer_until_settled() + def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None: validated_event = validate_execute_complete_event(event) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/redshift_cluster.py b/providers/amazon/src/airflow/providers/amazon/aws/triggers/redshift_cluster.py index ebd32a2b42388..42fa1e3911d61 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/redshift_cluster.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/redshift_cluster.py @@ -28,6 +28,28 @@ if TYPE_CHECKING: from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook +# Cluster lifecycle states during which a ``delete_cluster`` call is rejected with an +# ``InvalidClusterStateFault`` because an operation is already in flight. Once the cluster leaves all +# of these states it has "settled" into a deletable lifecycle (e.g. ``available`` or ``paused``) and the +# delete can be re-issued. +REDSHIFT_TRANSITIONAL_CLUSTER_STATES = frozenset( + { + "creating", + "deleting", + "modifying", + "rebooting", + "renaming", + "resizing", + "resuming", + "pausing", + "rotating-keys", + "updating-hsm", + "cancelling-resize", + "prep-for-resize", + "final-snapshot", + } +) + class RedshiftCreateClusterTrigger(AwsBaseWaiterTrigger): """ @@ -343,3 +365,90 @@ async def run(self) -> AsyncIterator[TriggerEvent]: await asyncio.sleep(self.poke_interval) except Exception as e: yield TriggerEvent({"status": "error", "message": str(e)}) + + +class RedshiftClusterSettledTrigger(BaseTrigger): + """ + Poll a Redshift cluster until it settles into a non-transitional (deletable) lifecycle. + + A ``delete_cluster`` call is rejected with ``InvalidClusterStateFault`` while an operation is in + flight (e.g. a pause or resize). Because a busy cluster can settle into *different* terminal states + depending on the in-flight operation (a ``pausing`` cluster becomes ``paused``; a ``resizing`` cluster + becomes ``available``), this trigger fires as soon as the status leaves every transitional state + rather than waiting for a single hardcoded ``target_status``. The operator then re-issues the delete. + + :param aws_conn_id: Reference to AWS connection id for redshift + :param cluster_identifier: unique identifier of a cluster + :param poke_interval: polling period in seconds to check for the status + :param max_attempts: maximum number of polls before emitting an error event + :param region_name: The AWS region where the cluster is. Used to build the hook. + :param verify: Whether or not to verify SSL certificates. Used to build the hook. + :param botocore_config: Configuration dictionary for the botocore client. Used to build the hook. + """ + + def __init__( + self, + *, + aws_conn_id: str | None, + cluster_identifier: str, + poke_interval: float = 30, + max_attempts: int = 30, + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, + ): + super().__init__() + self.aws_conn_id = aws_conn_id + self.cluster_identifier = cluster_identifier + self.poke_interval = poke_interval + self.max_attempts = max_attempts + self.region_name = region_name + self.verify = verify + self.botocore_config = botocore_config + + def serialize(self) -> tuple[str, dict[str, Any]]: + """Serialize RedshiftClusterSettledTrigger arguments and classpath.""" + return ( + "airflow.providers.amazon.aws.triggers.redshift_cluster.RedshiftClusterSettledTrigger", + { + "aws_conn_id": self.aws_conn_id, + "cluster_identifier": self.cluster_identifier, + "poke_interval": self.poke_interval, + "max_attempts": self.max_attempts, + "region_name": self.region_name, + "verify": self.verify, + "botocore_config": self.botocore_config, + }, + ) + + @cached_property + def hook(self) -> RedshiftHook: + return RedshiftHook( + aws_conn_id=self.aws_conn_id, + region_name=self.region_name, + verify=self.verify, + config=self.botocore_config, + ) + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Run async until the cluster leaves every transitional state (or is already gone).""" + try: + for _ in range(self.max_attempts): + status = await self.hook.cluster_status_async(self.cluster_identifier) + if status is None or status.lower() not in REDSHIFT_TRANSITIONAL_CLUSTER_STATES: + yield TriggerEvent( + {"status": "success", "message": "Cluster settled", "cluster_state": status} + ) + return + await asyncio.sleep(self.poke_interval) + yield TriggerEvent( + { + "status": "error", + "message": ( + f"Cluster {self.cluster_identifier} did not settle into a deletable state " + f"within {self.max_attempts} attempts." + ), + } + ) + except Exception as e: + yield TriggerEvent({"status": "error", "message": str(e)}) diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_cluster.py b/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_cluster.py index 586d5ef55799e..a185eb5d7742f 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_cluster.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_cluster.py @@ -34,6 +34,7 @@ RedshiftResumeClusterOperator, ) from airflow.providers.amazon.aws.triggers.redshift_cluster import ( + RedshiftClusterSettledTrigger, RedshiftCreateClusterSnapshotTrigger, RedshiftDeleteClusterTrigger, RedshiftPauseClusterTrigger, @@ -853,9 +854,9 @@ def test_delete_cluster_succeeds_on_second_attempt(self, _, mock_conn, mock_dele @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.cluster_status") @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.delete_cluster") def test_delete_cluster_deferrable_mode(self, mock_delete_cluster, mock_cluster_status): - """Test delete cluster operator with defer when deferrable param is true""" + """When the delete is accepted, deferrable mode waits for deletion to complete.""" mock_delete_cluster.return_value = True - mock_cluster_status.return_value = "available" + mock_cluster_status.return_value = "deleting" delete_cluster = RedshiftDeleteClusterOperator( task_id="task_test", cluster_identifier="test_cluster", @@ -869,16 +870,59 @@ def test_delete_cluster_deferrable_mode(self, mock_delete_cluster, mock_cluster_ assert isinstance(exc.value.trigger, RedshiftDeleteClusterTrigger), ( "Trigger is not a RedshiftDeleteClusterTrigger" ) + # Delete is attempted exactly once (no synchronous busy-retry loop in deferrable mode). + mock_delete_cluster.assert_called_once() - @mock.patch("airflow.providers.amazon.aws.operators.redshift_cluster.RedshiftDeleteClusterOperator.defer") @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.cluster_status") @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.delete_cluster") - def test_delete_cluster_deferrable_mode_in_paused_state( - self, mock_delete_cluster, mock_cluster_status, mock_defer + def test_delete_cluster_deferrable_mode_already_gone(self, mock_delete_cluster, mock_cluster_status): + """When the cluster is already gone after the delete, deferrable mode completes without deferring.""" + mock_delete_cluster.return_value = True + mock_cluster_status.return_value = "cluster_not_found" + delete_cluster = RedshiftDeleteClusterOperator( + task_id="task_test", + cluster_identifier="test_cluster", + deferrable=True, + wait_for_completion=False, + ) + + # No TaskDeferred is raised; the operator returns normally. + delete_cluster.execute(context=None) + mock_delete_cluster.assert_called_once() + + @mock.patch.object(RedshiftHook, "conn") + @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.delete_cluster") + def test_delete_cluster_deferrable_mode_busy_defers_to_settle_trigger( + self, mock_delete_cluster, mock_conn ): - """Test delete cluster operator with defer when deferrable param is true""" + """A busy cluster (InvalidClusterStateFault) defers to the settle-wait trigger, not a sync loop.""" + exception = boto3.client("redshift").exceptions.InvalidClusterStateFault({}, "test") + mock_conn.exceptions.InvalidClusterStateFault = type(exception) + mock_delete_cluster.side_effect = exception + + delete_cluster = RedshiftDeleteClusterOperator( + task_id="task_test", + cluster_identifier="test_cluster", + deferrable=True, + wait_for_completion=False, + ) + + with pytest.raises(TaskDeferred) as exc: + delete_cluster.execute(context=None) + + assert isinstance(exc.value.trigger, RedshiftClusterSettledTrigger), ( + "Trigger is not a RedshiftClusterSettledTrigger" + ) + assert exc.value.method_name == "_retry_delete_when_settled" + # Delete attempted once; the synchronous busy-retry loop never runs in deferrable mode. + mock_delete_cluster.assert_called_once() + + @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.cluster_status") + @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.delete_cluster") + def test_retry_delete_when_settled_reissues_delete(self, mock_delete_cluster, mock_cluster_status): + """The settle-wait callback re-issues the delete and defers to the delete-complete trigger.""" mock_delete_cluster.return_value = True - mock_cluster_status.return_value = "creating" + mock_cluster_status.return_value = "deleting" delete_cluster = RedshiftDeleteClusterOperator( task_id="task_test", cluster_identifier="test_cluster", @@ -886,10 +930,73 @@ def test_delete_cluster_deferrable_mode_in_paused_state( wait_for_completion=False, ) + with pytest.raises(TaskDeferred) as exc: + delete_cluster._retry_delete_when_settled( + context=None, event={"status": "success", "message": "Cluster settled"} + ) + + assert isinstance(exc.value.trigger, RedshiftDeleteClusterTrigger) + mock_delete_cluster.assert_called_once() + + @mock.patch.object(RedshiftHook, "conn") + @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.delete_cluster") + def test_retry_delete_when_settled_redefers_on_race(self, mock_delete_cluster, mock_conn): + """If the cluster re-enters a transitional state (race), the callback re-defers to settle-wait.""" + exception = boto3.client("redshift").exceptions.InvalidClusterStateFault({}, "test") + mock_conn.exceptions.InvalidClusterStateFault = type(exception) + mock_delete_cluster.side_effect = exception + + delete_cluster = RedshiftDeleteClusterOperator( + task_id="task_test", + cluster_identifier="test_cluster", + deferrable=True, + wait_for_completion=False, + ) + + with pytest.raises(TaskDeferred) as exc: + delete_cluster._retry_delete_when_settled( + context=None, event={"status": "success", "message": "Cluster settled"} + ) + + assert isinstance(exc.value.trigger, RedshiftClusterSettledTrigger) + assert exc.value.method_name == "_retry_delete_when_settled" + + def test_retry_delete_when_settled_error_event_raises(self): + """A non-success event from the settle-wait trigger raises AirflowException.""" + delete_cluster = RedshiftDeleteClusterOperator( + task_id="task_test", + cluster_identifier="test_cluster", + deferrable=True, + wait_for_completion=False, + ) with pytest.raises(AirflowException): - delete_cluster.execute(context=None) + delete_cluster._retry_delete_when_settled( + context=None, event={"status": "error", "message": "timed out"} + ) - assert not mock_defer.called + @mock.patch.object(RedshiftHook, "delete_cluster") + @mock.patch.object(RedshiftHook, "conn") + @mock.patch("time.sleep", return_value=None) + def test_delete_cluster_sync_mode_still_uses_busy_retry_loop( + self, mock_sleep, mock_conn, mock_delete_cluster + ): + """Sync mode (deferrable=False) must keep the synchronous busy-retry loop unchanged.""" + exception = boto3.client("redshift").exceptions.InvalidClusterStateFault({}, "test") + mock_conn.exceptions.InvalidClusterStateFault = type(exception) + mock_delete_cluster.side_effect = [exception, exception, True] + + redshift_operator = RedshiftDeleteClusterOperator( + task_id="task_test", + cluster_identifier="test_cluster", + aws_conn_id="aws_conn_test", + deferrable=False, + wait_for_completion=False, + ) + redshift_operator.execute(None) + + # Three synchronous attempts, and time.sleep was used between the retries. + assert mock_delete_cluster.call_count == 3 + assert mock_sleep.called def test_delete_cluster_execute_complete_success(self): """Asserts that logging occurs as expected""" diff --git a/providers/amazon/tests/unit/amazon/aws/triggers/test_redshift_cluster.py b/providers/amazon/tests/unit/amazon/aws/triggers/test_redshift_cluster.py index a55494ce8ab7c..a0a0f7fcc7e73 100644 --- a/providers/amazon/tests/unit/amazon/aws/triggers/test_redshift_cluster.py +++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_redshift_cluster.py @@ -23,6 +23,7 @@ import pytest from airflow.providers.amazon.aws.triggers.redshift_cluster import ( + RedshiftClusterSettledTrigger, RedshiftClusterTrigger, RedshiftCreateClusterSnapshotTrigger, RedshiftCreateClusterTrigger, @@ -147,6 +148,125 @@ async def test_redshift_cluster_sensor_trigger_exception(self, mock_cluster_stat assert TriggerEvent({"status": "error", "message": "Test exception"}) in task +class TestRedshiftClusterSettledTrigger: + def test_serialization(self): + """Asserts that RedshiftClusterSettledTrigger serializes its arguments and classpath.""" + trigger = RedshiftClusterSettledTrigger( + aws_conn_id="test_redshift_conn_id", + cluster_identifier="mock_cluster_identifier", + poke_interval=POLLING_PERIOD_SECONDS, + max_attempts=42, + ) + classpath, kwargs = trigger.serialize() + assert classpath == ( + "airflow.providers.amazon.aws.triggers.redshift_cluster.RedshiftClusterSettledTrigger" + ) + assert kwargs == { + "aws_conn_id": "test_redshift_conn_id", + "cluster_identifier": "mock_cluster_identifier", + "poke_interval": POLLING_PERIOD_SECONDS, + "max_attempts": 42, + "region_name": None, + "verify": None, + "botocore_config": None, + } + + def test_serializes_generic_hook_params(self): + """Asserts the generic AWS hook params are serialized and used to build the hook.""" + trigger = RedshiftClusterSettledTrigger( + aws_conn_id="test_redshift_conn_id", + cluster_identifier="mock_cluster_identifier", + poke_interval=POLLING_PERIOD_SECONDS, + region_name="eu-west-1", + verify=False, + botocore_config={"read_timeout": 42}, + ) + _, kwargs = trigger.serialize() + assert kwargs["region_name"] == "eu-west-1" + assert kwargs["verify"] is False + assert kwargs["botocore_config"] == {"read_timeout": 42} + + hook = trigger.hook + assert hook.aws_conn_id == "test_redshift_conn_id" + assert hook._region_name == "eu-west-1" + assert hook._verify is False + assert hook._config.read_timeout == 42 + + @pytest.mark.asyncio + @pytest.mark.parametrize("settled_status", ["available", "paused", "cluster_not_found"]) + @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.cluster_status_async") + async def test_fires_when_settled(self, mock_cluster_status, settled_status): + """Fires success as soon as the cluster leaves every transitional state.""" + mock_cluster_status.return_value = settled_status + trigger = RedshiftClusterSettledTrigger( + aws_conn_id="test_redshift_conn_id", + cluster_identifier="mock_cluster_identifier", + poke_interval=POLLING_PERIOD_SECONDS, + ) + actual = await trigger.run().asend(None) + assert actual == TriggerEvent( + {"status": "success", "message": "Cluster settled", "cluster_state": settled_status} + ) + + @pytest.mark.asyncio + @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.cluster_status_async") + async def test_fires_when_cluster_gone(self, mock_cluster_status): + """A missing cluster (status None) counts as settled.""" + mock_cluster_status.return_value = None + trigger = RedshiftClusterSettledTrigger( + aws_conn_id="test_redshift_conn_id", + cluster_identifier="mock_cluster_identifier", + poke_interval=POLLING_PERIOD_SECONDS, + ) + actual = await trigger.run().asend(None) + assert actual.payload["status"] == "success" + + @pytest.mark.asyncio + @mock.patch("asyncio.sleep") + @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.cluster_status_async") + async def test_keeps_polling_while_transitional(self, mock_cluster_status, mock_sleep): + """Keeps polling while the cluster is transitional, then fires when it settles.""" + mock_cluster_status.side_effect = ["pausing", "pausing", "paused"] + trigger = RedshiftClusterSettledTrigger( + aws_conn_id="test_redshift_conn_id", + cluster_identifier="mock_cluster_identifier", + poke_interval=POLLING_PERIOD_SECONDS, + ) + actual = await trigger.run().asend(None) + assert actual.payload["status"] == "success" + assert actual.payload["cluster_state"] == "paused" + assert mock_cluster_status.call_count == 3 + + @pytest.mark.asyncio + @mock.patch("asyncio.sleep") + @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.cluster_status_async") + async def test_errors_after_max_attempts(self, mock_cluster_status, mock_sleep): + """Emits an error event if the cluster never settles within max_attempts.""" + mock_cluster_status.return_value = "resizing" + trigger = RedshiftClusterSettledTrigger( + aws_conn_id="test_redshift_conn_id", + cluster_identifier="mock_cluster_identifier", + poke_interval=POLLING_PERIOD_SECONDS, + max_attempts=3, + ) + actual = await trigger.run().asend(None) + assert actual.payload["status"] == "error" + assert mock_cluster_status.call_count == 3 + + @pytest.mark.asyncio + @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.cluster_status_async") + async def test_error_on_exception(self, mock_cluster_status): + """Emits an error event when polling raises.""" + mock_cluster_status.side_effect = Exception("boom") + trigger = RedshiftClusterSettledTrigger( + aws_conn_id="test_redshift_conn_id", + cluster_identifier="mock_cluster_identifier", + poke_interval=POLLING_PERIOD_SECONDS, + ) + actual = await trigger.run().asend(None) + assert actual == TriggerEvent({"status": "error", "message": "boom"}) + + WAITER_TRIGGER_PARAMS = [ pytest.param( RedshiftCreateClusterTrigger,