Skip to content

Commit 0d85b5f

Browse files
committed
feat(bigquery): add bigquery_read_biglake_gcs_storage_api sample to read a BigLake Cloud Storage table via BigQuery Storage Read API
Problem: The BigQuery documentation (https://cloud.google.com/bigquery/docs/query-cloud-storage-using-biglake) needs a standalone, tested Python code sample demonstrating how to read a BigLake Cloud Storage table via BigQuery Storage Read API (region tag `bigquery_read_biglake_gcs_storage_api`). Solution: - Add `read_biglake_gcs_storage_api(project_id: str, dataset_id: str, table_id: str) -> None` in `bigquery/biglake/bigquery_read_biglake_gcs_storage_api.py` (Stream rows from a BigLake table backed by Cloud Storage using the BigQuery Storage Read API client.) - Add unit tests (`bigquery/biglake/bigquery_read_biglake_gcs_storage_api_test.py`) covering both the success path and API error handling (`GoogleAPICallError`)
1 parent 186be98 commit 0d85b5f

5 files changed

Lines changed: 188 additions & 0 deletions

File tree

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# [START bigquery_read_biglake_gcs_storage_api]
16+
import sys
17+
18+
from google.api_core import exceptions
19+
from google.cloud import bigquery_storage_v1
20+
21+
22+
def read_biglake_gcs_storage_api(
23+
project_id: str,
24+
dataset_id: str,
25+
table_id: str,
26+
) -> None:
27+
"""Reads projected columns from a BigLake table using the BigQuery Storage Read API.
28+
29+
Args:
30+
project_id: Google Cloud project ID (for example, 'example-project').
31+
dataset_id: BigQuery dataset ID containing the BigLake table
32+
(for example, 'example_dataset').
33+
table_id: BigLake table ID to read (for example, 'example_biglake_table').
34+
"""
35+
client = bigquery_storage_v1.BigQueryReadClient()
36+
parent = f"projects/{project_id}"
37+
table_ref = f"projects/{project_id}/datasets/{dataset_id}/tables/{table_id}"
38+
39+
# Project specific columns in the ReadSession to minimize network transfer from Cloud Storage.
40+
read_options = bigquery_storage_v1.types.ReadSession.TableReadOptions(
41+
selected_fields=["id", "label", "score"]
42+
)
43+
requested_session = bigquery_storage_v1.types.ReadSession(
44+
table=table_ref,
45+
data_format=bigquery_storage_v1.types.DataFormat.AVRO,
46+
read_options=read_options,
47+
)
48+
request = bigquery_storage_v1.types.CreateReadSessionRequest(
49+
parent=parent,
50+
read_session=requested_session,
51+
max_stream_count=1,
52+
)
53+
54+
try:
55+
print(f"Creating Storage Read API session for BigLake table {table_ref}")
56+
session = client.create_read_session(request=request)
57+
rows_read = 0
58+
if session.streams:
59+
reader = client.read_rows(session.streams[0].name)
60+
for response in reader:
61+
rows_read += response.row_count
62+
except exceptions.NotFound as err:
63+
print(
64+
f"error: BigLake table '{table_id}' not found in dataset '{dataset_id}' "
65+
f"for project '{project_id}': {err.message}",
66+
file=sys.stderr,
67+
)
68+
return
69+
except exceptions.GoogleAPICallError as err:
70+
print(
71+
f"error: failed to read BigLake table '{table_id}' in dataset "
72+
f"'{dataset_id}' for project '{project_id}': {err.message}",
73+
file=sys.stderr,
74+
)
75+
return
76+
77+
print(f"Resource ID: {table_ref}")
78+
print(f"Read Session: {session.name}")
79+
print(f"Rows Read: {rows_read}")
80+
print(f"Stream Count: {len(session.streams)}")
81+
82+
83+
# [END bigquery_read_biglake_gcs_storage_api]
Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from unittest import mock
16+
17+
import bigquery_read_biglake_gcs_storage_api
18+
from google.api_core import exceptions
19+
import pytest
20+
21+
22+
def test_read_biglake_gcs_storage_api_success(
23+
monkeypatch: pytest.MonkeyPatch,
24+
capsys: pytest.CaptureFixture[str],
25+
) -> None:
26+
mock_client = mock.MagicMock()
27+
mock_session = mock.MagicMock()
28+
mock_session.name = "projects/example-project/locations/us/sessions/s1"
29+
mock_stream = mock.MagicMock()
30+
mock_stream.name = "projects/example-project/locations/us/sessions/s1/streams/0"
31+
mock_session.streams = [mock_stream]
32+
mock_client.create_read_session.return_value = mock_session
33+
34+
mock_response = mock.MagicMock()
35+
mock_response.row_count = 3
36+
mock_client.read_rows.return_value = [mock_response]
37+
38+
monkeypatch.setattr(
39+
bigquery_read_biglake_gcs_storage_api.bigquery_storage_v1,
40+
"BigQueryReadClient",
41+
lambda: mock_client,
42+
)
43+
44+
bigquery_read_biglake_gcs_storage_api.read_biglake_gcs_storage_api(
45+
project_id="example-project",
46+
dataset_id="example_dataset",
47+
table_id="example_biglake_table",
48+
)
49+
50+
captured = capsys.readouterr()
51+
assert (
52+
"Resource ID: projects/example-project/datasets/example_dataset/tables/example_biglake_table"
53+
in captured.out
54+
)
55+
assert "Rows Read: 3" in captured.out
56+
assert "Stream Count: 1" in captured.out
57+
58+
59+
def test_read_biglake_gcs_storage_api_error(
60+
monkeypatch: pytest.MonkeyPatch,
61+
capsys: pytest.CaptureFixture[str],
62+
) -> None:
63+
mock_client = mock.MagicMock()
64+
mock_client.create_read_session.side_effect = exceptions.NotFound("table not found")
65+
monkeypatch.setattr(
66+
bigquery_read_biglake_gcs_storage_api.bigquery_storage_v1,
67+
"BigQueryReadClient",
68+
lambda: mock_client,
69+
)
70+
71+
bigquery_read_biglake_gcs_storage_api.read_biglake_gcs_storage_api(
72+
project_id="example-project",
73+
dataset_id="example_dataset",
74+
table_id="example_biglake_table",
75+
)
76+
77+
captured = capsys.readouterr()
78+
assert "error:" in captured.err
79+
assert "example_biglake_table" in captured.err

‎bigquery/biglake/noxfile_config.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
TEST_CONFIG_OVERRIDE = {
16+
# You can opt out from the test for specific Python versions.
17+
"ignored_versions": ["2.7", "3.7", "3.8", "3.9", "3.10", "3.11"],
18+
# Enforce type hints on new samples.
19+
"enforce_type_hints": True,
20+
"gcloud_project_env": "GOOGLE_CLOUD_PROJECT",
21+
"pip_version_override": None,
22+
"envs": {},
23+
}
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
pytest==9.0.3; python_version >= "3.10"

‎bigquery/biglake/requirements.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
google-cloud-bigquery==3.47.0
2+
google-cloud-bigquery-storage==2.43.0

0 commit comments

Comments
 (0)