Repository: jina-ai/serve
Pull request: fix: fix client error 4391
fixing https://github.com/jina-ai/jina/issues/4391


New behavior:
if an exception is raised then :

- [x] if there is an on_error callback then the callback is called and the original exception is passed to the on_error
- [x] Otherwise we just throw an ClientError which wrap the actual exception
- [x] docs


Problem statement:
`on_error` broken for Client with unreachable server
```python
from jina import Client, DocumentArray


def on_error(response):
    print('on_error called')
    print(response.to_dict())


for protocol in ['http', 'grpc', 'websocket']:
    print(f'\n\n\n\nprotocol: {protocol}')
    try:
        Client(host='0.0.0.0', protocol=protocol, port=12345).post(
            '/blah',
            inputs=DocumentArray.empty(10),
            on_error=on_error,
        )
    except Exception as e:
        print(f'protocol {protocol} raised \n{e}')
```

**HTTP** 
- No exceptions raised. 
- `on_error` is not invoked.

**gRPC**
- Exception raised.
- `on_error` is not invoked.

**Websocket**
- No exceptions raised. 
- `on_error` is not invoked.

----

```
protocol: http
     HTTPClient@1909863[E]:Error while fetching response from HTTP server ClientConnectorError(ConnectionKey(host='0.0.0.0', port=12345, is_ssl=False, ssl=None, proxy=None, proxy_auth=None, proxy_headers_hash=None), ConnectionRefusedError(111, 'Connection refused'))




protocol: grpc
     GRPCClient@1909863[E]:gRPC error: StatusCode.UNAVAILABLE failed to connect to all addresses
the ongoing request is terminated as the server is not available or closed already
protocol grpc raised 
<AioRpcError of RPC that terminated with:
        status = StatusCode.UNAVAILABLE
        details = "failed to connect to all addresses"
        debug_error_string = "{"created":"@1645593433.556564109","description":"Failed to pick subchannel","file":"src/core/ext/filters/client_channel/client_channel.cc","file_line":3134,"referenced_errors":[{"created":"@1645593433.556563329","description":"failed to connect to all addresses","file":"src/core/lib/transport/error_utils.cc","file_line":163,"grpc_status":14}]}"
>




protocol: websocket
WebSocketClient@1909863[E]:Error while streaming response from websocket server ClientConnectorError(ConnectionKey(host='0.0.0.0', port=12345, is_ssl=False, ssl=None, proxy=None, proxy_auth=None, proxy_headers_hash=None), ConnectionRefusedError(111, 'Connection refused'))
Unclosed client session
client_session: <aiohttp.client.ClientSession object at 0x7f9159c2e710>
```

Hints from issue discussion:

Commit message: fix: on_error depcration warning called only onced

Patch to review:
diff --git a/jina/clients/base/grpc.py b/jina/clients/base/grpc.py
index b81ad4e310bef..134af1f59e468 100644
--- a/jina/clients/base/grpc.py
+++ b/jina/clients/base/grpc.py
@@ -5,8 +5,8 @@
 import grpc
 
 from jina.clients.base import BaseClient
-from jina.clients.helper import callback_exec
-from jina.excepts import BadClient, BadClientInput
+from jina.clients.helper import callback_exec, callback_exec_on_error
+from jina.excepts import BadClient, BadClientInput, BaseJinaException
 from jina.logging.profile import ProgressBar
 from jina.proto import jina_pb2_grpc
 from jina.serve.networking import GrpcConnectionPool
