From 260d2cb612f07b019fc9a8e4e39e3cfef098dedb Mon Sep 17 00:00:00 2001 From: Jennifer Davis Date: Tue, 3 Nov 2026 23:31:56 +0000 Subject: [PATCH] fix(bigquery): restore IF NOT EXISTS in create_search_index and clean up test mocks Problem: Following review feedback across the BigQuery sample PRs (#14714, #14715, #14717): - `bigquery_create_search_index.py` should use `CREATE SEARCH INDEX IF NOT EXISTS` and handle `exceptions.NotFound` when the parent table or dataset is missing. - `bigquery_drop_vector_index.py` `NotFound` error message should reflect that `DROP VECTOR INDEX IF EXISTS` raises `NotFound` when the parent table or dataset is missing. - Unit tests in `bigquery_create_search_index_test.py`, `bigquery_drop_vector_index_test.py`, and `bigquery_query_manifest_external_table_test.py` defined unused mock attributes (`mock_row` on DDL statements and `mock_job.statement_type` on SELECT queries). Solution: - Restore `CREATE SEARCH INDEX IF NOT EXISTS` and catch `exceptions.NotFound` in `bigquery/search/bigquery_create_search_index.py` and its unit test. - Update the `NotFound` error message in `bigquery/vector_search/bigquery_drop_vector_index.py`. - Remove unused `mock_row` and `mock_job.statement_type` setup in the three corresponding unit test files. --- .../bigquery_query_manifest_external_table_test.py | 5 +++-- bigquery/search/bigquery_create_search_index.py | 8 ++++---- bigquery/search/bigquery_create_search_index_test.py | 7 ++++--- bigquery/vector_search/bigquery_drop_vector_index.py | 4 ++-- bigquery/vector_search/bigquery_drop_vector_index_test.py | 7 ++++--- 5 files changed, 17 insertions(+), 14 deletions(-) diff --git a/bigquery/biglake/bigquery_query_manifest_external_table_test.py b/bigquery/biglake/bigquery_query_manifest_external_table_test.py index cfe10fcdfa..b6f06dc9c0 100644 --- a/bigquery/biglake/bigquery_query_manifest_external_table_test.py +++ b/bigquery/biglake/bigquery_query_manifest_external_table_test.py @@ -26,7 +26,6 @@ def test_query_manifest_external_table_success( mock_client = mock.MagicMock() mock_job = mock.MagicMock() mock_job.job_id = "example-job-123" - mock_job.statement_type = "SCRIPT" mock_row = mock.MagicMock() mock_job.result.return_value = [mock_row] mock_client.query.return_value = mock_job @@ -54,7 +53,9 @@ def test_query_manifest_external_table_error( capsys: pytest.CaptureFixture[str], ) -> None: mock_client = mock.MagicMock() - mock_client.query.side_effect = exceptions.NotFound("resource error") + mock_job = mock.MagicMock() + mock_job.result.side_effect = exceptions.NotFound("resource error") + mock_client.query.return_value = mock_job monkeypatch.setattr( bigquery_query_manifest_external_table.bigquery, "Client", diff --git a/bigquery/search/bigquery_create_search_index.py b/bigquery/search/bigquery_create_search_index.py index 6aeff84cc6..ff690d0019 100644 --- a/bigquery/search/bigquery_create_search_index.py +++ b/bigquery/search/bigquery_create_search_index.py @@ -38,7 +38,7 @@ def create_search_index( table_ref = f"{project_id}.{dataset_id}.{table_id}" # Indexing ALL COLUMNS enables full-text SEARCH queries across every STRING or JSON field. sql = f""" - CREATE SEARCH INDEX `{index_name}` + CREATE SEARCH INDEX IF NOT EXISTS `{index_name}` ON `{table_ref}`(ALL COLUMNS); """ @@ -46,10 +46,10 @@ def create_search_index( print(f"Creating search index {index_name} on table: {table_ref}") query_job = client.query(sql) query_job.result() - except exceptions.AlreadyExists as err: + except exceptions.NotFound as err: print( - f"error: search index '{index_name}' already exists on table " - f"'{table_id}' in dataset '{dataset_id}' for project '{project_id}': {err.message}", + f"error: table '{table_id}' or dataset '{dataset_id}' not found " + f"for project '{project_id}': {err.message}", file=sys.stderr, ) return diff --git a/bigquery/search/bigquery_create_search_index_test.py b/bigquery/search/bigquery_create_search_index_test.py index f1a77df6bb..48242a0501 100644 --- a/bigquery/search/bigquery_create_search_index_test.py +++ b/bigquery/search/bigquery_create_search_index_test.py @@ -27,8 +27,7 @@ def test_create_search_index_success( mock_job = mock.MagicMock() mock_job.job_id = "example-job-123" mock_job.statement_type = "SCRIPT" - mock_row = mock.MagicMock() - mock_job.result.return_value = [mock_row] + mock_job.result.return_value = [] mock_client.query.return_value = mock_job monkeypatch.setattr( bigquery_create_search_index.bigquery, @@ -53,7 +52,9 @@ def test_create_search_index_error( capsys: pytest.CaptureFixture[str], ) -> None: mock_client = mock.MagicMock() - mock_client.query.side_effect = exceptions.AlreadyExists("resource error") + mock_job = mock.MagicMock() + mock_job.result.side_effect = exceptions.NotFound("resource error") + mock_client.query.return_value = mock_job monkeypatch.setattr( bigquery_create_search_index.bigquery, "Client", diff --git a/bigquery/vector_search/bigquery_drop_vector_index.py b/bigquery/vector_search/bigquery_drop_vector_index.py index 91c9c45e94..974d712e97 100644 --- a/bigquery/vector_search/bigquery_drop_vector_index.py +++ b/bigquery/vector_search/bigquery_drop_vector_index.py @@ -45,8 +45,8 @@ def drop_vector_index( query_job.result() except exceptions.NotFound as err: print( - f"error: table '{table_id}' or vector index '{index_name}' not found in " - f"dataset '{dataset_id}' for project '{project_id}': {err.message}", + f"error: table '{table_id}' or dataset '{dataset_id}' not found in " + f"project '{project_id}': {err.message}", file=sys.stderr, ) return diff --git a/bigquery/vector_search/bigquery_drop_vector_index_test.py b/bigquery/vector_search/bigquery_drop_vector_index_test.py index bf64f1f808..f441584f66 100644 --- a/bigquery/vector_search/bigquery_drop_vector_index_test.py +++ b/bigquery/vector_search/bigquery_drop_vector_index_test.py @@ -27,8 +27,7 @@ def test_drop_vector_index_success( mock_job = mock.MagicMock() mock_job.job_id = "example-job-123" mock_job.statement_type = "SCRIPT" - mock_row = mock.MagicMock() - mock_job.result.return_value = [mock_row] + mock_job.result.return_value = [] mock_client.query.return_value = mock_job monkeypatch.setattr( bigquery_drop_vector_index.bigquery, @@ -53,7 +52,9 @@ def test_drop_vector_index_error( capsys: pytest.CaptureFixture[str], ) -> None: mock_client = mock.MagicMock() - mock_client.query.side_effect = exceptions.NotFound("resource error") + mock_job = mock.MagicMock() + mock_job.result.side_effect = exceptions.NotFound("resource error") + mock_client.query.return_value = mock_job monkeypatch.setattr( bigquery_drop_vector_index.bigquery, "Client",