diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_data.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_data.py index 04dc8c3216297..810e5a839e195 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_data.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_data.py @@ -61,6 +61,9 @@ class RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]): :param session_id: the session identifier of the query :param session_keep_alive_seconds: duration in seconds to keep the session alive after the query finishes. The maximum time a session can keep alive is 24 hours + :param cancel_on_kill: If True (default), cancel the running Redshift statement when the task is + killed. This applies both while the operator is running and, for a deferred task, while it + waits in the triggerer. :param aws_conn_id: The Airflow connection used for AWS credentials. If this is ``None`` or empty then the default boto3 behaviour is used. If running Airflow in a distributed manner and aws_conn_id is None or @@ -104,6 +107,7 @@ def __init__( deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), session_id: str | None = None, session_keep_alive_seconds: int | None = None, + cancel_on_kill: bool = True, **kwargs, ) -> None: super().__init__(**kwargs) @@ -125,6 +129,7 @@ def __init__( self.deferrable = deferrable self.session_id = session_id self.session_keep_alive_seconds = session_keep_alive_seconds + self.cancel_on_kill = cancel_on_kill if self.deferrable and not self.wait_for_completion: self.log.warning( "deferrable=True and wait_for_completion=False are set; deferrable will be " @@ -175,6 +180,7 @@ def execute(self, context: Context) -> list[GetStatementResultResponseTypeDef] | region_name=self.region_name, verify=self.verify, botocore_config=self.botocore_config, + cancel_on_kill=self.cancel_on_kill, ), method_name="execute_complete", ) @@ -229,6 +235,8 @@ def get_sql_results( def on_kill(self) -> None: """Cancel the submitted redshift query.""" + if not self.cancel_on_kill: + return if hasattr(self, "statement_id"): self.log.info("Received a kill signal.") self.log.info("Stopping Query with statementId - %s", self.statement_id) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/redshift_data.py b/providers/amazon/src/airflow/providers/amazon/aws/triggers/redshift_data.py index c011fe75104e1..29ac1aceb6913 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/redshift_data.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/redshift_data.py @@ -41,6 +41,9 @@ class RedshiftDataTrigger(BaseTrigger): :param poll_interval: polling period in seconds to check for the status :param aws_conn_id: AWS connection ID for redshift :param region_name: aws region to use + :param cancel_on_kill: If True (default), cancel the running Redshift statement when the user + kills the deferred task (mark failed, clear, or mark success). Requires a version of + ``apache-airflow`` with ``BaseTrigger.on_kill()`` support; on older versions it is inert. """ def __init__( @@ -52,6 +55,7 @@ def __init__( region_name: str | None = None, verify: bool | str | None = None, botocore_config: dict | None = None, + cancel_on_kill: bool = True, ): super().__init__() self.statement_id = statement_id @@ -62,6 +66,7 @@ def __init__( self.region_name = region_name self.verify = verify self.botocore_config = botocore_config + self.cancel_on_kill = cancel_on_kill def serialize(self) -> tuple[str, dict[str, Any]]: """Serialize RedshiftDataTrigger arguments and classpath.""" @@ -75,6 +80,7 @@ def serialize(self) -> tuple[str, dict[str, Any]]: "region_name": self.region_name, "verify": self.verify, "botocore_config": self.botocore_config, + "cancel_on_kill": self.cancel_on_kill, }, ) @@ -87,6 +93,31 @@ def hook(self) -> RedshiftDataHook: config=self.botocore_config, ) + async def on_kill(self) -> None: + """Cancel the running Redshift statement when the user kills the deferred task.""" + # The triggerer invokes on_kill only on a user action (mark failed/success or clear), never on + # triggerer restart, redistribution, timeout, or normal completion, so cancelling here is safe. + if not self.cancel_on_kill or not self.statement_id: + return + self.log.info("Cancelling Redshift statement %s.", self.statement_id) + try: + async with await self.hook.get_async_conn() as client: + response = await client.cancel_statement(Id=self.statement_id) + # CancelStatement returns Status=False when Redshift declined the cancel (typically the + # statement already finished); surface that instead of claiming a successful cancel. + if response.get("Status"): + self.log.info("Redshift statement %s cancelled.", self.statement_id) + else: + self.log.warning( + "Redshift declined to cancel statement %s; it may have already finished.", + self.statement_id, + ) + except Exception: + self.log.exception( + "Failed to cancel Redshift statement %s. The query may still be running.", + self.statement_id, + ) + async def run(self) -> AsyncIterator[TriggerEvent]: try: while True: diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_data.py b/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_data.py index cb2f08db7e74c..255476ca8c036 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_data.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_data.py @@ -292,6 +292,23 @@ def test_on_kill_with_query(self, mock_conn): Id=STATEMENT_ID, ) + @mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.conn") + def test_on_kill_respects_cancel_on_kill_false(self, mock_conn): + mock_conn.execute_statement.return_value = {"Id": STATEMENT_ID, "SessionId": SESSION_ID} + operator = RedshiftDataOperator( + aws_conn_id=CONN_ID, + task_id=TASK_ID, + cluster_identifier="cluster_identifier", + sql=SQL, + database=DATABASE, + wait_for_completion=False, + cancel_on_kill=False, + ) + mock_ti = mock.MagicMock(name="MockedTaskInstance") + operator.execute({"ti": mock_ti}) + operator.on_kill() + mock_conn.cancel_statement.assert_not_called() + @mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.conn") def test_return_sql_result(self, mock_conn): expected_result = [{"Result": True}] @@ -397,6 +414,30 @@ def test_execute_defer(self, mock_exec_query, check_query_is_finished, deferrabl deferrable_operator.execute({"ti": mock_ti}) assert isinstance(exc.value.trigger, RedshiftDataTrigger) + assert exc.value.trigger.cancel_on_kill is True + + @mock.patch( + "airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.check_query_is_finished", + return_value=False, + ) + @mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.execute_query") + def test_execute_defer_propagates_cancel_on_kill_false(self, mock_exec_query, check_query_is_finished): + """cancel_on_kill=False on the operator reaches the trigger, so the deferred kill is a no-op.""" + operator = RedshiftDataOperator( + aws_conn_id=CONN_ID, + task_id=TASK_ID, + sql=SQL, + database=DATABASE, + cluster_identifier="cluster_identifier", + wait_for_completion=True, + poll_interval=5, + deferrable=True, + cancel_on_kill=False, + ) + with pytest.raises(TaskDeferred) as exc: + operator.execute({"ti": mock.MagicMock(name="MockedTaskInstance")}) + + assert exc.value.trigger.cancel_on_kill is False def test_execute_complete_failure(self, deferrable_operator): """Tests that an AirflowException is raised in case of error event""" diff --git a/providers/amazon/tests/unit/amazon/aws/triggers/test_redshift_data.py b/providers/amazon/tests/unit/amazon/aws/triggers/test_redshift_data.py index aa10a8f5ae1e1..d8bbcaccd3d45 100644 --- a/providers/amazon/tests/unit/amazon/aws/triggers/test_redshift_data.py +++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_redshift_data.py @@ -24,6 +24,7 @@ from airflow.providers.amazon.aws.hooks.redshift_data import ( ABORTED_STATE, FAILED_STATE, + RedshiftDataHook, RedshiftDataQueryAbortedError, RedshiftDataQueryFailedError, ) @@ -57,8 +58,21 @@ def test_redshift_data_trigger_serialization(self): "region_name": None, "botocore_config": None, "verify": None, + "cancel_on_kill": True, } + def test_redshift_data_trigger_serialization_cancel_on_kill_false(self): + """cancel_on_kill=False round-trips through serialization.""" + trigger = RedshiftDataTrigger( + statement_id="uuid", + task_id=TEST_TASK_ID, + aws_conn_id=TEST_CONN_ID, + poll_interval=POLL_INTERVAL, + cancel_on_kill=False, + ) + _, kwargs = trigger.serialize() + assert kwargs["cancel_on_kill"] is False + @pytest.mark.asyncio @pytest.mark.parametrize( ("return_value", "response"), @@ -151,3 +165,70 @@ async def test_redshift_data_trigger_exception( task = [i async for i in trigger.run()] assert len(task) == 1 assert TriggerEvent(expected_response) in task + + @pytest.mark.asyncio + @pytest.mark.parametrize("cancel_status", [True, False]) + @mock.patch.object(RedshiftDataHook, "get_async_conn") + async def test_on_kill_cancels_the_statement(self, mock_get_async_conn, cancel_status): + """on_kill() issues CancelStatement and consumes its Status flag for both outcomes. + + CancelStatement returns ``{"Status": bool}`` (False when Redshift declined, e.g. the + statement already finished); on_kill must read that shape without error either way. + """ + mock_client = mock.AsyncMock() + mock_client.cancel_statement.return_value = {"Status": cancel_status} + mock_cm = mock.AsyncMock() + mock_cm.__aenter__.return_value = mock_client + mock_get_async_conn.return_value = mock_cm + + trigger = RedshiftDataTrigger( + statement_id="uuid", + task_id=TEST_TASK_ID, + poll_interval=POLL_INTERVAL, + aws_conn_id=TEST_CONN_ID, + ) + await trigger.on_kill() + + mock_client.cancel_statement.assert_awaited_once_with(Id="uuid") + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("cancel_on_kill", "statement_id"), + [ + pytest.param(False, "uuid", id="disabled"), + pytest.param(True, "", id="no-statement-id"), + ], + ) + @mock.patch.object(RedshiftDataHook, "get_async_conn") + async def test_on_kill_does_not_cancel(self, mock_get_async_conn, cancel_on_kill, statement_id): + """on_kill() is a no-op (no connection opened) when disabled or without a statement_id.""" + trigger = RedshiftDataTrigger( + statement_id=statement_id, + task_id=TEST_TASK_ID, + poll_interval=POLL_INTERVAL, + aws_conn_id=TEST_CONN_ID, + cancel_on_kill=cancel_on_kill, + ) + await trigger.on_kill() + + mock_get_async_conn.assert_not_called() + + @pytest.mark.asyncio + @mock.patch.object(RedshiftDataHook, "get_async_conn") + async def test_on_kill_swallows_cancel_errors(self, mock_get_async_conn): + """on_kill() logs and swallows exceptions raised while cancelling.""" + mock_client = mock.AsyncMock() + mock_client.cancel_statement.side_effect = Exception("AWS API error") + mock_cm = mock.AsyncMock() + mock_cm.__aenter__.return_value = mock_client + mock_get_async_conn.return_value = mock_cm + + trigger = RedshiftDataTrigger( + statement_id="uuid", + task_id=TEST_TASK_ID, + poll_interval=POLL_INTERVAL, + aws_conn_id=TEST_CONN_ID, + ) + await trigger.on_kill() + + mock_client.cancel_statement.assert_awaited_once_with(Id="uuid")