Skip to content

Commit

Permalink
feat: support refresh callable on google.oauth2.credentials.Credentia…
Browse files Browse the repository at this point in the history
…ls (#812)

This is an optional parameter that can be set via the constructor.
It is used to provide the credentials with new tokens and their
expiration time on `refresh()` call.

```
def refresh_handler(request, scopes):
    # Generate a new token for the requested scopes by calling
    # an external process.
    return (
        "ACCESS_TOKEN",
        _helpers.utcnow() + datetime.timedelta(seconds=3600))

creds = google.oauth2.credentials.Credentials(
    scopes=scopes,
    refresh_handler=refresh_handler)
creds.refresh(request)
```

It is useful in the following cases:
- Useful in general when tokens are obtained by calling some
  external process on demand.
- Useful in particular for retrieving downscoped tokens from a
  token broker.

This should have no impact on existing behavior. Refresh tokens
will still have higher priority over refresh handlers.

A getter and setter is exposed to make it easy to set the callable
on unpickled credentials as the callable may not be easily serialized.

```
unpickled = pickle.loads(pickle.dumps(oauth_creds))
unpickled.refresh_handler = refresh_handler
```
  • Loading branch information
bojeil-google committed Jul 22, 2021
1 parent 63ac08a commit ec2fb18
Show file tree
Hide file tree
Showing 2 changed files with 353 additions and 3 deletions.
71 changes: 68 additions & 3 deletions google/oauth2/credentials.py
Expand Up @@ -74,6 +74,7 @@ def __init__(
quota_project_id=None,
expiry=None,
rapt_token=None,
refresh_handler=None,
):
"""
Args:
Expand Down Expand Up @@ -103,6 +104,13 @@ def __init__(
This project may be different from the project used to
create the credentials.
rapt_token (Optional[str]): The reauth Proof Token.
refresh_handler (Optional[Callable[[google.auth.transport.Request, Sequence[str]], [str, datetime]]]):
A callable which takes in the HTTP request callable and the list of
OAuth scopes and when called returns an access token string for the
requested scopes and its expiry datetime. This is useful when no
refresh tokens are provided and tokens are obtained by calling
some external process on demand. It is particularly useful for
retrieving downscoped tokens from a token broker.
"""
super(Credentials, self).__init__()
self.token = token
Expand All @@ -116,13 +124,20 @@ def __init__(
self._client_secret = client_secret
self._quota_project_id = quota_project_id
self._rapt_token = rapt_token
self.refresh_handler = refresh_handler

def __getstate__(self):
"""A __getstate__ method must exist for the __setstate__ to be called
This is identical to the default implementation.
See https://docs.python.org/3.7/library/pickle.html#object.__setstate__
"""
return self.__dict__
state_dict = self.__dict__.copy()
# Remove _refresh_handler function as there are limitations pickling and
# unpickling certain callables (lambda, functools.partial instances)
# because they need to be importable.
# Instead, the refresh_handler setter should be used to repopulate this.
del state_dict["_refresh_handler"]
return state_dict

def __setstate__(self, d):
"""Credentials pickled with older versions of the class do not have
Expand All @@ -138,6 +153,8 @@ def __setstate__(self, d):
self._client_secret = d.get("_client_secret")
self._quota_project_id = d.get("_quota_project_id")
self._rapt_token = d.get("_rapt_token")
# The refresh_handler setter should be used to repopulate this.
self._refresh_handler = None

@property
def refresh_token(self):
Expand Down Expand Up @@ -187,6 +204,31 @@ def rapt_token(self):
"""Optional[str]: The reauth Proof Token."""
return self._rapt_token

@property
def refresh_handler(self):
"""Returns the refresh handler if available.
Returns:
Optional[Callable[[google.auth.transport.Request, Sequence[str]], [str, datetime]]]:
The current refresh handler.
"""
return self._refresh_handler

@refresh_handler.setter
def refresh_handler(self, value):
"""Updates the current refresh handler.
Args:
value (Optional[Callable[[google.auth.transport.Request, Sequence[str]], [str, datetime]]]):
The updated value of the refresh handler.
Raises:
TypeError: If the value is not a callable or None.
"""
if not callable(value) and value is not None:
raise TypeError("The provided refresh_handler is not a callable or None.")
self._refresh_handler = value

@_helpers.copy_docstring(credentials.CredentialsWithQuotaProject)
def with_quota_project(self, quota_project_id):

Expand All @@ -205,6 +247,31 @@ def with_quota_project(self, quota_project_id):

@_helpers.copy_docstring(credentials.Credentials)
def refresh(self, request):
scopes = self._scopes if self._scopes is not None else self._default_scopes
# Use refresh handler if available and no refresh token is
# available. This is useful in general when tokens are obtained by calling
# some external process on demand. It is particularly useful for retrieving
# downscoped tokens from a token broker.
if self._refresh_token is None and self.refresh_handler:
token, expiry = self.refresh_handler(request, scopes=scopes)
# Validate returned data.
if not isinstance(token, str):
raise exceptions.RefreshError(
"The refresh_handler returned token is not a string."
)
if not isinstance(expiry, datetime):
raise exceptions.RefreshError(
"The refresh_handler returned expiry is not a datetime object."
)
if _helpers.utcnow() >= expiry - _helpers.CLOCK_SKEW:
raise exceptions.RefreshError(
"The credentials returned by the refresh_handler are "
"already expired."
)
self.token = token
self.expiry = expiry
return

if (
self._refresh_token is None
or self._token_uri is None
Expand All @@ -217,8 +284,6 @@ def refresh(self, request):
"token_uri, client_id, and client_secret."
)

scopes = self._scopes if self._scopes is not None else self._default_scopes

(
access_token,
refresh_token,
Expand Down

0 comments on commit ec2fb18

Please sign in to comment.