PYCO-69: Update error handling within retry logic

Changes
=======
* Handle wider range of httpx errors
* Allow BaseException to be raised if not a timeout or cancel scenario

Author: thejcfactor

acouchbase_analytics/protocol/_core/retries.py

@@ -19,7 +19,7 @@
 from functools import wraps
 from typing import TYPE_CHECKING, Any, Callable, Coroutine, Optional, Union
 
-from httpx import ConnectError, ConnectTimeout, ReadTimeout, WriteError, WriteTimeout
+from httpx import ConnectError, ConnectTimeout, CookieConflict, HTTPError, InvalidURL, ReadTimeout, StreamError
 
 from acouchbase_analytics.protocol._core.anyio_utils import sleep
 from couchbase_analytics.common.errors import AnalyticsError, InternalSDKError, TimeoutError
@@ -39,7 +39,7 @@ class AsyncRetryHandler:
 
     @staticmethod
     async def handle_httpx_retry(
-        ex: Union[ConnectError, ConnectTimeout, WriteError, WriteTimeout], ctx: AsyncRequestContext
+        ex: Union[ConnectError, ConnectTimeout], ctx: AsyncRequestContext
     ) -> Optional[Exception]:
         err_str = str(ex)
         if 'SSL:' in err_str:
@@ -107,7 +107,7 @@ async def wrapped_fn(self: AsyncHttpStreamingResponse) -> None:  # noqa: C901
                         continue
                     await self._request_context.shutdown(type(ex), ex, ex.__traceback__)
                     raise err from None
-                except (ConnectError, ConnectTimeout, WriteError, WriteTimeout) as ex:
+                except (ConnectError, ConnectTimeout) as ex:
                     err = await AsyncRetryHandler.handle_httpx_retry(ex, self._request_context)
                     if err is None:
                         continue
@@ -120,6 +120,12 @@ async def wrapped_fn(self: AsyncHttpStreamingResponse) -> None:  # noqa: C901
                     raise TimeoutError(
                         message='Request timed out.', context=str(self._request_context.error_context)
                     ) from None
+                except (CookieConflict, HTTPError, StreamError, InvalidURL) as ex:
+                    # these are not retriable errors, so we just shutdown the request context and raise the error
+                    await self._request_context.shutdown(type(ex), ex, ex.__traceback__)
+                    raise AnalyticsError(
+                        cause=ex, message=str(ex), context=str(self._request_context.error_context)
+                    ) from None
                 except AnalyticsError:
                     # if an AnalyticsError is raised, we have already shut down the request context
                     raise
@@ -142,9 +148,15 @@ async def wrapped_fn(self: AsyncHttpStreamingResponse) -> None:  # noqa: C901
                         raise CancelledError('Request was cancelled.') from None
                     if self._request_context.request_error is not None:
                         raise self._request_context.request_error from None
-                    raise InternalSDKError(
-                        cause=ex, message=str(ex), context=str(self._request_context.error_context)
-                    ) from None
+                    if isinstance(ex, Exception):
+                        # If the exception is an Exception, we raise it as an InternalSDKError as this is
+                        # an unexpected error in the SDK
+                        raise InternalSDKError(
+                            cause=ex, message=str(ex), context=str(self._request_context.error_context)
+                        ) from None
+                    # we should have handled CancelledError and TimeoutError above, so if we get here,
+                    # raise the BaseException as is (most likely a KeyboardInterrupt)
+                    raise ex
                 finally:
                     if not RequestState.is_okay(self._request_context.request_state):
                         await self.close()

couchbase_analytics/protocol/_core/retries.py

@@ -20,7 +20,7 @@
 from time import sleep
 from typing import TYPE_CHECKING, Callable, Optional, Union
 
-from httpx import ConnectError, ConnectTimeout, ReadTimeout, WriteError, WriteTimeout
+from httpx import ConnectError, ConnectTimeout, CookieConflict, HTTPError, InvalidURL, ReadTimeout, StreamError
 
 from couchbase_analytics.common.errors import AnalyticsError, InternalSDKError, TimeoutError
 from couchbase_analytics.common.logging import LogLevel
@@ -38,9 +38,7 @@ class RetryHandler:
     """
 
     @staticmethod
-    def handle_httpx_retry(
-        ex: Union[ConnectError, ConnectTimeout, WriteError, WriteTimeout], ctx: RequestContext
-    ) -> Optional[Exception]:
+    def handle_httpx_retry(ex: Union[ConnectError, ConnectTimeout], ctx: RequestContext) -> Optional[Exception]:
         err_str = str(ex)
         if 'SSL:' in err_str:
             message = 'TLS connection error occurred.'
@@ -105,7 +103,7 @@ def wrapped_fn(self: HttpStreamingResponse) -> None:  # noqa: C901
                         continue
                     self._request_context.shutdown(ex)
                     raise err from None
-                except (ConnectError, ConnectTimeout, WriteError, WriteTimeout) as ex:
+                except (ConnectError, ConnectTimeout) as ex:
                     err = RetryHandler.handle_httpx_retry(ex, self._request_context)
                     if err is None:
                         continue
@@ -118,6 +116,12 @@ def wrapped_fn(self: HttpStreamingResponse) -> None:  # noqa: C901
                     raise TimeoutError(
                         message='Request timed out.', context=str(self._request_context.error_context)
                     ) from None
+                except (CookieConflict, HTTPError, StreamError, InvalidURL) as ex:
+                    # these are not retriable errors, so we just shutdown the request context and raise the error
+                    self._request_context.shutdown(ex)
+                    raise AnalyticsError(
+                        cause=ex, message=str(ex), context=str(self._request_context.error_context)
+                    ) from None
                 except AnalyticsError:
                     # if an AnalyticsError is raised, we have already shut down the request context
                     raise
@@ -138,9 +142,15 @@ def wrapped_fn(self: HttpStreamingResponse) -> None:  # noqa: C901
                         ) from None
                     if self._request_context.cancelled:
                         raise CancelledError('Request was cancelled.') from None
-                    raise InternalSDKError(
-                        cause=ex, message=str(ex), context=str(self._request_context.error_context)
-                    ) from None
+                    if isinstance(ex, Exception):
+                        # If the exception is an Exception, we raise it as an InternalSDKError as this is
+                        # an unexpected error in the SDK
+                        raise InternalSDKError(
+                            cause=ex, message=str(ex), context=str(self._request_context.error_context)
+                        ) from None
+                    # we should have handled CancelledError and TimeoutError above, so if we get here,
+                    # raise the BaseException as is (most likely a KeyboardInterrupt)
+                    raise ex
                 finally:
                     if not RequestState.is_okay(self._request_context.request_state):
                         self.close()