feat: supported attach hint

CHANGELOG.md

@@ -1,3 +1,4 @@
+* Query session attach stream handles `NodeShutdown` and `SessionShutdown` session hints: on `NodeShutdown` the session's node connection is pessimized and the session is retired, on `SessionShutdown` the session is retired without touching the node
 * Support the timezone-carrying types `TzDate`, `TzDatetime` and `TzTimestamp` in query results and parameters (reading them previously raised `AttributeError`)
 * Drop support for Python 3.8 and 3.9; the minimum supported version is now Python 3.10
 * Accept native `datetime.datetime` values for `Datetime` and `Datetime64` query parameters (previously only integer seconds since the epoch were accepted)

tests/query/test_query_session_hints.py

@@ -0,0 +1,199 @@
+import asyncio
+from unittest import mock
+
+import pytest
+
+from ydb._grpc.common.protos import ydb_query_pb2
+from ydb._grpc.grpcwrapper import common_utils
+from ydb.aio.pool import ConnectionPool as AsyncConnectionPool
+from ydb.pool import ConnectionPool, ConnectionsCache
+from ydb.query.session import QuerySession
+
+
+def _make_session(node_id=42):
+    driver = mock.Mock()
+    driver._pessimize_node = mock.Mock()
+    session = QuerySession(driver)
+    session._session_id = "test-session"
+    session._node_id = node_id
+    return session, driver
+
+
+class TestQuerySessionAttachHints:
+    def test_node_shutdown_pessimizes_node_and_invalidates_session(self):
+        session, driver = _make_session(node_id=42)
+
+        session._handle_attach_session_state(
+            ydb_query_pb2.SessionState(
+                status=0,
+                node_shutdown=ydb_query_pb2.NodeShutdownHint(),
+            )
+        )
+
+        driver._pessimize_node.assert_called_once_with(42)
+        assert session._invalidated
+        assert session._closed
+
+    def test_session_shutdown_invalidates_without_pessimizing_node(self):
+        session, driver = _make_session(node_id=42)
+
+        session._handle_attach_session_state(
+            ydb_query_pb2.SessionState(
+                status=0,
+                session_shutdown=ydb_query_pb2.SessionShutdownHint(),
+            )
+        )
+
+        driver._pessimize_node.assert_not_called()
+        assert session._invalidated
+        assert session._closed
+
+    def test_node_shutdown_with_zero_node_id_delegates_to_driver(self):
+        session, driver = _make_session(node_id=0)
+
+        session._handle_attach_session_state(
+            ydb_query_pb2.SessionState(
+                status=0,
+                node_shutdown=ydb_query_pb2.NodeShutdownHint(),
+            )
+        )
+
+        driver._pessimize_node.assert_called_once_with(0)
+        assert session._invalidated
+
+    def test_node_shutdown_without_node_id_skips_pessimization(self):
+        session, driver = _make_session(node_id=None)
+
+        session._handle_attach_session_state(
+            ydb_query_pb2.SessionState(
+                status=0,
+                node_shutdown=ydb_query_pb2.NodeShutdownHint(),
+            )
+        )
+
+        driver._pessimize_node.assert_not_called()
+        assert session._invalidated
+        assert session._closed
+
+    def test_no_hint_does_not_invalidate(self):
+        session, driver = _make_session()
+
+        session._handle_attach_session_state(
+            ydb_query_pb2.SessionState(status=0),
+        )
+
+        driver._pessimize_node.assert_not_called()
+        assert not session._invalidated
+        assert not session._closed
+
+    def test_none_response_does_not_invalidate(self):
+        session, driver = _make_session()
+
+        session._handle_attach_session_state(None)
+
+        driver._pessimize_node.assert_not_called()
+        assert not session._invalidated
+        assert not session._closed
+
+    def test_attach_stream_wrapper_handles_hint_and_returns_status(self):
+        session, driver = _make_session(node_id=42)
+        response = ydb_query_pb2.SessionState(
+            status=0,
+            session_shutdown=ydb_query_pb2.SessionShutdownHint(),
+        )
+
+        status = session._attach_stream_wrapper(response)
+
+        driver._pessimize_node.assert_not_called()
+        assert session._invalidated
+        assert isinstance(status, common_utils.ServerStatus)
+
+    def test_attach_stream_wrapper_returns_status_without_hint(self):
+        session, driver = _make_session()
+        response = ydb_query_pb2.SessionState(status=0)
+
+        status = session._attach_stream_wrapper(response)
+
+        driver._pessimize_node.assert_not_called()
+        assert not session._invalidated
+        assert isinstance(status, common_utils.ServerStatus)
+
+
+class TestConnectionPoolAttachHintPessimization:
+    def test_connections_cache_get_connection_by_node_id(self):
+        cache = ConnectionsCache()
+        connection = mock.Mock()
+        connection.endpoint = "grpc://localhost:2135"
+        connection.node_id = 42
+        connection.add_cleanup_callback = mock.Mock()
+
+        cache.add(connection)
+
+        assert cache.get_connection_by_node_id(42) is connection
+        assert cache.get_connection_by_node_id(99) is None
+
+    def test_sync_pool_pessimizes_node_connection(self):
+        pool = ConnectionPool.__new__(ConnectionPool)
+        connection = mock.Mock()
+        pool._store = mock.Mock()
+        pool._store.get_connection_by_node_id.return_value = connection
+        pool._on_disconnected = mock.Mock()
+
+        pool._pessimize_node(42)
+
+        pool._store.get_connection_by_node_id.assert_called_once_with(42)
+        pool._on_disconnected.assert_called_once_with(connection)
+
+    def test_sync_pool_ignores_missing_node_connection(self):
+        pool = ConnectionPool.__new__(ConnectionPool)
+        pool._store = mock.Mock()
+        pool._store.get_connection_by_node_id.return_value = None
+        pool._on_disconnected = mock.Mock()
+
+        pool._pessimize_node(42)
+        pool._pessimize_node(0)
+        pool._pessimize_node(-1)
+
+        pool._store.get_connection_by_node_id.assert_called_once_with(42)
+        pool._on_disconnected.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_async_pool_ignores_non_positive_node_id(self):
+        pool = AsyncConnectionPool.__new__(AsyncConnectionPool)
+        pool._store = mock.Mock()
+        pool._on_disconnected = mock.Mock()
+
+        pool._pessimize_node(0)
+        pool._pessimize_node(-1)
+
+        pool._store.get_connection_by_node_id.assert_not_called()
+        pool._on_disconnected.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_async_pool_ignores_missing_node_connection(self):
+        pool = AsyncConnectionPool.__new__(AsyncConnectionPool)
+        pool._store = mock.Mock()
+        pool._store.get_connection_by_node_id.return_value = None
+        pool._on_disconnected = mock.Mock()
+
+        pool._pessimize_node(42)
+        await asyncio.sleep(0)
+
+        pool._store.get_connection_by_node_id.assert_called_once_with(42)
+        pool._on_disconnected.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_async_pool_pessimizes_node_connection(self):
+        pool = AsyncConnectionPool.__new__(AsyncConnectionPool)
+        connection = mock.Mock()
+        disconnect = mock.AsyncMock()
+        pool._store = mock.Mock()
+        pool._store.get_connection_by_node_id.return_value = connection
+        pool._on_disconnected = mock.Mock(return_value=disconnect)
+
+        pool._pessimize_node(42)
+        await asyncio.sleep(0)
+
+        pool._store.get_connection_by_node_id.assert_called_once_with(42)
+        pool._on_disconnected.assert_called_once_with(connection)
+        disconnect.assert_awaited_once_with()

