From a742e90bda8c763ebd0b984f78b299c914811c7a Mon Sep 17 00:00:00 2001 From: Steve Ahn <38049807+steveahnahn@users.noreply.github.com> Date: Wed, 8 Jul 2026 13:56:18 -0700 Subject: [PATCH] Cancel the Redshift statement when a user kills the deferred task A deferred RedshiftDataOperator parks its query in the triggerer, so the operator's own on_kill no longer runs once the task is deferred. When a user marks that task failed, clears it, or marks it success, the trigger had no on_kill hook, so the Redshift statement kept running against the cluster or workgroup even though the operator already cancels the statement on kill in the non-deferred path. On Redshift Serverless a runaway statement keeps billing RPU-hours, and on a provisioned cluster it holds a WLM slot, until it finishes on its own. This adds on_kill to RedshiftDataTrigger to cancel the running statement when the user acts on the deferred task, matching the behaviour already shipped for the EMR, Dataproc, BigQuery, and Dataflow triggers. A cancel_on_kill flag on both the operator and the trigger lets users opt out. --- .../amazon/aws/operators/redshift_data.py | 8 ++ .../amazon/aws/triggers/redshift_data.py | 21 ++++++ .../aws/operators/test_redshift_data.py | 18 +++++ .../amazon/aws/triggers/test_redshift_data.py | 75 +++++++++++++++++++ 4 files changed, 122 insertions(+) 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 3168b9f69c425..b31516d95ead1 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 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) 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..6b404d8b8bf50 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,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: 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 19182eb4c86b2..cddd80f06c1cc 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 @@ -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""" 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..6c3a61f117598 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,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")