WIP: cache credentials.

Author: tswast

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,65 @@ 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()
+
+
 class GbqConnector(object):
     def __init__(
         self,
@@ -173,6 +232,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 +245,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 +268,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/test_context.py

@@ -0,0 +1,38 @@
+# -*- coding: utf-8 -*-
+
+try:
+    from unittest import mock
+except ImportError:  # pragma: NO COVER
+    import mock
+
+import pytest
+
+import pandas_gbq
+
+
+@pytest.fixture(autouse=True)
+def mock_get_credentials(monkeypatch):
+    from pandas_gbq import auth
+
+    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):
+    assert pandas_gbq.context.credentials is not None
+    assert pandas_gbq.context.project is not None
+
+    # run read_gbq
+
+    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
+
+    # run read_gbq
+    mock_get_credentials.assert_never_called()  # Not called again.