Thanks to visit codestin.com
Credit goes to github.com

Skip to content

feat: Improve error message for drive based BQ table reads #344

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Jan 26, 2024
Merged

feat: Improve error message for drive based BQ table reads #344

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
ff00eea
feat: Improve error message for drive based BQ table reads
shobsi Jan 24, 2024
3834256
Merge remote-tracking branch 'refs/remotes/github/main' into shobs-gd…
shobsi Jan 25, 2024
535f2af
move exception handling deeper to apply to read_gbq*
shobsi Jan 25, 2024
ccebd05
add unit tests
shobsi Jan 25, 2024
c3c87e1
Merge remote-tracking branch 'refs/remotes/github/main' into shobs-gd…
shobsi Jan 25, 2024
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
21 changes: 17 additions & 4 deletions bigframes/session/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -597,9 +597,16 @@ def _get_snapshot_sql_and_primary_key(
).result()
)[0][0]
self._df_snapshot[table_ref] = snapshot_timestamp
table_expression = self.ibis_client.sql(
bigframes_io.create_snapshot_sql(table_ref, snapshot_timestamp)
)

try:
table_expression = self.ibis_client.sql(
bigframes_io.create_snapshot_sql(table_ref, snapshot_timestamp)
)
except google.api_core.exceptions.Forbidden as ex:
if "Drive credentials" in ex.message:
ex.message += "\nCheck https://cloud.google.com/bigquery/docs/query-drive-data#Google_Drive_permissions."
raise

return table_expression, primary_keys

def _read_gbq_table(
Expand Down Expand Up @@ -1451,7 +1458,13 @@ def _start_query(
job_config.labels = bigframes_io.create_job_configs_labels(
job_configs_labels=job_config.labels, api_methods=api_methods
)
query_job = self.bqclient.query(sql, job_config=job_config)

try:
query_job = self.bqclient.query(sql, job_config=job_config)
except google.api_core.exceptions.Forbidden as ex:
if "Drive credentials" in ex.message:
ex.message += "\nCheck https://cloud.google.com/bigquery/docs/query-drive-data#Google_Drive_permissions."
raise

opts = bigframes.options.display
if opts.progress_bar is not None and not query_job.configuration.dry_run:
Expand Down
30 changes: 29 additions & 1 deletion tests/unit/session/test_session.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ def test_read_gbq_missing_parts(missing_parts_table_id):
"not_found_table_id",
[("unknown.dataset.table"), ("project.unknown.table"), ("project.dataset.unknown")],
)
def test_read_gdb_not_found_tables(not_found_table_id):
def test_read_gbq_not_found_tables(not_found_table_id):
bqclient = mock.create_autospec(google.cloud.bigquery.Client, instance=True)
bqclient.project = "test-project"
bqclient.get_table.side_effect = google.api_core.exceptions.NotFound(
Expand All @@ -47,6 +47,34 @@ def test_read_gdb_not_found_tables(not_found_table_id):
session.read_gbq(not_found_table_id)


@pytest.mark.parametrize(
("api_name", "query_or_table"),
[
("read_gbq", "project.dataset.table"),
("read_gbq_table", "project.dataset.table"),
("read_gbq", "SELECT * FROM project.dataset.table"),
("read_gbq_query", "SELECT * FROM project.dataset.table"),
],
ids=[
"read_gbq_on_table",
"read_gbq_table",
"read_gbq_on_query",
"read_gbq_query",
],
)
def test_read_gbq_external_table_no_drive_access(api_name, query_or_table):
bqclient = mock.create_autospec(google.cloud.bigquery.Client, instance=True)
bqclient.project = "test-project"
bqclient.get_table.side_effect = google.api_core.exceptions.Forbidden(
"Access Denied: BigQuery BigQuery: Permission denied while getting Drive credentials."
)
session = resources.create_bigquery_session(bqclient=bqclient)

api = getattr(session, api_name)
with pytest.raises(google.api_core.exceptions.Forbidden):
api(query_or_table)


@mock.patch.dict(os.environ, {}, clear=True)
def test_session_init_fails_with_no_project():
with pytest.raises(
Expand Down