Fix stream next error handling

Author: vgvoleg

ydb/_utilities.py

@@ -149,9 +149,10 @@ def __next__(self):
 
 
 class SyncResponseIterator(object):
-    def __init__(self, it, wrapper):
+    def __init__(self, it, wrapper, error_converter=None):
         self.it = it
         self.wrapper = wrapper
+        self.error_converter = error_converter
 
     def cancel(self):
         self.it.cancel()
@@ -161,9 +162,16 @@ def __iter__(self):
         return self
 
     def _next(self):
-        res = self.wrapper(next(self.it))
+        try:
+            res = self.wrapper(next(self.it))
+        except BaseException as e:
+            if self.error_converter:
+                raise self.error_converter(e) from e
+            raise e
+
         if res is not None:
             return res
+
         return self._next()
 
     def next(self):

ydb/query/session.py

@@ -19,6 +19,7 @@
 from .transaction import QueryTxContext
 
 from .._constants import DEFAULT_INITIAL_RESPONSE_TIMEOUT, DEFAULT_LONG_STREAM_TIMEOUT
+from .._errors import stream_error_converter
 
 
 logger = logging.getLogger(__name__)
@@ -362,12 +363,13 @@ def execute(
         )
 
         return base.SyncResponseContextIterator(
-            stream_it,
-            lambda resp: base.wrap_execute_query_response(
+            it=stream_it,
+            wrapper=lambda resp: base.wrap_execute_query_response(
                 rpc_state=None,
                 response_pb=resp,
                 session_state=self._state,
                 session=self,
                 settings=self._settings,
             ),
+            error_converter=stream_error_converter,
         )

ydb/query/transaction.py

@@ -16,6 +16,7 @@
 from ..connection import _RpcState as RpcState
 
 from . import base
+from .._errors import stream_error_converter
 from ..settings import BaseRequestSettings
 
 logger = logging.getLogger(__name__)
@@ -500,14 +501,15 @@ def execute(
         )
 
         self._prev_stream = base.SyncResponseContextIterator(
-            stream_it,
-            lambda resp: base.wrap_execute_query_response(
+            it=stream_it,
+            wrapper=lambda resp: base.wrap_execute_query_response(
                 rpc_state=None,
                 response_pb=resp,
                 session_state=self._session_state,
                 tx=self,
                 commit_tx=commit_tx,
                 settings=self.session._settings,
             ),
+            error_converter=stream_error_converter,
         )
         return self._prev_stream