feat: Introduce environment variable to set refresh threshold (#1147)

google/auth/_helpers.py

@@ -17,13 +17,32 @@
 import base64
 import calendar
 import datetime
+import os
 import sys
 
 import six
 from six.moves import urllib
 
+from google.auth import environment_vars
 
-REFRESH_THRESHOLD = datetime.timedelta(seconds=300)
+
+# Token server doesn't provide a new a token when doing refresh unless the
+# token is expiring within 30 seconds, so refresh threshold should not be
+# more than 30 seconds. Otherwise auth lib will send tons of refresh requests
+# until 30 seconds before the expiration, and cause a spike of CPU usage.
+REFRESH_THRESHOLD = datetime.timedelta(seconds=20)
+
+
+def get_refresh_threshold():
+    env_threshold = os.environ.get(
+        environment_vars.GOOGLE_CLOUD_TOKEN_REFRESH_THRESHOLD
+    )
+
+    if env_threshold is not None:
+        # If the value is not an int, we can throw exception and let it bubble up
+        return datetime.timedelta(seconds=int(env_threshold))
+
+    return REFRESH_THRESHOLD
 
 
 def copy_docstring(source_class):

google/auth/credentials.py

@@ -65,7 +65,7 @@ def expired(self):
 
         # Remove some threshold from expiry to err on the side of reporting
         # expiration early so that we avoid the 401-refresh-retry loop.
-        skewed_expiry = self.expiry - _helpers.REFRESH_THRESHOLD
+        skewed_expiry = self.expiry - _helpers.get_refresh_threshold()
         return _helpers.utcnow() >= skewed_expiry
 
     @property

google/auth/environment_vars.py

@@ -29,6 +29,10 @@
 situations (such as Google App Engine).
 """
 
+GOOGLE_CLOUD_TOKEN_REFRESH_THRESHOLD = "GOOGLE_CLOUD_TOKEN_REFRESH_THRESHOLD"
+"""Environment variable defines a time in seconds. If a token is going to
+expire within this time, then it will be refreshed"""
+
 CREDENTIALS = "GOOGLE_APPLICATION_CREDENTIALS"
 """Environment variable defining the location of Google application default
 credentials."""

google/oauth2/credentials.py

@@ -272,7 +272,7 @@ def refresh(self, request):
                 raise exceptions.RefreshError(
                     "The refresh_handler returned expiry is not a datetime object."
                 )
-            if _helpers.utcnow() >= expiry - _helpers.REFRESH_THRESHOLD:
+            if _helpers.utcnow() >= expiry - _helpers.get_refresh_threshold():
                 raise exceptions.RefreshError(
                     "The credentials returned by the refresh_handler are "
                     "already expired."
@@ -361,7 +361,7 @@ def from_authorized_user_info(cls, info, scopes=None):
                 expiry.rstrip("Z").split(".")[0], "%Y-%m-%dT%H:%M:%S"
             )
         else:
-            expiry = _helpers.utcnow() - _helpers.REFRESH_THRESHOLD
+            expiry = _helpers.utcnow() - _helpers.get_refresh_threshold()
 
         # process scopes, which needs to be a seq
         if scopes is None and "scopes" in info:

tests/compute_engine/test_credentials.py

@@ -64,7 +64,7 @@ def test_default_state(self):
 
     @mock.patch(
         "google.auth._helpers.utcnow",
-        return_value=datetime.datetime.min + _helpers.REFRESH_THRESHOLD,
+        return_value=datetime.datetime.min + _helpers.get_refresh_threshold(),
     )
     @mock.patch("google.auth.compute_engine._metadata.get", autospec=True)
     def test_refresh_success(self, get, utcnow):
@@ -98,7 +98,7 @@ def test_refresh_success(self, get, utcnow):
 
     @mock.patch(
         "google.auth._helpers.utcnow",
-        return_value=datetime.datetime.min + _helpers.REFRESH_THRESHOLD,
+        return_value=datetime.datetime.min + _helpers.get_refresh_threshold(),
     )
     @mock.patch("google.auth.compute_engine._metadata.get", autospec=True)
     def test_refresh_success_with_scopes(self, get, utcnow):

tests/oauth2/test_credentials.py

@@ -115,7 +115,7 @@ def test_invalid_refresh_handler(self):
     @mock.patch("google.oauth2.reauth.refresh_grant", autospec=True)
     @mock.patch(
         "google.auth._helpers.utcnow",
-        return_value=datetime.datetime.min + _helpers.REFRESH_THRESHOLD,
+        return_value=datetime.datetime.min + _helpers.get_refresh_threshold(),
     )
     def test_refresh_success(self, unused_utcnow, refresh_grant):
         token = "token"
@@ -175,7 +175,7 @@ def test_refresh_no_refresh_token(self):
     @mock.patch("google.oauth2.reauth.refresh_grant", autospec=True)
     @mock.patch(
         "google.auth._helpers.utcnow",
-        return_value=datetime.datetime.min + _helpers.REFRESH_THRESHOLD,
+        return_value=datetime.datetime.min + _helpers.get_refresh_threshold(),
     )
     def test_refresh_with_refresh_token_and_refresh_handler(
         self, unused_utcnow, refresh_grant
@@ -361,7 +361,7 @@ def test_refresh_with_refresh_handler_invalid_expiry(self):
 
     @mock.patch("google.auth._helpers.utcnow", return_value=datetime.datetime.min)
     def test_refresh_with_refresh_handler_expired_token(self, unused_utcnow):
-        expected_expiry = datetime.datetime.min + _helpers.REFRESH_THRESHOLD
+        expected_expiry = datetime.datetime.min + _helpers.get_refresh_threshold()
         # Simulate refresh handler returns an expired token.
         refresh_handler = mock.Mock(return_value=("TOKEN", expected_expiry))
         scopes = ["email", "profile"]
@@ -391,7 +391,7 @@ def test_refresh_with_refresh_handler_expired_token(self, unused_utcnow):
     @mock.patch("google.oauth2.reauth.refresh_grant", autospec=True)
     @mock.patch(
         "google.auth._helpers.utcnow",
-        return_value=datetime.datetime.min + _helpers.REFRESH_THRESHOLD,
+        return_value=datetime.datetime.min + _helpers.get_refresh_threshold(),
     )
     def test_credentials_with_scopes_requested_refresh_success(
         self, unused_utcnow, refresh_grant
@@ -457,7 +457,7 @@ def test_credentials_with_scopes_requested_refresh_success(
     @mock.patch("google.oauth2.reauth.refresh_grant", autospec=True)
     @mock.patch(
         "google.auth._helpers.utcnow",
-        return_value=datetime.datetime.min + _helpers.REFRESH_THRESHOLD,
+        return_value=datetime.datetime.min + _helpers.get_refresh_threshold(),
     )
     def test_credentials_with_only_default_scopes_requested(
         self, unused_utcnow, refresh_grant
@@ -521,7 +521,7 @@ def test_credentials_with_only_default_scopes_requested(
     @mock.patch("google.oauth2.reauth.refresh_grant", autospec=True)
     @mock.patch(
         "google.auth._helpers.utcnow",
-        return_value=datetime.datetime.min + _helpers.REFRESH_THRESHOLD,
+        return_value=datetime.datetime.min + _helpers.get_refresh_threshold(),
     )
     def test_credentials_with_scopes_returned_refresh_success(
         self, unused_utcnow, refresh_grant
@@ -588,7 +588,7 @@ def test_credentials_with_scopes_returned_refresh_success(
     @mock.patch("google.oauth2.reauth.refresh_grant", autospec=True)
     @mock.patch(
         "google.auth._helpers.utcnow",
-        return_value=datetime.datetime.min + _helpers.REFRESH_THRESHOLD,
+        return_value=datetime.datetime.min + _helpers.get_refresh_threshold(),
     )
     def test_credentials_with_scopes_refresh_failure_raises_refresh_error(
         self, unused_utcnow, refresh_grant

tests/test__helpers.py

@@ -13,18 +13,29 @@
 # limitations under the License.
 
 import datetime
+import os
 
+import mock
 import pytest  # type: ignore
 from six.moves import urllib
 
 from google.auth import _helpers
+from google.auth import environment_vars
 
 
 class SourceClass(object):
     def func(self):  # pragma: NO COVER
         """example docstring"""
 
 
+@mock.patch.dict(os.environ)
+def test_get_refresh_threshold():
+    assert _helpers.get_refresh_threshold() == _helpers.REFRESH_THRESHOLD
+
+    os.environ[environment_vars.GOOGLE_CLOUD_TOKEN_REFRESH_THRESHOLD] = "300"
+    assert _helpers.get_refresh_threshold() == datetime.timedelta(seconds=300)
+
+
 def test_copy_docstring_success():
     def func():  # pragma: NO COVER
         pass

tests/test_credentials.py

@@ -47,7 +47,7 @@ def test_expired_and_valid():
     # accomodation. These credentials should be valid.
     credentials.expiry = (
         datetime.datetime.utcnow()
-        + _helpers.REFRESH_THRESHOLD
+        + _helpers.get_refresh_threshold()
         + datetime.timedelta(seconds=1)
     )
 

tests/test_downscoped.py

@@ -670,7 +670,7 @@ def test_before_request_expired(self, utcnow):
         # accommodation. These credentials should be valid.
         credentials.expiry = (
             datetime.datetime.min
-            + _helpers.REFRESH_THRESHOLD
+            + _helpers.get_refresh_threshold()
             + datetime.timedelta(seconds=1)
         )
 

tests/test_external_account.py

@@ -1457,7 +1457,7 @@ def test_before_request_expired(self, utcnow):
         # accomodation. These credentials should be valid.
         credentials.expiry = (
             datetime.datetime.min
-            + _helpers.REFRESH_THRESHOLD
+            + _helpers.get_refresh_threshold()
             + datetime.timedelta(seconds=1)
         )
 
@@ -1510,7 +1510,7 @@ def test_before_request_impersonation_expired(self, utcnow):
         # accomodation. These credentials should be valid.
         credentials.expiry = (
             datetime.datetime.min
-            + _helpers.REFRESH_THRESHOLD
+            + _helpers.get_refresh_threshold()
             + datetime.timedelta(seconds=1)
         )
 

tests/test_iam.py

@@ -45,7 +45,7 @@ def __init__(self):
             super(CredentialsImpl, self).__init__()
             self.token = "token"
             # Force refresh
-            self.expiry = datetime.datetime.min + _helpers.REFRESH_THRESHOLD
+            self.expiry = datetime.datetime.min + _helpers.get_refresh_threshold()
 
         def refresh(self, request):
             pass

tests/test_impersonated_credentials.py

@@ -216,11 +216,11 @@ def test_refresh_source_credentials(self, time_skew):
         credentials = self.make_credentials(lifetime=None)
 
         # Source credentials is refreshed only if it is expired within
-        # _helpers.REFRESH_THRESHOLD from now. We add a time_skew to the expiry, so
+        # _helpers.get_refresh_threshold() from now. We add a time_skew to the expiry, so
         # source credentials is refreshed only if time_skew <= 0.
         credentials._source_credentials.expiry = (
             _helpers.utcnow()
-            + _helpers.REFRESH_THRESHOLD
+            + _helpers.get_refresh_threshold()
             + datetime.timedelta(seconds=time_skew)
         )
         credentials._source_credentials.token = "Token"
@@ -243,7 +243,7 @@ def test_refresh_source_credentials(self, time_skew):
             assert not credentials.expired
 
             # Source credentials is refreshed only if it is expired within
-            # _helpers.REFRESH_THRESHOLD
+            # _helpers.get_refresh_threshold()
             if time_skew > 0:
                 source_cred_refresh.assert_not_called()
             else:

tests/transport/test_grpc.py

@@ -80,7 +80,7 @@ def test_call_no_refresh(self):
 
     def test_call_refresh(self):
         credentials = CredentialsStub()
-        credentials.expiry = datetime.datetime.min + _helpers.REFRESH_THRESHOLD
+        credentials.expiry = datetime.datetime.min + _helpers.get_refresh_threshold()
         request = mock.create_autospec(transport.Request)
 
         plugin = google.auth.transport.grpc.AuthMetadataPlugin(credentials, request)

tests_async/oauth2/test_credentials_async.py

@@ -62,7 +62,7 @@ def test_default_state(self):
     @mock.patch("google.oauth2._reauth_async.refresh_grant", autospec=True)
     @mock.patch(
         "google.auth._helpers.utcnow",
-        return_value=datetime.datetime.min + _helpers.REFRESH_THRESHOLD,
+        return_value=datetime.datetime.min + _helpers.get_refresh_threshold(),
     )
     @pytest.mark.asyncio
     async def test_refresh_success(self, unused_utcnow, refresh_grant):
@@ -124,7 +124,7 @@ async def test_refresh_no_refresh_token(self):
     @mock.patch("google.oauth2._reauth_async.refresh_grant", autospec=True)
     @mock.patch(
         "google.auth._helpers.utcnow",
-        return_value=datetime.datetime.min + _helpers.REFRESH_THRESHOLD,
+        return_value=datetime.datetime.min + _helpers.get_refresh_threshold(),
     )
     @pytest.mark.asyncio
     async def test_credentials_with_scopes_requested_refresh_success(
@@ -188,7 +188,7 @@ async def test_credentials_with_scopes_requested_refresh_success(
     @mock.patch("google.oauth2._reauth_async.refresh_grant", autospec=True)
     @mock.patch(
         "google.auth._helpers.utcnow",
-        return_value=datetime.datetime.min + _helpers.REFRESH_THRESHOLD,
+        return_value=datetime.datetime.min + _helpers.get_refresh_threshold(),
     )
     @pytest.mark.asyncio
     async def test_credentials_with_scopes_returned_refresh_success(
@@ -251,7 +251,7 @@ async def test_credentials_with_scopes_returned_refresh_success(
     @mock.patch("google.oauth2._reauth_async.refresh_grant", autospec=True)
     @mock.patch(
         "google.auth._helpers.utcnow",
-        return_value=datetime.datetime.min + _helpers.REFRESH_THRESHOLD,
+        return_value=datetime.datetime.min + _helpers.get_refresh_threshold(),
     )
     @pytest.mark.asyncio
     async def test_credentials_with_scopes_refresh_failure_raises_refresh_error(

tests_async/test_credentials_async.py

@@ -47,7 +47,7 @@ def test_expired_and_valid():
     # accomodation. These credentials should be valid.
     credentials.expiry = (
         datetime.datetime.utcnow()
-        + _helpers.REFRESH_THRESHOLD
+        + _helpers.get_refresh_threshold()
         + datetime.timedelta(seconds=1)
     )