Move retry logic to client module

Author: mdesmet

tests/unit/test_client.py

@@ -28,7 +28,8 @@
     SERVER_ADDRESS
 from trino import constants
 from trino.auth import KerberosAuthentication, _OAuth2TokenBearer
-from trino.client import TrinoQuery, TrinoRequest, TrinoResult, ClientSession
+from trino.client import TrinoQuery, TrinoRequest, TrinoResult, ClientSession, _DelayExponential, _retry_with, \
+    _RetryWithExponentialBackoff
 
 
 @mock.patch("trino.client.TrinoRequest.http")
@@ -947,3 +948,57 @@ def json(self):
 
         # Validate the result is an instance of TrinoResult
         assert isinstance(result, TrinoResult)
+
+
+def test_delay_exponential_without_jitter():
+    max_delay = 1200.0
+    get_delay = _DelayExponential(base=5, jitter=False, max_delay=max_delay)
+    results = [
+        10.0,
+        20.0,
+        40.0,
+        80.0,
+        160.0,
+        320.0,
+        640.0,
+        max_delay,  # rather than 1280.0
+        max_delay,  # rather than 2560.0
+    ]
+    for i, result in enumerate(results, start=1):
+        assert get_delay(i) == result
+
+
+def test_delay_exponential_with_jitter():
+    max_delay = 120.0
+    get_delay = _DelayExponential(base=10, jitter=False, max_delay=max_delay)
+    for i in range(10):
+        assert get_delay(i) <= max_delay
+
+
+class SomeException(Exception):
+    pass
+
+
+def test_retry_with():
+    max_attempts = 3
+    with_retry = _retry_with(
+        handle_retry=_RetryWithExponentialBackoff(),
+        handled_exceptions=[SomeException],
+        conditions={},
+        max_attempts=max_attempts,
+    )
+
+    class FailerUntil(object):
+        def __init__(self, until=1):
+            self.attempt = 0
+            self._until = until
+
+        def __call__(self):
+            self.attempt += 1
+            if self.attempt > self._until:
+                return
+            raise SomeException(self.attempt)
+
+    with_retry(FailerUntil(2).__call__)()
+    with pytest.raises(SomeException):
+        with_retry(FailerUntil(3).__call__)()

tests/unit/test_exceptions.py

@@ -34,9 +34,12 @@
 """
 
 import copy
+import functools
 import os
+import random
 import re
 import threading
+import time
 import urllib.parse
 from datetime import datetime, timedelta, timezone
 from decimal import Decimal
@@ -227,6 +230,34 @@ def __repr__(self):
         )
 
 
+class _DelayExponential(object):
+    def __init__(
+            self, base=0.1, exponent=2, jitter=True, max_delay=2 * 3600  # 100ms  # 2 hours
+    ):
+        self._base = base
+        self._exponent = exponent
+        self._jitter = jitter
+        self._max_delay = max_delay
+
+    def __call__(self, attempt):
+        delay = float(self._base) * (self._exponent ** attempt)
+        if self._jitter:
+            delay *= random.random()
+        delay = min(float(self._max_delay), delay)
+        return delay
+
+
+class _RetryWithExponentialBackoff(object):
+    def __init__(
+            self, base=0.1, exponent=2, jitter=True, max_delay=2 * 3600  # 100ms  # 2 hours
+    ):
+        self._get_delay = _DelayExponential(base, exponent, jitter, max_delay)
+
+    def retry(self, func, args, kwargs, err, attempt):
+        delay = self._get_delay(attempt)
+        time.sleep(delay)
+
+
 class TrinoRequest(object):
     """
     Manage the HTTP requests of a Trino query.
@@ -286,7 +317,7 @@ def __init__(
         redirect_handler: Any = None,
         max_attempts: int = MAX_ATTEMPTS,
         request_timeout: Union[float, Tuple[float, float]] = constants.DEFAULT_REQUEST_TIMEOUT,
-        handle_retry=exceptions.RetryWithExponentialBackoff(),
+        handle_retry=_RetryWithExponentialBackoff(),
         verify: bool = True,
     ) -> None:
         self._client_session = client_session
@@ -383,9 +414,9 @@ def max_attempts(self, value) -> None:
             self._delete = self._http_session.delete
             return
 
-        with_retry = exceptions.retry_with(
+        with_retry = _retry_with(
             self._handle_retry,
-            exceptions=self._exceptions,
+            handled_exceptions=self._exceptions,
             conditions=(
                 # need retry when there is no exception but the status code is 502, 503, or 504
                 lambda response: getattr(response, "status_code", None)
@@ -779,3 +810,32 @@ def cancelled(self) -> bool:
     @property
     def response_headers(self):
         return self._response_headers
+
+
+def _retry_with(handle_retry, handled_exceptions, conditions, max_attempts):
+    def wrapper(func):
+        @functools.wraps(func)
+        def decorated(*args, **kwargs):
+            error = None
+            result = None
+            for attempt in range(1, max_attempts + 1):
+                try:
+                    result = func(*args, **kwargs)
+                    if any(guard(result) for guard in conditions):
+                        handle_retry.retry(func, args, kwargs, None, attempt)
+                        continue
+                    return result
+                except Exception as err:
+                    error = err
+                    if any(isinstance(err, exc) for exc in handled_exceptions):
+                        handle_retry.retry(func, args, kwargs, err, attempt)
+                        continue
+                    break
+            logger.info("failed after %s attempts", attempt)
+            if error is not None:
+                raise error
+            return result
+
+        return decorated
+
+    return wrapper

trino/client.py

@@ -16,10 +16,6 @@
 """
 
 
-import functools
-import random
-import time
-
 import trino.logging
 
 logger = trino.logging.get_logger(__name__)
@@ -120,63 +116,6 @@ class TrinoUserError(TrinoQueryError):
     pass
 
 
-def retry_with(handle_retry, exceptions, conditions, max_attempts):
-    def wrapper(func):
-        @functools.wraps(func)
-        def decorated(*args, **kwargs):
-            error = None
-            result = None
-            for attempt in range(1, max_attempts + 1):
-                try:
-                    result = func(*args, **kwargs)
-                    if any(guard(result) for guard in conditions):
-                        handle_retry.retry(func, args, kwargs, None, attempt)
-                        continue
-                    return result
-                except Exception as err:
-                    error = err
-                    if any(isinstance(err, exc) for exc in exceptions):
-                        handle_retry.retry(func, args, kwargs, err, attempt)
-                        continue
-                    break
-            logger.info("failed after %s attempts", attempt)
-            if error is not None:
-                raise error
-            return result
-
-        return decorated
-
-    return wrapper
-
-
-class DelayExponential(object):
-    def __init__(
-        self, base=0.1, exponent=2, jitter=True, max_delay=2 * 3600  # 100ms  # 2 hours
-    ):
-        self._base = base
-        self._exponent = exponent
-        self._jitter = jitter
-        self._max_delay = max_delay
-
-    def __call__(self, attempt):
-        delay = float(self._base) * (self._exponent ** attempt)
-        if self._jitter:
-            delay *= random.random()
-        delay = min(float(self._max_delay), delay)
-        return delay
-
-
-class RetryWithExponentialBackoff(object):
-    def __init__(
-        self, base=0.1, exponent=2, jitter=True, max_delay=2 * 3600  # 100ms  # 2 hours
-    ):
-        self._get_delay = DelayExponential(base, exponent, jitter, max_delay)
-
-    def retry(self, func, args, kwargs, err, attempt):
-        delay = self._get_delay(attempt)
-        time.sleep(delay)
-
-
 # PEP 249
 class Error(Exception):
     pass