From 202667113d0b3daa7455593d019d61b0dd9ea84f Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Fri, 31 Aug 2018 12:01:19 -0700 Subject: [PATCH 1/3] [BUG] skip _try_credentials check for to_gbq Don't do a query read when all you need to write is write credentials. --- docs/source/changelog.rst | 1 + pandas_gbq/auth.py | 27 ++++++++++++++++++++------- pandas_gbq/gbq.py | 5 +++++ tests/unit/test_auth.py | 4 +++- tests/unit/test_gbq.py | 11 +++++++++++ 5 files changed, 40 insertions(+), 8 deletions(-) diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index b4b53488..8681de27 100644 --- a/docs/source/changelog.rst +++ b/docs/source/changelog.rst @@ -12,6 +12,7 @@ Changelog (:issue:`128`) - Reduced verbosity of logging from ``read_gbq``, particularly for short queries. (:issue:`201`) +- Avoid ``SELECT 1`` query when running ``to_gbq``. (:issue:`202`) .. _changelog-0.6.0: diff --git a/pandas_gbq/auth.py b/pandas_gbq/auth.py index c27d342e..2bc2efea 100644 --- a/pandas_gbq/auth.py +++ b/pandas_gbq/auth.py @@ -16,20 +16,28 @@ def get_credentials( - private_key=None, project_id=None, reauth=False, auth_local_webserver=False + private_key=None, + project_id=None, + reauth=False, + auth_local_webserver=False, + try_credentials=None, ): + if try_credentials is None: + try_credentials = _try_credentials + if private_key: return get_service_account_credentials(private_key) # Try to retrieve Application Default Credentials credentials, default_project = get_application_default_credentials( - project_id=project_id + try_credentials, project_id=project_id ) if credentials: return credentials, default_project credentials = get_user_account_credentials( + try_credentials, project_id=project_id, reauth=reauth, auth_local_webserver=auth_local_webserver, @@ -79,7 +87,7 @@ def get_service_account_credentials(private_key): ) -def get_application_default_credentials(project_id=None): +def get_application_default_credentials(try_credentials, project_id=None): """ This method tries to retrieve the "default application credentials". This could be useful for running code on Google Cloud Platform. @@ -111,10 +119,11 @@ def get_application_default_credentials(project_id=None): # used with BigQuery. For example, we could be running on a GCE instance # that does not allow the BigQuery scopes. billing_project = project_id or default_project - return _try_credentials(billing_project, credentials), billing_project + return try_credentials(billing_project, credentials), billing_project def get_user_account_credentials( + try_credentials, project_id=None, reauth=False, auth_local_webserver=False, @@ -151,7 +160,9 @@ def get_user_account_credentials( os.rename("bigquery_credentials.dat", credentials_path) credentials = load_user_account_credentials( - project_id=project_id, credentials_path=credentials_path + try_credentials, + project_id=project_id, + credentials_path=credentials_path, ) client_config = { @@ -187,7 +198,9 @@ def get_user_account_credentials( return credentials -def load_user_account_credentials(project_id=None, credentials_path=None): +def load_user_account_credentials( + try_credentials, project_id=None, credentials_path=None +): """ Loads user account credentials from a local file. @@ -230,7 +243,7 @@ def load_user_account_credentials(project_id=None, credentials_path=None): request = google.auth.transport.requests.Request() credentials.refresh(request) - return _try_credentials(project_id, credentials) + return try_credentials(project_id, credentials) def get_default_credentials_path(): diff --git a/pandas_gbq/gbq.py b/pandas_gbq/gbq.py index eba08009..c45384e4 100644 --- a/pandas_gbq/gbq.py +++ b/pandas_gbq/gbq.py @@ -171,6 +171,7 @@ def __init__( auth_local_webserver=False, dialect="legacy", location=None, + try_credentials=None, ): from google.api_core.exceptions import GoogleAPIError from google.api_core.exceptions import ClientError @@ -189,6 +190,7 @@ def __init__( project_id=project_id, reauth=reauth, auth_local_webserver=auth_local_webserver, + try_credentials=try_credentials, ) if self.project_id is None: @@ -804,6 +806,9 @@ def to_gbq( private_key=private_key, auth_local_webserver=auth_local_webserver, location=location, + # Avoid reads when writing tables. + # https://github.com/pydata/pandas-gbq/issues/202 + try_credentials=lambda project, creds: creds, ) dataset_id, table_id = destination_table.rsplit(".", 1) diff --git a/tests/unit/test_auth.py b/tests/unit/test_auth.py index 50c71b27..d8107a40 100644 --- a/tests/unit/test_auth.py +++ b/tests/unit/test_auth.py @@ -95,7 +95,9 @@ def mock_default_credentials(scopes=None, request=None): google.auth.credentials.Credentials ) - def mock_load_credentials(project_id=None, credentials_path=None): + def mock_load_credentials( + try_credentials, project_id=None, credentials_path=None + ): return mock_user_credentials monkeypatch.setattr( diff --git a/tests/unit/test_gbq.py b/tests/unit/test_gbq.py index df80d5d2..7557145a 100644 --- a/tests/unit/test_gbq.py +++ b/tests/unit/test_gbq.py @@ -210,6 +210,17 @@ def test_to_gbq_with_verbose_old_pandas_no_warnings(recwarn, min_bq_version): assert len(recwarn) == 0 +def test_to_gbq_doesnt_run_query(recwarn, min_bq_version): + try: + gbq.to_gbq( + DataFrame([[1]]), "dataset.tablename", project_id="my-project" + ) + except gbq.TableCreationError: + pass + + gbq.GbqConnector.get_client(None).query.assert_not_called() + + def test_read_gbq_with_no_project_id_given_should_fail(monkeypatch): from pandas_gbq import auth From 906ffdf4fe81ce0da034f6288b4b15955a1ae53f Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Fri, 31 Aug 2018 12:15:50 -0700 Subject: [PATCH 2/3] Fix auth tests. --- tests/system/test_auth.py | 25 +++++++++++++++++++------ 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/tests/system/test_auth.py b/tests/system/test_auth.py index 7ccc79c2..f7cdf014 100644 --- a/tests/system/test_auth.py +++ b/tests/system/test_auth.py @@ -66,9 +66,13 @@ def test_get_application_default_credentials_does_not_throw_error(): with mock.patch( "google.auth.default", side_effect=DefaultCredentialsError() ): - credentials, _ = auth.get_application_default_credentials() + credentials, _ = auth.get_application_default_credentials( + try_credentials=auth._try_credentials + ) else: - credentials, _ = auth.get_application_default_credentials() + credentials, _ = auth.get_application_default_credentials( + try_credentials=auth._try_credentials + ) assert credentials is None @@ -77,7 +81,9 @@ def test_get_application_default_credentials_returns_credentials(): pytest.skip("Cannot get default_credentials " "from the environment!") from google.auth.credentials import Credentials - credentials, default_project = auth.get_application_default_credentials() + credentials, default_project = auth.get_application_default_credentials( + try_credentials=auth._try_credentials + ) assert isinstance(credentials, Credentials) assert default_project is not None @@ -88,7 +94,9 @@ def test_get_user_account_credentials_bad_file_returns_credentials(): from google.auth.credentials import Credentials with mock.patch("__main__.open", side_effect=IOError()): - credentials = auth.get_user_account_credentials() + credentials = auth.get_user_account_credentials( + try_credentials=auth._try_credentials + ) assert isinstance(credentials, Credentials) @@ -97,7 +105,9 @@ def test_get_user_account_credentials_returns_credentials(project_id): from google.auth.credentials import Credentials credentials = auth.get_user_account_credentials( - project_id=project_id, auth_local_webserver=True + project_id=project_id, + auth_local_webserver=True, + try_credentials=auth._try_credentials, ) assert isinstance(credentials, Credentials) @@ -107,6 +117,9 @@ def test_get_user_account_credentials_reauth_returns_credentials(project_id): from google.auth.credentials import Credentials credentials = auth.get_user_account_credentials( - project_id=project_id, auth_local_webserver=True, reauth=True + project_id=project_id, + auth_local_webserver=True, + reauth=True, + try_credentials=auth._try_credentials, ) assert isinstance(credentials, Credentials) From 3607d63588bc7f1151e456fe47c62bc7093f0cf7 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Fri, 31 Aug 2018 13:36:45 -0700 Subject: [PATCH 3/3] Fix mock for Python 2.7. --- tests/unit/test_gbq.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_gbq.py b/tests/unit/test_gbq.py index 7557145a..4a42e057 100644 --- a/tests/unit/test_gbq.py +++ b/tests/unit/test_gbq.py @@ -46,6 +46,7 @@ def mock_bigquery_client(monkeypatch): # Mock table creation. mock_client.get_table.side_effect = NotFound("nope") monkeypatch.setattr(gbq.GbqConnector, "get_client", lambda _: mock_client) + return mock_client def mock_none_credentials(*args, **kwargs): @@ -210,7 +211,9 @@ def test_to_gbq_with_verbose_old_pandas_no_warnings(recwarn, min_bq_version): assert len(recwarn) == 0 -def test_to_gbq_doesnt_run_query(recwarn, min_bq_version): +def test_to_gbq_doesnt_run_query( + recwarn, mock_bigquery_client, min_bq_version +): try: gbq.to_gbq( DataFrame([[1]]), "dataset.tablename", project_id="my-project" @@ -218,7 +221,7 @@ def test_to_gbq_doesnt_run_query(recwarn, min_bq_version): except gbq.TableCreationError: pass - gbq.GbqConnector.get_client(None).query.assert_not_called() + mock_bigquery_client.query.assert_not_called() def test_read_gbq_with_no_project_id_given_should_fail(monkeypatch):