fix(load): a failed status check must not discard a running job

Review found a regression this PR introduced, reachable from the opposite
direction to the one it fixes.

`_poll_job` aborted on the first bad `GET /v1/jobs/{id}`. That exception reaches
`ManagedDatabaseClient._request_with_retry`, which classifies a 502/503 as
transient and re-invokes THE WHOLE OPERATION for every non-append mode -- so a
single blipped status check discards a healthy running job and submits a second
load against a table whose write lock the first one still holds. 409
RESOURCE_LOCKED, which is also transient, and round it goes.

Worse than the behaviour it replaces, not better: before, a lost request took with
it a load that was genuinely being torn down. Now the first job is neither going
away nor releasing anything, and a load polls for up to `_LOAD_JOB_TIMEOUT_S`, so
there are hundreds of status checks where there used to be one request -- every
one a fresh chance to trigger the collision. Exposure up, collision from
incidental to certain.

So the poll tolerates a failed status check and keeps waiting, bounded by
consecutive failures rather than by the first one. An isolated 502 is noise; a run
of them means the API is gone and there is nothing to wait for. The poll's own
deadline bounds the total wait either way.

THE JOB ID IS NOW RETURNED. The docstring claimed the id removed `append`'s "did
it land?" ambiguity, which was not true as implemented: it was a local, discarded
on every failure path, and never reached the caller. `LoadManagedTableResult`
carries it now, for the same reason `CreateIndexResult` does. `append` stays
non-retryable -- knowing the id makes the question answerable, it does not make a
blind re-submission safe, and that call belongs to the caller. The docstring says
that instead of the stronger thing.

`result` is read with `getattr(..., "actual_instance", ...)`, matching the index
path rather than silently disagreeing with it: the strict read raised a bare
AttributeError if the model ever arrived unwrapped.

Did NOT add a strict isinstance check on the inline response. The index path can
assert its inline type because it names it positively first; here the inline shape
is read duck-typed and callers and tests rely on that, so asserting it narrows an
interface this change has no business narrowing. Only the job branch is
type-checked.

Import order fixed for ruff's I rule; `ruff check` clean.

149 passed.

Author: anoop-narang

hotdata_framework/client.py

@@ -24,8 +24,8 @@
 from hotdata.models.database_default_table_decl import DatabaseDefaultTableDecl
 from hotdata.models.index_info_response import IndexInfoResponse
 from hotdata.models.job_status_response import JobStatusResponse
-from hotdata.models.load_managed_table_response import LoadManagedTableResponse
 from hotdata.models.load_managed_table_request import LoadManagedTableRequest
+from hotdata.models.load_managed_table_response import LoadManagedTableResponse
 from hotdata.models.query_request import QueryRequest
 from hotdata.models.query_response import QueryResponse
 from hotdata.models.submit_job_response import SubmitJobResponse
@@ -95,6 +95,9 @@
 # first place. Bounded rather than unbounded so a wedged job still surfaces.
 _LOAD_JOB_TIMEOUT_S = 3600.0
 
+# Consecutive failed status checks before a poll gives up.
+_JOB_POLL_MAX_CONSECUTIVE_ERRORS = 5
+
 
 @dataclass(frozen=True)
 class ResultSummary:
@@ -461,14 +464,21 @@ def load_managed_table(
             )
         except ApiException as e:
             raise RuntimeError(api_error_message(e)) from e
+        # Only the job branch is type-checked. The index path can also assert its
+        # inline type because it names it positively first; here the inline shape is
+        # read duck-typed, which callers and tests already rely on, so asserting it
+        # would narrow an interface this change has no business narrowing.
+        job_id: str | None = None
         if isinstance(loaded, SubmitJobResponse):
