3030- `OAuthServiceFailureError`: OAuthService exception
3131"""
3232
33+ import contextvars
3334import json
3435import os
3536import six
5051
5152
5253
54+ _OAUTH_AUTH_DOMAIN = contextvars .ContextVar ('OAUTH_AUTH_DOMAIN' )
55+ _OAUTH_EMAIL = contextvars .ContextVar ('OAUTH_EMAIL' )
56+ _OAUTH_USER_ID = contextvars .ContextVar ('OAUTH_USER_ID' )
57+ _OAUTH_CLIENT_ID = contextvars .ContextVar ('OAUTH_CLIENT_ID' )
58+ _OAUTH_IS_ADMIN = contextvars .ContextVar ('OAUTH_IS_ADMIN' )
59+ _OAUTH_ERROR_CODE = contextvars .ContextVar ('OAUTH_ERROR_CODE' )
60+ _OAUTH_ERROR_DETAIL = contextvars .ContextVar ('OAUTH_ERROR_DETAIL' )
61+ _OAUTH_LAST_SCOPE = contextvars .ContextVar ('OAUTH_LAST_SCOPE' )
62+ _OAUTH_AUTHORIZED_SCOPES = contextvars .ContextVar ('OAUTH_AUTHORIZED_SCOPES' )
63+
64+ _TESTBED_RESET_TOKENS = dict ()
65+
5366
5467class Error (Exception ):
5568 """Base error class for this module."""
@@ -117,7 +130,7 @@ def is_current_user_admin(_scope=None):
117130 """
118131
119132 _maybe_call_get_oauth_user (_scope )
120- return os . environ . get ('OAUTH_IS_ADMIN' , '0' ) == '1'
133+ return _OAUTH_IS_ADMIN . get (None )
121134
122135
123136def get_oauth_consumer_key ():
@@ -165,14 +178,14 @@ def get_authorized_scopes(scope):
165178
166179
167180def _maybe_call_get_oauth_user (scope ):
168- """Makes an GetOAuthUser RPC and stores the results in os.environ .
181+ """Makes an GetOAuthUser RPC and stores the results in context .
169182
170183 This method will only make the RPC if 'OAUTH_ERROR_CODE' has not already
171184 been set or 'OAUTH_LAST_SCOPE' is different to str(_scopes).
172185
173186 Args:
174- scope: The custom OAuth scope or an iterable of scopes at least one of
175- which is accepted.
187+ scope: The custom OAuth scope or an iterable of scopes at least one of which
188+ is accepted.
176189 """
177190
178191 if not scope :
@@ -181,8 +194,8 @@ def _maybe_call_get_oauth_user(scope):
181194 scope_str = scope
182195 else :
183196 scope_str = str (sorted (scope ))
184- if ('OAUTH_ERROR_CODE' not in os . environ or
185- os . environ . get ('OAUTH_LAST_SCOPE' , None ) != scope_str or
197+ if (_OAUTH_ERROR_CODE . get ( None ) is None or
198+ _OAUTH_LAST_SCOPE . get (None ) != scope_str or
186199 os .environ .get ('TESTONLY_OAUTH_SKIP_CACHE' )):
187200 req = user_service_pb2 .GetOAuthUserRequest ()
188201 if scope :
@@ -194,35 +207,39 @@ def _maybe_call_get_oauth_user(scope):
194207 resp = user_service_pb2 .GetOAuthUserResponse ()
195208 try :
196209 apiproxy_stub_map .MakeSyncCall ('user' , 'GetOAuthUser' , req , resp )
197- os .environ ['OAUTH_EMAIL' ] = resp .email
198- os .environ ['OAUTH_AUTH_DOMAIN' ] = resp .auth_domain
199- os .environ ['OAUTH_USER_ID' ] = resp .user_id
200- os .environ ['OAUTH_CLIENT_ID' ] = resp .client_id
201-
202- os .environ ['OAUTH_AUTHORIZED_SCOPES' ] = json .dumps (list (resp .scopes ))
203- if resp .is_admin :
204- os .environ ['OAUTH_IS_ADMIN' ] = '1'
205- else :
206- os .environ ['OAUTH_IS_ADMIN' ] = '0'
207- os .environ ['OAUTH_ERROR_CODE' ] = ''
210+ token = _OAUTH_EMAIL .set (resp .email )
211+ _TESTBED_RESET_TOKENS [_OAUTH_EMAIL ] = token
212+ token = _OAUTH_AUTH_DOMAIN .set (resp .auth_domain )
213+ _TESTBED_RESET_TOKENS [_OAUTH_AUTH_DOMAIN ] = token
214+ token = _OAUTH_USER_ID .set (resp .user_id )
215+ _TESTBED_RESET_TOKENS [_OAUTH_USER_ID ] = token
216+ token = _OAUTH_CLIENT_ID .set (resp .client_id )
217+ _TESTBED_RESET_TOKENS [_OAUTH_CLIENT_ID ] = token
218+ token = _OAUTH_AUTHORIZED_SCOPES .set (json .dumps (list (resp .scopes )))
219+ _TESTBED_RESET_TOKENS [_OAUTH_AUTHORIZED_SCOPES ] = token
220+ token = _OAUTH_IS_ADMIN .set (resp .is_admin )
221+ _TESTBED_RESET_TOKENS [_OAUTH_IS_ADMIN ] = token
222+ token = _OAUTH_ERROR_CODE .set ('' )
223+ _TESTBED_RESET_TOKENS [_OAUTH_ERROR_CODE ] = token
208224 except apiproxy_errors .ApplicationError as e :
209- os .environ ['OAUTH_ERROR_CODE' ] = str (e .application_error )
210- os .environ ['OAUTH_ERROR_DETAIL' ] = e .error_detail
211- os .environ ['OAUTH_LAST_SCOPE' ] = scope_str
225+ token = _OAUTH_ERROR_CODE .set (str (e .application_error ))
226+ _TESTBED_RESET_TOKENS [_OAUTH_ERROR_CODE ] = token
227+ token = _OAUTH_ERROR_DETAIL .set (e .error_detail )
228+ _TESTBED_RESET_TOKENS [_OAUTH_ERROR_DETAIL ] = token
229+ token = _OAUTH_LAST_SCOPE .set (scope_str )
230+ _TESTBED_RESET_TOKENS [_OAUTH_LAST_SCOPE ] = token
212231 _maybe_raise_exception ()
213232
214233
215234def _maybe_raise_exception ():
216- """Raises an error if one has been stored in os.environ .
235+ """Raises an error if one has been stored in context .
217236
218237 This method requires that 'OAUTH_ERROR_CODE' has already been set (an empty
219238 string indicates that there is no actual error).
220239 """
221- assert 'OAUTH_ERROR_CODE' in os .environ
222- error = os .environ ['OAUTH_ERROR_CODE' ]
240+ error = _OAUTH_ERROR_CODE .get ()
223241 if error :
224- assert 'OAUTH_ERROR_DETAIL' in os .environ
225- error_detail = os .environ ['OAUTH_ERROR_DETAIL' ]
242+ error_detail = _OAUTH_ERROR_DETAIL .get ()
226243 if error == str (user_service_pb2 .UserServiceError .NOT_ALLOWED ):
227244 raise NotAllowedError (error_detail )
228245 elif error == str (user_service_pb2 .UserServiceError .OAUTH_INVALID_REQUEST ):
@@ -236,42 +253,38 @@ def _maybe_raise_exception():
236253
237254
238255def _get_user_from_environ ():
239- """Returns a User based on values stored in os.environ .
256+ """Returns a User based on values stored in context .
240257
241258 This method requires that 'OAUTH_EMAIL', 'OAUTH_AUTH_DOMAIN', and
242259 'OAUTH_USER_ID' have already been set.
243260
244261 Returns:
245262 User
246263 """
247- assert 'OAUTH_EMAIL' in os .environ
248- assert 'OAUTH_AUTH_DOMAIN' in os .environ
249- assert 'OAUTH_USER_ID' in os .environ
250- return users .User (email = os .environ ['OAUTH_EMAIL' ],
251- _auth_domain = os .environ ['OAUTH_AUTH_DOMAIN' ],
252- _user_id = os .environ ['OAUTH_USER_ID' ])
264+ return users .User (
265+ email = _OAUTH_EMAIL .get (),
266+ _auth_domain = _OAUTH_AUTH_DOMAIN .get (),
267+ _user_id = _OAUTH_USER_ID .get ())
253268
254269
255270def _get_client_id_from_environ ():
256- """Returns Client ID based on values stored in os.environ .
271+ """Returns Client ID based on values stored in context .
257272
258273 This method requires that 'OAUTH_CLIENT_ID' has already been set.
259274
260275 Returns:
261276 string: the value of Client ID.
262277 """
263- assert 'OAUTH_CLIENT_ID' in os .environ
264- return os .environ ['OAUTH_CLIENT_ID' ]
278+ return _OAUTH_CLIENT_ID .get ()
265279
266280
267281def _get_authorized_scopes_from_environ ():
268- """Returns authorized scopes based on values stored in os.environ .
282+ """Returns authorized scopes based on values stored in context .
269283
270284 This method requires that 'OAUTH_AUTHORIZED_SCOPES' has already been set.
271285
272286 Returns:
273287 list: the list of OAuth scopes.
274288 """
275- assert 'OAUTH_AUTHORIZED_SCOPES' in os .environ
276289
277- return json .loads (os . environ [ 'OAUTH_AUTHORIZED_SCOPES' ] )
290+ return json .loads (_OAUTH_AUTHORIZED_SCOPES . get () )
0 commit comments