ydb/aio/pool.py

@@ -3,7 +3,7 @@
 import asyncio
 import logging
 import random
-from typing import Any, Callable, Optional, Tuple, TYPE_CHECKING
+from typing import Any, Callable, Optional, Tuple, TYPE_CHECKING, cast
 
 from ydb import issues
 from ydb.opentelemetry.tracing import SpanName, create_ydb_span
@@ -311,6 +311,15 @@ async def __wrapper__() -> None:
 
         return __wrapper__
 
+    def _pessimize_node(self, node_id: int) -> None:
+        """Deprioritize the connection attached to the given YDB node."""
+        if node_id <= 0:
+            return
+
+        connection = cast(Optional[Connection], self._store.get_connection_by_node_id(node_id))
+        if connection is not None:
+            asyncio.get_running_loop().create_task(self._on_disconnected(connection)())
+
     async def wait(self, timeout: Optional[float] = 7.0, fail_fast: bool = False) -> None:  # type: ignore[override]  # async override of sync method
         with create_ydb_span(SpanName.DRIVER_INITIALIZE, self._driver_config, kind="internal").attach_context():
             await self._store.get(fast_fail=fail_fast, wait_timeout=timeout if timeout is not None else 7.0)

ydb/aio/query/session.py

@@ -14,7 +14,6 @@
 from .. import _utilities
 from ... import issues
 from ...settings import BaseRequestSettings
