Skip to content

Commit b2e0bcb

Browse files
dauregtswast
authored andcommitted
read project_id from credentials if available (#258)
* read project_id from credentials if available * Move project inference from credentials to shared GbqConnector To more easily test the project ID inference, this commit also sends the project ID as part of the query method call. Even though this is redundant because the client object also contains the project ID, we aren't exercising the bigquery.Client constructor in the unit tests to avoid making real API calls.
1 parent dd79b4f commit b2e0bcb

File tree

3 files changed

+59
-0
lines changed

3 files changed

+59
-0
lines changed

docs/source/changelog.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ Enhancements
2525
- Allow ``table_schema`` in :func:`to_gbq` to contain only a subset of columns,
2626
with the rest being populated using the DataFrame dtypes (:issue:`218`)
2727
(contributed by @johnpaton)
28+
- Read ``project_id`` in :func:`to_gbq` from provided ``credentials`` if
29+
available (contributed by @daureg)
2830

2931
.. _changelog-0.9.0:
3032

pandas_gbq/gbq.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -318,6 +318,11 @@ def __init__(
318318
self.credentials = credentials
319319
default_project = None
320320

321+
# Service account credentials have a project associated with them.
322+
# Prefer that project if none was supplied.
323+
if self.project_id is None and hasattr(self.credentials, "project_id"):
324+
self.project_id = credentials.project_id
325+
321326
# Load credentials from cache.
322327
if not self.credentials:
323328
self.credentials = context.credentials
@@ -421,6 +426,7 @@ def run_query(self, query, **kwargs):
421426
query,
422427
job_config=bigquery.QueryJobConfig.from_api_repr(job_config),
423428
location=self.location,
429+
project=self.project_id,
424430
)
425431
logger.debug("Query running...")
426432
except (RefreshError, ValueError):

tests/unit/test_gbq.py

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,27 @@ def mock_get_credentials(*args, **kwargs):
4747
return mock_credentials, "default-project"
4848

4949

50+
@pytest.fixture
51+
def mock_service_account_credentials():
52+
import google.oauth2.service_account
53+
54+
mock_credentials = mock.create_autospec(
55+
google.oauth2.service_account.Credentials
56+
)
57+
return mock_credentials
58+
59+
60+
@pytest.fixture
61+
def mock_compute_engine_credentials():
62+
import google.auth.compute_engine
63+
64+
mock_credentials = mock.create_autospec(
65+
google.auth.compute_engine.Credentials
66+
)
67+
return mock_credentials
68+
69+
70+
@pytest.fixture
5071
def mock_get_user_credentials(*args, **kwargs):
5172
import google.auth.credentials
5273

@@ -260,6 +281,36 @@ def test_read_gbq_with_inferred_project_id(monkeypatch):
260281
assert df is not None
261282

262283

284+
def test_read_gbq_with_inferred_project_id_from_service_account_credentials(
285+
mock_bigquery_client, mock_service_account_credentials
286+
):
287+
mock_service_account_credentials.project_id = "service_account_project_id"
288+
df = gbq.read_gbq(
289+
"SELECT 1",
290+
dialect="standard",
291+
credentials=mock_service_account_credentials,
292+
)
293+
assert df is not None
294+
mock_bigquery_client.query.assert_called_once_with(
295+
"SELECT 1",
296+
job_config=mock.ANY,
297+
location=None,
298+
project="service_account_project_id",
299+
)
300+
301+
302+
def test_read_gbq_without_inferred_project_id_from_compute_engine_credentials(
303+
mock_compute_engine_credentials
304+
):
305+
with pytest.raises(ValueError) as exception:
306+
gbq.read_gbq(
307+
"SELECT 1",
308+
dialect="standard",
309+
credentials=mock_compute_engine_credentials,
310+
)
311+
assert "Could not determine project ID" in str(exception)
312+
313+
263314
def test_read_gbq_with_invalid_private_key_json_should_fail():
264315
with pytest.raises(pandas_gbq.exceptions.InvalidPrivateKeyFormat):
265316
gbq.read_gbq(

0 commit comments

Comments
 (0)