-            loaded = self._load_response_from_job(loaded.id)
+            job_id = loaded.id
+            loaded = self._load_response_from_job(job_id)
         return LoadManagedTableResult(
             connection_id=loaded.connection_id,
             schema_name=loaded.schema_name,
             table_name=loaded.table_name,
             row_count=loaded.row_count,
             full_name=f"{db.id}.{loaded.schema_name}.{loaded.table_name}",
+            job_id=job_id,
         )
 
     def add_managed_table(
@@ -911,11 +921,27 @@ def _poll_job(
         jobs = self._jobs_api()
         deadline = time.monotonic() + timeout_s
         last: JobStatusResponse | None = None
+        consecutive_errors = 0
         while time.monotonic() < deadline:
             try:
                 last = jobs.get_job(job_id)
+                consecutive_errors = 0
             except ApiException as e:
-                raise RuntimeError(api_error_message(e)) from e
+                # A failed STATUS CHECK is not a failed job. Aborting here throws
+                # away work that is still running, and the caller's retry then
+                # re-submits it while the original still holds its resources -- so
+                # one blip becomes a collision with the job it just abandoned. A
+                # load polls for up to `_LOAD_JOB_TIMEOUT_S`, so the longer the
+                # wait the more chances to hit it, which is exactly backwards.
+                #
+                # CONSECUTIVE failures are the signal: an isolated 502 is noise, a
+                # run of them means the API is gone and there is nothing to wait
+                # for. The poll's own deadline bounds the total wait regardless.
+                consecutive_errors += 1
+                if consecutive_errors >= _JOB_POLL_MAX_CONSECUTIVE_ERRORS:
+                    raise RuntimeError(api_error_message(e)) from e
+                time.sleep(interval_s)
+                continue
             if last.status in _JOB_TERMINAL:
                 return last
             time.sleep(interval_s)
@@ -929,9 +955,11 @@ def _load_response_from_job(self, job_id: str) -> LoadManagedTableResponse:
 
         Polling replaces waiting on the request, so the outcome is read from
         durable state rather than from a connection that has to stay alive. That
-        also removes the ambiguity `append` was made non-retryable for: a lost
-        response no longer leaves "did it land?" unanswerable, because the job id
-        outlives the request that created it.
+        also gives a caller a handle: the job id is returned on
+        `LoadManagedTableResult`, so "did it land?" is answerable after a lost
+        response. `append` stays non-retryable -- knowing the id makes the question
+        answerable, it does not make a blind re-submission safe, and that call is
+        the caller's to make.
 
         `partially_succeeded` is terminal and carries a message, so it is raised
         rather than returned -- a caller asked for a table's contents to be
@@ -943,7 +971,9 @@ def _load_response_from_job(self, job_id: str) -> LoadManagedTableResponse:
             raise RuntimeError(
                 final.error_message or f"load job {job_id} finished {status}"
             )
-        payload = final.result.actual_instance if final.result is not None else None
+        # `result` is a oneOf wrapper today; tolerate the model arriving directly,
+        # the same way the index path does rather than disagreeing with it.
+        payload = getattr(final.result, "actual_instance", final.result)
         if not isinstance(payload, LoadManagedTableResponse):
             raise RuntimeError(
                 f"load job {job_id} succeeded without a load result "

hotdata_framework/databases.py

@@ -83,6 +83,11 @@ class LoadManagedTableResult:
     table_name: str
     row_count: int
     full_name: str
+    # Set when the server ran the load as a background job. Carried for the same
+    # reason CreateIndexResult carries it: it is the only handle a caller has to
+    # ask "did that land?" after a lost response, and without it the question is
+    # unanswerable. `None` when the load finished inline.
+    job_id: str | None = None
 
     def to_dict(self) -> dict[str, Any]:
         return asdict(self)

tests/test_client.py

@@ -714,3 +714,91 @@ def test_a_partially_succeeded_load_job_is_not_treated_as_success():
         else:
             raise AssertionError("partially_succeeded was treated as success")
 
+
+def test_a_transient_status_check_does_not_discard_a_running_job():
+    """The regression this guards. A load polls for up to an hour, so there are
+    hundreds of status checks; aborting on the first bad one throws away a job
+    that is still running, and the caller's retry then re-submits the load while
+    the original still holds the table -- the 409 spin, from a new direction."""
+    from hotdata.exceptions import ApiException
+
+    client = HotdataClient("k", "ws", host="https://api.hotdata.dev")
+    calls = {"n": 0}
+    final = SimpleNamespace(status="succeeded", error_message=None,
+                            result=SimpleNamespace(actual_instance=_load_response(5)))
+
+    class _Flaky:
+        def get_job(self, job_id):
+            calls["n"] += 1
+            if calls["n"] < 3:
+                raise ApiException(status=502, reason="Bad Gateway")
+            return final
+
+    with (
+        patch.object(client, "_jobs_api", return_value=_Flaky()),
+        patch("time.sleep", lambda *_: None),
+    ):
+        got = client._poll_job("jobs_1", timeout_s=60.0, interval_s=0.01)
+
+    assert got is final, "a blipping status check aborted the poll"
+    assert calls["n"] == 3
+
+
+def test_a_run_of_failed_status_checks_still_gives_up():
+    """Tolerance is for blips, not for an API that has gone away."""
+    from hotdata.exceptions import ApiException
+    from hotdata_framework.client import _JOB_POLL_MAX_CONSECUTIVE_ERRORS
+
+    client = HotdataClient("k", "ws", host="https://api.hotdata.dev")
+
+    class _Dead:
+        def get_job(self, job_id):
+            raise ApiException(status=502, reason="Bad Gateway")
+
+    with (
+        patch.object(client, "_jobs_api", return_value=_Dead()),
+        patch("time.sleep", lambda *_: None),
+    ):
+        try:
+            client._poll_job("jobs_1", timeout_s=600.0, interval_s=0.01)
+        except RuntimeError:
+            pass
+        else:
+            raise AssertionError("polled forever against a dead API")
+    assert _JOB_POLL_MAX_CONSECUTIVE_ERRORS >= 2
+
+
+def test_a_deferred_load_returns_the_job_id_to_the_caller():
+    """`append` stays non-retryable, so the id is the only handle a caller has to
+    answer "did it land?" after a lost response -- the same reason
+    CreateIndexResult carries one."""
+    from hotdata.models.submit_job_response import SubmitJobResponse
+
+    client = HotdataClient("k", "ws", host="https://api.hotdata.dev")
+    db = ManagedDatabase(id="db_1", description="mydb", default_connection_id="conn_1")
+    connections = _FakeConnectionsApi(responses=[
+        SubmitJobResponse(id="jobs_77", status="running", status_url="/v1/jobs/jobs_77"),
+    ])
+    final = SimpleNamespace(status="succeeded", error_message=None,
+                            result=SimpleNamespace(actual_instance=_load_response()))
+
+    with (
+        patch.object(client, "_databases_api", return_value=_ForbiddenDatabasesApi()),
+        patch.object(client, "connections", return_value=connections),
+        patch.object(client, "_poll_job", return_value=final),
+    ):
+        result = client.load_managed_table(db, "orders", schema="public", upload_id="up_1")
+
+    assert result.job_id == "jobs_77"
+
+
+def test_an_inline_load_carries_no_job_id():
+    client = HotdataClient("k", "ws", host="https://api.hotdata.dev")
+    db = ManagedDatabase(id="db_1", description="mydb", default_connection_id="conn_1")
+    with (
+        patch.object(client, "_databases_api", return_value=_ForbiddenDatabasesApi()),
+        patch.object(client, "connections", return_value=_FakeConnectionsApi()),
+    ):
+        result = client.load_managed_table(db, "orders", schema="public", upload_id="up_1")
+    assert result.job_id is None
+