Repository: googleapis/python-bigquery

Issue:
bug: client doesn't retry "Job exceeded rate limits" for DDL query jobs that exceed quota for table update operations 
In https://github.com/googleapis/python-bigquery-sqlalchemy/pull/1009#discussion_r1457644849 it seems that the query in https://btx-internal.corp.google.com/invocations/ffafb866-6bc0-423f-a86b-df69fb270d57/targets/cloud-devrel%2Fclient-libraries%2Fpython%2Fgoogleapis%2Fpython-bigquery-sqlalchemy%2Fpresubmit%2Fprerelease-deps;config=default/log with rate limits exceeded errors are not retried.

#### Environment details

  - OS type and version:
  - Python version: `python --version`
  - pip version: `pip --version`
  - `google-cloud-bigquery` version: `pip show google-cloud-bigquery`

#### Steps to reproduce

Run a DDL query more than 5 times in 10 seconds, violating the five table metadata update operations per 10 seconds per table limit (https://cloud.google.com/bigquery/quotas#standard_tables).

#### Code example

```python
import google.cloud.bigquery
bqclient = google.cloud.bigquery.Client()
sql = "ALTER TABLE `swast-scratch.my_dataset.my_table` ADD COLUMN IF NOT EXISTS my_string_col STRING"
for _ in range(100):
    bqclient.query_and_wait(sql)
```

#### Stack trace

```
BadRequest                                Traceback (most recent call last)
Input In [4], in <cell line: 1>()
      1 for _ in range(100):
----> 2     bqclient.query_and_wait(sql)

File ~/src/github.com/googleapis/python-bigquery/google/cloud/bigquery/client.py:3503, in Client.query_and_wait(self, query, job_config, location, project, api_timeout, wait_timeout, retry, job_retry, page_size, max_results)
   3497     _verify_job_config_type(job_config, QueryJobConfig)
   3499 job_config = _job_helpers.job_config_with_defaults(
   3500     job_config, self._default_query_job_config
   3501 )
-> 3503 return _job_helpers.query_and_wait(
   3504     self,
   3505     query,
   3506     job_config=job_config,
   3507     location=location,
   3508     project=project,
   3509     api_timeout=api_timeout,
   3510     wait_timeout=wait_timeout,
   3511     retry=retry,
   3512     job_retry=job_retry,
   3513     page_size=page_size,
   3514     max_results=max_results,
   3515 )

File ~/src/github.com/googleapis/python-bigquery/google/cloud/bigquery/_job_helpers.py:498, in query_and_wait(client, query, job_config, location, project, api_timeout, wait_timeout, retry, job_retry, page_size, max_results)
    481     return table.RowIterator(
    482         client=client,
    483         api_request=functools.partial(client._call_api, retry, timeout=api_timeout),
   (...)
    494         num_dml_affected_rows=query_results.num_dml_affected_rows,
    495     )
    497 if job_retry is not None:
--> 498     return job_retry(do_query)()
    499 else:
    500     return do_query()

File /opt/miniconda3/envs/dev-3.10/lib/python3.10/site-packages/google/api_core/retry.py:349, in Retry.__call__.<locals>.retry_wrapped_func(*args, **kwargs)
    345 target = functools.partial(func, *args, **kwargs)
    346 sleep_generator = exponential_sleep_generator(
    347     self._initial, self._maximum, multiplier=self._multiplier
    348 )
--> 349 return retry_target(
    350     target,
    351     self._predicate,
    352     sleep_generator,
    353     self._timeout,
    354     on_error=on_error,
    355 )

File /opt/miniconda3/envs/dev-3.10/lib/python3.10/site-packages/google/api_core/retry.py:191, in retry_target(target, predicate, sleep_generator, timeout, on_error, **kwargs)
    189 for sleep in sleep_generator:
    190     try:
--> 191         return target()
    193     # pylint: disable=broad-except
    194     # This function explicitly must deal with broad exceptions.
    195     except Exception as exc:

File ~/src/github.com/googleapis/python-bigquery/google/cloud/bigquery/_job_helpers.py:439, in query_and_wait.<locals>.do_query()
    437 # For easier testing, handle the retries ourselves.
    438 if retry is not None:
--> 439     response = retry(client._call_api)(
    440         retry=None,  # We're calling the retry decorator ourselves.
    441         span_name="BigQuery.query",
    442         span_attributes=span_attributes,
    443         method="POST",
    444         path=path,
    445         data=request_body,
    446         timeout=api_timeout,
    447     )
    448 else:
    449     response = client._call_api(
    450         retry=None,
    451         span_name="BigQuery.query",
   (...)
    456         timeout=api_timeout,
    457     )

File /opt/miniconda3/envs/dev-3.10/lib/python3.10/site-packages/google/api_core/retry.py:349, in Retry.__call__.<locals>.retry_wrapped_func(*args, **kwargs)
    345 target = functools.partial(func, *args, **kwargs)
    346 sleep_generator = exponential_sleep_generator(
    347     self._initial, self._maximum, multiplier=self._multiplier
    348 )
--> 349 return retry_target(
    350     target,
    351     self._predicate,
    352     sleep_generator,
    353     self._timeout,
    354     on_error=on_error,
    355 )

File /opt/miniconda3/envs/dev-3.10/lib/python3.10/site-packages/google/api_core/retry.py:191, in retry_target(target, predicate, sleep_generator, timeout, on_error, **kwargs)
    189 for sleep in sleep_generator:
    190     try:
--> 191         return target()
    193     # pylint: disable=broad-except
    194     # This function explicitly must deal with broad exceptions.
    195     except Exception as exc:

File ~/src/github.com/googleapis/python-bigquery/google/cloud/bigquery/client.py:827, in Client._call_api(self, retry, span_name, span_attributes, job_ref, headers, **kwargs)
    823 if span_name is not None:
    824     with create_span(
    825         name=span_name, attributes=span_attributes, client=self, job_ref=job_ref
    826     ):
--> 827         return call()
    829 return call()

File /opt/miniconda3/envs/dev-3.10/lib/python3.10/site-packages/google/cloud/_http/__init__.py:494, in JSONConnection.api_request(self, method, path, query_params, data, content_type, headers, api_base_url, api_version, expect_json, _target_object, timeout, extra_api_info)
    482 response = self._make_request(
    483     method=method,
    484     url=url,
   (...)
    490     extra_api_info=extra_api_info,
    491 )
    493 if not 200 <= response.status_code < 300:
--> 494     raise exceptions.from_http_response(response)
    496 if expect_json and response.content:
    497     return response.json()

BadRequest: 400 POST https://bigquery.googleapis.com/bigquery/v2/projects/swast-scratch/queries?prettyPrint=false: Job exceeded rate limits: Your table exceeded quota for table update operations. For more information, see https://cloud.google.com/bigquery/docs/troubleshoot-quotas

In [5]: import sys

In [6]: exc = sys.last_value

In [7]: exc
Out[7]: google.api_core.exceptions.BadRequest('POST https://bigquery.googleapis.com/bigquery/v2/projects/swast-scratch/queries?prettyPrint=false: Job exceeded rate limits: Your table exceeded quota for table update operations. For more information, see https://cloud.google.com/bigquery/docs/troubleshoot-quotas')

In [8]: exc.reason

In [9]: exc.errors
Out[9]: 
[{'message': 'Job exceeded rate limits: Your table exceeded quota for table update operations. For more information, see https://cloud.google.com/bigquery/docs/troubleshoot-quotas',
  'domain': 'global',
  'reason': 'jobRateLimitExceeded'}]

In [10]: exc.errors[0]["reason"]
Out[10]: 'jobRateLimitExceeded'
```


Discussion:
A fix should just require adding "jobRateLimitExceeded" to https://github.com/googleapis/python-bigquery/blob/main/google/cloud/bigquery/retry.py#L76 but I'd also like to see a unit test where the raised exception looks like the one above.

Resolving patch:
diff --git a/google/cloud/bigquery/retry.py b/google/cloud/bigquery/retry.py
--- a/google/cloud/bigquery/retry.py
+++ b/google/cloud/bigquery/retry.py
@@ -73,7 +73,7 @@ def _should_retry(exc):
 deadline on the retry object.
 """
 
-job_retry_reasons = "rateLimitExceeded", "backendError"
+job_retry_reasons = "rateLimitExceeded", "backendError", "jobRateLimitExceeded"
 
 
 def _job_should_retry(exc):


Test patch:
diff --git a/tests/unit/test_job_retry.py b/tests/unit/test_job_retry.py
--- a/tests/unit/test_job_retry.py
+++ b/tests/unit/test_job_retry.py
@@ -22,6 +22,10 @@
 import google.api_core.retry
 import freezegun
 
+from google.cloud.bigquery.client import Client
+from google.cloud.bigquery import _job_helpers
+from google.cloud.bigquery.retry import DEFAULT_JOB_RETRY
+
 from .helpers import make_connection
 
 
@@ -240,3 +244,79 @@ def test_raises_on_job_retry_on_result_with_non_retryable_jobs(client):
         ),
     ):
         job.result(job_retry=google.api_core.retry.Retry())
+
+
+def test_query_and_wait_retries_job_for_DDL_queries():
+    """
+    Specific test for retrying DDL queries with "jobRateLimitExceeded" error:
+    https://github.com/googleapis/python-bigquery/issues/1790
+    """
+    freezegun.freeze_time(auto_tick_seconds=1)
+    client = mock.create_autospec(Client)
+    client._call_api.__name__ = "_call_api"
+    client._call_api.__qualname__ = "Client._call_api"
+    client._call_api.__annotations__ = {}
+    client._call_api.__type_params__ = ()
+    client._call_api.side_effect = (
+        {
+            "jobReference": {
+                "projectId": "response-project",
+                "jobId": "abc",
+                "location": "response-location",
+            },
+            "jobComplete": False,
+        },
+        google.api_core.exceptions.InternalServerError(
+            "job_retry me", errors=[{"reason": "jobRateLimitExceeded"}]
+        ),
+        google.api_core.exceptions.BadRequest(
+            "retry me", errors=[{"reason": "jobRateLimitExceeded"}]
+        ),
+        {
+            "jobReference": {
+                "projectId": "response-project",
+                "jobId": "abc",
+                "location": "response-location",
+            },
+            "jobComplete": True,
+            "schema": {
+                "fields": [
+                    {"name": "full_name", "type": "STRING", "mode": "REQUIRED"},
+                    {"name": "age", "type": "INT64", "mode": "NULLABLE"},
+                ],
+            },
+            "rows": [
+                {"f": [{"v": "Whillma Phlyntstone"}, {"v": "27"}]},
+                {"f": [{"v": "Bhetty Rhubble"}, {"v": "28"}]},
+                {"f": [{"v": "Phred Phlyntstone"}, {"v": "32"}]},
+                {"f": [{"v": "Bharney Rhubble"}, {"v": "33"}]},
+            ],
+        },
+    )
+    rows = _job_helpers.query_and_wait(
+        client,
+        query="SELECT 1",
+        location="request-location",
+        project="request-project",
+        job_config=None,
+        page_size=None,
+        max_results=None,
+        retry=DEFAULT_JOB_RETRY,
+        job_retry=DEFAULT_JOB_RETRY,
+    )
+    assert len(list(rows)) == 4
+
+    # Relevant docs for the REST API path: https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query
+    # and https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/getQueryResults
+    query_request_path = "/projects/request-project/queries"
+
+    calls = client._call_api.call_args_list
+    _, kwargs = calls[0]
+    assert kwargs["method"] == "POST"
+    assert kwargs["path"] == query_request_path
+
+    # TODO: Add assertion statements for response paths after PR#1797 is fixed
+
+    _, kwargs = calls[3]
+    assert kwargs["method"] == "POST"
+    assert kwargs["path"] == query_request_path
diff --git a/tests/unit/test_retry.py b/tests/unit/test_retry.py
--- a/tests/unit/test_retry.py
+++ b/tests/unit/test_retry.py
@@ -129,3 +129,30 @@ def test_DEFAULT_JOB_RETRY_deadline():
 
     # Make sure we can retry the job at least once.
     assert DEFAULT_JOB_RETRY._deadline > DEFAULT_RETRY._deadline
+
+
+def test_DEFAULT_JOB_RETRY_job_rate_limit_exceeded_retry_predicate():
+    """Tests the retry predicate specifically for jobRateLimitExceeded."""
+    from google.cloud.bigquery.retry import DEFAULT_JOB_RETRY
+    from google.api_core.exceptions import ClientError
+
+    # Non-ClientError exceptions should never trigger a retry
+    assert not DEFAULT_JOB_RETRY._predicate(TypeError())
+
+    # ClientError without specific reason shouldn't trigger a retry
+    assert not DEFAULT_JOB_RETRY._predicate(ClientError("fail"))
+
+    # ClientError with generic reason "idk" shouldn't trigger a retry
+    assert not DEFAULT_JOB_RETRY._predicate(
+        ClientError("fail", errors=[dict(reason="idk")])
+    )
+
+    # ClientError with reason "jobRateLimitExceeded" should trigger a retry
+    assert DEFAULT_JOB_RETRY._predicate(
+        ClientError("fail", errors=[dict(reason="jobRateLimitExceeded")])
+    )
+
+    # Other retryable reasons should still work as expected
+    assert DEFAULT_JOB_RETRY._predicate(
+        ClientError("fail", errors=[dict(reason="backendError")])
+    )