[ENH] cache credentials in-memory

This commit adds a global pandas_gbq.context variable which caches the
project ID and credentials across calls to read_gbq and to_gbq.

Author: tswast

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

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
 

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

tests/unit/conftest.py

@@ -0,0 +1,40 @@
+# -*- 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

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")
+
+    mock_get_credentials.assert_called_once()
+    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()

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