Skip to content

Commit

Permalink
feat: Improve error message for drive based BQ table reads (#344)
Browse files Browse the repository at this point in the history
* feat: Improve error message for drive based BQ table reads

* move exception handling deeper to apply to read_gbq*

* add unit tests
  • Loading branch information
shobsi authored Jan 26, 2024
1 parent 9c34d83 commit 0794788
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 5 deletions.
21 changes: 17 additions & 4 deletions bigframes/session/__init__.py
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
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

0 comments on commit 0794788

Please sign in to comment.