Add durable execution to SnowflakeSqlApiOperator

[amoghrajesh]

---

Was generative AI tooling used to co-author this PR?

• Yes (please specify the tool below)

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:

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:

Later try repeats the same effort which was already done:

Duplicate effort:

First try suceeded in snowflake at 5:37:43 even when airflow task failed at: 5:37:17

After changes

Initial try where worker crashed

Worker is back up but job has already completed:

Since I had a custom backend configured, data stored is:

[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 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) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

[kaxil]

On the failure path query_ids is never pushed to xcom. When a statement errors, poll_on_queries() reports the error and this method raises RuntimeError at line 567, before it reaches this xcom_push. That is the case where the handles matter most, since you need them to look the failed statements up in Snowflake.

The pre-PR non-deferrable path pushed query_ids immediately after execute_query (before the poll loop), so the handles were recorded even when a query failed. This moves the push after the loop, so a failure now loses them.

Consider pushing the xcom near the top of poll_until_complete (right after self.query_ids = ..., before the poll loop) so it lands on both the fresh and reconnect paths and survives an error. Non-blocking.