SCRUM-1160-adapters(bigquery): retry once on the exact bytes-billed shortfall

CHANGELOG.md

@@ -11,6 +11,27 @@ tag releases both in lockstep, so entries below are keyed by the engine version.
 
 ### Fixed
 
+- **`explore query` no longer fails inside its own confirmed budget on
+  BigQuery** ([#320]). The server-side `maximum_bytes_billed` cap was set to
+  this command's own reservation against the cumulative session ceiling
+  (sized to the dry-run estimate), not the wider per-command budget the
+  operator actually confirmed. BigQuery's own execution-time rounding of
+  bytes billed can exceed any dry-run estimate regardless of how accurate
+  that estimate was, so a multi-table statement confirmed at a budget six
+  times its estimate still failed with `bytesBilledLimitExceeded`, a
+  self-imposed cap the error message never named as the cause.
+
+  BigQuery's own refusal already states the exact byte count it needed
+  (`"163595928. 164626432 or higher required."`), so a statement that hits
+  the cap now widens its charge to precisely that number and retries once,
+  rather than guessing at a margin or handing the warehouse a cap wider than
+  what this command actually reserved. A retry that still can't fit the
+  confirmed ceiling refuses on the real number, exactly as it would have
+  without the retry; the concurrency guarantee a cumulative session ceiling
+  depends on (two commands sharing one ceiling can never jointly overspend
+  it, [#159]) is unaffected, since the widening goes through the same locked
+  admission path an estimate drifting past its booking already used.
+
 - **`get_dialect` now raises on an unrecognized connector instead of
   silently parsing every subsequent statement as DuckDB** ([#319]). A
   hyphenated BigQuery project id, the shape BigQuery itself hands out and

packages/dex-core/src/exmergo_dex_core/adapters/bigquery.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import re
 from collections.abc import Callable
 from dataclasses import dataclass
 from typing import Any
@@ -66,6 +67,16 @@
 # the math instead of letting the server fail the job after the fact.
 _MIN_BILLED_BYTES = 10 * 1024 * 1024
 
+# BigQuery's own refusal names the exact byte count it needed (issue #320):
+# "Query exceeded limit for bytes billed: <cap>. <required> or higher
+# required." No dry run predicts this number, since it reflects the server's
+# own execution-time rounding of bytes billed (observed per table scanned,
+# not per query), which is exactly the gap between a dry-run-based estimate
+# and what the job actually needed. Parsed rather than guessed at with a
+# margin, so the retry below asks for precisely what BigQuery says it needs,
+# whatever the underlying rounding rule turns out to be.
+_BYTES_BILLED_REQUIRED_RE = re.compile(r"(\d+) or higher required")
+
 # Field types whose values are nested or non-scalar: no approx-distinct, no
 # min/max, and non-null counting via COUNTIF (COUNT DISTINCT is invalid on
 # them and plain COUNT is not supported for every one of these types).
@@ -705,7 +716,7 @@ def exact_distinct_counts(
                 "not cover the extra scan; uniqueness verdicts stay approximate",
             )
             return {}
-        _job, iterator = self._run(sql)
+        _job, iterator = self._run(sql, floored)
         rows = list(iterator)
         return {name: int(rows[0][f"d_{i}"]) for i, name in enumerate(columns)}
 
@@ -737,7 +748,7 @@ def distinct_combination_counts(
                 "cover the extra scan; grain stays unknown",
             )
             return {}
-        _job, iterator = self._run(sql)
+        _job, iterator = self._run(sql, floored)
         rows = list(iterator)
         return {
             tuple(combo): int(rows[0][f"d_{i}"]) for i, combo in enumerate(combinations)
@@ -779,7 +790,7 @@ def value_domain_counts(
                 "cover the extra scan; no value domain reported",
             )
             return {}
-        _job, iterator = self._run(sql)
+        _job, iterator = self._run(sql, floored)
         rows = list(iterator)
         return {
             name: ValueDomainSample(
@@ -1054,19 +1065,30 @@ def _execute(
         """SELECT-only guard, free dry-run, gate charge, then the capped run."""
 
         assert_select_only(sql, dialect=self.dialect)
-        self.cost_gate.charge(self._dry_run(sql))
-        return self._run(sql, timeout_seconds=timeout_seconds, max_results=max_results)
+        estimate = self._dry_run(sql)
+        self.cost_gate.charge(estimate)
+        return self._run(
+            sql, estimate, timeout_seconds=timeout_seconds, max_results=max_results
+        )
 
     def _run(
         self,
         sql: str,
+        dry_run_estimate: float,
         *,
         timeout_seconds: float | None = None,
         max_results: int | None = None,
+        _retried: bool = False,
     ) -> tuple[Any, Any]:
         """The single billed door past the gate: run with the server-side byte
         cap, wait for completion (bounded when a timeout is given), account the
-        actual billed bytes, and return (job, row iterator)."""
+        actual billed bytes, and return (job, row iterator).
+
+        ``dry_run_estimate`` is what :meth:`_execute` already charged for this
+        exact statement; a bytes-billed refusal below widens the charge by the
+        gap between that estimate and what BigQuery says it actually needed,
+        rather than charging the full requirement a second time on top of it.
+        """
 
         cap = self.cost_gate.remaining_for_statement()
         if cap is not None and cap < _MIN_BILLED_BYTES:
@@ -1087,6 +1109,27 @@ def _run(
             iterator = job.result(timeout=timeout_seconds, max_results=max_results)
         except self._api_exceptions.BadRequest as exc:
             if "bytes billed" in str(exc) or "bytesBilledLimitExceeded" in str(exc):
+                required = _parse_bytes_billed_required(str(exc))
+                if required is not None and not _retried:
+                    # No dry run predicts BigQuery's own execution-time
+                    # rounding of bytes billed (issue #320), so the estimate
+                    # this command already charged for the statement can be
+                    # a genuine underestimate even though nothing was wrong
+                    # with it at dry-run time. BigQuery's own refusal names
+                    # exactly what it needed; widen the charge by the gap and
+                    # retry once with that as the new cap. Raises the same
+                    # OverCeilingError/ConfirmationRequiredError this would
+                    # raise anyway if the confirmed ceiling itself can't
+                    # cover the real requirement, so a genuine over-budget
+                    # query still refuses, correctly, on the real number.
+                    self.cost_gate.charge(required - dry_run_estimate)
+                    return self._run(
+                        sql,
+                        required,
+                        timeout_seconds=timeout_seconds,
+                        max_results=max_results,
+                        _retried=True,
+                    )
                 raise OverCeilingError(
                     "the query would bill more than the remaining budget "
                     "(server-side maximum_bytes_billed); raise --budget or "
@@ -1155,3 +1198,12 @@ def _quote_ident(name: str) -> str:
 
     escaped = name.replace("`", "\\`")
     return f"`{escaped}`"
+
+
+def _parse_bytes_billed_required(message: str) -> float | None:
+    """The byte count BigQuery's own bytes-billed refusal names as required,
+    or ``None`` if the message doesn't have the expected shape (a future
+    wording change should degrade to the old flat refusal, not a crash)."""
+
+    match = _BYTES_BILLED_REQUIRED_RE.search(message)
+    return float(match.group(1)) if match else None

packages/dex-core/tests/adapters/test_connect_bigquery.py

@@ -335,6 +335,81 @@ def test_server_side_cap_translates_when_the_estimate_drifts(fake_bq_client):
     assert "budget" in str(exc_info.value)
 
 
+# --- issue #320: the server-side cap must not be pinned below the confirmed
+# budget by BigQuery's own execution-time billing rounding, which no dry run
+# predicts. The bug's actual trigger is a `session_ceiling`: without one,
+# `remaining_for_statement()` already equals the confirmed ceiling directly
+# (see `test_every_executed_job_carries_maximum_bytes_billed` above); with
+# one, the per-statement cap is additionally bounded by this command's own
+# reservation (booked at confirm time, sized to the dry-run estimate), which
+# a dry run cannot widen for a gap execution alone reveals.
+
+
+def test_billing_rounding_retries_once_within_the_confirmed_budget(fake_bq_client):
+    # Dry run is accurate to within 1%: nothing was wrong with the estimate at
+    # dry-run time, but the real bill still exceeds the tight reservation a
+    # session_ceiling books it against, the way BigQuery's own execution-time
+    # rounding can regardless of how accurate the dry run was.
+    fake_bq_client.dry_run_underestimate = 0.99
+    fake_bq_client.tables["test-proj.shop.customers"].num_bytes = 160 * MB
+    adapter = make_adapter(
+        fake_bq_client,
+        ceiling=1_000 * MB,  # roughly 6x the estimate, like the reported budget
+        session_ceiling=10_000 * MB,
+    )
+    sql = "SELECT COUNT(*) FROM `test-proj`.`shop`.`customers`"
+    adapter.cost_gate.preflight_command(
+        adapter.query_estimate(sql)
+    )  # confirm handshake
+
+    adapter.run_query(sql, max_rows=10, timeout_seconds=30)
+
+    non_dry = [c for c in fake_bq_client.query_calls if not c.dry_run]
+    assert len(non_dry) == 2  # the too-tight attempt, then the retry
+    first_cap = non_dry[0].job_config.maximum_bytes_billed
+    second_cap = non_dry[1].job_config.maximum_bytes_billed
+    assert first_cap < 160 * MB  # pinned to the reservation, not the ceiling
+    assert second_cap == 160 * MB  # widened to exactly what BigQuery required
+
+
+def test_billing_rounding_retry_still_refuses_past_the_confirmed_ceiling(
+    fake_bq_client,
+):
+    # The real bill (200 MB) exceeds not just the reservation but the
+    # confirmed per-command ceiling itself (120 MB): the retry's own widening
+    # must refuse here, on the real number, rather than silently granting more
+    # than the operator confirmed.
+    fake_bq_client.dry_run_underestimate = 0.5
+    fake_bq_client.tables["test-proj.shop.customers"].num_bytes = 200 * MB
+    adapter = make_adapter(fake_bq_client, ceiling=120 * MB, session_ceiling=1_000 * MB)
+    sql = "SELECT COUNT(*) FROM `test-proj`.`shop`.`customers`"
+    adapter.cost_gate.preflight_command(adapter.query_estimate(sql))
+
+    with pytest.raises(OverCeilingError) as exc_info:
+        adapter.run_query(sql, max_rows=10, timeout_seconds=30)
+    assert "budget" in str(exc_info.value)
+    # No second server-side attempt: the widening itself refused before a
+    # retry could even be issued.
+    non_dry = [c for c in fake_bq_client.query_calls if not c.dry_run]
+    assert len(non_dry) == 1
+
+
+def test_parse_bytes_billed_required_reads_bigquerys_own_number():
+    from exmergo_dex_core.adapters.bigquery import _parse_bytes_billed_required
+
+    message = (
+        "Query exceeded limit for bytes billed: 163595928. "
+        "164626432 or higher required."
+    )
+    assert _parse_bytes_billed_required(message) == 164626432.0
+
+
+def test_parse_bytes_billed_required_degrades_to_none_on_an_unexpected_message():
+    from exmergo_dex_core.adapters.bigquery import _parse_bytes_billed_required
+
+    assert _parse_bytes_billed_required("bytesBilledLimitExceeded") is None
+
+
 def test_timeout_cancels_the_job(fake_bq_client):
     fake_bq_client.result_error = TimeoutError("deadline")
     adapter = make_adapter(fake_bq_client)