Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
83 changes: 83 additions & 0 deletions bigquery/vector_search/bigquery_get_vector_index_status.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,83 @@
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# [START bigquery_get_vector_index_status]
import sys

from google.api_core import exceptions
from google.cloud import bigquery


def get_vector_index_status(
project_id: str,
dataset_id: str,
table_id: str,
) -> None:
"""Queries INFORMATION_SCHEMA.VECTOR_INDEXES to inspect vector index build status.

Args:
project_id: Google Cloud project ID (for example, 'example-project').
dataset_id: BigQuery dataset ID containing the indexed embeddings table
(for example, 'example_dataset').
table_id: BigQuery table ID to inspect vector indexes for
(for example, 'example_embeddings').
"""
client = bigquery.Client(project=project_id)
table_ref = f"{project_id}.{dataset_id}.{table_id}"
# Parameterize the table_name filter when querying INFORMATION_SCHEMA.VECTOR_INDEXES.
job_config = bigquery.QueryJobConfig(
query_parameters=[
bigquery.ScalarQueryParameter("table_id", "STRING", table_id),
]
)
sql = f"""
SELECT
index_name,
index_status,
coverage_percentage
FROM `{project_id}.{dataset_id}.INFORMATION_SCHEMA.VECTOR_INDEXES`
WHERE table_name = @table_id;
"""

try:
print(f"Checking vector index status for table: {table_ref}")
query_job = client.query(sql, job_config=job_config)
rows = list(query_job.result())
except exceptions.NotFound as err:
print(
f"error: dataset '{dataset_id}' or table '{table_id}' not found in "
f"project '{project_id}': {err.message}",
file=sys.stderr,
)
return
except exceptions.GoogleAPICallError as err:
print(
f"error: failed to get vector index status for table '{table_id}' in "
f"dataset '{dataset_id}' for project '{project_id}': {err.message}",
file=sys.stderr,
)
return

index_name = rows[0].index_name if rows else "NONE"
index_status = rows[0].index_status if rows else "NONE"
coverage_percentage = rows[0].coverage_percentage if rows else 0
print(f"Resource ID: {table_ref}")
print(f"Index Count: {len(rows)}")
print(f"Index Name: {index_name}")
print(f"Index Status: {index_status}")
print(f"Coverage Percentage: {coverage_percentage}%")
print(f"Job ID: {query_job.job_id}")


# [END bigquery_get_vector_index_status]
71 changes: 71 additions & 0 deletions bigquery/vector_search/bigquery_get_vector_index_status_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from unittest import mock

import bigquery_get_vector_index_status
from google.api_core import exceptions
import pytest


def test_get_vector_index_status_success(
capsys: pytest.CaptureFixture[str],
) -> None:
with mock.patch("google.cloud.bigquery.Client") as mock_client_cls:
mock_client = mock_client_cls.return_value
mock_job = mock.MagicMock()
mock_job.job_id = "example-job-123"
mock_row = mock.MagicMock()
mock_row.index_name = "example_vector_index"
mock_row.index_status = "ACTIVE"
mock_row.coverage_percentage = 100
mock_job.result.return_value = [mock_row]
mock_client.query.return_value = mock_job

bigquery_get_vector_index_status.get_vector_index_status(
project_id="example-project",
dataset_id="example_dataset",
table_id="example_embeddings",
)

mock_client.query.assert_called_once()
captured = capsys.readouterr()
assert (
"Resource ID: example-project.example_dataset.example_embeddings"
in captured.out
)
assert "Index Count: 1" in captured.out
assert "Index Name: example_vector_index" in captured.out
assert "Index Status: ACTIVE" in captured.out
Comment on lines +43 to +50

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Update the test assertions to verify that the coverage percentage is correctly printed in the output.

Suggested change
captured = capsys.readouterr()
assert (
"Resource ID: example-project.example_dataset.example_embeddings"
in captured.out
)
assert "Index Count: 1" in captured.out
assert "Index Name: example_vector_index" in captured.out
assert "Index Status: ACTIVE" in captured.out
captured = capsys.readouterr()
assert (
"Resource ID: example-project.example_dataset.example_embeddings"
in captured.out
)
assert "Index Count: 1" in captured.out
assert "Index Name: example_vector_index" in captured.out
assert "Index Status: ACTIVE" in captured.out
assert "Coverage Percentage: 100" in captured.out

assert "Coverage Percentage: 100%" in captured.out


def test_get_vector_index_status_not_found(
capsys: pytest.CaptureFixture[str],
) -> None:
with mock.patch("google.cloud.bigquery.Client") as mock_client_cls:
mock_client = mock_client_cls.return_value
mock_job = mock.MagicMock()
mock_job.result.side_effect = exceptions.NotFound("Dataset not found")
mock_client.query.return_value = mock_job

bigquery_get_vector_index_status.get_vector_index_status(
project_id="example-project",
dataset_id="example_dataset",
table_id="example_embeddings",
)

captured = capsys.readouterr()
assert "error:" in captured.err
assert "example_embeddings" in captured.err
23 changes: 23 additions & 0 deletions bigquery/vector_search/noxfile_config.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

TEST_CONFIG_OVERRIDE = {
# You can opt out from the test for specific Python versions.
"ignored_versions": ["2.7", "3.7", "3.8", "3.9", "3.10", "3.11"],
# Enforce type hints on new samples.
"enforce_type_hints": True,
"gcloud_project_env": "GOOGLE_CLOUD_PROJECT",
"pip_version_override": None,
"envs": {},
}
1 change: 1 addition & 0 deletions bigquery/vector_search/requirements-test.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
pytest==9.0.3; python_version >= "3.10"
1 change: 1 addition & 0 deletions bigquery/vector_search/requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
google-cloud-bigquery==3.47.0
Loading