[Python] Reset BigQuery retries for each batch

[bvolpato]

BigQuery streaming inserts can send a recoverable row to failed_rows because an earlier successful batch consumed its retry allowance. The pipeline can finish successfully while that row is missing from BigQuery.

The retry iterator currently belongs to the bundle and advances after every request, including successful requests. This also increases the backoff for later batches: after 20 successful batches, their first transient retry can sleep for 750 to 1500 seconds.

Create the iterator per flushed batch, matching the documented per-group retry limit. Failed rows and their insert IDs retain the existing retry behavior.

Reproduction with BigQuery

I ran a public WriteToBigQuery pipeline against a temporary table in my own GCP project. After a successful real insertAll request, the client injects one row-level backendError before sending the next request. This makes the transient failure reproducible without trying to induce a BigQuery outage.

Run	Rows verified in BigQuery	Failed-row output
Current upstream	order-200	order-100
This PR	order-100, order-200	Empty
Upstream without fault injection	order-100, order-200	Empty

Every run finished DONE. With the fix, the rejected batch retries using the same insert ID and the real request succeeds. Shuffle order can differ, so the injected error targets the second request rather than a specific order ID. The error is controlled; successful requests and table contents were verified against BigQuery.

The reduced pipeline below uses an existing table with schema run_mode:STRING,order_id:STRING,cents:INTEGER. Run it with the SDK being compared installed, GCP dependencies available, and gcloud authenticated:

uv run --active --no-project python retry_budget_repro.py PROJECT:DATASET.TABLE upstream

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.12"
# dependencies = ["apache-beam[gcp]", "requests"]
# ///
import json
from pathlib import Path
import subprocess
import sys
import tempfile

import apache_beam as beam
import requests
from apache_beam.io.gcp import bigquery_tools
from apache_beam.io.gcp.bigquery import BigQueryDisposition
from apache_beam.io.gcp.bigquery import WriteToBigQuery
from apache_beam.options.pipeline_options import PipelineOptions


class ControlledInsertClient:
    def __init__(self, trace):
        self.trace = trace

    def insert_rows_json(self, table, json_rows, row_ids, **kwargs):
        calls = Path(self.trace).read_text().splitlines()
        if len(calls) == 1:
            errors = [{"index": 0, "errors": [{"reason": "backendError"}]}]
            response = "injected backendError"
        else:
            project, dataset, table_id = table.split(".")
            token = subprocess.run(
                ["gcloud", "auth", "print-access-token"],
                check=True, capture_output=True, text=True).stdout.strip()
            result = requests.post(
                f"https://bigquery.googleapis.com/bigquery/v2/projects/{project}"
                f"/datasets/{dataset}/tables/{table_id}/insertAll",
                headers={"Authorization": "Bearer " + token},
                json={
                    "skipInvalidRows": True,
                    "rows": [{"insertId": row_id, "json": row}
                             for row, row_id in zip(json_rows, row_ids)]},
                timeout=30)
            result.raise_for_status()
            errors = result.json().get("insertErrors", [])
            response = "real insertAll success" if not errors else errors
        with open(self.trace, "a") as handle:
            handle.write(json.dumps({
                "rows": json_rows, "insert_ids": row_ids,
                "response": response}) + "\n")
        return errors


def print_failed(row):
    print("FAILED_ROW=" + json.dumps(row))


table, run_mode = sys.argv[1:3]
with tempfile.TemporaryDirectory() as directory:
    trace = Path(directory) / "requests.jsonl"
    trace.touch()
    pipeline = beam.Pipeline(options=PipelineOptions([
        "--runner=DirectRunner", "--direct_num_workers=1",
        "--direct_running_mode=in_memory"]))
    result = (
        pipeline
        | beam.Create([
            {"run_mode": run_mode, "order_id": "order-100", "cents": 100},
            {"run_mode": run_mode, "order_id": "order-200", "cents": 200}])
        | WriteToBigQuery(
            table,
            method=WriteToBigQuery.Method.STREAMING_INSERTS,
            create_disposition=BigQueryDisposition.CREATE_NEVER,
            write_disposition=BigQueryDisposition.WRITE_APPEND,
            batch_size=1,
            num_streaming_keys=1,
            max_retries=1,
            insert_retry_strategy=bigquery_tools.RetryStrategy.RETRY_ON_TRANSIENT_ERROR,
            test_client=ControlledInsertClient(str(trace))))
    result.failed_rows | beam.Map(print_failed)
    print("PIPELINE_STATE=" + pipeline.run().wait_until_finish())
    print(trace.read_text())

Validation

Compared the complete BigQuery modules from current upstream and this PR over the same SDK dependency environment. The SDK selected Prism for the public pipeline. This comparison did not rebuild the entire SDK.

The focused regressions fail on upstream and pass with the fix. They cover successful prior batches, separate destinations, an earlier exhausted batch, and initial backoff. The zero-retry control passes on both implementations. The affected streaming transform, error-handling, and pipeline suites pass, as do repository-pinned YAPF, Ruff, and whitespace checks.

From sdks/python, with the SDK and GCP test dependencies installed:

python -m pytest apache_beam/io/gcp/bigquery_test.py -q \
  -k 'successful_batches or exhausted_batch or zero_retries'

Credential discovery was stubbed only for the mocked local test suites. The live table reproduction used authenticated BigQuery requests.

Downsides

Each failing batch now receives its configured retry allowance independently. Several failing batches can therefore make more total attempts than the previous accidental bundle-wide limit. The configured per-batch limit is unchanged.

---

• If this contribution is large, please file an Apache Individual Contributor License Agreement.
• Replace these instructions with a description of your change and why you have made it.
• Mention the appropriate issue in your description (for example: addresses #123 or fixes #456).
• Update CHANGES.md with noteworthy changes.

[github-actions]

Assigning reviewers:

R: @claudevdm for label python.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 56.23%. Comparing base (76f6e6d) to head (3f16104).
⚠️ Report is 111 commits behind head on master.

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #39995      +/-   ##
============================================
+ Coverage     56.16%   56.23%   +0.06%     
  Complexity     2288     2288              
============================================
  Files          1124     1124              
  Lines        178336   178359      +23     
  Branches       1489     1489              
============================================
+ Hits         100170   100293     +123     
+ Misses        75645    75545     -100     
  Partials       2521     2521              

Flag	Coverage Δ
python	79.68% <100.00%> (+0.12%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[claudevdm]

Thanks!