diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index da39ab43ce..bd813c8c6b 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -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( @@ -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: diff --git a/tests/unit/session/test_session.py b/tests/unit/session/test_session.py index d38a393f27..5fc8996993 100644 --- a/tests/unit/session/test_session.py +++ b/tests/unit/session/test_session.py @@ -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( @@ -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(