Repository: googleapis/python-bigquery (python)

Issue:
`Client.query()` sometimes raises 404, despite the query job being successfully started
#### Environment details

  - OS type and version: Google Colab
  - Python version: `python --version` 3.11
  - pip version: `pip --version` N/A
  - `google-cloud-bigquery` version: `pip show google-cloud-bigquery` 3.25.0

#### Steps to reproduce

It's a bit difficult to reproduce, but it seems to happen most often with creating a vector index via [CREATE VECTOR INDEX DDL](https://cloud.google.com/bigquery/docs/vector-index#create_a_vector_index) as generated by [bigframes.bigquery.create_vector_index](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.bigquery#bigframes_bigquery_create_vector_index). 

#### Code example

<details><summary>Long Python sample, sorry</summary>

```python
PROJECT_ID = "my-project"  # @param {type:"string"}
DATASET_ID = "my_dataset"  # @param {type:"string"}
CONNECTION_ID = "bqml_llm_conn"  # @param {type:"string"}

# Set the project id
! gcloud config set project {PROJECT_ID}

REGION = "US"  # @param {type: "string"}

## Text embedding variables

TEXT_TABLE_ID = "PATENT_TABLE_BF"
TEXT_MODEL_ID = "llm_embedding_model"
TEXT_EMBEDDING_TABLE_ID = "patent_embedding_BF_test"
TEXT_EMBEDDING_TABLE_ID2 = "patent_embedding"
TEXT_MODEL_ENDPOINT = "text-embedding-004"
TEXT_VECTOR_INDEX_ID = "patent_index_BF"

import bigframes.pandas as bf
import bigframes.ml as bf_ml
import bigframes.bigquery as bf_bq
import bigframes.ml.llm as bf_llm


from google.cloud import bigquery
from google.cloud import storage

# Construct a BigQuery client object.
client = bigquery.Client()

import pandas as pd
from IPython.display import Image, display
from PIL import Image as PILImage
import io

bf.options.bigquery.project = PROJECT_ID
bf.options.bigquery.location = REGION
bf.options.bigquery.ordering_mode = "partial"

text_model = bf_llm.TextEmbeddingGenerator() ## which endpoint this uses??

publications = bf.read_gbq('patents-public-data.google_patents_research.publications')
keep = (publications.embedding_v1.str.len() > 0) & (publications.title.str.len() > 0) & (publications.abstract.str.len() > 0)
publications = publications[["publication_number", "title", "abstract"]].rename(columns={'abstract': 'content'})

## creating embeddings for 10000 rows only for testing, change in the line below for embedding higher number of rows
publications_subset = publications[publications["content"].str.len() > 30].peek(10000)
## publications_subset = publications_subset[["publication_number", "title", "abstract"]].rename(columns={'abstract': 'content'})


embedding = text_model.predict(publications_subset)[["publication_number", "title", "content", "ml_generate_embedding_result","ml_generate_embedding_status"]]

## filter out rows where the embedding generation failed. the embedding status value is empty if the embedding generation was successful
embedding = embedding[~embedding["ml_generate_embedding_status"].isnull()]

# store embeddings in a BQ table
embedding.to_gbq(f"{DATASET_ID}.{TEXT_EMBEDDING_TABLE_ID}", if_exists='replace')

## python code to create index
bf_bq.create_vector_index(
    table_id = f"{DATASET_ID}.{TEXT_EMBEDDING_TABLE_ID}",
    column_name = "ml_generate_embedding_result",
    replace= True,
    index_name = "bf_python_index",
    distance_type="cosine",
    index_type= "ivf"
)
```

</details>

Unfortunately, even this is not a guarantee, but it does seem pretty reproducible on one of the projects my colleague is using.

#### Stack trace

