Cancel the Redshift statement when a user kills the deferred task

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
 
     def execute(self, context: Context) -> list[GetStatementResultResponseTypeDef] | list[str]:
         """Execute a statement against Amazon Redshift."""
@@ -170,6 +175,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",
                 )
@@ -224,6 +230,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)

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,21 @@ 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."""
+        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:
+                await client.cancel_statement(Id=self.statement_id)
+            self.log.info("Redshift statement %s cancelled.", 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:

providers/amazon/tests/unit/amazon/aws/operators/test_redshift_data.py

@@ -277,6 +277,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}]
@@ -382,6 +399,7 @@ 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
 
     def test_execute_complete_failure(self, deferrable_operator):
         """Tests that an AirflowException is raised in case of error event"""

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,64 @@ 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
+    @mock.patch.object(RedshiftDataHook, "get_async_conn")
+    async def test_on_kill_cancels_the_statement(self, mock_get_async_conn):
+        """on_kill() cancels the running statement when enabled and a statement_id is set."""
+        mock_client = mock.AsyncMock()
+        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")