This is an automated email from the ASF dual-hosted git repository.
amoghrajesh 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 5d1e2043793 Add durable execution to `RedshiftDataOperator` (#69530)
5d1e2043793 is described below
commit 5d1e2043793f5b7076aaebbe888265a11c8f9ffb
Author: Amogh Desai <[email protected]>
AuthorDate: Mon Jul 13 12:30:56 2026 +0530
Add durable execution to `RedshiftDataOperator` (#69530)
---
.../docs/operators/redshift/redshift_data.rst | 42 +++++
.../amazon/aws/operators/redshift_data.py | 113 ++++++++++-
.../amazon/aws/operators/test_redshift_data.py | 209 ++++++++++++++++++++-
3 files changed, 350 insertions(+), 14 deletions(-)
diff --git a/providers/amazon/docs/operators/redshift/redshift_data.rst
b/providers/amazon/docs/operators/redshift/redshift_data.rst
index 9486e98fb25..7720e304b76 100644
--- a/providers/amazon/docs/operators/redshift/redshift_data.rst
+++ b/providers/amazon/docs/operators/redshift/redshift_data.rst
@@ -66,6 +66,48 @@ the XCom and pass it to the ``session_id`` parameter. This
is useful when you wo
:start-after: [START howto_operator_redshift_data_session_reuse]
:end-before: [END howto_operator_redshift_data_session_reuse]
+Durable execution
+==================
+
+``RedshiftDataOperator`` submits a statement and then polls it to completion
on the worker. By
+default the operator runs in a *durable* mode that makes this crash-safe: the
Redshift statement
+id is persisted to task state before polling begins, so if the worker crashes
or is preempted and
+the task is retried, the operator reconnects to the statement that is already
executing in
+Redshift instead of resubmitting the SQL.
+
+On retry the operator checks the prior statement's state:
+
+* if it is still running, the operator reconnects and continues polling
+* if it already succeeded, the operator returns immediately without
resubmitting
+* if it failed terminally, or its id has expired and is no longer found, the
operator submits the
+ SQL fresh
+
+This protection also applies when ``wait_for_completion=False`` -- even though
that task attempt
+never polls at all, a retry after a successful submission still reconnects
rather than
+resubmitting, since the statement id is persisted immediately after submission
regardless of
+whether the task waits for it to finish.
+
+Durable execution requires Airflow 3.3 or newer, since it relies on the task
state store. On
+earlier Airflow versions the flag is a no-op and the operator always submits
fresh SQL on retry,
+exactly as before. If the task state store is unavailable at runtime, the
operator logs that
+crash recovery is disabled and behaves the same way.
+
+To opt out and always submit fresh SQL on retry, set ``durable=False``:
+
+.. code-block:: python
+
+ statement = RedshiftDataOperator(
+ task_id="redshift_data",
+ database="dev",
+ sql="SELECT * FROM table",
+ cluster_identifier="cluster_identifier",
+ durable=False,
+ )
+
+Durable execution applies to the synchronous path. When ``deferrable=True`` is
set, the Triggerer
+already tracks the statement across the wait, so deferrable mode takes
precedence and ``durable``
+has no effect.
+
Reference
---------
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 04dc8c32162..fddc921fbdd 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
@@ -17,25 +17,47 @@
# under the License.
from __future__ import annotations
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
-from airflow.providers.amazon.aws.hooks.redshift_data import RedshiftDataHook
+import botocore.exceptions
+
+from airflow.providers.amazon.aws.hooks.redshift_data import FAILURE_STATES,
FINISHED_STATE, RedshiftDataHook
from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
from airflow.providers.amazon.aws.triggers.redshift_data import
RedshiftDataTrigger
from airflow.providers.amazon.aws.utils import validate_execute_complete_event
from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
from airflow.providers.common.compat.sdk import AirflowException, conf
+try:
+ from airflow.sdk import ResumableJobMixin
+except ImportError:
+
+ class ResumableJobMixin: # type: ignore[no-redef]
+ """Airflow <3.3 stub, task_state_store unavailable, always submits
fresh."""
+
+ external_id_key: str = "redshift_statement_id"
+
+ def __init__(self, *, durable: bool = True, **kwargs: Any) -> None:
+ super().__init__(**kwargs)
+ self.durable = durable
+
+ def execute_resumable(self, context):
+ external_id = self.submit_job(context)
+ self.poll_until_complete(external_id, context)
+ return self.get_job_result(external_id, context)
+
+
if TYPE_CHECKING:
from mypy_boto3_redshift_data.type_defs import (
DescribeStatementResponseTypeDef,
GetStatementResultResponseTypeDef,
)
+ from pydantic import JsonValue
from airflow.sdk import Context
-class RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
+class RedshiftDataOperator(ResumableJobMixin,
AwsBaseOperator[RedshiftDataHook]):
"""
Executes SQL Statements against an Amazon Redshift cluster using Redshift
Data.
@@ -71,9 +93,14 @@ class
RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
:param botocore_config: Configuration dictionary (key-values) for botocore
client. See:
https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :param durable: When ``True`` (the default), the Redshift statement id is
persisted to task
+ state before polling begins. A worker crash on retry reconnects to the
existing statement
+ instead of resubmitting the SQL. Set to ``False`` to always submit
fresh SQL on retry.
+ Requires Airflow 3.3+; ignored silently on earlier versions.
"""
aws_hook_class = RedshiftDataHook
+ external_id_key = "redshift_statement_id"
template_fields = aws_template_fields(
"cluster_identifier",
"database",
@@ -133,13 +160,13 @@ class
RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
def execute(self, context: Context) ->
list[GetStatementResultResponseTypeDef] | list[str]:
"""Execute a statement against Amazon Redshift."""
- self.log.info("Executing statement: %s", self.sql)
+ if not self.deferrable:
+ self.execute_resumable(context)
+ return self._sql_results
- # Set wait_for_completion to False so that it waits for the status in
the deferred task.
- wait_for_completion = self.wait_for_completion
- if self.deferrable:
- wait_for_completion = False
+ self.log.info("Executing statement: %s", self.sql)
+ # wait_for_completion=False so that it waits for the status in the
deferred task instead.
query_execution_output = self.hook.execute_query(
database=self.database,
sql=self.sql,
@@ -150,7 +177,7 @@ class
RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
secret_arn=self.secret_arn,
statement_name=self.statement_name,
with_event=self.with_event,
- wait_for_completion=wait_for_completion,
+ wait_for_completion=False,
poll_interval=self.poll_interval,
session_id=self.session_id,
session_keep_alive_seconds=self.session_keep_alive_seconds,
@@ -162,7 +189,7 @@ class
RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
if query_execution_output.session_id:
context["ti"].xcom_push(key="session_id",
value=query_execution_output.session_id)
- if self.deferrable and self.wait_for_completion:
+ if self.wait_for_completion:
is_finished: bool =
self.hook.check_query_is_finished(self.statement_id)
if not is_finished:
self.defer(
@@ -237,3 +264,69 @@ class
RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
self.hook.conn.cancel_statement(Id=self.statement_id)
except Exception as ex:
self.log.error("Unable to cancel query. Exiting. %s", ex)
+
+ def submit_job(self, context: Context) -> str:
+ """Submit the statement for execution and return its statement id."""
+ output = self.hook.execute_query(
+ database=self.database,
+ sql=self.sql,
+ cluster_identifier=self.cluster_identifier,
+ workgroup_name=self.workgroup_name,
+ db_user=self.db_user,
+ parameters=self.parameters,
+ secret_arn=self.secret_arn,
+ statement_name=self.statement_name,
+ with_event=self.with_event,
+ wait_for_completion=False,
+ poll_interval=self.poll_interval,
+ session_id=self.session_id,
+ session_keep_alive_seconds=self.session_keep_alive_seconds,
+ )
+ # Set immediately (before any polling) so on_kill can cancel even if
the worker
+ # dies before poll_until_complete runs.
+ self.statement_id = output.statement_id
+ if output.session_id:
+ context["ti"].xcom_push(key="session_id", value=output.session_id)
+ return self.statement_id
+
+ def get_job_status(self, external_id: JsonValue, context: Context) -> str:
+ """Query the raw statement status; a missing/expired statement
degrades to NOT_FOUND."""
+ statement_id = cast("str", external_id)
+ try:
+ resp = self.hook.conn.describe_statement(Id=statement_id)
+ except botocore.exceptions.ClientError as e:
+ if e.response["Error"]["Code"] in ("ResourceNotFoundException",
"ValidationException"):
+ return "NOT_FOUND"
+ raise
+ return resp["Status"]
+
+ def is_job_active(self, status: str) -> bool:
+ return status not in (FINISHED_STATE, *FAILURE_STATES, "NOT_FOUND")
+
+ def is_job_succeeded(self, status: str) -> bool:
+ return status == FINISHED_STATE
+
+ def poll_until_complete(self, external_id: JsonValue, context: Context) ->
None:
+ statement_id = cast("str", external_id)
+ self.statement_id = statement_id
+ if self.wait_for_completion:
+ self.hook.wait_for_results(statement_id,
poll_interval=self.poll_interval)
+ # Fetch here (not just in get_job_result): on reconnect the mixin
calls only
+ # poll_until_complete, never get_job_result, so execute() reads
self._sql_results
+ # afterward rather than relying on either method's return value.
+ self._sql_results = self.get_sql_results(
+ statement_id=statement_id, return_sql_result=self.return_sql_result
+ )
+ self._result_fetched = True
+
+ def get_job_result(
+ self, external_id: JsonValue, context: Context
+ ) -> list[GetStatementResultResponseTypeDef] | list[str]:
+ statement_id = cast("str", external_id)
+ self.statement_id = statement_id
+ if not getattr(self, "_result_fetched", False):
+ self._sql_results = self.get_sql_results(
+ statement_id=statement_id,
return_sql_result=self.return_sql_result
+ )
+ self._result_fetched = True
+ return self._sql_results
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 cb2f08db7e7..6058fdd190b 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
@@ -19,6 +19,7 @@ from __future__ import annotations
from unittest import mock
+import botocore.exceptions
import pytest
from airflow.providers.amazon.aws.hooks.redshift_data import
QueryExecutionOutput
@@ -26,6 +27,7 @@ from airflow.providers.amazon.aws.operators.redshift_data
import RedshiftDataOpe
from airflow.providers.amazon.aws.triggers.redshift_data import
RedshiftDataTrigger
from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS
from unit.amazon.aws.utils.test_template_fields import validate_template_fields
CONN_ID = "aws_conn_test"
@@ -33,6 +35,7 @@ TASK_ID = "task_id"
SQL = "sql"
DATABASE = "database"
STATEMENT_ID = "statement_id"
+NEW_STATEMENT_ID = "new_statement_id"
SESSION_ID = "session_id"
@@ -145,7 +148,7 @@ class TestRedshiftDataOperator:
statement_name=statement_name,
parameters=parameters,
with_event=False,
- wait_for_completion=True, # Matches above
+ wait_for_completion=False, # submit_job always submits async; the
mixin owns waiting
poll_interval=poll_interval,
session_id=None,
session_keep_alive_seconds=None,
@@ -201,7 +204,7 @@ class TestRedshiftDataOperator:
statement_name=statement_name,
parameters=parameters,
with_event=False,
- wait_for_completion=True,
+ wait_for_completion=False, # submit_job always submits async; the
mixin owns waiting
poll_interval=poll_interval,
session_id=None,
session_keep_alive_seconds=None,
@@ -217,7 +220,6 @@ class TestRedshiftDataOperator:
statement_name = "statement_name"
parameters = [{"name": "id", "value": "1"}]
poll_interval = 5
- wait_for_completion = True
# Like before, return a statement ID and a status
mock_conn.execute_statement.return_value = {"Id": STATEMENT_ID}
@@ -253,7 +255,7 @@ class TestRedshiftDataOperator:
statement_name=statement_name,
parameters=parameters,
with_event=False,
- wait_for_completion=wait_for_completion,
+ wait_for_completion=False, # submit_job always submits async; the
mixin owns waiting
poll_interval=poll_interval,
session_id=None,
session_keep_alive_seconds=123,
@@ -479,3 +481,202 @@ class TestRedshiftDataOperator:
wait_for_completion=False,
)
validate_template_fields(operator)
+
+
[email protected](
+ not AIRFLOW_V_3_3_PLUS, reason="task_state_store (durable execution)
requires Airflow 3.3+"
+)
+class TestRedshiftDataOperatorDurable:
+ @staticmethod
+ def _context(task_store=None):
+ ctx: dict = {"ti": mock.MagicMock(stats_tags={})}
+ if task_store is not None:
+ ctx["task_state_store"] = task_store
+ return ctx
+
+ @staticmethod
+ def _make_operator(**kwargs):
+ return RedshiftDataOperator(
+ task_id=TASK_ID,
+ database=DATABASE,
+ sql=SQL,
+ cluster_identifier="cluster_identifier",
+ **kwargs,
+ )
+
+ @staticmethod
+ def _not_found_error():
+ return botocore.exceptions.ClientError(
+ error_response={"Error": {"Code": "ResourceNotFoundException",
"Message": "not found"}},
+ operation_name="DescribeStatement",
+ )
+
+
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.execute_query")
+
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.conn")
+ def test_persists_statement_id_to_task_state_store_on_fresh_submit(self,
mock_conn, mock_exec_query):
+ operator = self._make_operator()
+ mock_exec_query.return_value =
QueryExecutionOutput(statement_id=STATEMENT_ID, session_id=None)
+ mock_conn.describe_statement.return_value = {"Status": "FINISHED"}
+ task_store = mock.MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = None
+
+ operator.execute(self._context(task_store))
+
+ mock_exec_query.assert_called_once()
+ task_store.set.assert_called_once_with("redshift_statement_id",
STATEMENT_ID)
+
+
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.execute_query")
+
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.conn")
+ def test_reconnects_to_running_statement_without_resubmitting(self,
mock_conn, mock_exec_query):
+ operator = self._make_operator()
+ # get_job_status sees it still running, wait_for_results then sees it
finish,
+ # get_sql_results makes one more describe_statement call to check for
sub-statements.
+ mock_conn.describe_statement.side_effect = [
+ {"Status": "STARTED"},
+ {"Status": "FINISHED"},
+ {"Status": "FINISHED"},
+ ]
+ task_store = mock.MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = STATEMENT_ID
+
+ result = operator.execute(self._context(task_store))
+
+ mock_exec_query.assert_not_called()
+ task_store.set.assert_not_called()
+ assert result == [STATEMENT_ID]
+
+
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.execute_query")
+
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.conn")
+ def test_already_succeeded_returns_result_without_polling(self, mock_conn,
mock_exec_query):
+ operator = self._make_operator()
+ mock_conn.describe_statement.return_value = {"Status": "FINISHED"}
+ task_store = mock.MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = STATEMENT_ID
+
+ result = operator.execute(self._context(task_store))
+
+ mock_exec_query.assert_not_called()
+ assert result == [STATEMENT_ID]
+ # Only get_job_status + get_sql_results should have called
describe_statement -- no poll loop.
+ assert mock_conn.describe_statement.call_count == 2
+
+
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.execute_query")
+
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.conn")
+ def test_resubmits_when_stored_statement_in_terminal_failure(self,
mock_conn, mock_exec_query):
+ operator = self._make_operator()
+ mock_exec_query.return_value =
QueryExecutionOutput(statement_id=NEW_STATEMENT_ID, session_id=None)
+ # get_job_status sees the stored statement failed; after fresh
resubmit, polling succeeds.
+ mock_conn.describe_statement.side_effect = [
+ {"Status": "FAILED"},
+ {"Status": "FINISHED"},
+ {"Status": "FINISHED"},
+ ]
+ task_store = mock.MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = STATEMENT_ID
+
+ result = operator.execute(self._context(task_store))
+
+ mock_exec_query.assert_called_once()
+ task_store.set.assert_called_once_with("redshift_statement_id",
NEW_STATEMENT_ID)
+ assert result == [NEW_STATEMENT_ID]
+
+
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.execute_query")
+
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.conn")
+ def test_resubmits_when_stored_statement_not_found(self, mock_conn,
mock_exec_query):
+ operator = self._make_operator()
+ mock_exec_query.return_value =
QueryExecutionOutput(statement_id=NEW_STATEMENT_ID, session_id=None)
+ mock_conn.describe_statement.side_effect = [
+ self._not_found_error(),
+ {"Status": "FINISHED"},
+ {"Status": "FINISHED"},
+ ]
+ task_store = mock.MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = STATEMENT_ID
+
+ result = operator.execute(self._context(task_store))
+
+ mock_exec_query.assert_called_once()
+ task_store.set.assert_called_once_with("redshift_statement_id",
NEW_STATEMENT_ID)
+ assert result == [NEW_STATEMENT_ID]
+
+
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.execute_query")
+
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.conn")
+ def test_durable_false_never_touches_task_state_store(self, mock_conn,
mock_exec_query):
+ operator = self._make_operator(durable=False)
+ mock_exec_query.return_value =
QueryExecutionOutput(statement_id=STATEMENT_ID, session_id=None)
+ mock_conn.describe_statement.return_value = {"Status": "FINISHED"}
+ task_store = mock.MagicMock(spec_set=["get", "set"])
+
+ operator.execute(self._context(task_store))
+
+ mock_exec_query.assert_called_once()
+ task_store.get.assert_not_called()
+ task_store.set.assert_not_called()
+
+
@mock.patch("airflow.providers.amazon.aws.operators.redshift_data.RedshiftDataOperator.defer")
+
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.execute_query")
+
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.conn")
+ def test_deferrable_unaffected_by_durable(self, mock_conn,
mock_exec_query, mock_defer):
+ operator = self._make_operator(deferrable=True)
+ mock_exec_query.return_value =
QueryExecutionOutput(statement_id=STATEMENT_ID, session_id=None)
+ mock_conn.describe_statement.return_value = {"Status": "STARTED"}
+ task_store = mock.MagicMock(spec_set=["get", "set"])
+
+ operator.execute(self._context(task_store))
+
+ assert mock_defer.called
+ task_store.get.assert_not_called()
+ task_store.set.assert_not_called()
+
+
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.execute_query")
+
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.conn")
+ def
test_wait_for_completion_false_still_protects_against_duplicate_submission(
+ self, mock_conn, mock_exec_query
+ ):
+ """A prior attempt already submitted successfully; a retry must not
resubmit even
+ though wait_for_completion=False means this attempt never polls."""
+ operator = self._make_operator(wait_for_completion=False)
+ mock_conn.describe_statement.return_value = {"Status": "FINISHED"}
+ task_store = mock.MagicMock(spec_set=["get", "set"])
+ task_store.get.return_value = STATEMENT_ID
+
+ result = operator.execute(self._context(task_store))
+
+ mock_exec_query.assert_not_called()
+ assert result == [STATEMENT_ID]
+
+
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.conn")
+ def test_get_job_status_returns_raw_status(self, mock_conn):
+ operator = self._make_operator()
+ mock_conn.describe_statement.return_value = {"Status": "STARTED"}
+
+ assert operator.get_job_status(STATEMENT_ID, context={}) == "STARTED"
+
+
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.conn")
+ def test_get_job_status_not_found_on_resource_not_found(self, mock_conn):
+ operator = self._make_operator()
+ mock_conn.describe_statement.side_effect = self._not_found_error()
+
+ assert operator.get_job_status(STATEMENT_ID, context={}) == "NOT_FOUND"
+
+
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_data.RedshiftDataHook.conn")
+ def test_get_job_status_reraises_other_client_errors(self, mock_conn):
+ operator = self._make_operator()
+ mock_conn.describe_statement.side_effect =
botocore.exceptions.ClientError(
+ error_response={"Error": {"Code": "InternalServerException",
"Message": "boom"}},
+ operation_name="DescribeStatement",
+ )
+
+ with pytest.raises(botocore.exceptions.ClientError):
+ operator.get_job_status(STATEMENT_ID, context={})
+
+ def test_is_job_active_and_is_job_succeeded_predicates(self):
+ operator = self._make_operator()
+
+ for status in ("PICKED", "STARTED", "SUBMITTED"):
+ assert operator.is_job_active(status) is True
+ for status in ("FINISHED", "FAILED", "ABORTED", "NOT_FOUND"):
+ assert operator.is_job_active(status) is False
+
+ assert operator.is_job_succeeded("FINISHED") is True
+ assert operator.is_job_succeeded("STARTED") is False
