1fanwang opened a new pull request, #72422:
URL: https://github.com/apache/airflow/pull/72422
When a worker dies after submitting a Databricks SQL statement, a
synchronous task retry submits the statement again. The original statement can
still be running, so users can get duplicate work and cost.
Before this change, the retry submits `statement-2` while `statement-1`
remains active. After this change, Airflow 3.3+ stores `statement-1` in task
state and reconnects to it. A retry also returns an already successful
statement without resubmitting, and starts fresh only after failure,
cancellation, closure, or a missing statement.
This uses the AIP-103 `ResumableJobMixin` contract only when
`wait_for_termination=True` and `deferrable=False`. Fire-and-forget and
deferrable execution keep their existing paths. Reconnect restores the operator
field and statement-ID XCom used by cancellation, links, and OpenLineage.
## Testing
The live proof below reproduces the duplicate submission against
`upstream/main`, then shows the retry reconnecting on this branch. The complete
SQL statements test file and Databricks provider suite also passed.
<details>
<summary>Live task-state red/green proof</summary>
The local proof driver called `execute()` twice with the real
`TaskStateStoreAccessor` and supervisor message types. Its Databricks hook
stand-in raised after the first submission to simulate a worker crash. The
driver and extracted pre-change source were kept out of the commit.
```python
from __future__ import annotations
import importlib
import importlib.util
import os
import sys
from pathlib import Path
from typing import Any
from unittest import mock
from uuid import UUID
from airflow.providers.databricks.hooks.databricks import SQLStatementState
from airflow.sdk._shared.state import TaskScope
from airflow.sdk.execution_time import task_runner
from airflow.sdk.execution_time.comms import GetTaskStateStore,
SetTaskStateStore, TaskStateStoreResult
from airflow.sdk.execution_time.context import TaskStateStoreAccessor
MODULE_NAME = "airflow.providers.databricks.operators.databricks"
if source_path := os.environ.get("DATABRICKS_OPERATOR_SOURCE"):
module_spec = importlib.util.spec_from_file_location(MODULE_NAME,
Path(source_path))
if module_spec is None or module_spec.loader is None:
raise RuntimeError(f"cannot load {source_path}")
databricks_module = importlib.util.module_from_spec(module_spec)
sys.modules[MODULE_NAME] = databricks_module
module_spec.loader.exec_module(databricks_module)
else:
databricks_module = importlib.import_module(MODULE_NAME)
DatabricksSQLStatementsOperator =
databricks_module.DatabricksSQLStatementsOperator
class WorkerCrash(Exception):
pass
class LocalDatabricksHook:
def __init__(self) -> None:
self.submissions = 0
self.status_checks: dict[str, int] = {}
def post_sql_statement(self, json: dict[str, Any]) -> str:
self.submissions += 1
return f"statement-{self.submissions}"
def get_sql_statement_state(self, statement_id: str) ->
SQLStatementState:
checks = self.status_checks.get(statement_id, 0) + 1
self.status_checks[statement_id] = checks
if statement_id == "statement-1" and checks == 1:
raise WorkerCrash("simulated worker crash after submission")
if statement_id == "statement-1" and checks == 2:
return SQLStatementState("RUNNING")
return SQLStatementState("SUCCEEDED")
def main(expected_submissions: int) -> None:
stored: dict[str, Any] = {}
def send(message: Any) -> TaskStateStoreResult | None:
if isinstance(message, GetTaskStateStore):
value = stored.get(message.key)
return TaskStateStoreResult(value=value) if value is not None
else None
if isinstance(message, SetTaskStateStore):
stored[message.key] = message.value
return None
task_state_store = TaskStateStoreAccessor(
ti_id=UUID("00000000-0000-0000-0000-000000000001"),
scope=TaskScope(dag_id="proof", run_id="run", task_id="sql"),
)
context = {
"task_state_store": task_state_store,
"ti": mock.MagicMock(stats_tags={}),
}
hook = LocalDatabricksHook()
with (
mock.patch.object(task_runner, "SUPERVISOR_COMMS",
mock.Mock(send=mock.Mock(side_effect=send)), create=True),
mock.patch.object(DatabricksSQLStatementsOperator, "_get_hook",
return_value=hook),
):
first = DatabricksSQLStatementsOperator(
task_id="sql",
statement="SELECT 1",
warehouse_id="warehouse",
polling_period_seconds=0,
)
try:
first.execute(context)
except WorkerCrash as error:
print(f"first_attempt={error}")
retry = DatabricksSQLStatementsOperator(
task_id="sql",
statement="SELECT 1",
warehouse_id="warehouse",
polling_period_seconds=0,
)
retry.execute(context)
print(f"operator_source={databricks_module.__file__}")
print(f"stored={stored}")
print(f"submissions={hook.submissions}")
print(f"retry_statement_id={retry.statement_id}")
if hook.submissions != expected_submissions:
raise RuntimeError(f"expected {expected_submissions} submissions")
if __name__ == "__main__":
main(expected_submissions=int(sys.argv[1]))
```
Save the driver as `dev/databricks_resumability_e2e.py`, then run:
```console
mkdir -p dev/databricks-resumability-prechange
git archive upstream/main
providers/databricks/src/airflow/providers/databricks/operators/databricks.py \
| tar -x -C dev/databricks-resumability-prechange
DATABRICKS_OPERATOR_SOURCE="$PWD/dev/databricks-resumability-prechange/providers/databricks/src/airflow/providers/databricks/operators/databricks.py"
\
AIRFLOW_HOME="$PWD/dev/databricks-resumability-airflow-home" \
AIRFLOW__CORE__LOAD_EXAMPLES=False \
.venv/bin/uv run --project providers/databricks python
dev/databricks_resumability_e2e.py 2
```
Pre-change log, with the local source path redacted:
```text
2026-09-02T04:41:45.028098Z [warning ]
DatabricksSQLStatementsOperator.execute cannot be called outside of the Task
Runner!
[airflow.task.operators.airflow.providers.databricks.operators.databricks.DatabricksSQLStatementsOperator]
loc=operator.py:439
2026-09-02T04:41:45.031377Z [info ] SQL Statement submitted with
statement_id: statement-1
[airflow.task.operators.airflow.providers.databricks.operators.databricks.DatabricksSQLStatementsOperator]
loc=databricks.py:1624
first_attempt=simulated worker crash after submission
2026-09-02T04:41:45.031676Z [warning ]
DatabricksSQLStatementsOperator.execute cannot be called outside of the Task
Runner!
[airflow.task.operators.airflow.providers.databricks.operators.databricks.DatabricksSQLStatementsOperator]
loc=operator.py:439
2026-09-02T04:41:45.031852Z [info ] SQL Statement submitted with
statement_id: statement-2
[airflow.task.operators.airflow.providers.databricks.operators.databricks.DatabricksSQLStatementsOperator]
loc=databricks.py:1624
2026-09-02T04:41:45.032015Z [info ] sql completed successfully.
[airflow.task.operators.airflow.providers.databricks.operators.databricks.DatabricksSQLStatementsOperator]
loc=mixins.py:107
operator_source=<pre-change source>
stored={}
submissions=2
retry_statement_id=statement-2
```
Post-change command:
```console
AIRFLOW_HOME="$PWD/dev/databricks-resumability-airflow-home" \
AIRFLOW__CORE__LOAD_EXAMPLES=False \
.venv/bin/uv run --project providers/databricks python
dev/databricks_resumability_e2e.py 1
```
Post-change log, with the local source path redacted:
```text
2026-09-02T05:23:04.722106Z [warning ]
DatabricksSQLStatementsOperator.execute cannot be called outside of the Task
Runner!
[airflow.task.operators.airflow.providers.databricks.operators.databricks.DatabricksSQLStatementsOperator]
loc=operator.py:439
2026-09-02T05:23:04.727990Z [info ] SQL Statement submitted with
statement_id: statement-1
[airflow.task.operators.airflow.providers.databricks.operators.databricks.DatabricksSQLStatementsOperator]
loc=databricks.py:1648
first_attempt=simulated worker crash after submission
2026-09-02T05:23:04.737522Z [warning ]
DatabricksSQLStatementsOperator.execute cannot be called outside of the Task
Runner!
[airflow.task.operators.airflow.providers.databricks.operators.databricks.DatabricksSQLStatementsOperator]
loc=operator.py:439
2026-09-02T05:23:04.748079Z [info ] Reconnecting to existing job
[airflow.task.operators.airflow.providers.databricks.operators.databricks.DatabricksSQLStatementsOperator]
external_id=statement-1 external_id_key=databricks_sql_statement_id
loc=resumablejobmixin.py:161 status=RUNNING
2026-09-02T05:23:04.748360Z [info ] sql completed successfully.
[airflow.task.operators.airflow.providers.databricks.operators.databricks.DatabricksSQLStatementsOperator]
loc=mixins.py:107
operator_source=<working-tree source>
stored={'databricks_sql_statement_id': 'statement-1'}
submissions=1
retry_statement_id=statement-1
```
</details>
---
Please check the type of change your PR introduces:
- [ ] Bugfix
- [x] Feature
- [ ] New Provider
- [ ] Improvement
- [ ] Documentation Fix
- [ ] Refactoring
- [ ] Other
### Relevant Issue(s)
None.
### Description
See above.
### How did you test it?
See the live proof above.
### Did you add documentation?
Yes. The SQL statements guide covers the recovery behavior and version floor.
### Does this introduce a breaking change?
No.
### Checklist
- [x] I have checked that there are no existing pull requests for the same
change.
- [x] I have added a commit message that describes my changes.
- [x] I have added tests that prove my fix is effective or that my feature
works.
- [x] I have updated the documentation accordingly.
### Do you use Generative AI to contribute to this PR?
- [ ] No
- [x] Yes, I used GitHub Copilot CLI (GPT-5.6 Sol) to assist with this pull
request.
### AI-assisted review checklist
- [x] I have manually reviewed all AI-generated code and verified its
correctness.
- [x] I have added comprehensive tests for AI-generated code paths.
- [x] I have verified that no sensitive data or secrets were included in AI
prompts.
- [x] I understand that I am responsible for the entire content of this PR.
### PR Checklist
- [x] I have locally rebased my branch onto the latest main.
- [x] I have run relevant tests and they pass.
- [x] I have run pre-commit checks and they pass.
- [x] I have updated documentation where applicable.
--
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]
