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/biglake/bigquery_read_biglake_gcs_storage_api.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_read_biglake_gcs_storage_api]
import sys

from google.api_core import exceptions
from google.cloud import bigquery_storage_v1


def read_biglake_gcs_storage_api(
project_id: str,
dataset_id: str,
table_id: str,
) -> None:
"""Reads projected columns from a BigLake table using the BigQuery Storage Read API.

Args:
project_id: Google Cloud project ID (for example, 'example-project').
dataset_id: BigQuery dataset ID containing the BigLake table
(for example, 'example_dataset').
table_id: BigLake table ID to read (for example, 'example_biglake_table').
"""
client = bigquery_storage_v1.BigQueryReadClient()
parent = f"projects/{project_id}"
table_ref = f"projects/{project_id}/datasets/{dataset_id}/tables/{table_id}"

# Project specific columns in the ReadSession to minimize network transfer from Cloud Storage.
read_options = bigquery_storage_v1.types.ReadSession.TableReadOptions(
selected_fields=["id", "label", "score"]
)
requested_session = bigquery_storage_v1.types.ReadSession(
table=table_ref,
data_format=bigquery_storage_v1.types.DataFormat.AVRO,
read_options=read_options,
)
request = bigquery_storage_v1.types.CreateReadSessionRequest(
parent=parent,
read_session=requested_session,
max_stream_count=1,
)

try:
print(f"Creating Storage Read API session for BigLake table {table_ref}")
session = client.create_read_session(request=request)
rows_read = 0
if session.streams:
reader = client.read_rows(session.streams[0].name)
for response in reader:
rows_read += response.row_count
except exceptions.NotFound as err:
print(
f"error: BigLake table '{table_id}' not found in dataset '{dataset_id}' "
f"for project '{project_id}': {err.message}",
file=sys.stderr,
)
return
except exceptions.GoogleAPICallError as err:
print(
f"error: failed to read BigLake table '{table_id}' in dataset "
f"'{dataset_id}' for project '{project_id}': {err.message}",
file=sys.stderr,
)
return

print(f"Resource ID: {table_ref}")
print(f"Read Session: {session.name}")
print(f"Rows Read: {rows_read}")
print(f"Stream Count: {len(session.streams)}")


# [END bigquery_read_biglake_gcs_storage_api]
79 changes: 79 additions & 0 deletions bigquery/biglake/bigquery_read_biglake_gcs_storage_api_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
# 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_read_biglake_gcs_storage_api
from google.api_core import exceptions
import pytest


def test_read_biglake_gcs_storage_api_success(
monkeypatch: pytest.MonkeyPatch,
capsys: pytest.CaptureFixture[str],
) -> None:
mock_client = mock.MagicMock()
mock_session = mock.MagicMock()
mock_session.name = "projects/example-project/locations/us/sessions/s1"
mock_stream = mock.MagicMock()
mock_stream.name = "projects/example-project/locations/us/sessions/s1/streams/0"
mock_session.streams = [mock_stream]
mock_client.create_read_session.return_value = mock_session

mock_response = mock.MagicMock()
mock_response.row_count = 3
mock_client.read_rows.return_value = [mock_response]

monkeypatch.setattr(
bigquery_read_biglake_gcs_storage_api.bigquery_storage_v1,
"BigQueryReadClient",
lambda: mock_client,
)

bigquery_read_biglake_gcs_storage_api.read_biglake_gcs_storage_api(
project_id="example-project",
dataset_id="example_dataset",
table_id="example_biglake_table",
)

captured = capsys.readouterr()
assert (
"Resource ID: projects/example-project/datasets/example_dataset/tables/example_biglake_table"
in captured.out
)
assert "Rows Read: 3" in captured.out
assert "Stream Count: 1" in captured.out


def test_read_biglake_gcs_storage_api_error(
monkeypatch: pytest.MonkeyPatch,
capsys: pytest.CaptureFixture[str],
) -> None:
mock_client = mock.MagicMock()
mock_client.create_read_session.side_effect = exceptions.NotFound("table not found")
monkeypatch.setattr(
bigquery_read_biglake_gcs_storage_api.bigquery_storage_v1,
"BigQueryReadClient",
lambda: mock_client,
)

bigquery_read_biglake_gcs_storage_api.read_biglake_gcs_storage_api(
project_id="example-project",
dataset_id="example_dataset",
table_id="example_biglake_table",
)

captured = capsys.readouterr()
assert "error:" in captured.err
assert "example_biglake_table" in captured.err
23 changes: 23 additions & 0 deletions bigquery/biglake/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/biglake/requirements-test.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
pytest==9.0.3; python_version >= "3.10"
2 changes: 2 additions & 0 deletions bigquery/biglake/requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
google-cloud-bigquery==3.47.0
google-cloud-bigquery-storage==2.43.0
Loading