refactor: use google.auth TokenState for credentials validity (#1089)

Author: jackwotherspoon

google/cloud/sql/connector/client.py

@@ -21,7 +21,6 @@
 import aiohttp
 from cryptography.hazmat.backends import default_backend
 from cryptography.x509 import load_pem_x509_certificate
-import google.auth.transport.requests
 
 from google.cloud.sql.connector.refresh_utils import _downscope_credentials
 from google.cloud.sql.connector.version import __version__ as version
@@ -113,9 +112,6 @@ async def _get_metadata(
             addresses and their type and a string representing the
             certificate authority.
         """
-        if not self._credentials.valid:
-            request = google.auth.transport.requests.Request()
-            self._credentials.refresh(request)
 
         headers = {
             "Authorization": f"Bearer {self._credentials.token}",
@@ -177,10 +173,6 @@ async def _get_ephemeral(
 
         logger.debug(f"['{instance}']: Requesting ephemeral certificate")
 
-        if not self._credentials.valid:
-            request = google.auth.transport.requests.Request()
-            self._credentials.refresh(request)
-
         headers = {
             "Authorization": f"Bearer {self._credentials.token}",
         }

google/cloud/sql/connector/instance.py

@@ -26,6 +26,8 @@
 from typing import Any, Dict, Tuple, TYPE_CHECKING
 
 import aiohttp
+from google.auth.credentials import TokenState
+from google.auth.transport import requests
 
 from google.cloud.sql.connector.client import CloudSQLClient
 from google.cloud.sql.connector.exceptions import AutoIAMAuthNotSupported
@@ -231,6 +233,10 @@ async def _perform_refresh(self) -> ConnectionInfo:
 
             logger.debug(f"['{self._instance_connection_string}']: Creating context")
 
+            # before making Cloud SQL Admin API calls, refresh creds
+            if not self._client._credentials.token_state == TokenState.FRESH:
+                self._client._credentials.refresh(requests.Request())
+
             metadata_task = asyncio.create_task(
                 self._client._get_metadata(
                     self._project,

tests/unit/mocks.py

@@ -13,13 +13,14 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 # file containing all mocks used for Cloud SQL Python Connector unit tests
 
 import datetime
 import json
 import ssl
 from tempfile import TemporaryDirectory
-from typing import Any, Callable, Dict, Optional, Tuple
+from typing import Any, Callable, Dict, Literal, Optional, Tuple
 
 from aiohttp import web
 from cryptography import x509
@@ -28,7 +29,9 @@
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
 from cryptography.x509.oid import NameOID
+from google.auth import _helpers
 from google.auth.credentials import Credentials
+from google.auth.credentials import TokenState
 
 from google.cloud.sql.connector.connector import _DEFAULT_UNIVERSE_DOMAIN
 from google.cloud.sql.connector.utils import generate_keys
@@ -48,7 +51,7 @@ def __class__(self) -> Credentials:
         # set class type to google auth Credentials
         return Credentials
 
-    def refresh(self, request: Callable) -> None:
+    def refresh(self, _: Callable) -> None:
         """Refreshes the access token."""
         self.token = "12345"
         self.expiry = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(
@@ -75,13 +78,33 @@ def universe_domain(self) -> str:
         return self._universe_domain
 
     @property
-    def valid(self) -> bool:
-        """Checks the validity of the credentials.
-
-        This is True if the credentials have a token and the token
-        is not expired.
+    def token_state(
+        self,
+    ) -> Literal[TokenState.FRESH, TokenState.STALE, TokenState.INVALID]:
         """
-        return self.token is not None and not self.expired
+        Tracks the state of a token.
+        FRESH: The token is valid. It is not expired or close to expired, or the token has no expiry.
+        STALE: The token is close to expired, and should be refreshed. The token can be used normally.
+        INVALID: The token is expired or invalid. The token cannot be used for a normal operation.
+        """
+        if self.token is None:
+            return TokenState.INVALID
+
+        # Credentials that can't expire are always treated as fresh.
+        if self.expiry is None:
+            return TokenState.FRESH
+
+        expired = datetime.datetime.now(datetime.timezone.utc) >= self.expiry
+        if expired:
+            return TokenState.INVALID
+
+        is_stale = datetime.datetime.now(datetime.timezone.utc) >= (
+            self.expiry - _helpers.REFRESH_THRESHOLD
+        )
+        if is_stale:
+            return TokenState.STALE
+
+        return TokenState.FRESH
 
 
 def generate_cert(

tests/unit/test_refresh_utils.py

@@ -13,12 +13,14 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 import asyncio
 import datetime
 
 from conftest import SCOPES  # type: ignore
 import google.auth
 from google.auth.credentials import Credentials
+from google.auth.credentials import TokenState
 import google.oauth2.credentials
 from mock import Mock
 from mock import patch
@@ -32,7 +34,7 @@
 @pytest.fixture
 def credentials() -> Credentials:
     credentials = Mock(spec=Credentials)
-    credentials.valid = True
+    credentials.token_state = TokenState.FRESH
     credentials.token = "12345"
     return credentials