CM-46371 - Add retry behavior for HTTP requests

Author: MarshalX

cycode/cli/apps/auth/auth_manager.py

@@ -2,8 +2,6 @@
 import webbrowser
 from typing import TYPE_CHECKING, Tuple
 
-from requests import Request
-
 from cycode.cli.exceptions.custom_exceptions import AuthProcessError
 from cycode.cli.user_settings.configuration_manager import ConfigurationManager
 from cycode.cli.user_settings.credentials_manager import CredentialsManager
@@ -53,7 +51,7 @@ def start_session(self, code_challenge: str) -> str:
         return auth_session.session_id
 
     def redirect_to_login_page(self, code_challenge: str, session_id: str) -> None:
-        login_url = self._build_login_url(code_challenge, session_id)
+        login_url = self.auth_client.build_login_url(code_challenge, session_id)
         webbrowser.open(login_url)
 
     def get_api_token(self, session_id: str, code_verifier: str) -> 'ApiToken':
@@ -75,19 +73,11 @@ def get_api_token_polling(self, session_id: str, code_verifier: str) -> 'ApiToke
                 raise AuthProcessError('Error while obtaining API token')
             time.sleep(self.POLLING_WAIT_INTERVAL_IN_SECONDS)
 
-        raise AuthProcessError('session expired')
+        raise AuthProcessError('Timeout while obtaining API token (session expired)')
 
     def save_api_token(self, api_token: 'ApiToken') -> None:
         self.credentials_manager.update_credentials(api_token.client_id, api_token.secret)
 
-    def _build_login_url(self, code_challenge: str, session_id: str) -> str:
-        app_url = self.configuration_manager.get_cycode_app_url()
-        login_url = f'{app_url}/account/sign-in'
-        query_params = {'source': 'cycode_cli', 'code_challenge': code_challenge, 'session_id': session_id}
-        # TODO(MarshalX). Use auth_client instead and don't depend on "requests" lib here
-        request = Request(url=login_url, params=query_params)
-        return request.prepare().url
-
     def _generate_pkce_code_pair(self) -> Tuple[str, str]:
         code_verifier = generate_random_string(self.CODE_VERIFIER_LENGTH)
         code_challenge = hash_string_to_sha256(code_verifier)

cycode/cyclient/auth_client.py

@@ -1,9 +1,9 @@
 from typing import Optional
 
-from requests import Response
+from requests import Request, Response
 
 from cycode.cli.exceptions.custom_exceptions import HttpUnauthorizedError, RequestHttpError
-from cycode.cyclient import models
+from cycode.cyclient import config, models
 from cycode.cyclient.cycode_client import CycodeClient
 
 
@@ -13,6 +13,11 @@ class AuthClient:
     def __init__(self) -> None:
         self.cycode_client = CycodeClient()
 
+    @staticmethod
+    def build_login_url(code_challenge: str, session_id: str) -> str:
+        query_params = {'source': 'cycode_cli', 'code_challenge': code_challenge, 'session_id': session_id}
+        return Request(url=f'{config.cycode_app_url}/account/sign-in', params=query_params).prepare().url
+
     def start_session(self, code_challenge: str) -> models.AuthenticationSession:
         path = f'{self.AUTH_CONTROLLER_PATH}/start'
         body = {'code_challenge': code_challenge}

cycode/cyclient/config.py

@@ -14,6 +14,14 @@
     cycode_api_url = consts.DEFAULT_CYCODE_API_URL
 
 
+cycode_app_url = configuration_manager.get_cycode_app_url()
+if not is_valid_url(cycode_app_url):
+    logger.warning(
+        'Invalid Cycode APP URL: %s, using default value (%s)', cycode_app_url, consts.DEFAULT_CYCODE_APP_URL
+    )
+    cycode_app_url = consts.DEFAULT_CYCODE_APP_URL
+
+
 def _is_on_premise_installation(cycode_domain: str) -> bool:
     return not cycode_api_url.endswith(cycode_domain)
 

