diff --git a/docs/source/api.rst b/docs/source/api.rst index f5bcf957..a189bae5 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -14,6 +14,10 @@ API Reference read_gbq to_gbq + context + Context .. autofunction:: read_gbq .. autofunction:: to_gbq +.. autodata:: context +.. autoclass:: Context diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index 8681de27..c192578d 100644 --- a/docs/source/changelog.rst +++ b/docs/source/changelog.rst @@ -1,6 +1,14 @@ Changelog ========= +.. _changelog-0.7.0: + +0.7.0 / [unreleased] +-------------------- + +- Add :class:`pandas_gbq.Context` to cache credentials in-memory, across + calls to ``read_gbq`` and ``to_gbq``. (:issue:`198`, :issue:`208`) + .. _changelog-0.6.1: 0.6.1 / [unreleased] diff --git a/pandas_gbq/__init__.py b/pandas_gbq/__init__.py index febda7c6..401d114a 100644 --- a/pandas_gbq/__init__.py +++ b/pandas_gbq/__init__.py @@ -1,4 +1,4 @@ -from .gbq import to_gbq, read_gbq # noqa +from .gbq import to_gbq, read_gbq, Context, context # noqa from ._version import get_versions diff --git a/pandas_gbq/gbq.py b/pandas_gbq/gbq.py index c45384e4..6add4cdc 100644 --- a/pandas_gbq/gbq.py +++ b/pandas_gbq/gbq.py @@ -162,6 +162,72 @@ class TableCreationError(ValueError): pass +class Context(object): + """Storage for objects to be used throughout a session. + + A Context object is initialized when the ``pandas_gbq`` module is + imported, and can be found at :attr:`pandas_gbq.context`. + """ + + def __init__(self): + self._credentials = None + self._project = None + + @property + def credentials(self): + """google.auth.credentials.Credentials: Credentials to use for Google + APIs. + + Note: + These credentials are automatically cached in memory by calls to + :func:`pandas_gbq.read_gbq` and :func:`pandas_gbq.to_gbq`. To + manually set the credentials, construct an + :class:`google.auth.credentials.Credentials` object and set it as + the context credentials as demonstrated in the example below. See + `auth docs`_ for more information on obtaining credentials. + + Example: + Manually setting the context credentials: + >>> import pandas_gbq + >>> from google.oauth2 import service_account + >>> credentials = (service_account + ... .Credentials.from_service_account_file( + ... '/path/to/key.json')) + >>> pandas_gbq.context.credentials = credentials + .. _auth docs: http://google-auth.readthedocs.io + /en/latest/user-guide.html#obtaining-credentials + """ + return self._credentials + + @credentials.setter + def credentials(self, value): + self._credentials = value + + @property + def project(self): + """str: Default project to use for calls to Google APIs. + + Example: + Manually setting the context project: + >>> import pandas_gbq + >>> pandas_gbq.context.project = 'my-project' + """ + return self._project + + @project.setter + def project(self, value): + self._project = value + + +# Create an empty context, used to cache credentials. +context = Context() +"""A :class:`pandas_gbq.Context` object used to cache credentials. + +Credentials automatically are cached in-memory by :func:`pandas_gbq.read_gbq` +and :func:`pandas_gbq.to_gbq`. +""" + + class GbqConnector(object): def __init__( self, @@ -173,6 +239,7 @@ def __init__( location=None, try_credentials=None, ): + global context from google.api_core.exceptions import GoogleAPIError from google.api_core.exceptions import ClientError from pandas_gbq import auth @@ -185,13 +252,20 @@ def __init__( self.auth_local_webserver = auth_local_webserver self.dialect = dialect self.credentials_path = _get_credentials_file() - self.credentials, default_project = auth.get_credentials( - private_key=private_key, - project_id=project_id, - reauth=reauth, - auth_local_webserver=auth_local_webserver, - try_credentials=try_credentials, - ) + + # Load credentials from cache. + self.credentials = context.credentials + default_project = context.project + + # Credentials were explicitly asked for, so don't use the cache. + if private_key or reauth or not self.credentials: + self.credentials, default_project = auth.get_credentials( + private_key=private_key, + project_id=project_id, + reauth=reauth, + auth_local_webserver=auth_local_webserver, + try_credentials=try_credentials, + ) if self.project_id is None: self.project_id = default_project @@ -201,6 +275,12 @@ def __init__( "Could not determine project ID and one was not supplied." ) + # Cache the credentials if they haven't been set yet. + if context.credentials is None: + context.credentials = self.credentials + if context.project is None: + context.project = self.project_id + self.client = self.get_client() # BQ Queries costs $5 per TB. First 1 TB per month is free diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py new file mode 100644 index 00000000..ece8421d --- /dev/null +++ b/tests/unit/conftest.py @@ -0,0 +1,41 @@ +# -*- coding: utf-8 -*- + +try: + from unittest import mock +except ImportError: # pragma: NO COVER + import mock + +import pytest + + +@pytest.fixture(autouse=True, scope="function") +def reset_context(): + import pandas_gbq + + pandas_gbq.context.credentials = None + pandas_gbq.context.project = None + + +@pytest.fixture(autouse=True) +def mock_bigquery_client(monkeypatch): + from pandas_gbq import gbq + from google.api_core.exceptions import NotFound + import google.cloud.bigquery + import google.cloud.bigquery.table + + mock_client = mock.create_autospec(google.cloud.bigquery.Client) + mock_schema = [google.cloud.bigquery.SchemaField("_f0", "INTEGER")] + # Mock out SELECT 1 query results. + mock_query = mock.create_autospec(google.cloud.bigquery.QueryJob) + mock_query.job_id = "some-random-id" + mock_query.state = "DONE" + mock_rows = mock.create_autospec(google.cloud.bigquery.table.RowIterator) + mock_rows.total_rows = 1 + mock_rows.schema = mock_schema + mock_rows.__iter__.return_value = [(1,)] + mock_query.result.return_value = mock_rows + mock_client.query.return_value = mock_query + # Mock table creation. + mock_client.get_table.side_effect = NotFound("nope") + monkeypatch.setattr(gbq.GbqConnector, "get_client", lambda _: mock_client) + return mock_client diff --git a/tests/unit/test_context.py b/tests/unit/test_context.py new file mode 100644 index 00000000..352ece7e --- /dev/null +++ b/tests/unit/test_context.py @@ -0,0 +1,38 @@ +# -*- coding: utf-8 -*- + +try: + from unittest import mock +except ImportError: # pragma: NO COVER + import mock + +import pytest + + +@pytest.fixture(autouse=True) +def mock_get_credentials(monkeypatch): + from pandas_gbq import auth + import google.auth.credentials + + mock_credentials = mock.MagicMock(google.auth.credentials.Credentials) + mock_get_credentials = mock.Mock() + mock_get_credentials.return_value = (mock_credentials, "my-project") + + monkeypatch.setattr(auth, "get_credentials", mock_get_credentials) + return mock_get_credentials + + +def test_read_gbq_should_save_credentials(mock_get_credentials): + import pandas_gbq + + assert pandas_gbq.context.credentials is None + assert pandas_gbq.context.project is None + + pandas_gbq.read_gbq("SELECT 1", dialect="standard") + + assert mock_get_credentials.call_count == 1 + mock_get_credentials.reset_mock() + assert pandas_gbq.context.credentials is not None + assert pandas_gbq.context.project is not None + + pandas_gbq.read_gbq("SELECT 1", dialect="standard") + mock_get_credentials.assert_not_called() diff --git a/tests/unit/test_gbq.py b/tests/unit/test_gbq.py index 4a42e057..1f3ec9a4 100644 --- a/tests/unit/test_gbq.py +++ b/tests/unit/test_gbq.py @@ -25,30 +25,6 @@ def min_bq_version(): return pkg_resources.parse_version("0.32.0") -@pytest.fixture(autouse=True) -def mock_bigquery_client(monkeypatch): - from google.api_core.exceptions import NotFound - import google.cloud.bigquery - import google.cloud.bigquery.table - - mock_client = mock.create_autospec(google.cloud.bigquery.Client) - mock_schema = [google.cloud.bigquery.SchemaField("_f0", "INTEGER")] - # Mock out SELECT 1 query results. - mock_query = mock.create_autospec(google.cloud.bigquery.QueryJob) - mock_query.job_id = "some-random-id" - mock_query.state = "DONE" - mock_rows = mock.create_autospec(google.cloud.bigquery.table.RowIterator) - mock_rows.total_rows = 1 - mock_rows.schema = mock_schema - mock_rows.__iter__.return_value = [(1,)] - mock_query.result.return_value = mock_rows - mock_client.query.return_value = mock_query - # 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): return None, None