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 ea0d1272db8c6..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,6 +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: (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( @@ -856,6 +865,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 +875,26 @@ 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: + 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: self.hook.delete_cluster( cluster_identifier=self.cluster_identifier, @@ -882,51 +903,98 @@ 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 - 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 db6466125a035..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, @@ -785,14 +786,77 @@ 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") 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", @@ -806,16 +870,82 @@ 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", + 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, 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", @@ -823,10 +953,50 @@ 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, 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,