amoghrajesh opened a new pull request, #69530:
URL: https://github.com/apache/airflow/pull/69530
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull
request.
Write a good git commit message following this guide:
http://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the
documentation.
For user-facing UI changes, please attach before/after screenshots (or a
short
screen recording) so reviewers can assess the visual impact.
Feel free to ping (in general) for the review if you do not see reaction for
a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we
sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR,
please
change below checkbox to `[X]` followed by the name of the tool, uncomment
the "Generated-by".
-->
- [ ] Yes (please specify the tool below)
<!--
Generated-by: [Tool Name] following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
### What
`RedshiftDataOperator` submits a statement to the Redshift Data API and
polls it to completion on the worker. On a worker crash or preemption mid-poll,
Airflow retries the task by calling `execute()` again which resubmits the SQL
from scratch, since nothing about the in-flight statement id is persisted
across attempts.
### Current behaviour
A retry after a crash always resubmits the full SQL, even if the original
statement is still running (or already finished) in Redshift.
- For non-idempotent SQL (`INSERT`, `COPY`, `UPDATE`, `CREATE TABLE`) this
risks duplicate writes;
- For expensive queries it's wasted compute, since the orphaned original
execution keeps running with nobody polling it.
This also applies, in a narrower window, to `wait_for_completion=False`:
even though that mode does not poll at all, a retry after a successful
submission still resubmits today, since the id is never persisted regardless of
whether the task waits.
### Proposed change
Adds `ResumableJobMixin` support (Airflow 3.3+) to `RedshiftDataOperator`,
following the same pattern already done for the Databricks operators and
`SnowflakeSqlApiOperator`. Before polling begins, the submitted statement id is
persisted to `task_state_store`. On retry, the operator reads it back and:
- reconnects and keeps polling if the statement is still running
- returns immediately without resubmitting if it already succeeded
- submits the SQL fresh if the statement failed, or its id has expired and
is no longer found (`ClientError`)
`durable=True` is the default; set `durable=False` to keep the old "always
submit fresh on retry" behavior. `deferrable=True` takes precedence over
`durable`, the Triggerer already tracks the statement across the wait in that
mode.
Critically, the persistence-on-submit design means the
`wait_for_completion=False` case is protected too, not just the blocking-wait
case & the id is written to `task_state_store` immediately after submission
regardless of whether the task polls for it.
### Changes of Note
- `submit_job` now always calls `hook.execute_query(...,
wait_for_completion=False, ...)` and hence the mixin, not the hook, owns
polling. The statement id is set on the operator immediately after submission
(not after the wait completes), which also fixes a pre-existing `on_kill` gap:
today, killing the task mid-wait doesn't cancel the statement because
`self.statement_id` isn't set yet.
- `get_job_status` calls `describe_statement` directly rather than reusing
`check_query_is_finished`/`parse_statement_response`, since those raise on
failure states rather than returning a status string, which would break the
mixin's status-in/decision-out contract.
- `is_job_active` uses a negative test (`not in (FINISHED, *FAILURE_STATES,
NOT_FOUND)`) rather than a positive running-states allowlist, deliberately
avoiding a real pitfall found in `RedshiftDataTrigger.is_still_running`, which
uses a positive allowlist and would misclassify any future unlisted Redshift
status as "done."
- `get_sql_results` (the actual return value fetch) was moved into
`poll_until_complete` rather than left solely in `get_job_result`, because the
mixin calls only `poll_until_complete` on reconnect - `get_job_result` is
skipped entirely on that path. A first pass missed returning the fetched value
from `poll_until_complete` itself, which silently turned `execute()`'s return
value into `None` on reconnect; caught via the durable test suite and fixed.
- Verified there was no missing-state prerequisite bugfix needed here
(unlike the Databricks port, which needed `BLOCKED`/`WAITING_FOR_RETRY` added
to a state allowlist first) - Redshift's
`RUNNING_STATES`/`FINISHED_STATE`/`FAILURE_STATES` are already complete against
the real AWS API's `Status` enum.
### User implications / backcompat
No breaking change. `durable` defaults to `True` on Airflow 3.3+; on earlier
versions it's a no-op stub and the operator always submits fresh, exactly as
before. If `task_state_store` isn't available at runtime, the operator logs
that crash recovery is disabled and falls back to the same fresh-submit
behavior.
### Testing
<TODO>
---
* Read the **[Pull Request
Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)**
for more information. Note: commit author/co-author name and email in commits
become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal
([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals))
is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License
Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment:
`{pr_number}.significant.rst`, in
[airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
You can add this file in a follow-up commit after the PR is created so you
know the PR number.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
