-
-
Notifications
You must be signed in to change notification settings - Fork 893
refactor 401 retrying into both auth methods #1910
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
gmishkin
wants to merge
3
commits into
pycontribs:main
Choose a base branch
from
gmishkin:jira-session-retry
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change | ||||
---|---|---|---|---|---|---|
|
@@ -310,53 +310,35 @@ def _sort_and_quote_values(self, values): | |||||
return [quote(value, safe="~") for value in ordered_values] | ||||||
|
||||||
|
||||||
class JiraCookieAuth(AuthBase): | ||||||
"""Jira Cookie Authentication. | ||||||
|
||||||
Allows using cookie authentication as described by `jira api docs <https://developer.atlassian.com/server/jira/platform/cookie-based-authentication/>`_ | ||||||
""" | ||||||
class RetryingJiraAuth(AuthBase): | ||||||
"""Base class for Jira authentication handlers that need to retry requests on 401 responses.""" | ||||||
|
||||||
def __init__( | ||||||
self, session: ResilientSession, session_api_url: str, auth: tuple[str, str] | ||||||
): | ||||||
"""Cookie Based Authentication. | ||||||
|
||||||
Args: | ||||||
session (ResilientSession): The Session object to communicate with the API. | ||||||
session_api_url (str): The session api url to use. | ||||||
auth (Tuple[str, str]): The username, password tuple. | ||||||
""" | ||||||
def __init__(self, session: ResilientSession | None = None): | ||||||
self._session = session | ||||||
self._session_api_url = session_api_url # e.g ."/rest/auth/1/session" | ||||||
self.__auth = auth | ||||||
self._retry_counter_401 = 0 | ||||||
self._max_allowed_401_retries = 1 # 401 aren't recoverable with retries really | ||||||
|
||||||
def init_session(self): | ||||||
"""Auth mechanism specific code to re-initialize the Jira session.""" | ||||||
raise NotImplementedError() | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
let's make this an abstract base class and define this as an abstract method. To allow us to pick up errors in static analysis rather than runtime. |
||||||
|
||||||
@property | ||||||
def cookies(self): | ||||||
"""Return the cookies from the session.""" | ||||||
assert ( | ||||||
self._session is not None | ||||||
) # handle_401 should've caught this before attempting retry | ||||||
return self._session.cookies | ||||||
|
||||||
def _increment_401_retry_counter(self): | ||||||
self._retry_counter_401 += 1 | ||||||
|
||||||
def _reset_401_retry_counter(self): | ||||||
self._retry_counter_401 = 0 | ||||||
|
||||||
def __call__(self, request: requests.PreparedRequest): | ||||||
request.register_hook("response", self.handle_401) | ||||||
return request | ||||||
|
||||||
def init_session(self): | ||||||
"""Initialise the Session object's cookies, so we can use the session cookie. | ||||||
def _increment_401_retry_counter(self): | ||||||
self._retry_counter_401 += 1 | ||||||
|
||||||
Raises HTTPError if the post returns an erroring http response | ||||||
""" | ||||||
username, password = self.__auth | ||||||
authentication_data = {"username": username, "password": password} | ||||||
r = self._session.post( # this also goes through the handle_401() hook | ||||||
self._session_api_url, data=json.dumps(authentication_data) | ||||||
) | ||||||
r.raise_for_status() | ||||||
def _reset_401_retry_counter(self): | ||||||
self._retry_counter_401 = 0 | ||||||
|
||||||
def handle_401(self, response: requests.Response, **kwargs) -> requests.Response: | ||||||
"""Refresh cookies if the session cookie has expired. Then retry the request. | ||||||
|
@@ -367,43 +349,99 @@ def handle_401(self, response: requests.Response, **kwargs) -> requests.Response | |||||
Returns: | ||||||
requests.Response | ||||||
""" | ||||||
if ( | ||||||
is_retryable_401 = ( | ||||||
response.status_code == 401 | ||||||
and self._retry_counter_401 < self._max_allowed_401_retries | ||||||
): | ||||||
) | ||||||
|
||||||
if is_retryable_401 and self._session is not None: | ||||||
LOG.info("Trying to refresh the cookie auth session...") | ||||||
self._increment_401_retry_counter() | ||||||
self.init_session() | ||||||
response = self.process_original_request(response.request.copy()) | ||||||
elif is_retryable_401 and self._session is None: | ||||||
LOG.warning("No session was passed to constructor, can't refresh cookies.") | ||||||
|
||||||
self._reset_401_retry_counter() | ||||||
return response | ||||||
|
||||||
def process_original_request(self, original_request: requests.PreparedRequest): | ||||||
self.update_cookies(original_request) | ||||||
return self.send_request(original_request) | ||||||
|
||||||
def update_cookies(self, original_request: requests.PreparedRequest): | ||||||
"""Auth mechanism specific cookie handling prior to retrying.""" | ||||||
raise NotImplementedError() | ||||||
|
||||||
def send_request(self, request: requests.PreparedRequest): | ||||||
if self._session is not None: | ||||||
request.prepare_cookies(self.cookies) # post-update re-prepare | ||||||
return self._session.send(request) | ||||||
|
||||||
|
||||||
class JiraCookieAuth(RetryingJiraAuth): | ||||||
"""Jira Cookie Authentication. | ||||||
|
||||||
Allows using cookie authentication as described by `jira api docs <https://developer.atlassian.com/server/jira/platform/cookie-based-authentication/>`_ | ||||||
""" | ||||||
|
||||||
def __init__( | ||||||
self, session: ResilientSession, session_api_url: str, auth: tuple[str, str] | ||||||
): | ||||||
"""Cookie Based Authentication. | ||||||
|
||||||
Args: | ||||||
session (ResilientSession): The Session object to communicate with the API. | ||||||
session_api_url (str): The session api url to use. | ||||||
auth (Tuple[str, str]): The username, password tuple. | ||||||
""" | ||||||
super().__init__(session) | ||||||
self._session_api_url = session_api_url # e.g ."/rest/auth/1/session" | ||||||
self.__auth = auth | ||||||
|
||||||
def init_session(self): | ||||||
"""Initialise the Session object's cookies, so we can use the session cookie. | ||||||
|
||||||
Raises HTTPError if the post returns an erroring http response | ||||||
""" | ||||||
assert ( | ||||||
self._session is not None | ||||||
) # Constructor for this subclass always takes a session | ||||||
username, password = self.__auth | ||||||
authentication_data = {"username": username, "password": password} | ||||||
r = self._session.post( # this also goes through the handle_401() hook | ||||||
self._session_api_url, data=json.dumps(authentication_data) | ||||||
) | ||||||
r.raise_for_status() | ||||||
|
||||||
def update_cookies(self, original_request: requests.PreparedRequest): | ||||||
# Cookie header needs first to be deleted for the header to be updated using the | ||||||
# prepare_cookies method. See request.PrepareRequest.prepare_cookies | ||||||
if "Cookie" in original_request.headers: | ||||||
del original_request.headers["Cookie"] | ||||||
original_request.prepare_cookies(self.cookies) | ||||||
|
||||||
def send_request(self, request: requests.PreparedRequest): | ||||||
return self._session.send(request) | ||||||
|
||||||
|
||||||
class TokenAuth(AuthBase): | ||||||
class TokenAuth(RetryingJiraAuth): | ||||||
"""Bearer Token Authentication.""" | ||||||
|
||||||
def __init__(self, token: str): | ||||||
def __init__(self, token: str, session: ResilientSession | None = None): | ||||||
super().__init__(session) | ||||||
# setup any auth-related data here | ||||||
self._token = token | ||||||
|
||||||
def __call__(self, r: requests.PreparedRequest): | ||||||
# modify and return the request | ||||||
r.headers["authorization"] = f"Bearer {self._token}" | ||||||
return r | ||||||
return super().__call__(r) | ||||||
|
||||||
def init_session(self): | ||||||
pass # token should still work, only thing needed is to clear session cookies which happens next | ||||||
|
||||||
def update_cookies(self, _): | ||||||
assert ( | ||||||
self._session is not None | ||||||
) # handle_401 on the superclass should've caught this before attempting retry | ||||||
self._session.cookies.clear_session_cookies() | ||||||
|
||||||
|
||||||
class JIRA: | ||||||
|
@@ -4499,7 +4537,7 @@ def _create_token_session(self, token_auth: str): | |||||
|
||||||
Header structure: "authorization": "Bearer <token_auth>". | ||||||
""" | ||||||
self._session.auth = TokenAuth(token_auth) | ||||||
self._session.auth = TokenAuth(token_auth, session=self._session) | ||||||
|
||||||
def _set_avatar(self, params, url, avatar): | ||||||
data = {"id": avatar} | ||||||
|
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Let's remove the default of
None
here, and then remove as many if not all the plain asserts.If there do end up being situations where the session is None when passing to the constructor of the subclasses then let us move to a subclass
__init__
method override, and rather than asserts, either use warnings, or if it should really fail, let's use an Exception with suitable message.I want to avoid
asserts
in the production codebase.