diff --git a/providers/amazon/docs/operators/redshift/redshift_data.rst b/providers/amazon/docs/operators/redshift/redshift_data.rst index 9486e98fb25ab..7720e304b76bc 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 04dc8c3216297..fddc921fbddca 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 @@ def __init__( 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 @@ def execute(self, context: Context) -> list[GetStatementResultResponseTypeDef] | 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 @@ def execute(self, context: Context) -> list[GetStatementResultResponseTypeDef] | 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 @@ def on_kill(self) -> None: 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 cb2f08db7e74c..6058fdd190b92 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 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.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 @@ SQL = "sql" DATABASE = "database" STATEMENT_ID = "statement_id" +NEW_STATEMENT_ID = "new_statement_id" SESSION_ID = "session_id" @@ -145,7 +148,7 @@ def test_execute(self, mock_conn, mock_exec_query): 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 @@ def test_execute_with_workgroup_name(self, mock_conn, mock_exec_query): 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 @@ def test_execute_new_session(self, mock_conn, mock_exec_query): 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 @@ def test_execute_new_session(self, mock_conn, mock_exec_query): 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 @@ def test_template_fields(self): wait_for_completion=False, ) validate_template_fields(operator) + + +@pytest.mark.skipif( + 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