```
/usr/local/lib/python3.10/dist-packages/bigframes/session/_io/bigquery/__init__.py in start_query_with_client(bq_client, sql, job_config, max_results, timeout, api_name, metrics)
    229 
    230     try:
--> 231         query_job = bq_client.query(sql, job_config=job_config, timeout=timeout)
    232     except google.api_core.exceptions.Forbidden as ex:
    233         if "Drive credentials" in ex.message:

/usr/local/lib/python3.10/dist-packages/google/cloud/bigquery/client.py in query(self, query, job_config, job_id, job_id_prefix, location, project, retry, timeout, job_retry, api_method)
   3490             )
   3491         elif api_method == enums.QueryApiMethod.INSERT:
-> 3492             return _job_helpers.query_jobs_insert(
   3493                 self,
   3494                 query,

/usr/local/lib/python3.10/dist-packages/google/cloud/bigquery/_job_helpers.py in query_jobs_insert(client, query, job_config, job_id, job_id_prefix, location, project, retry, timeout, job_retry)
    157             return query_job
    158 
--> 159     future = do_query()
    160     # The future might be in a failed state now, but if it's
    161     # unrecoverable, we'll find out when we ask for it's result, at which

/usr/local/lib/python3.10/dist-packages/google/cloud/bigquery/_job_helpers.py in do_query()
    143 
    144             try:
--> 145                 query_job = client.get_job(
    146                     job_id,
    147                     project=project,

/usr/local/lib/python3.10/dist-packages/google/cloud/bigquery/client.py in get_job(self, job_id, project, location, retry, timeout)
   2199         span_attributes = {"path": path, "job_id": job_id, "location": location}
   2200 
-> 2201         resource = self._call_api(
   2202             retry,
   2203             span_name="BigQuery.getJob",

/usr/local/lib/python3.10/dist-packages/google/cloud/bigquery/client.py in _call_api(self, retry, span_name, span_attributes, job_ref, headers, **kwargs)
    831                 name=span_name, attributes=span_attributes, client=self, job_ref=job_ref
    832             ):
--> 833                 return call()
    834 
    835         return call()

/usr/local/lib/python3.10/dist-packages/google/api_core/retry/retry_unary.py in retry_wrapped_func(*args, **kwargs)
    291                 self._initial, self._maximum, multiplier=self._multiplier
    292             )
--> 293             return retry_target(
    294                 target,
    295                 self._predicate,

/usr/local/lib/python3.10/dist-packages/google/api_core/retry/retry_unary.py in retry_target(target, predicate, sleep_generator, timeout, on_error, exception_factory, **kwargs)
    151         except Exception as exc:
    152             # defer to shared logic for handling errors
--> 153             _retry_error_helper(
    154                 exc,
    155                 deadline,

/usr/local/lib/python3.10/dist-packages/google/api_core/retry/retry_base.py in _retry_error_helper(exc, deadline, next_sleep, error_list, predicate_fn, on_error_fn, exc_factory_fn, original_timeout)
    210             original_timeout,
    211         )
--> 212         raise final_exc from source_exc
    213     if on_error_fn is not None:
    214         on_error_fn(exc)

/usr/local/lib/python3.10/dist-packages/google/api_core/retry/retry_unary.py in retry_target(target, predicate, sleep_generator, timeout, on_error, exception_factory, **kwargs)
    142     for sleep in sleep_generator:
    143         try:
--> 144             result = target()
    145             if inspect.isawaitable(result):
    146                 warnings.warn(_ASYNC_RETRY_WARNING)

/usr/local/lib/python3.10/dist-packages/google/cloud/_http/__init__.py in api_request(self, method, path, query_params, data, content_type, headers, api_base_url, api_version, expect_json, _target_object, timeout, extra_api_info)
    492 
    493         if not 200 <= response.status_code < 300:
--> 494             raise exceptions.from_http_response(response)
    495 
    496         if expect_json and response.content:

NotFound: 404 GET https://bigquery.googleapis.com/bigquery/v2/projects/my-project/jobs/c540a00e-05ed-4caf-b76c-b23d16b9617f?projection=full&location=US&prettyPrint=false: Not found: Job my-project:US.c540a00e-05ed-4caf-b76c-b23d16b9617f
```

Note that it is getting a 404 here: https://github.com/googleapis/python-bigquery/blob/b03a2afabde7f42be45f62fabd3dc0e6a9a493e1/google/cloud/bigquery/_job_helpers.py#L145

This is after a Conflict exception, meaning the job was created successfully, the retry failed because it already exists, but then getting that already existing job failed with 404. Something funky is going on with the networking I think for this case to happen, but it should be possible to simulate in unit tests, at least.


Pull request description:
fix: retry 404 errors in `Client.query(...)`
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/python-bigquery/issues/new/choose) before writing your code!  That way we can discuss the change, evaluate designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)

Fixes #2134 🦕