-from ..._grpc.grpcwrapper import common_utils
 from ..._grpc.grpcwrapper import ydb_query_public_types as _ydb_query_public
 
 from ...query import base
@@ -53,7 +52,7 @@ async def _attach(self) -> None:
         self._stream = await self._attach_call()
         self._status_stream = _utilities.AsyncResponseIterator(
             self._stream,
-            lambda response: common_utils.ServerStatus.from_proto(response),
+            self._attach_stream_wrapper,
         )
 
         try:

ydb/pool.py

@@ -160,6 +160,10 @@ def remove(self, connection: Connection) -> None:
             self.connections.pop(connection.endpoint, None)
             self.outdated.pop(connection.endpoint, None)
 
+    def get_connection_by_node_id(self, node_id: Optional[int]) -> Optional[Connection]:
+        with self.lock:
+            return self.connections_by_node_id.get(node_id)
+
 
 class Discovery(threading.Thread):
     def __init__(self, store: ConnectionsCache, driver_config: "DriverConfig") -> None:
@@ -494,6 +498,15 @@ def _on_disconnected(self, connection: Connection) -> None:
         if self._discovery_thread:
             self._discovery_thread.notify_disconnected()
 
+    def _pessimize_node(self, node_id: int) -> None:
+        """Deprioritize the connection attached to the given YDB node."""
+        if node_id <= 0:
+            return
+
+        connection = self._store.get_connection_by_node_id(node_id)
+        if connection is not None:
+            self._on_disconnected(connection)
+
     def discovery_debug_details(self) -> str:
         """
         Returns debug string about last errors

ydb/query/session.py

@@ -156,6 +156,24 @@ def _check_session_ready_to_use(self) -> None:
         if self._closed:
             raise RuntimeError(f"Session is not active, session_id: {self._session_id}, closed: {self._closed}")
 
+    def _attach_stream_wrapper(self, response_pb):
+        """Map attach-stream protobuf frames to ServerStatus and handle session hints."""
+        self._handle_attach_session_state(response_pb)
+        return common_utils.ServerStatus.from_proto(response_pb)
+
+    def _handle_attach_session_state(self, response_pb) -> None:
+        """Retire the session when the server sends a shutdown hint on the attach stream."""
+        if response_pb is None:
+            return
+
+        hint = response_pb.WhichOneof("session_hint")
+        if hint == "node_shutdown":
+            if self._node_id is not None:
+                self._driver._pessimize_node(self._node_id)
+            self._close_session(invalidate=True)
+        elif hint == "session_shutdown":
+            self._close_session(invalidate=True)
+
     def _close_session(self, invalidate: bool = False) -> None:
         if self._closed:
             return
@@ -362,7 +380,7 @@ def _attach(self, first_resp_timeout: int = DEFAULT_INITIAL_RESPONSE_TIMEOUT) ->
         self._stream = self._attach_call()
         status_stream = _utilities.SyncResponseIterator(
             self._stream,
-            lambda response: common_utils.ServerStatus.from_proto(response),
+            self._attach_stream_wrapper,
         )
 
         try: