feat: try colab credentials in get_user_credentials if client_id not specified (#76)

* feat: try colab credentials in `get_user_credentials` if `client_id` not specified

* fix AttributeError: 'tuple' object has no attribute 'valid'

* Update pydata_google_auth/auth.py

* Update pydata_google_auth/auth.py

* add unit tests

* add unit tests, docs, and always try colab

* update readthedocs config

* fix indentation

* try again

* fix requirements

Author: tswast

.readthedocs.yml

@@ -1,6 +1,16 @@
-requirements_file: docs/requirements-docs.txt
+version: 2
+
 build:
-  image: latest
+  os: ubuntu-24.04
+  tools:
+    python: "3.10"
+
 python:
-  pip_install: true
-  version: 3.6
+  install:
+    - requirements: docs/requirements-docs.txt
+    # Install our python package before building the docs
+    - method: pip
+      path: .
+
+sphinx:
+  configuration: docs/source/conf.py

docs/requirements-docs.txt

@@ -1,5 +1,12 @@
-ipython
-matplotlib
-numpydoc
-sphinx==4.0.2
-sphinx_rtd_theme
+# We need to pin to specific versions of the `sphinxcontrib-*` packages
+# which still support sphinx 4.x.
+# See https://github.com/googleapis/sphinx-docfx-yaml/issues/344
+# and https://github.com/googleapis/sphinx-docfx-yaml/issues/345.
+sphinxcontrib-applehelp==1.0.4
+sphinxcontrib-devhelp==1.0.2
+sphinxcontrib-htmlhelp==2.0.1
+sphinxcontrib-qthelp==1.0.3
+sphinxcontrib-serializinghtml==1.1.5
+sphinx==4.5.0
+alabaster
+recommonmark

docs/source/conf.py

@@ -37,8 +37,6 @@
     "sphinx.ext.extlinks",
     "sphinx.ext.todo",
     "sphinx.ext.napoleon",
-    "IPython.sphinxext.ipython_console_highlighting",
-    "IPython.sphinxext.ipython_directive",
     "sphinx.ext.intersphinx",
     "sphinx.ext.coverage",
     "sphinx.ext.ifconfig",
@@ -129,20 +127,9 @@
 
 # -- Options for HTML output ----------------------------------------------
 
-# Taken from docs.readthedocs.io:
-# on_rtd is whether we are on readthedocs.io
-on_rtd = os.environ.get("READTHEDOCS", None) == "True"
-
-if not on_rtd:  # only import and set the theme if we're building docs locally
-    import sphinx_rtd_theme
-
-    html_theme = "sphinx_rtd_theme"
-    html_theme_path = [sphinx_rtd_theme.get_html_theme_path()]
-
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes.
-#
-# html_theme = 'alabaster'
+html_theme = "alabaster"
 
 # Theme options are theme-specific and customize the look and feel of a theme
 # further.  For a list of options available for each theme, see the

noxfile.py

@@ -27,7 +27,9 @@
 BLACK_VERSION = "black==22.12.0"
 BLACK_PATHS = ["docs", "pydata_google_auth", "tests", "noxfile.py", "setup.py"]
 
-DEFAULT_PYTHON_VERSION = "3.8"
+SPHINX_VERSION = "sphinx==4.5.0"
+
+DEFAULT_PYTHON_VERSION = "3.10"
 SYSTEM_TEST_PYTHON_VERSIONS = ["3.9", "3.13"]
 UNIT_TEST_PYTHON_VERSIONS = ["3.9", "3.10", "3.11", "3.12", "3.13"]
 
@@ -191,8 +193,21 @@ def cover(session):
 def docs(session):
     """Build the docs for this library."""
 
-    session.install("sphinx==4.0.2", "sphinx_rtd_theme", "ipython")
     session.install("-e", ".")
+    session.install(
+        # We need to pin to specific versions of the `sphinxcontrib-*` packages
+        # which still support sphinx 4.x.
+        # See https://github.com/googleapis/sphinx-docfx-yaml/issues/344
+        # and https://github.com/googleapis/sphinx-docfx-yaml/issues/345.
+        "sphinxcontrib-applehelp==1.0.4",
+        "sphinxcontrib-devhelp==1.0.2",
+        "sphinxcontrib-htmlhelp==2.0.1",
+        "sphinxcontrib-qthelp==1.0.3",
+        "sphinxcontrib-serializinghtml==1.1.5",
+        SPHINX_VERSION,
+        "alabaster",
+        "recommonmark",
+    )
 
     shutil.rmtree(os.path.join("docs", "source", "_build"), ignore_errors=True)
     session.run(

pydata_google_auth/auth.py

@@ -163,26 +163,48 @@ def default(
     return credentials, None
 
 
-def _ensure_application_default_credentials_in_colab_environment():
-    # This is a special handling for google colab environment where we want to
-    # use the colab specific authentication flow
-    # https://github.com/googlecolab/colabtools/blob/3c8772efd332289e1c6d1204826b0915d22b5b95/google/colab/auth.py#L209
+def try_colab_auth_import():
     try:
         from google.colab import auth
 
-        auth.authenticate_user()
+        return auth
     except Exception:
         # We are catching a broad exception class here because we want to be
         # agnostic to anything that could internally go wrong in the google
         # colab auth. Some of the known exception we want to pass on are:
         #
         # ModuleNotFoundError: No module named 'google.colab'
         # ImportError: cannot import name 'auth' from 'google.cloud'
+        return None
+
+
+def get_colab_default_credentials(scopes):
+    """This is a special handling for google colab environment where we want to
+    use the colab specific authentication flow.
+
+    See:
+    https://github.com/googlecolab/colabtools/blob/3c8772efd332289e1c6d1204826b0915d22b5b95/google/colab/auth.py#L209
+    """
+    auth = try_colab_auth_import()
+    if auth is None:
+        return None, None
+
+    try:
+        auth.authenticate_user()
+
+        # authenticate_user() sets the default credentials, but we
+        # still need to get the token from those default credentials.
+        return get_application_default_credentials(scopes=scopes)
+    except Exception:
+        # We are catching a broad exception class here because we want to be
+        # agnostic to anything that could internally go wrong in the google
+        # colab auth. Some of the known exception we want to pass on are:
+        #
         # MessageError: Error: credential propagation was unsuccessful
         #
         # The MessageError happens on Vertex Colab when it fails to resolve auth
         # from the Compute Engine Metadata server.
-        pass
+        return None, None
 
 
 def get_application_default_credentials(scopes):
@@ -205,9 +227,6 @@ def get_application_default_credentials(scopes):
         from the environment. Or, the retrieved credentials do not
         have access to the project (project_id) on BigQuery.
     """
-
-    _ensure_application_default_credentials_in_colab_environment()
-
     try:
         credentials, project = google.auth.default(scopes=scopes)
     except (google.auth.exceptions.DefaultCredentialsError, IOError) as exc:
@@ -240,8 +259,12 @@ def get_user_credentials(
     """
     Gets user account credentials.
 
-    This function authenticates using user credentials, either loading saved
-    credentials from the cache or by going through the OAuth 2.0 flow.
+    This function authenticates using user credentials, by trying to
+
+    1. Authenticate using ``google.colab.authenticate_user()``
+    2. Load saved credentials from the ``credentials_cache``
+    3. Go through the OAuth 2.0 flow (with provided ``client_id`` and
+       ``client_secret``)
 
     The default read-write cache attempts to read credentials from a file on
     disk. If these credentials are not found or are invalid, it begins an
@@ -311,6 +334,20 @@ def get_user_credentials(
     pydata_google_auth.exceptions.PyDataCredentialsError
         If unable to get valid user credentials.
     """
+
+    # Try to authenticate the user with Colab-based credentials, if possible.
+    # The default_project ignored for colab credentials. It's not usually set,
+    # anyway.
+    credentials, _ = get_colab_default_credentials(scopes)
+
+    # Break early to avoid trying to fetch any other kinds of credentials.
+    # Prefer Colab credentials over any credentials based on the default
+    # client ID.
+    if credentials:
+        # Make sure to exit early since we don't want to try to save these
+        # credentials to a cache file.
+        return credentials
+
     if auth_local_webserver is not None:
         use_local_webserver = auth_local_webserver
 

tests/unit/test_auth.py

@@ -12,11 +12,18 @@
 
 from google.oauth2 import service_account
 from pydata_google_auth import exceptions
+import pydata_google_auth.cache
 
 
 TEST_SCOPES = ["https://www.googleapis.com/auth/cloud-platform"]
 
 
+class FakeCredentials(object):
+    @property
+    def valid(self):
+        return True
+
+
 @pytest.fixture
 def module_under_test():
     from pydata_google_auth import auth
@@ -57,10 +64,51 @@ def mock_default_credentials(scopes=None, request=None):
     assert credentials is mock_user_credentials
 
 
-class FakeCredentials(object):
-    @property
-    def valid(self):
-        return True
+def test_get_user_credentials_tries_colab_first(monkeypatch, module_under_test):
+    colab_auth_module = mock.Mock()
+    try_colab_auth_import = mock.Mock(return_value=colab_auth_module)
+    monkeypatch.setattr(
+        module_under_test, "try_colab_auth_import", try_colab_auth_import
+    )
+    default_credentials = mock.create_autospec(google.auth.credentials.Credentials)
+    default_call_times = 0
+
+    # Can't use a Mock because we want to check authenticate_user.
+    def mock_default(scopes=None, request=None):
+        nonlocal default_call_times
+        default_call_times += 1
+
+        # Make sure colab auth is called first.
+        colab_auth_module.authenticate_user.assert_called_once_with()
+
+        return (
+            default_credentials,
+            "colab-project",  # In reality, often None.
+        )
+
+    monkeypatch.setattr(google.auth, "default", mock_default)
+
+    credentials = module_under_test.get_user_credentials(TEST_SCOPES)
+
+    assert credentials is default_credentials
+    assert default_call_times == 1
+
+
+def test_get_user_credentials_skips_colab_if_no_colab(monkeypatch, module_under_test):
+    try_colab_auth_import = mock.Mock(return_value=None)
+    monkeypatch.setattr(
+        module_under_test, "try_colab_auth_import", try_colab_auth_import
+    )
+    credentials_cache = mock.create_autospec(pydata_google_auth.cache.CredentialsCache)
+    loaded_credentials = mock.Mock()
+    credentials_cache.load.return_value = loaded_credentials
+
+    credentials = module_under_test.get_user_credentials(
+        TEST_SCOPES,
+        credentials_cache=credentials_cache,
+    )
+
+    assert credentials is loaded_credentials
 
 
 def test_load_service_account_credentials(monkeypatch, tmp_path, module_under_test):