Patch:
diff --git a/google/cloud/bigquery/_job_helpers.py b/google/cloud/bigquery/_job_helpers.py
index e66ab276..b028cd35 100644
--- a/google/cloud/bigquery/_job_helpers.py
+++ b/google/cloud/bigquery/_job_helpers.py
@@ -47,6 +47,7 @@ from google.api_core import retry as retries
 from google.cloud.bigquery import job
 import google.cloud.bigquery.query
 from google.cloud.bigquery import table
+import google.cloud.bigquery.retry
 from google.cloud.bigquery.retry import POLLING_DEFAULT_VALUE
 
 # Avoid circular imports
@@ -142,12 +143,28 @@ def query_jobs_insert(
                 raise create_exc
 
             try:
+                # Sometimes we get a 404 after a Conflict. In this case, we
+                # have pretty high confidence that by retrying the 404, we'll
+                # (hopefully) eventually recover the job.
+                # https://github.com/googleapis/python-bigquery/issues/2134
+                #
+                # Allow users who want to completely disable retries to
+                # continue to do so by setting retry to None.
+                get_job_retry = retry
+                if retry is not None:
+                    # TODO(tswast): Amend the user's retry object with allowing
+                    # 404 to retry when there's a public way to do so.
+                    # https://github.com/googleapis/python-api-core/issues/796
+                    get_job_retry = (
+                        google.cloud.bigquery.retry._DEFAULT_GET_JOB_CONFLICT_RETRY
+                    )
+
                 query_job = client.get_job(
                     job_id,
                     project=project,
                     location=location,
-                    retry=retry,
-                    timeout=timeout,
+                    retry=get_job_retry,
+                    timeout=google.cloud.bigquery.retry.DEFAULT_GET_JOB_TIMEOUT,
                 )
             except core_exceptions.GoogleAPIError:  # (includes RetryError)
                 raise
@@ -156,7 +173,13 @@ def query_jobs_insert(
         else:
             return query_job
 
+    # Allow users who want to completely disable retries to
+    # continue to do so by setting job_retry to None.
+    if job_retry is not None:
+        do_query = google.cloud.bigquery.retry._DEFAULT_QUERY_JOB_INSERT_RETRY(do_query)
+
     future = do_query()
+
     # The future might be in a failed state now, but if it's
     # unrecoverable, we'll find out when we ask for it's result, at which
     # point, we may retry.
diff --git a/google/cloud/bigquery/retry.py b/google/cloud/bigquery/retry.py
index 10958980..999d0e85 100644
--- a/google/cloud/bigquery/retry.py
+++ b/google/cloud/bigquery/retry.py
@@ -82,6 +82,32 @@ on ``DEFAULT_RETRY``. For example, to change the deadline to 30 seconds,
 pass ``retry=bigquery.DEFAULT_RETRY.with_deadline(30)``.
 """
 
+
+def _should_retry_get_job_conflict(exc):
+    """Predicate for determining when to retry a jobs.get call after a conflict error.
+
+    Sometimes we get a 404 after a Conflict. In this case, we
+    have pretty high confidence that by retrying the 404, we'll
+    (hopefully) eventually recover the job.
+    https://github.com/googleapis/python-bigquery/issues/2134
+
+    Note: we may be able to extend this to user-specified predicates
+    after https://github.com/googleapis/python-api-core/issues/796
+    to tweak existing Retry object predicates.
+    """
+    return isinstance(exc, exceptions.NotFound) or _should_retry(exc)
+
+
+# Pick a deadline smaller than our other deadlines since we want to timeout
+# before those expire.
+_DEFAULT_GET_JOB_CONFLICT_DEADLINE = _DEFAULT_RETRY_DEADLINE / 3.0
+_DEFAULT_GET_JOB_CONFLICT_RETRY = retry.Retry(
+    predicate=_should_retry_get_job_conflict,
+    deadline=_DEFAULT_GET_JOB_CONFLICT_DEADLINE,
+)
+"""Private, may be removed in future."""
+
+
 # Note: Take care when updating DEFAULT_TIMEOUT to anything but None. We
 # briefly had a default timeout, but even setting it at more than twice the
 # theoretical server-side default timeout of 2 minutes was not enough for
@@ -142,6 +168,34 @@ DEFAULT_JOB_RETRY = retry.Retry(
 The default job retry object.
 """
 
+
+def _query_job_insert_should_retry(exc):
+    # Per https://github.com/googleapis/python-bigquery/issues/2134, sometimes
+    # we get a 404 error. In this case, if we get this far, assume that the job
+    # doesn't actually exist and try again. We can't add 404 to the default
+    # job_retry because that happens for errors like "this table does not
+    # exist", which probably won't resolve with a retry.
+    if isinstance(exc, exceptions.RetryError):
+        exc = exc.cause
+
+    if isinstance(exc, exceptions.NotFound):
+        message = exc.message
+        # Don't try to retry table/dataset not found, just job not found.
+        # The URL contains jobs, so use whitespace to disambiguate.
+        return message is not None and " job" in message.lower()
+
+    return _job_should_retry(exc)
+
+
+_DEFAULT_QUERY_JOB_INSERT_RETRY = retry.Retry(
+    predicate=_query_job_insert_should_retry,
+    # jobs.insert doesn't wait for the job to complete, so we don't need the
+    # long _DEFAULT_JOB_DEADLINE for this part.
+    deadline=_DEFAULT_RETRY_DEADLINE,
+)
+"""Private, may be removed in future."""
+
+
 DEFAULT_GET_JOB_TIMEOUT = 128
 """
 Default timeout for Client.get_job().


Test patch:
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
index 6897c255..4f13d6ec 100644
--- a/tests/unit/test_client.py
+++ b/tests/unit/test_client.py
@@ -28,9 +28,12 @@ import unittest
 from unittest import mock
 import warnings
 
-import requests
+import freezegun
 import packaging
 import pytest
+import requests
+
+import google.api
 
 
 try:
@@ -55,6 +58,8 @@ from google.api_core import client_info
 import google.cloud._helpers
 from google.cloud import bigquery
 
+from google.cloud.bigquery import job as bqjob
+import google.cloud.bigquery._job_helpers
 from google.cloud.bigquery.dataset import DatasetReference
 from google.cloud.bigquery import exceptions
 from google.cloud.bigquery import ParquetOptions
@@ -5308,6 +5313,173 @@ class TestClient(unittest.TestCase):
             with pytest.raises(DataLoss, match="we lost your job, sorry"):
                 client.query("SELECT 1;", job_id=None)
 
+    def test_query_job_rpc_fail_w_conflict_random_id_job_fetch_fails_no_retries(self):
+        from google.api_core.exceptions import Conflict
+        from google.api_core.exceptions import DataLoss
+        from google.cloud.bigquery.job import QueryJob
+
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
+
+        job_create_error = Conflict("Job already exists.")
+        job_begin_patcher = mock.patch.object(
+            QueryJob, "_begin", side_effect=job_create_error
+        )
+        get_job_patcher = mock.patch.object(
+            client, "get_job", side_effect=DataLoss("we lost your job, sorry")
+        )
+
+        with job_begin_patcher, get_job_patcher:
+            # If get job request fails but supposedly there does exist a job
+            # with this ID already, raise the exception explaining why we
+            # couldn't recover the job.
+            with pytest.raises(DataLoss, match="we lost your job, sorry"):
+                client.query(
+                    "SELECT 1;",
+                    job_id=None,
+                    # Explicitly test with no retries to make sure those branches are covered.
+                    retry=None,
+                    job_retry=None,
+                )
+
+    def test_query_job_rpc_fail_w_conflict_random_id_job_fetch_retries_404(self):
+        """Regression test for https://github.com/googleapis/python-bigquery/issues/2134
+
+        Sometimes after a Conflict, the fetch fails with a 404, but we know
+        because of the conflict that really the job does exist. Retry until we
+        get the job status (or timeout).
+        """
+        job_id = "abc123"
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
+        conn = client._connection = make_connection(
+            # We're mocking QueryJob._begin, so this is only going to be
+            # jobs.get requests and responses.
+            google.api_core.exceptions.TooManyRequests("this is retriable by default"),
+            google.api_core.exceptions.NotFound("we lost your job"),
+            google.api_core.exceptions.NotFound("we lost your job again, sorry"),
+            {
+                "jobReference": {
+                    "projectId": self.PROJECT,
+                    "location": "TESTLOC",
+                    "jobId": job_id,
+                }
+            },
+        )
+
+        job_create_error = google.api_core.exceptions.Conflict("Job already exists.")
+        job_begin_patcher = mock.patch.object(
+            bqjob.QueryJob, "_begin", side_effect=job_create_error
+        )
+        job_id_patcher = mock.patch.object(
+            google.cloud.bigquery._job_helpers,
+            "make_job_id",
+            return_value=job_id,
+        )
+
+        with job_begin_patcher, job_id_patcher:
+            # If get job request fails there does exist a job
+            # with this ID already, retry 404 until we get it (or fails for a
+            # non-retriable reason, see other tests).
+            result = client.query("SELECT 1;", job_id=None)
+
+        jobs_get_path = mock.call(
+            method="GET",
+            path=f"/projects/{self.PROJECT}/jobs/{job_id}",
+            query_params={
+                "projection": "full",
+            },
+            timeout=google.cloud.bigquery.retry.DEFAULT_GET_JOB_TIMEOUT,
+        )
+        conn.api_request.assert_has_calls(
+            # Double-check that it was jobs.get that was called for each of our
+            # mocked responses.
+            [jobs_get_path]
+            * 4,
+        )
+        assert result.job_id == job_id
+
+    def test_query_job_rpc_fail_w_conflict_random_id_job_fetch_retries_404_and_query_job_insert(
+        self,
+    ):
+        """Regression test for https://github.com/googleapis/python-bigquery/issues/2134
+
+        Sometimes after a Conflict, the fetch fails with a 404. If it keeps
+        failing with a 404, assume that the job actually doesn't exist.
+        """
+        job_id_1 = "abc123"
+        job_id_2 = "xyz789"
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
+
+        # We're mocking QueryJob._begin, so that the connection should only get
+        # jobs.get requests.
+        job_create_error = google.api_core.exceptions.Conflict("Job already exists.")
+        job_begin_patcher = mock.patch.object(
+            bqjob.QueryJob, "_begin", side_effect=job_create_error
+        )
+        conn = client._connection = make_connection(
+            google.api_core.exceptions.NotFound("we lost your job again, sorry"),
+            {
+                "jobReference": {
+                    "projectId": self.PROJECT,
+                    "location": "TESTLOC",
+                    "jobId": job_id_2,
+                }
+            },
+        )
+
+        # Choose a small deadline so the 404 retries give up.
+        retry = (
+            google.cloud.bigquery.retry._DEFAULT_GET_JOB_CONFLICT_RETRY.with_deadline(1)
+        )
+        job_id_patcher = mock.patch.object(
+            google.cloud.bigquery._job_helpers,
+            "make_job_id",
+            side_effect=[job_id_1, job_id_2],
+        )
+        retry_patcher = mock.patch.object(
+            google.cloud.bigquery.retry,
+            "_DEFAULT_GET_JOB_CONFLICT_RETRY",
+            retry,
+        )
+
+        with freezegun.freeze_time(
+            "2025-01-01 00:00:00",
+            # 10x the retry deadline to guarantee a timeout.
+            auto_tick_seconds=10,
+        ), job_begin_patcher, job_id_patcher, retry_patcher:
+            # If get job request fails there does exist a job
+            # with this ID already, retry 404 until we get it (or fails for a
+            # non-retriable reason, see other tests).
+            result = client.query("SELECT 1;", job_id=None)
+
+        jobs_get_path_1 = mock.call(
+            method="GET",
+            path=f"/projects/{self.PROJECT}/jobs/{job_id_1}",
+            query_params={
+                "projection": "full",
+            },
+            timeout=google.cloud.bigquery.retry.DEFAULT_GET_JOB_TIMEOUT,
+        )
+        jobs_get_path_2 = mock.call(
+            method="GET",
+            path=f"/projects/{self.PROJECT}/jobs/{job_id_2}",
+            query_params={
+                "projection": "full",
+            },
+            timeout=google.cloud.bigquery.retry.DEFAULT_GET_JOB_TIMEOUT,
+        )
+        conn.api_request.assert_has_calls(
+            # Double-check that it was jobs.get that was called for each of our
+            # mocked responses.
+            [jobs_get_path_1, jobs_get_path_2],
+        )
+        assert result.job_id == job_id_2
+
     def test_query_job_rpc_fail_w_conflict_random_id_job_fetch_succeeds(self):
         from google.api_core.exceptions import Conflict
         from google.cloud.bigquery.job import QueryJob