@@ -68,22 +68,43 @@ async def _get_results(
             my_code = rpc_ex.code()
             my_details = rpc_ex.details()
             msg = f'gRPC error: {my_code} {my_details}'
-            if my_code == grpc.StatusCode.UNAVAILABLE:
-                self.logger.error(
-                    f'{msg}\nthe ongoing request is terminated as the server is not available or closed already'
-                )
-                raise rpc_ex
-            elif my_code == grpc.StatusCode.INTERNAL:
-                self.logger.error(f'{msg}\ninternal error on the server side')
-                raise rpc_ex
-            elif (
-                my_code == grpc.StatusCode.UNKNOWN
-                and 'asyncio.exceptions.TimeoutError' in my_details
-            ):
-                raise BadClientInput(
-                    f'{msg}\n'
-                    'often the case is that you define/send a bad input iterator to jina, '
-                    'please double check your input iterator'
-                ) from rpc_ex
-            else:
-                raise BadClient(msg) from rpc_ex
+
+            try:
+                if my_code == grpc.StatusCode.UNAVAILABLE:
+                    self.logger.error(
+                        f'{msg}\nthe ongoing request is terminated as the server is not available or closed already'
+                    )
+                    raise rpc_ex
+                elif my_code == grpc.StatusCode.INTERNAL:
+                    self.logger.error(f'{msg}\ninternal error on the server side')
+                    raise rpc_ex
+                elif (
+                    my_code == grpc.StatusCode.UNKNOWN
+                    and 'asyncio.exceptions.TimeoutError' in my_details
+                ):
+                    raise BadClientInput(
+                        f'{msg}\n'
+                        'often the case is that you define/send a bad input iterator to jina, '
+                        'please double check your input iterator'
+                    ) from rpc_ex
+                else:
+                    raise BadClient(msg) from rpc_ex
+
+            except (
+                grpc.aio._call.AioRpcError,
+                BaseJinaException,
+            ) as e:  # depending on if there are callbacks we catch or not the exception
+                if on_error or on_always:
+                    if on_error:
+                        callback_exec_on_error(on_error, e, self.logger)
+                    if on_always:
+                        callback_exec(
+                            response=None,
+                            on_error=None,
+                            on_done=None,
+                            on_always=on_always,
+                            continue_on_error=self.continue_on_error,
+                            logger=self.logger,
+                        )
+                else:
+                    raise e
diff --git a/jina/clients/base/http.py b/jina/clients/base/http.py
index 28e9b4eda81af..bcadd1c63fe07 100644
--- a/jina/clients/base/http.py
+++ b/jina/clients/base/http.py
@@ -4,7 +4,7 @@
 
 from jina.clients.base.helper import HTTPClientlet
 from jina.clients.base import BaseClient
-from jina.clients.helper import callback_exec
+from jina.clients.helper import callback_exec, callback_exec_on_error
 from jina.excepts import BadClient
 from jina.importer import ImportExtensions
 from jina.logging.profile import ProgressBar
@@ -108,3 +108,18 @@ def _result_handler(result):
                 self.logger.error(
                     f'Error while fetching response from HTTP server {e!r}'
                 )
+
+                if on_error or on_always:
+                    if on_error:
+                        callback_exec_on_error(on_error, e, self.logger)
+                    if on_always:
+                        callback_exec(
+                            response=None,
+                            on_error=None,
+                            on_done=None,
+                            on_always=on_always,
+                            continue_on_error=self.continue_on_error,
+                            logger=self.logger,
+                        )
+                else:
+                    raise e
diff --git a/jina/clients/base/websocket.py b/jina/clients/base/websocket.py
index dbcf920d83755..1f348949f6939 100644
--- a/jina/clients/base/websocket.py
+++ b/jina/clients/base/websocket.py
@@ -6,7 +6,8 @@
 
 from jina.clients.base.helper import WebsocketClientlet
 from jina.clients.base import BaseClient
-from jina.clients.helper import callback_exec
+from jina.clients.helper import callback_exec, callback_exec_on_error
+from jina.excepts import BadClient
 from jina.importer import ImportExtensions
 from jina.logging.profile import ProgressBar
 from jina.serve.stream import RequestStreamer
@@ -134,3 +135,18 @@ def _request_handler(request: 'Request') -> 'asyncio.Future':
                 self.logger.error(
                     f'Error while streaming response from websocket server {e!r}'
                 )
+
+                if on_error or on_always:
+                    if on_error:
+                        callback_exec_on_error(on_error, e, self.logger)
+                    if on_always:
+                        callback_exec(
+                            response=None,
+                            on_error=None,
+                            on_done=None,
+                            on_always=on_always,
+                            continue_on_error=self.continue_on_error,
+                            logger=self.logger,
+                        )
+                else:
+                    raise e
diff --git a/jina/clients/helper.py b/jina/clients/helper.py
index c29f1f3650d11..20844fe09477c 100644
--- a/jina/clients/helper.py
+++ b/jina/clients/helper.py
@@ -1,7 +1,9 @@
 """Helper functions for clients in Jina."""
 
 from functools import wraps
-from typing import Callable
+from inspect import signature
+from typing import Callable, Optional
+import warnings
 
 from jina.excepts import BadClientCallback
 from jina.logging.logger import JinaLogger
@@ -81,3 +83,23 @@ def callback_exec(
         _safe_callback(on_done, continue_on_error, logger)(response)
     if on_always:
         _safe_callback(on_always, continue_on_error, logger)(response)
+
+
+def callback_exec_on_error(
+    on_error: Callable,
+    exception: Exception,
+    logger: JinaLogger,
+    response: Optional = None,
+) -> None:
+    """Execute the on_error callback with the response, Use when an error outside the response status was thrown.
+    :param on_error: the on_error callback
+    :param exception: the exception with was thrown and led to the call of on_error
+    :param logger: a logger instance
+    :param response: the response
+    """
+
+    @wraps(on_error)
+    def on_error_wrap(resp):
+        on_error(resp, exception)
+
+    _safe_callback(on_error_wrap, False, logger)(response)
diff --git a/jina/clients/mixin.py b/jina/clients/mixin.py
index f800573edb5d1..d7b3898f83ef9 100644
--- a/jina/clients/mixin.py
+++ b/jina/clients/mixin.py
@@ -1,6 +1,7 @@
-from functools import partialmethod
+from functools import partialmethod, wraps
 from typing import Optional, Dict, List, AsyncGenerator, TYPE_CHECKING, Union
 import warnings
+from inspect import signature
 
 from jina.helper import run_async
 
@@ -90,6 +91,7 @@ async def _get_results(*args, **kwargs):
             return_results = True
 
         parameters = _include_results_field_in_param(parameters)
+        on_error = _wrap_on_error(on_error) if on_error is not None else on_error
 
         return run_async(
             _get_results,
@@ -148,6 +150,7 @@ async def post(
         c.continue_on_error = continue_on_error
 
         parameters = _include_results_field_in_param(parameters)
+        on_error = _wrap_on_error(on_error) if on_error is not None else on_error
 
         async for r in c._get_results(
             inputs=inputs,
@@ -167,3 +170,23 @@ async def post(
     search = partialmethod(post, '/search')
     update = partialmethod(post, '/update')
     delete = partialmethod(post, '/delete')
+
+
+def _wrap_on_error(on_error):
+
+    num_args = len(signature(on_error).parameters)
+    if num_args == 1:
+        warnings.warn(
+            "on_error callback taking only the response parameters is deprecated. Please add one parameter "
+            "to handle additional optional Exception as well",
+            DeprecationWarning,
+        )
+
+        @wraps(on_error)
+        def _fn(resp, exception):  # just skip the exception
+            return on_error(resp)
+
+        return _fn
+
+    else:
+        return on_error
diff --git a/jina/excepts.py b/jina/excepts.py
index d2ee7be7912c5..ec69e1a8fe9ca 100644
--- a/jina/excepts.py
+++ b/jina/excepts.py
@@ -1,93 +1,97 @@
 """This modules defines all kinds of exceptions raised in Jina."""
 
 
-class BaseJinaExeception:
+class BaseJinaException(BaseException):
     """A base class for all exceptions raised by Jina"""
 
 
-class ExecutorFailToLoad(SystemError, BaseJinaExeception):
+class ExecutorFailToLoad(SystemError, BaseJinaException):
     """When the executor can not be loaded in pod/deployment."""
 
 
-class RuntimeFailToStart(SystemError, BaseJinaExeception):
+class RuntimeFailToStart(SystemError, BaseJinaException):
     """When pod/deployment is failed to started."""
 
 
-class ScalingFails(SystemError, BaseJinaExeception):
+class ScalingFails(SystemError, BaseJinaException):
     """When scaling is unsuccessful for an Executor."""
 
 
-class RuntimeTerminated(KeyboardInterrupt, BaseJinaExeception):
+class RuntimeTerminated(KeyboardInterrupt, BaseJinaException):
     """The event loop of BasePod ends."""
 
 
-class FlowTopologyError(Exception, BaseJinaExeception):
+class FlowTopologyError(Exception, BaseJinaException):
     """Flow exception when the topology is ambiguous."""
 
 
-class FlowMissingDeploymentError(Exception, BaseJinaExeception):
+class FlowMissingDeploymentError(Exception, BaseJinaException):
     """Flow exception when a deployment can not be found in the flow."""
 
 
-class FlowBuildLevelError(Exception, BaseJinaExeception):
+class FlowBuildLevelError(Exception, BaseJinaException):
     """Flow exception when required build level is higher than the current build level."""
 
 
-class BadConfigSource(FileNotFoundError, BaseJinaExeception):
+class BadConfigSource(FileNotFoundError, BaseJinaException):
     """The yaml config file is bad, not loadable or not exist."""
 
 
-class BadClient(Exception, BaseJinaExeception):
+class BadClient(Exception, BaseJinaException):
     """A wrongly defined client, can not communicate with jina server correctly."""
 
 
-class BadClientCallback(BadClient, BaseJinaExeception):
+class BadClientCallback(BadClient, BaseJinaException):
     """Error in the callback function on the client side."""
 
 
-class BadClientInput(BadClient, BaseJinaExeception):
+class BadClientInput(BadClient, BaseJinaException):
     """Error in the request generator function on the client side."""
 
 
-class BadRequestType(TypeError, BaseJinaExeception):
+class BadRequestType(TypeError, BaseJinaException):
     """Exception when can not construct a request object from given data."""
 
 
-class BadImageNameError(Exception, BaseJinaExeception):
+class BadClientResponse(Exception, BaseJinaException):
+    """Exception when can not construct a request object from given data."""
+
+
+class BadImageNameError(Exception, BaseJinaException):
     """Exception when an image name can not be found either local & remote"""
 
 
-class BadYAMLVersion(Exception, BaseJinaExeception):
+class BadYAMLVersion(Exception, BaseJinaException):
     """Exception when YAML config specifies a wrong version number."""
 
 
-class DaemonConnectivityError(Exception, BaseJinaExeception):
+class DaemonConnectivityError(Exception, BaseJinaException):
     """Exception to raise when jina daemon is not reachable."""
 
 
-class DaemonWorkspaceCreationFailed(Exception, BaseJinaExeception):
+class DaemonWorkspaceCreationFailed(Exception, BaseJinaException):
     """Exception to raise when jina daemon is not connectable."""
 
 
-class DaemonPodCreationFailed(Exception, BaseJinaExeception):
+class DaemonPodCreationFailed(Exception, BaseJinaException):
     """Exception to raise when jina daemon is not connectable."""
 
 
-class NotSupportedError(Exception, BaseJinaExeception):
+class NotSupportedError(Exception, BaseJinaException):
     """Exception when user accidentally using a retired argument."""
 
 
-class RuntimeRunForeverEarlyError(Exception, BaseJinaExeception):
+class RuntimeRunForeverEarlyError(Exception, BaseJinaException):
     """Raised when an error occurs when starting the run_forever of Runtime"""
 
 
-class DockerVersionError(SystemError, BaseJinaExeception):
+class DockerVersionError(SystemError, BaseJinaException):
     """Raised when the docker version is incompatible"""
 
 
-class DaemonInvalidDockerfile(FileNotFoundError, BaseJinaExeception):
+class DaemonInvalidDockerfile(FileNotFoundError, BaseJinaException):
     """Raised when invalid dockerfile is passed to JinaD"""
 
 
-class NoContainerizedError(Exception, BaseJinaExeception):
+class NoContainerizedError(Exception, BaseJinaException):
     """Raised when trying to use non-containerized Executor in K8s or Docker Compose"""
diff --git a/tests/unit/clients/python/test_on_err.py b/tests/unit/clients/python/test_on_err.py
index e49cb0f339e69..f7347571cc0f5 100644
--- a/tests/unit/clients/python/test_on_err.py
+++ b/tests/unit/clients/python/test_on_err.py
@@ -1,8 +1,14 @@
-import numpy as np
-import pytest
+from typing import Optional
+
+import aiohttp
+import grpc
 
 from jina.excepts import BadClientCallback
-from jina import Flow
+from jina import Flow, Client
+
+import numpy as np
+import pytest
+from docarray import DocumentArray
 from docarray.document.generators import from_ndarray
 
 
@@ -43,3 +49,88 @@ def validate(x):
             continue_on_error=True,
         )
         assert t == 1
+
+
+@pytest.mark.parametrize(
+    'protocol,exception',
+    [
+        ('websocket', aiohttp.ClientError),
+        ('grpc', grpc.aio._call.AioRpcError),
+        ('http', aiohttp.ClientError),
+    ],
+)
+def test_client_on_error_call(protocol, exception):
+
+    with pytest.raises(exception):
+        Client(host='0.0.0.0', protocol=protocol, port=12345).post(
+            '/blah',
+            inputs=DocumentArray.empty(10),
+        )
+
+
+@pytest.mark.parametrize(
+    'protocol,exception',
+    [
+        ('websocket', aiohttp.client_exceptions.ClientConnectorError),
+        ('grpc', grpc.aio._call.AioRpcError),
+        ('http', aiohttp.client_exceptions.ClientConnectorError),
+    ],
+)
+def test_client_on_error_raise_exception(protocol, exception):
+    class OnError:
+        def __init__(self):
+            self.is_called = False
+
+        def __call__(self, response, exception_param: Optional[Exception] = None):
+            self.is_called = True
+            assert type(exception_param) == exception
+
+    on_error = OnError()
+
+    Client(host='0.0.0.0', protocol=protocol, port=12345).post(
+        '/blah',
+        inputs=DocumentArray.empty(10),
+        on_error=on_error,
+    )
+
+    assert on_error.is_called
+
+
+@pytest.mark.parametrize('protocol', ['websocket', 'grpc', 'http'])
+def test_client_on_error_deprecation(protocol):
+    class OnError:
+        def __init__(self):
+            self.is_called = False
+
+        def __call__(self, response):  # this is deprecated
+            self.is_called = True
+
+    on_error = OnError()
+
+    Client(host='0.0.0.0', protocol=protocol, port=12345).post(
+        '/blah',
+        inputs=DocumentArray.empty(10),
+        on_error=on_error,
+    )
+
+    assert on_error.is_called
+
+
+@pytest.mark.parametrize('protocol', ['websocket', 'grpc', 'http'])
+def test_client_on_always_after_exception(protocol):
+    class OnAlways:
+        def __init__(self):
+            self.is_called = False
+
+        def __call__(self, response):
+            self.is_called = True
+
+    on_always = OnAlways()
+
+    Client(host='0.0.0.0', protocol=protocol, port=12345).post(
+        '/blah',
+        inputs=DocumentArray.empty(10),
+        on_always=on_always,
+    )
+
+    assert on_always.is_called
diff --git a/tests/unit/serve/runtimes/gateway/http/test_app.py b/tests/unit/serve/runtimes/gateway/http/test_app.py
index 33c38a1fcf4ba..8eefaa577cea3 100644
--- a/tests/unit/serve/runtimes/gateway/http/test_app.py
+++ b/tests/unit/serve/runtimes/gateway/http/test_app.py
@@ -2,6 +2,7 @@
 import ssl
 from tempfile import NamedTemporaryFile
 
+import aiohttp
 import pytest
 import requests as req
 from fastapi.testclient import TestClient
@@ -202,11 +203,11 @@ def test_uvicorn_ssl_with_flow(cert_pem, key_pem, protocol, capsys):
         ],
     ) as f:
         os.environ['JINA_LOG_LEVEL'] = 'ERROR'
-        Client(protocol=protocol, port=f.port_expose, https=True).index([Document()])
-        assert (
-            '''certificate verify failed: self signed certificate'''
-            in capsys.readouterr().out
-        )
+
+        with pytest.raises(aiohttp.ClientConnectorCertificateError):
+            Client(protocol=protocol, port=f.port_expose, https=True).index(
+                [Document()]
+            )
 
 
 da = DocumentArray([Document(text='text_input')])


Code review comments:
diff_hunk: @@ -325,7 +325,7 @@ class MyExecutor(Executor):
 f = Flow().add(uses=MyExecutor)
 
 
-def print_why(resp):
+def print_why(resp, exception: None):
path: docs/fundamentals/executor/executor-in-flow.md
text: @user1:
exception: None seems wrong