cycode/cyclient/cycode_client_base.py

@@ -1,11 +1,12 @@
 import os
 import platform
 import ssl
-from typing import Callable, ClassVar, Dict, Optional
+from typing import TYPE_CHECKING, Callable, ClassVar, Dict, Optional
 
 import requests
 from requests import Response, exceptions
 from requests.adapters import HTTPAdapter
+from tenacity import retry, retry_if_exception, stop_after_attempt, wait_random_exponential
 
 from cycode.cli.exceptions.custom_exceptions import (
     HttpUnauthorizedError,
@@ -19,6 +20,9 @@
 from cycode.cyclient.headers import get_cli_user_agent, get_correlation_id
 from cycode.cyclient.logger import logger
 
+if TYPE_CHECKING:
+    from tenacity import RetryCallState
+
 
 class SystemStorageSslContext(HTTPAdapter):
     def init_poolmanager(self, *args, **kwargs) -> None:
@@ -45,6 +49,47 @@ def _get_request_function() -> Callable:
     return session.request
 
 
+_REQUEST_ERRORS_TO_RETRY = (
+    RequestTimeout,
+    RequestConnectionError,
+    exceptions.ChunkedEncodingError,
+    exceptions.ContentDecodingError,
+)
+_RETRY_MAX_ATTEMPTS = 3
+_RETRY_STOP_STRATEGY = stop_after_attempt(_RETRY_MAX_ATTEMPTS)
+_RETRY_WAIT_STRATEGY = wait_random_exponential(multiplier=1, min=2, max=10)
+
+
+def _retry_before_sleep(retry_state: 'RetryCallState') -> None:
+    exception_name = 'None'
+    if retry_state.outcome.failed:
+        exception = retry_state.outcome.exception()
+        exception_name = f'{exception.__class__.__name__}'
+
+    logger.debug(
+        'Retrying request after error: %s. Attempt %s of %s. Upcoming sleep: %s',
+        exception_name,
+        retry_state.attempt_number,
+        _RETRY_MAX_ATTEMPTS,
+        retry_state.upcoming_sleep,
+    )
+
+
+def _should_retry_exception(exception: BaseException) -> bool:
+    if 'PYTEST_CURRENT_TEST' in os.environ:
+        # We are running under pytest, don't retry
+        return False
+
+    # Don't retry client errors (400, 401, etc.)
+    if isinstance(exception, RequestHttpError):
+        return not exception.status_code < 500
+
+    is_request_error = isinstance(exception, _REQUEST_ERRORS_TO_RETRY)
+    is_server_error = isinstance(exception, RequestHttpError) and exception.status_code >= 500
+
+    return is_request_error or is_server_error
+
+
 class CycodeClientBase:
     MANDATORY_HEADERS: ClassVar[Dict[str, str]] = {
         'User-Agent': get_cli_user_agent(),
@@ -72,6 +117,13 @@ def put(self, url_path: str, body: Optional[dict] = None, headers: Optional[dict
     def get(self, url_path: str, headers: Optional[dict] = None, **kwargs) -> Response:
         return self._execute(method='get', endpoint=url_path, headers=headers, **kwargs)
 
+    @retry(
+        retry=retry_if_exception(_should_retry_exception),
+        stop=_RETRY_STOP_STRATEGY,
+        wait=_RETRY_WAIT_STRATEGY,
+        reraise=True,
+        before_sleep=_retry_before_sleep,
+    )
     def _execute(
         self,
         method: str,

poetry.lock

@@ -41,6 +41,7 @@ pyjwt = ">=2.8.0,<3.0"
 rich = ">=13.9.4, <14"
 patch-ng = "1.18.1"
 typer = "^0.15.2"
+tenacity = ">=9.0.0,<9.1.0"
 
 [tool.poetry.group.test.dependencies]
 mock = ">=4.0.3,<4.1.0"