Intermediate changes

commit_hash:fedb8a96cf9e063e9419c4449773b483370fe336

7 files changed, 46 insertions, 10 deletions

diff --git a/contrib/python/ydb/py3/.dist-info/METADATA b/contrib/python/ydb/py3/.dist-info/METADATA
index 09428b518f6..d754214e8f6 100644
--- a/contrib/python/ydb/py3/.dist-info/METADATA
+++ b/contrib/python/ydb/py3/.dist-info/METADATA
@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ydb
-Version: 3.21.7
+Version: 3.21.8
 Summary: YDB Python SDK
 Home-page: http://github.com/ydb-platform/ydb-python-sdk
 Author: Yandex LLC
diff --git a/contrib/python/ydb/py3/ya.make b/contrib/python/ydb/py3/ya.make
index 59a91bb5816..8958d7b97e4 100644
--- a/contrib/python/ydb/py3/ya.make
+++ b/contrib/python/ydb/py3/ya.make
@@ -2,7 +2,7 @@
 
 PY3_LIBRARY()
 
-VERSION(3.21.7)
+VERSION(3.21.8)
 
 LICENSE(Apache-2.0)
 
diff --git a/contrib/python/ydb/py3/ydb/_errors.py b/contrib/python/ydb/py3/ydb/_errors.py
index 1e2308ef394..b19de749cc5 100644
--- a/contrib/python/ydb/py3/ydb/_errors.py
+++ b/contrib/python/ydb/py3/ydb/_errors.py
@@ -1,5 +1,7 @@
 from dataclasses import dataclass
-from typing import Optional
+from typing import Optional, Union
+
+import grpc
 
 from . import issues
 
@@ -52,3 +54,26 @@ def check_retriable_error(err, retry_settings, attempt):
 class ErrorRetryInfo:
     is_retriable: bool
     sleep_timeout_seconds: Optional[float]
+
+
+def stream_error_converter(exc: BaseException) -> Union[issues.Error, BaseException]:
+    """Converts gRPC stream errors to appropriate YDB exception types.
+
+    This function takes a base exception and converts specific gRPC aio stream errors
+    to their corresponding YDB exception types for better error handling and semantic
+    clarity.
+
+    Args:
+        exc (BaseException): The original exception to potentially convert.
+
+    Returns:
+        BaseException: Either a converted YDB exception or the original exception
+                      if no specific conversion rule applies.
+    """
+    if isinstance(exc, (grpc.RpcError, grpc.aio.AioRpcError)):
+        if exc.code() == grpc.StatusCode.UNAVAILABLE:
+            return issues.Unavailable(exc.details() or "")
+        if exc.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
+            return issues.DeadlineExceed("Deadline exceeded on request")
+        return issues.Error("Stream has been terminated. Original exception: {}".format(str(exc.details())))
+    return exc
diff --git a/contrib/python/ydb/py3/ydb/aio/_utilities.py b/contrib/python/ydb/py3/ydb/aio/_utilities.py
index 296cd256f13..53a7d412f04 100644
--- a/contrib/python/ydb/py3/ydb/aio/_utilities.py
+++ b/contrib/python/ydb/py3/ydb/aio/_utilities.py
@@ -2,9 +2,10 @@ import asyncio
 
 
 class AsyncResponseIterator(object):
-    def __init__(self, it, wrapper):
+    def __init__(self, it, wrapper, error_converter=None):
         self.it = it.__aiter__()
         self.wrapper = wrapper
+        self.error_converter = error_converter
 
     def cancel(self):
         self.it.cancel()
@@ -17,7 +18,13 @@ class AsyncResponseIterator(object):
         return self
 
     async def _next(self):
-        res = self.wrapper(await self.it.__anext__())
+        try:
+            res = self.wrapper(await self.it.__anext__())
+        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 await self._next()
diff --git a/contrib/python/ydb/py3/ydb/aio/query/session.py b/contrib/python/ydb/py3/ydb/aio/query/session.py
index 7a7ba5baef2..83d527b7438 100644
--- a/contrib/python/ydb/py3/ydb/aio/query/session.py
+++ b/contrib/python/ydb/py3/ydb/aio/query/session.py
@@ -19,6 +19,7 @@ from ...query.session import (
 )
 
 from ..._constants import DEFAULT_INITIAL_RESPONSE_TIMEOUT
+from ..._errors import stream_error_converter
 
 
 class QuerySession(BaseQuerySession):
@@ -151,12 +152,13 @@ class QuerySession(BaseQuerySession):
         )
 
         return AsyncResponseContextIterator(
-            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,
         )
diff --git a/contrib/python/ydb/py3/ydb/aio/query/transaction.py b/contrib/python/ydb/py3/ydb/aio/query/transaction.py
index c9a6e445c93..2c313a4a55d 100644
--- a/contrib/python/ydb/py3/ydb/aio/query/transaction.py
+++ b/contrib/python/ydb/py3/ydb/aio/query/transaction.py
@@ -11,6 +11,7 @@ from ...query.transaction import (
     BaseQueryTxContext,
     QueryTxStateEnum,
 )
+from ..._errors import stream_error_converter
 
 logger = logging.getLogger(__name__)
 
@@ -181,8 +182,8 @@ class QueryTxContext(BaseQueryTxContext):
         )
 
         self._prev_stream = AsyncResponseContextIterator(
-            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,
@@ -190,5 +191,6 @@ class QueryTxContext(BaseQueryTxContext):
                 commit_tx=commit_tx,
                 settings=self.session._settings,
             ),
+            error_converter=stream_error_converter,
         )
         return self._prev_stream
diff --git a/contrib/python/ydb/py3/ydb/ydb_version.py b/contrib/python/ydb/py3/ydb/ydb_version.py
index d953c2ebcdd..68d3e8ca2bc 100644
--- a/contrib/python/ydb/py3/ydb/ydb_version.py
+++ b/contrib/python/ydb/py3/ydb/ydb_version.py
@@ -1 +1 @@
-VERSION = "3.21.7"
+VERSION = "3.21.8"