fix: retry 'job exceeded rate limits' for DDL queries (#1794)

* fix: retry 'job exceeded rate limits' for DDL queries * Fixed retry test logic to better align to library standards * added docstring for test * deleted extra coverage file * Update tests/unit/test_job_retry.py Co-authored-by: Tim Swast <swast@google.com> * requested changes to retry jobs test * slight change to assert statemet * added TODO statements and fixed default job retry * modify sleep time and path names --------- Co-authored-by: Tim Swast <swast@google.com>
googleapis · Jan 24, 2024 · 39f33b2 · 39f33b2
1 parent b402a6d
commit 39f33b2
Show file tree

Hide file tree

Showing 3 changed files with 108 additions and 1 deletion.
diff --git 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):

diff --git 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
@@ -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")])
+    )