Skip to content

Commit c304465

Browse files
author
Alejandro Casanovas
committed
Updated LockableFileSystemTokenBackend
1 parent 6c69c47 commit c304465

File tree

1 file changed

+35
-17
lines changed

1 file changed

+35
-17
lines changed

examples/token_backends.py

+35-17
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ def should_refresh_token(self, con: Optional[Connection] = None, username: Optio
6262
# 1) check if the token is already a new one:
6363
old_access_token = self.get_access_token(username=username)
6464
if old_access_token:
65-
self.load_token()
65+
self.load_token() # retrieve again the token from the backend
6666
new_access_token = self.get_access_token(username=username)
6767
if old_access_token["secret"] != new_access_token["secret"]:
6868
# The token is different so the refresh took part somewhere else.
@@ -130,8 +130,8 @@ class LockableFileSystemTokenBackend(FileSystemTokenBackend):
130130
"""
131131

132132
def __init__(self, *args, **kwargs):
133-
self.max_tries = kwargs.pop("max_tries")
134-
self.fs_wait = False
133+
self.max_tries: int = kwargs.pop("max_tries", 3)
134+
self.fs_wait: bool = False
135135
super().__init__(*args, **kwargs)
136136

137137
def should_refresh_token(self, con: Optional[Connection] = None, username: Optional[str] = None):
@@ -150,10 +150,10 @@ def should_refresh_token(self, con: Optional[Connection] = None, username: Optio
150150
unlocks the file. Since refreshing has been taken care of, the calling
151151
method does not need to refresh and we return None.
152152
153-
If we are blocked because the file is locked, that means another
153+
If we are blocked because the file is locked, that means another
154154
instance is using it. We'll change the backend's state to waiting,
155155
sleep for 2 seconds, reload a token into memory from the file (since
156-
another process is using it, we can assume it's being updated), and
156+
another process is using it, we can assume it's being updated), and
157157
loop again.
158158
159159
If this newly loaded token is not expired, the other instance loaded
@@ -165,27 +165,45 @@ def should_refresh_token(self, con: Optional[Connection] = None, username: Optio
165165
runtime exception
166166
"""
167167

168+
# 1) check if the token is already a new one:
169+
old_access_token = self.get_access_token(username=username)
170+
if old_access_token:
171+
self.load_token() # retrieve again the token from the backend
172+
new_access_token = self.get_access_token(username=username)
173+
if old_access_token["secret"] != new_access_token["secret"]:
174+
# The token is different so the refresh took part somewhere else.
175+
# Return False so the connection can update the token access from the backend into the session
176+
return False
177+
178+
# 2) Here the token stored in the token backend and in the token cache of this instance is the same
168179
for i in range(self.max_tries, 0, -1):
169180
try:
170-
with Lock(self.token_path, "r+",
171-
fail_when_locked=True, timeout=0):
172-
log.debug("Locked oauth token file")
173-
if con.get_refresh_token() is False:
181+
with Lock(self.token_path, "r+", fail_when_locked=True, timeout=0) as token_file:
182+
# we were able to lock the file ourselves so proceed to refresh the token
183+
# we have to do the refresh here as we must do it with the lock applied
184+
log.debug("Locked oauth token file. Refreshing the token now...")
185+
token_refreshed = con.refresh_token()
186+
if token_refreshed is False:
174187
raise RuntimeError("Token Refresh Operation not working")
175-
log.info("New oauth token fetched")
188+
189+
# we have refreshed the auth token ourselves to we must take care of
190+
# updating the header and save the token file
191+
con.update_session_auth_header()
192+
log.debug("New oauth token fetched. Saving the token data into the file")
193+
token_file.write(self.serialize())
176194
log.debug("Unlocked oauth token file")
177195
return None
178196
except LockException:
197+
# somebody else has adquired a lock so will be in the process of updating the token
179198
self.fs_wait = True
180-
log.warning(f"Oauth file locked. Sleeping for 2 seconds... retrying {i - 1} more times.")
199+
log.debug(f"Oauth file locked. Sleeping for 2 seconds... retrying {i - 1} more times.")
181200
time.sleep(2)
182201
log.debug("Waking up and rechecking token file for update from other instance...")
183-
self.token = self.load_token()
184-
# else:
185-
# log.info('Token was refreshed by another instance...')
186-
# self.fs_wait = False
187-
# return False
202+
# Assume the token has been already updated
203+
self.load_token()
204+
# Return False so the connection can update the token access from the backend into the session
205+
return False
188206

189207
# if we exit the loop, that means we were locked out of the file after
190208
# multiple retries give up and throw an error - something isn't right
191-
raise RuntimeError("Could not access locked token file after {self.max_tries}")
209+
raise RuntimeError(f"Could not access locked token file after {self.max_tries}")

0 commit comments

Comments
 (0)