amoghrajesh opened a new pull request, #69477:
URL: https://github.com/apache/airflow/pull/69477
<!-- 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
`SnowflakeSqlApiOperator` submits one or more SQL statements to the
Snowflake SQL API and polls their statement handles to completion on the
worker. On a worker crash or preemption mid-poll, Airflow retries the task by
calling `execute()` again which re-submits the SQL from scratch, since nothing
about the in-flight statement handles is persisted across attempts.
### Current behaviour
A retry after a crash always resubmits the full SQL, even if the original
statements are still running (or already finished) in Snowflake. For
non-idempotent SQL (`INSERT`, `UPDATE`, `CREATE TABLE`, etc.) this risks
duplicate writes; for expensive queries it's wasted warehouse compute, since
the orphaned original execution keeps running with nobody polling it.
### Proposed change
Adds `ResumableJobMixin` support (Airflow 3.3+) to
`SnowflakeSqlApiOperator`, following the same pattern already completed for
`DatabricksSubmitRunOperator`/`DatabricksRunNowOperator`. Before polling
begins, the submitted statement handles are persisted to `task_state_store`. On
retry, the operator reads them back and:
- reconnects and keeps polling if any handle is still running
(already-finished handles are not re-run, only the ones still in progress are
waited on)
- returns immediately without resubmitting if every handle already succeeded
- submits the SQL fresh if a handle failed, or its handle expired past
Snowflake's retention window (404)
`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 handles across the wait in that mode.
Includes a companion fix (already merged separately, #69450) that removed an
unconditional per-handle sleep in `poll_on_queries()`, which this port relies
on to stay latency-neutral on the already-resolved case.
### Changes of Note
- `get_job_status` aggregates across the list of statement handles into the
single status string the mixin's interface expects: any `error` wins over any
`running`, which wins over all-`success` — matching Snowflake's all-or-nothing
submission semantics (there's no per-statement repair, unlike Databricks'
`repair_run`).
- The mixin calls `poll_until_complete` alone on reconnect —
`get_job_result` is never invoked on that path. Since `get_job_result` is where
`check_query_output` (the actual result fetch/log) lives, `check_query_output`
was moved into `poll_until_complete` itself, guarded by a flag so the
fresh-submit path (which calls both methods) doesn't fetch/push twice.
- A 404 from an expired/unknown statement handle is caught in
`get_job_status` and treated as a `not_found` sentinel, degrading to a fresh
resubmit rather than failing the task.
### 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.
One minor, intentional behavior shift: a fresh submission that is still
running on its very first status check now incurs one `poll_interval` of
latency it didn't before (the old code had a sleep-free pre-check before
entering the poll loop; the durable path goes straight into the sleep-guarded
loop). One-time cost, not compounding, and documented in the test suite.
### Testing
Follow this to create a snowflake connection using RSA:
https://medium.com/@chik0di/using-the-snowflakesqlapi-operator-in-airflow-0206632db2a3
DAG used:
```python
from datetime import datetime, timedelta
from airflow.providers.snowflake.operators.snowflake import
SnowflakeSqlApiOperator
from airflow.sdk import DAG
with DAG(
dag_id="snowflake_sql_api",
schedule=None,
start_date=datetime(2024, 1, 1),
catchup=False,
) as dag:
multi_stmt = SnowflakeSqlApiOperator(
task_id="multi_statement",
snowflake_conn_id="snowflake_default",
sql=(
"SELECT CURRENT_TIMESTAMP(); "
"CALL SYSTEM$WAIT(30); "
"SELECT CURRENT_ACCOUNT();"
),
statement_count=3,
deferrable=False,
retries=2,
retry_delay=timedelta(seconds=5),
)
```
#### Before changes
Before:
Initial try where worker crashed:
<img width="1660" height="874" alt="image"
src="https://github.com/user-attachments/assets/74fe5f6e-c047-446b-a6c6-d65d05716d23";
/>
Later try repeats the same effort which was already done:
<img width="1660" height="917" alt="image"
src="https://github.com/user-attachments/assets/c2fbf8e7-2e39-467b-8928-16bf93865b9c";
/>
Duplicate effort:
<img width="1471" height="86" alt="image"
src="https://github.com/user-attachments/assets/a442b318-5456-4bc5-9303-3c96fd65fcdb";
/>
First try suceeded in snowflake at 5:37:43 even when airflow task failed at:
5:37:17
<img width="1471" height="794" alt="image"
src="https://github.com/user-attachments/assets/0b6d58d4-faa5-4ea7-9238-e606df389619";
/>
<img width="1471" height="794" alt="image"
src="https://github.com/user-attachments/assets/516bc3e2-e7a2-4d6f-b50a-9d8cafdaefd8";
/>
#### After changes
Initial try where worker crashed
<img width="1663" height="919" alt="image"
src="https://github.com/user-attachments/assets/02ece83b-4d9a-4c82-a8d9-a1dbec6223c9";
/>
Worker is back up but job has already completed:
<img width="1663" height="919" alt="image"
src="https://github.com/user-attachments/assets/507d32b5-dcf6-42d2-9129-d6355e5bfc14";
/>
<img width="1663" height="919" alt="image"
src="https://github.com/user-attachments/assets/36296e0d-a46e-494d-b0a8-9c8dd464e2e1";
/>
<img width="1663" height="919" alt="image"
src="https://github.com/user-attachments/assets/9548e343-6533-450d-8aef-db7408fe0c10";
/>
Since I had a custom backend configured, data stored is:
```shell
[Breeze:3.10.20] root@62f2e50bd27b:/opt/airflow$ cat
/tmp/airflow_state/ti_multi_statement/snowflake_query_ids.json
["01c5873f-061b-2e62-0068-2101308e710b"]
```
No duplicate resubmission.
---
* 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]
