Intermediate changes

commit_hash:aa7097ab35f6d94c4f7ae567c25849825b9199c7

17 files changed, 113 insertions, 50 deletions

diff --git a/contrib/python/ydb/py3/.dist-info/METADATA b/contrib/python/ydb/py3/.dist-info/METADATA
index 1158c004283..c7532ac6d14 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.9
+Version: 3.21.12
 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 634ea53ac19..811dfeb664a 100644
--- a/contrib/python/ydb/py3/ya.make
+++ b/contrib/python/ydb/py3/ya.make
@@ -2,7 +2,7 @@
 
 PY3_LIBRARY()
 
-VERSION(3.21.9)
+VERSION(3.21.12)
 
 LICENSE(Apache-2.0)
 
diff --git a/contrib/python/ydb/py3/ydb/_errors.py b/contrib/python/ydb/py3/ydb/_errors.py
index b19de749cc5..3f4263502cf 100644
--- a/contrib/python/ydb/py3/ydb/_errors.py
+++ b/contrib/python/ydb/py3/ydb/_errors.py
@@ -1,7 +1,5 @@
 from dataclasses import dataclass
-from typing import Optional, Union
-
-import grpc
+from typing import Optional
 
 from . import issues
 
@@ -15,6 +13,7 @@ _errors_retriable_slow_backoff_types = [
     issues.Overloaded,
     issues.SessionPoolEmpty,
     issues.ConnectionError,
+    issues.ConnectionLost,
 ]
 _errors_retriable_slow_backoff_idempotent_types = [
     issues.Undetermined,
@@ -22,6 +21,10 @@ _errors_retriable_slow_backoff_idempotent_types = [
 
 
 def check_retriable_error(err, retry_settings, attempt):
+    if isinstance(err, issues.Cancelled):
+        if retry_settings.retry_cancelled:
+            return ErrorRetryInfo(True, retry_settings.fast_backoff.calc_timeout(attempt))
+
     if isinstance(err, issues.NotFound):
         if retry_settings.retry_not_found:
             return ErrorRetryInfo(True, retry_settings.fast_backoff.calc_timeout(attempt))
@@ -54,26 +57,3 @@ 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/_grpc/grpcwrapper/common_utils.py b/contrib/python/ydb/py3/ydb/_grpc/grpcwrapper/common_utils.py
index 004faf1524f..569a75d00e8 100644
--- a/contrib/python/ydb/py3/ydb/_grpc/grpcwrapper/common_utils.py
+++ b/contrib/python/ydb/py3/ydb/_grpc/grpcwrapper/common_utils.py
@@ -6,6 +6,7 @@ import concurrent.futures
 import contextvars
 import datetime
 import functools
+import logging
 import typing
 from typing import (
     Optional,
@@ -36,6 +37,8 @@ from ... import issues, connection
 from ...settings import BaseRequestSettings
 from ..._constants import DEFAULT_LONG_STREAM_TIMEOUT
 
+logger = logging.getLogger(__name__)
+
 
 class IFromProto(abc.ABC):
     @staticmethod
diff --git a/contrib/python/ydb/py3/ydb/_topic_reader/topic_reader.py b/contrib/python/ydb/py3/ydb/_topic_reader/topic_reader.py
index d477c9ca1bb..38ee1be6598 100644
--- a/contrib/python/ydb/py3/ydb/_topic_reader/topic_reader.py
+++ b/contrib/python/ydb/py3/ydb/_topic_reader/topic_reader.py
@@ -85,7 +85,7 @@ class PublicReaderSettings:
         )
 
     def _retry_settings(self) -> RetrySettings:
-        return RetrySettings(idempotent=True)
+        return RetrySettings(idempotent=True, retry_cancelled=True)
 
 
 class RetryPolicy:
diff --git a/contrib/python/ydb/py3/ydb/_topic_reader/topic_reader_asyncio.py b/contrib/python/ydb/py3/ydb/_topic_reader/topic_reader_asyncio.py
index b855a80b99c..818eb1a9db4 100644
--- a/contrib/python/ydb/py3/ydb/_topic_reader/topic_reader_asyncio.py
+++ b/contrib/python/ydb/py3/ydb/_topic_reader/topic_reader_asyncio.py
@@ -248,12 +248,15 @@ class ReaderReconnector:
                 self._state_changed.set()
                 await self._stream_reader.wait_error()
             except BaseException as err:
+                logger.debug("reader %s, attempt %s connection loop error %s", self._id, attempt, err)
                 retry_info = check_retriable_error(err, self._settings._retry_settings(), attempt)
                 if not retry_info.is_retriable:
                     logger.debug("reader %s stop connection loop due to %s", self._id, err)
                     self._set_first_error(err)
                     return
 
+                logger.debug("sleep before retry for %s seconds", retry_info.sleep_timeout_seconds)
+
                 await asyncio.sleep(retry_info.sleep_timeout_seconds)
 
                 attempt += 1
diff --git a/contrib/python/ydb/py3/ydb/_topic_writer/topic_writer_asyncio.py b/contrib/python/ydb/py3/ydb/_topic_writer/topic_writer_asyncio.py
index d39606d1797..1c0e410b218 100644
--- a/contrib/python/ydb/py3/ydb/_topic_writer/topic_writer_asyncio.py
+++ b/contrib/python/ydb/py3/ydb/_topic_writer/topic_writer_asyncio.py
@@ -434,7 +434,7 @@ class WriterAsyncIOReconnector:
             raise self._stop_reason.exception()
 
     async def _connection_loop(self):
-        retry_settings = RetrySettings()  # todo
+        retry_settings = RetrySettings(retry_cancelled=True)  # todo
 
         while True:
             attempt = 0  # todo calc and reset
@@ -485,15 +485,16 @@ class WriterAsyncIOReconnector:
             except issues.Error as err:
                 err_info = check_retriable_error(err, retry_settings, attempt)
                 if not err_info.is_retriable or self._tx is not None:  # no retries in tx writer
+                    logger.debug("writer reconnector %s stop connection loop due to %s", self._id, err)
                     self._stop(err)
                     return
 
-                await asyncio.sleep(err_info.sleep_timeout_seconds)
                 logger.debug(
                     "writer reconnector %s retry in %s seconds",
                     self._id,
                     err_info.sleep_timeout_seconds,
                 )
+                await asyncio.sleep(err_info.sleep_timeout_seconds)
 
             except (asyncio.CancelledError, Exception) as err:
                 self._stop(err)
diff --git a/contrib/python/ydb/py3/ydb/aio/_utilities.py b/contrib/python/ydb/py3/ydb/aio/_utilities.py
index 53a7d412f04..062545d8fa6 100644
--- a/contrib/python/ydb/py3/ydb/aio/_utilities.py
+++ b/contrib/python/ydb/py3/ydb/aio/_utilities.py
@@ -2,10 +2,9 @@ import asyncio
 
 
 class AsyncResponseIterator(object):
-    def __init__(self, it, wrapper, error_converter=None):
+    def __init__(self, it, wrapper):
         self.it = it.__aiter__()
         self.wrapper = wrapper
-        self.error_converter = error_converter
 
     def cancel(self):
         self.it.cancel()
@@ -18,12 +17,7 @@ class AsyncResponseIterator(object):
         return self
 
     async def _next(self):
-        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
+        res = self.wrapper(await self.it.__anext__())
 
         if res is not None:
             return res
diff --git a/contrib/python/ydb/py3/ydb/aio/connection.py b/contrib/python/ydb/py3/ydb/aio/connection.py
index 21461b4bad0..d0e34a7f7ef 100644
--- a/contrib/python/ydb/py3/ydb/aio/connection.py
+++ b/contrib/python/ydb/py3/ydb/aio/connection.py
@@ -11,6 +11,7 @@ from ydb.connection import (
     _log_request,
     _log_response,
     _rpc_error_handler,
+    _is_disconnect_needed,
     _get_request_timeout,
     _set_server_timeouts,
     _RpcState as RpcState,
@@ -102,6 +103,33 @@ class _RpcState(RpcState):
         raise NotImplementedError
 
 
+class _SafeAsyncIterator:
+    def __init__(self, resp, rpc_state, on_disconnected_callback):
+        self.resp = resp
+        self.it = resp.__aiter__()
+        self.rpc_state = rpc_state
+        self.on_disconnected_callback = on_disconnected_callback
+
+    def cancel(self):
+        self.resp.cancel()
+        return self
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        try:
+            return await self.it.__anext__()
+        except grpc.RpcError as rpc_error:
+            ydb_error = _rpc_error_handler(self.rpc_state, rpc_error, use_unavailable=True)
+            if _is_disconnect_needed(ydb_error):
+                await self.on_disconnected_callback()
+            raise ydb_error
+
+    def __getattr__(self, item):
+        return getattr(self.resp, item)
+
+
 class Connection:
     __slots__ = (
         "endpoint",
@@ -191,6 +219,11 @@ class Connection:
 
             response = await feature
             _log_response(rpc_state, response)
+
+            if hasattr(response, "__aiter__"):
+                # NOTE(vgvoleg): for stream results we should also be able to handle disconnects
+                response = _SafeAsyncIterator(response, rpc_state, on_disconnected)
+
             return response if wrap_result is None else wrap_result(rpc_state, response, *wrap_args)
         except grpc.RpcError as rpc_error:
             if on_disconnected:
@@ -234,7 +267,7 @@ class Connection:
         been terminated are cancelled. If grace is None, this method will wait until all tasks are finished.
         :return: None
         """
-        logger.info("Closing channel for endpoint %s", self.endpoint)
+        logger.debug("Closing channel for endpoint %s", self.endpoint)
 
         self.closing = True
 
diff --git a/contrib/python/ydb/py3/ydb/aio/pool.py b/contrib/python/ydb/py3/ydb/aio/pool.py
index 99a3cfdb8da..0c75fa30426 100644
--- a/contrib/python/ydb/py3/ydb/aio/pool.py
+++ b/contrib/python/ydb/py3/ydb/aio/pool.py
@@ -270,7 +270,7 @@ class ConnectionPool(IConnectionPool):
                 self._discovery.notify_disconnected()
             raise
 
-        return await connection(
+        res = await connection(
             request,
             stub,
             rpc_name,
@@ -279,3 +279,5 @@ class ConnectionPool(IConnectionPool):
             wrap_args,
             self._on_disconnected(connection),
         )
+
+        return res
diff --git a/contrib/python/ydb/py3/ydb/aio/query/session.py b/contrib/python/ydb/py3/ydb/aio/query/session.py
index 13906164d60..98ea1849561 100644
--- a/contrib/python/ydb/py3/ydb/aio/query/session.py
+++ b/contrib/python/ydb/py3/ydb/aio/query/session.py
@@ -23,7 +23,6 @@ from ...query.session import (
 )
 
 from ..._constants import DEFAULT_INITIAL_RESPONSE_TIMEOUT
-from ..._errors import stream_error_converter
 
 
 class QuerySession(BaseQuerySession):
@@ -164,7 +163,6 @@ class QuerySession(BaseQuerySession):
                 session=self,
                 settings=self._settings,
             ),
-            error_converter=stream_error_converter,
         )
 
     async def explain(
diff --git a/contrib/python/ydb/py3/ydb/aio/query/transaction.py b/contrib/python/ydb/py3/ydb/aio/query/transaction.py
index 2c313a4a55d..9b2db2efb46 100644
--- a/contrib/python/ydb/py3/ydb/aio/query/transaction.py
+++ b/contrib/python/ydb/py3/ydb/aio/query/transaction.py
@@ -11,7 +11,6 @@ from ...query.transaction import (
     BaseQueryTxContext,
     QueryTxStateEnum,
 )
-from ..._errors import stream_error_converter
 
 logger = logging.getLogger(__name__)
 
@@ -191,6 +190,5 @@ 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/connection.py b/contrib/python/ydb/py3/ydb/connection.py
index d5b6ed50c69..0ab63069f6d 100644
--- a/contrib/python/ydb/py3/ydb/connection.py
+++ b/contrib/python/ydb/py3/ydb/connection.py
@@ -67,6 +67,7 @@ def _rpc_error_handler(
     rpc_state,
     rpc_error: typing.Union[grpc.RpcError, grpc.aio.AioRpcError, grpc.Call, grpc.aio.Call],
     on_disconnected: typing.Callable[[], None] = None,
+    use_unavailable: bool = False,
 ):
     """
     RPC call error handler, that translates gRPC error into YDB issue
@@ -74,7 +75,7 @@ def _rpc_error_handler(
     :param rpc_error: an underlying rpc error to handle
     :param on_disconnected: a handler to call on disconnected connection
     """
-    logger.info("%s: received error, %s", rpc_state, rpc_error)
+    logger.debug("%s: received error, %s", rpc_state, rpc_error)
     if isinstance(rpc_error, (grpc.RpcError, grpc.aio.AioRpcError, grpc.Call, grpc.aio.Call)):
         if rpc_error.code() == grpc.StatusCode.UNAUTHENTICATED:
             return issues.Unauthenticated(rpc_error.details())
@@ -82,6 +83,10 @@ def _rpc_error_handler(
             return issues.DeadlineExceed("Deadline exceeded on request")
         elif rpc_error.code() == grpc.StatusCode.UNIMPLEMENTED:
             return issues.Unimplemented("Method or feature is not implemented on server!")
+        elif rpc_error.code() == grpc.StatusCode.CANCELLED:
+            return issues.Cancelled(rpc_error.details())
+        elif use_unavailable and rpc_error.code() == grpc.StatusCode.UNAVAILABLE:
+            return issues.Unavailable(rpc_error.details())
 
     logger.debug("%s: unhandled rpc error, disconnecting channel", rpc_state)
     if on_disconnected is not None:
@@ -90,6 +95,16 @@ def _rpc_error_handler(
     return issues.ConnectionLost("Rpc error, reason %s" % str(rpc_error))
 
 
+def _is_disconnect_needed(error):
+    return isinstance(
+        error,
+        (
+            issues.ConnectionLost,
+            issues.Unavailable,
+        ),
+    )
+
+
 def _on_response_callback(rpc_state, call_state_unref, wrap_result=None, on_disconnected=None, wrap_args=()):
     """
     Callback to be executed on received RPC response
@@ -325,6 +340,33 @@ class EndpointKey(object):
         self.node_id = node_id
 
 
+class _SafeSyncIterator:
+    def __init__(self, resp, rpc_state, on_disconnected_callback):
+        self.resp = resp
+        self.it = resp.__iter__()
+        self.rpc_state = rpc_state
+        self.on_disconnected_callback = on_disconnected_callback
+
+    def cancel(self):
+        self.resp.cancel()
+        return self
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            return self.it.__next__()
+        except grpc.RpcError as rpc_error:
+            ydb_error = _rpc_error_handler(self.rpc_state, rpc_error, use_unavailable=True)
+            if _is_disconnect_needed(ydb_error):
+                self.on_disconnected_callback()
+            raise ydb_error
+
+    def __getattr__(self, item):
+        return getattr(self.resp, item)
+
+
 class Connection(object):
     __slots__ = (
         "endpoint",
@@ -466,6 +508,11 @@ class Connection(object):
                 compression=getattr(settings, "compression", None),
             )
             _log_response(rpc_state, response)
+
+            if hasattr(response, "__iter__"):
+                # NOTE(vgvoleg): for stream results we should also be able to handle disconnects
+                response = _SafeSyncIterator(response, rpc_state, on_disconnected)
+
             return response if wrap_result is None else wrap_result(rpc_state, response, *wrap_args)
         except grpc.RpcError as rpc_error:
             raise _rpc_error_handler(rpc_state, rpc_error, on_disconnected)
@@ -499,7 +546,7 @@ class Connection(object):
         Closes the underlying gRPC channel
         :return: None
         """
-        logger.info("Closing channel for endpoint %s", self.endpoint)
+        logger.debug("Closing channel for endpoint %s", self.endpoint)
         with self.lock:
             self.closing = True
 
diff --git a/contrib/python/ydb/py3/ydb/pool.py b/contrib/python/ydb/py3/ydb/pool.py
index 476ea674d9f..0fb3e86a3a3 100644
--- a/contrib/python/ydb/py3/ydb/pool.py
+++ b/contrib/python/ydb/py3/ydb/pool.py
@@ -470,7 +470,9 @@ class ConnectionPool(IConnectionPool):
             wrap_args,
             lambda: self._on_disconnected(connection),
         )
+
         tracing.trace(self.tracer, {"response": res}, trace_level=tracing.TraceLevel.DEBUG)
+
         return res
 
     @_utilities.wrap_async_call_exceptions
diff --git a/contrib/python/ydb/py3/ydb/retries.py b/contrib/python/ydb/py3/ydb/retries.py
index c9c23b1a918..5331f1b00f0 100644
--- a/contrib/python/ydb/py3/ydb/retries.py
+++ b/contrib/python/ydb/py3/ydb/retries.py
@@ -32,6 +32,7 @@ class RetrySettings(object):
         fast_backoff_settings=None,
         slow_backoff_settings=None,
         idempotent=False,
+        retry_cancelled=False,
     ):
         self.max_retries = max_retries
         self.max_session_acquire_timeout = max_session_acquire_timeout
@@ -45,6 +46,7 @@ class RetrySettings(object):
         self.retry_not_found = True
         self.idempotent = idempotent
         self.retry_internal_error = True
+        self.retry_cancelled = retry_cancelled
         self.unknown_error_handler = lambda e: None
         self.get_session_client_timeout = get_session_client_timeout
         if max_session_acquire_timeout is not None:
diff --git a/contrib/python/ydb/py3/ydb/scheme.py b/contrib/python/ydb/py3/ydb/scheme.py
index 263d1c65d33..27458dacd0c 100644
--- a/contrib/python/ydb/py3/ydb/scheme.py
+++ b/contrib/python/ydb/py3/ydb/scheme.py
@@ -124,7 +124,7 @@ class SchemeEntryType(enum.IntEnum):
         return entry == SchemeEntryType.EXTERNAL_DATA_SOURCE
 
     @staticmethod
-    def is_external_view(entry):
+    def is_view(entry):
         """
         :param entry: A scheme entry to check
         :return: True if scheme entry is a view and False otherwise
@@ -132,7 +132,7 @@ class SchemeEntryType(enum.IntEnum):
         return entry == SchemeEntryType.VIEW
 
     @staticmethod
-    def is_external_resource_pool(entry):
+    def is_resource_pool(entry):
         """
         :param entry: A scheme entry to check
         :return: True if scheme entry is a resource pool and False otherwise
diff --git a/contrib/python/ydb/py3/ydb/ydb_version.py b/contrib/python/ydb/py3/ydb/ydb_version.py
index 70f9f0a38ed..1493a7be5d3 100644
--- a/contrib/python/ydb/py3/ydb/ydb_version.py
+++ b/contrib/python/ydb/py3/ydb/ydb_version.py
@@ -1 +1 @@
-VERSION = "3.21.9"
+VERSION = "3.21.12"