Reconnect to the running Livy batch on retry instead of resubmitting

[1fanwang]

Why

LivyOperator waiting synchronously (deferrable=False with polling_interval > 0) holds the
Spark batch on the worker. If the worker is lost mid-poll, the retry posts a brand-new Livy
batch — the original Spark application keeps running and the work is duplicated.

ResumableJobMixin (Airflow 3.3, AIP-103) exists to make exactly this synchronous-wait path
crash-safe: persist the external job id before polling, and on retry reconnect to the running job
instead of resubmitting. Livy is a clean fit — a synchronous submit-then-poll operator, the same
shape the mixin was built for and the SparkSubmitOperator already uses.

What

LivyOperator now subclasses ResumableJobMixin and routes its synchronous-poll path through
execute_resumable:

• submit_job posts the batch and returns its id; get_job_status / is_job_active /
  is_job_succeeded classify Livy BatchState; poll_until_complete reuses the existing
  poll_for_termination; get_job_result pushes the app_id XCom.
• The batch id is persisted to task_state_store before polling, so a retry reads it back and
  reconnects to the running batch.
• Deferrable (Triggerer owns the wait) and fire-and-forget (polling_interval=0, nothing
  to reconnect to) paths are untouched. An Airflow-2 stub keeps the provider importable on 2.x.

Crash-safety is opt-in through the mixin's durable flag (default on); set durable=False to keep
the always-resubmit behaviour.

Tests

A new TestLivyOperatorResumable suite (gated on Airflow 3.3+) covers fresh-submit-persists-before-poll,
the three retry decisions (reconnect / return / resubmit) across real BatchState values, graceful
degradation without a task_state_store, and durable=False. The existing LivyOperator suite is
unchanged.

End-to-end (live, Breeze)

A real worker crash during the synchronous wait, against an in-memory Livy stand-in that counts
POST /batches. A LivyOperator(durable=True, deferrable=False, polling_interval=3) submits a
batch; the worker is SIGKILLed mid-poll; the scheduler retries. Attempt 2 reads the persisted
batch id back, reconnects to the still-running batch, and finishes it — with no second submit.

Raw

attempt 1: POST /batches -> batch id 1, worker polling      (POST count = 1)
SIGKILL worker (pid 406)
attempt 2 (try_number=2):
   "Reconnecting to existing job"
   "Batch with id 1 terminated with state: success"
task run_batch: success
POST /batches count after both attempts: 1   <- single submit, reconnected; no duplicate batch

Risk

Only the synchronous-poll path changes; deferrable and fire-and-forget are byte-for-byte the same.
The reconnect logic is the shared mixin core, already covered by its own tests.

---

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

• Yes (please specify the tool below)

---

• 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.

[1fanwang]

cc @amoghrajesh @vikramkoka

[potiuk]

Real duplication bug: with deferrable=False and polling_interval > 0 the batch is held on the worker, so losing the worker mid-poll meant the retry posted a brand-new Livy batch while the original Spark application kept running. Duplicated work, doubled cluster cost, and nothing looks wrong afterwards because the retry succeeds.

I stacked this against the six providers that already adopted AIP-103 — redshift_data, spark_submit, databricks, bigquery, snowflake — and it matches the house pattern: same try/except ImportError fallback shim, external_id_key following the <system>_<id-noun> convention, and execute_resumable gated on a condition rather than applied unconditionally. The comment above the guard explaining why it exists and when it can go is a genuine improvement over the others, which just have the shim.

The thing I most expected to be wrong is handled by the mixin rather than here: a legitimately failed Spark job does not get reconnected to on retry — execute_resumable falls through and resubmits fresh — so ordinary retry semantics still work.

Two details worth crediting. Excluding the non-polling path is right, since there is nothing to reconnect to when the operator does not wait. And re-setting self._batch_id inside poll_until_complete so on_kill() can still delete the batch after a reconnect, where submit_job was skipped, is the kind of thing that usually only surfaces in production.

I pushed one small change: the fallback stub's docstring and external_id_key now match the other five providers, which ship a byte-identical shim. Keeping the wording aligned means the seven copies stay greppable as one set — which matters, because that is now seven copies of the same shim with no shared home in common.compat. Worth someone folding them together, though not in this PR.

---

Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting