Intermediate changes

commit_hash:baf7b0a6d4bac3a6190dbe6f1ef4cf241820805d

6 files changed, 71 insertions, 13 deletions

diff --git a/contrib/python/ydb/py3/.dist-info/METADATA b/contrib/python/ydb/py3/.dist-info/METADATA
index 51d9ebee2ed..c9e4bb99485 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.28.1
+Version: 3.28.2
 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 5bb09ac2d2e..87991083694 100644
--- a/contrib/python/ydb/py3/ya.make
+++ b/contrib/python/ydb/py3/ya.make
@@ -2,7 +2,7 @@
 
 PY3_LIBRARY()
 
-VERSION(3.28.1)
+VERSION(3.28.2)
 
 LICENSE(Apache-2.0)
 
diff --git a/contrib/python/ydb/py3/ydb/aio/query/base.py b/contrib/python/ydb/py3/ydb/aio/query/base.py
index 66df3703236..a1d6e5d74fc 100644
--- a/contrib/python/ydb/py3/ydb/aio/query/base.py
+++ b/contrib/python/ydb/py3/ydb/aio/query/base.py
@@ -12,12 +12,30 @@ class AsyncResponseContextIterator(_utilities.AsyncResponseIterator):
     async def _next(self):
         try:
             return await super()._next()
-        except Exception as e:
+        except StopAsyncIteration:
+            # Normal stream termination is not an error and must not invalidate
+            # the session.
+            raise
+        except BaseException as e:
+            # BaseException (not Exception) because asyncio.CancelledError
+            # inherits from BaseException in Python 3.8+. A stream interrupted
+            # by a cancel must also be reported to _on_error so the session can
+            # be invalidated; otherwise the next caller that picks this session
+            # out of the pool races the undrained stream and the server can
+            # reply with SessionBusy.
             if self._on_error:
                 self._on_error(e)
-            raise e
+            raise
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        #  To close stream on YDB it is necessary to scroll through it to the end
-        async for _ in self:
+        #  To close stream on YDB it is necessary to scroll through it to the end.
+        # Errors that happen during the cleanup drain have already been reported
+        # to _on_error inside _next, so swallow them here — re-raising from
+        # __aexit__ would mask whatever exception is already propagating out of
+        # the `async with` body and would leave callers (e.g. the tx __aexit__)
+        # unable to run their own cleanup (rollback).
+        try:
+            async for _ in self:
+                pass
+        except BaseException:
             pass
diff --git a/contrib/python/ydb/py3/ydb/query/base.py b/contrib/python/ydb/py3/ydb/query/base.py
index e7764e1c7eb..bf0d80b9dad 100644
--- a/contrib/python/ydb/py3/ydb/query/base.py
+++ b/contrib/python/ydb/py3/ydb/query/base.py
@@ -83,14 +83,29 @@ class SyncResponseContextIterator(_utilities.SyncResponseIterator):
     def _next(self):
         try:
             return super()._next()
-        except Exception as e:
+        except StopIteration:
+            # Normal stream termination is not an error and must not invalidate
+            # the session.
+            raise
+        except BaseException as e:
+            # BaseException (not Exception) for parity with the async iterator:
+            # KeyboardInterrupt / SystemExit should still invalidate the session
+            # before they propagate, otherwise the next caller that reuses the
+            # session races the undrained stream and the server can reply with
+            # SessionBusy.
             if self._on_error:
                 self._on_error(e)
-            raise e
+            raise
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        #  To close stream on YDB it is necessary to scroll through it to the end
-        for _ in self:
+        #  To close stream on YDB it is necessary to scroll through it to the end.
+        # Errors during the cleanup drain have already been reported to _on_error
+        # inside _next; swallow them here so __exit__ does not mask a primary
+        # exception and the caller's own cleanup (e.g. tx rollback) can still run.
+        try:
+            for _ in self:
+                pass
+        except BaseException:
             pass
 
 
diff --git a/contrib/python/ydb/py3/ydb/query/session.py b/contrib/python/ydb/py3/ydb/query/session.py
index af4b7ec6b06..b28cba8baed 100644
--- a/contrib/python/ydb/py3/ydb/query/session.py
+++ b/contrib/python/ydb/py3/ydb/query/session.py
@@ -136,8 +136,33 @@ class BaseQuerySession(abc.ABC, Generic[DriverT]):
             except Exception:
                 pass
 
-    def _on_execute_stream_error(self, e: Exception) -> None:
-        if isinstance(e, issues.DeadlineExceed):
+    def _on_execute_stream_error(self, e: BaseException) -> None:
+        # The execute stream is a single gRPC call that carries all of a
+        # query's response parts. If any of these errors surface while reading
+        # it, the server-side stream is either known-dead (BadSession,
+        # ConnectionError, DeadlineExceed) or undrained and un-resumable
+        # (SessionBusy: server thinks this session still has the previous
+        # query running; Cancelled: the call has been torn down mid-flight),
+        # which means a subsequent query on the same session can race the
+        # stragglers and get a spurious SessionBusy back. Drop the session so
+        # the pool creates a fresh one on the next acquire.
+        #
+        # Accepts BaseException so that asyncio.CancelledError (not an
+        # issues.Error subclass) — the case documented in the bug report —
+        # also invalidates here.
+        if isinstance(e, issues.Error):
+            if isinstance(
+                e,
+                (
+                    issues.DeadlineExceed,
+                    issues.SessionBusy,
+                    issues.BadSession,
+                    issues.ConnectionError,
+                    issues.Cancelled,
+                ),
+            ):
+                self._invalidate()
+        else:
             self._invalidate()
 
     # Overloads for _create_call
diff --git a/contrib/python/ydb/py3/ydb/ydb_version.py b/contrib/python/ydb/py3/ydb/ydb_version.py
index fb99e7eeda8..7804556fba4 100644
--- a/contrib/python/ydb/py3/ydb/ydb_version.py
+++ b/contrib/python/ydb/py3/ydb/ydb_version.py
@@ -1 +1 @@
-VERSION = "3.28.1"
+VERSION = "3.28.2"