From 4be9af0f54032f57780cfc2d77bfb98717174ff1 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Tue, 7 Jul 2026 12:40:37 +0530 Subject: [PATCH 1/2] Add durable execution to RedshiftDataOperator --- .../docs/operators/redshift/redshift_data.rst | 42 ++++ .../amazon/aws/operators/redshift_data.py | 105 ++++++++- .../aws/operators/test_redshift_data.py | 209 +++++++++++++++++- 3 files changed, 343 insertions(+), 13 deletions(-) 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 3168b9f69c425..93e3873cd5d1c 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 @@ -19,13 +19,34 @@ from typing import TYPE_CHECKING, Any -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, @@ -35,7 +56,7 @@ 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 +92,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", @@ -128,13 +154,12 @@ 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: + return self.execute_resumable(context) - # 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, @@ -145,7 +170,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, @@ -157,7 +182,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( @@ -232,3 +257,65 @@ 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: str, context: Context) -> str: + """Query the raw statement status; a missing/expired statement degrades to NOT_FOUND.""" + try: + resp = self.hook.conn.describe_statement(Id=external_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: str, context: Context + ) -> list[GetStatementResultResponseTypeDef] | list[str]: + self.statement_id = external_id + if self.wait_for_completion: + self.hook.wait_for_results(external_id, poll_interval=self.poll_interval) + # Fetch here (not just in get_job_result): on reconnect the mixin calls only + # poll_until_complete and returns its value directly, never calling get_job_result. + self._sql_results = self.get_sql_results( + statement_id=external_id, return_sql_result=self.return_sql_result + ) + self._result_fetched = True + return self._sql_results + + def get_job_result( + self, external_id: str, context: Context + ) -> list[GetStatementResultResponseTypeDef] | list[str]: + self.statement_id = external_id + if getattr(self, "_result_fetched", False): + return self._sql_results + return self.get_sql_results(statement_id=external_id, return_sql_result=self.return_sql_result) 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..90cdee9c8b9d3 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" @@ -130,7 +133,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, @@ -186,7 +189,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, @@ -202,7 +205,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} @@ -238,7 +240,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, @@ -464,3 +466,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 From 929929d11d708ce6c9305fc74a25120c28999846 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Fri, 10 Jul 2026 11:27:54 +0530 Subject: [PATCH 2/2] fixing mypy errors --- .../amazon/aws/operators/redshift_data.py | 40 +++++++++++-------- 1 file changed, 23 insertions(+), 17 deletions(-) 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 55782f2b217ca..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,7 +17,7 @@ # under the License. from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import botocore.exceptions @@ -52,6 +52,7 @@ def execute_resumable(self, context): DescribeStatementResponseTypeDef, GetStatementResultResponseTypeDef, ) + from pydantic import JsonValue from airflow.sdk import Context @@ -160,7 +161,8 @@ def __init__( def execute(self, context: Context) -> list[GetStatementResultResponseTypeDef] | list[str]: """Execute a statement against Amazon Redshift.""" if not self.deferrable: - return self.execute_resumable(context) + self.execute_resumable(context) + return self._sql_results self.log.info("Executing statement: %s", self.sql) @@ -287,10 +289,11 @@ def submit_job(self, context: Context) -> str: context["ti"].xcom_push(key="session_id", value=output.session_id) return self.statement_id - def get_job_status(self, external_id: str, context: Context) -> str: + 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=external_id) + 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" @@ -303,24 +306,27 @@ def is_job_active(self, status: str) -> bool: def is_job_succeeded(self, status: str) -> bool: return status == FINISHED_STATE - def poll_until_complete( - self, external_id: str, context: Context - ) -> list[GetStatementResultResponseTypeDef] | list[str]: - self.statement_id = external_id + 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(external_id, poll_interval=self.poll_interval) + 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 and returns its value directly, never calling get_job_result. + # 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=external_id, return_sql_result=self.return_sql_result + statement_id=statement_id, return_sql_result=self.return_sql_result ) self._result_fetched = True - return self._sql_results def get_job_result( - self, external_id: str, context: Context + self, external_id: JsonValue, context: Context ) -> list[GetStatementResultResponseTypeDef] | list[str]: - self.statement_id = external_id - if getattr(self, "_result_fetched", False): - return self._sql_results - return self.get_sql_results(statement_id=external_id, return_sql_result=self.return_sql_result) + 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