This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new d9b7eb0f8a4 Cancel the Redshift statement when a user kills the
deferred task (#69676)
d9b7eb0f8a4 is described below
commit d9b7eb0f8a48c2059ef630786628ed01ba73741f
Author: Steve Ahn <[email protected]>
AuthorDate: Sat Aug 1 11:01:30 2026 -0700
Cancel the Redshift statement when a user kills the deferred task (#69676)
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 +++
.../providers/amazon/aws/triggers/redshift_data.py | 31 +++++++++
.../amazon/aws/operators/test_redshift_data.py | 41 +++++++++++
.../unit/amazon/aws/triggers/test_redshift_data.py | 81 ++++++++++++++++++++++
4 files changed, 161 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 fddc921fbdd..b81403be214 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
@@ -83,6 +83,9 @@ class RedshiftDataOperator(ResumableJobMixin,
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
@@ -131,6 +134,7 @@ class RedshiftDataOperator(ResumableJobMixin,
AwsBaseOperator[RedshiftDataHook])
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)
@@ -152,6 +156,7 @@ class RedshiftDataOperator(ResumableJobMixin,
AwsBaseOperator[RedshiftDataHook])
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 "
@@ -202,6 +207,7 @@ class RedshiftDataOperator(ResumableJobMixin,
AwsBaseOperator[RedshiftDataHook])
region_name=self.region_name,
verify=self.verify,
botocore_config=self.botocore_config,
+ cancel_on_kill=self.cancel_on_kill,
),
method_name="execute_complete",
)
@@ -256,6 +262,8 @@ class RedshiftDataOperator(ResumableJobMixin,
AwsBaseOperator[RedshiftDataHook])
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 c011fe75104..29ac1aceb69 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 @@ class RedshiftDataTrigger(BaseTrigger):
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 @@ class RedshiftDataTrigger(BaseTrigger):
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 @@ class RedshiftDataTrigger(BaseTrigger):
"region_name": self.region_name,
"verify": self.verify,
"botocore_config": self.botocore_config,
+ "cancel_on_kill": self.cancel_on_kill,
},
)
@@ -87,6 +93,31 @@ class RedshiftDataTrigger(BaseTrigger):
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 6058fdd190b..1b33c7ab183 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
@@ -294,6 +294,23 @@ class TestRedshiftDataOperator:
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}]
@@ -399,6 +416,30 @@ class TestRedshiftDataOperator:
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 aa10a8f5ae1..d8bbcaccd3d 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 @@ import pytest
from airflow.providers.amazon.aws.hooks.redshift_data import (
ABORTED_STATE,
FAILED_STATE,
+ RedshiftDataHook,
RedshiftDataQueryAbortedError,
RedshiftDataQueryFailedError,
)
@@ -57,8 +58,21 @@ class TestRedshiftDataTrigger:
"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 @@ class TestRedshiftDataTrigger:
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")
