Fix Copilot review issues for attach hint PR

Look up node connections under ConnectionsCache lock when
pessimizing after NodeShutdown hints, and strip grpcio-tools
version gates from v6 gRPC stubs so imports work with grpcio>=1.42.0.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: KirillKurdyukov

CHANGELOG.md

@@ -1,6 +1,7 @@
 * 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
 * Bumped `ydb-api-protos` and regenerated gRPC/protobuf stubs (v3–v6) to include query service session hints
 * Fix incompatibility with protobuf 6.30–6.31.0: regenerate v6 stubs with the lowest 6.x gencode floor (6.30.0) instead of 6.31.1
+* 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
 
 ## 3.29.4 ##
 * Fix leaked topic reader stream when close interrupts stream creation during reconnect

generate_protoc.py

@@ -1,11 +1,50 @@
 import os
 import pathlib
+import re
 import shutil
 
 from typing import List
 from argparse import ArgumentParser
 
-from grpc_tools import command
+_GRPC_VERSION_GATE_RE = re.compile(
+    r"GRPC_GENERATED_VERSION = '[^']+'\n"
+    r"GRPC_VERSION = grpc\.__version__\n"
+    r"_version_not_supported = False\n\n"
+    r"try:\n"
+    r"    from grpc\._utilities import first_version_is_lower\n"
+    r"    _version_not_supported = first_version_is_lower\(GRPC_VERSION, GRPC_GENERATED_VERSION\)\n"
+    r"except ImportError:\n"
+    r"    _version_not_supported = True\n\n"
+    r"if _version_not_supported:\n"
+    r"    raise RuntimeError\(\n"
+    r"(?:        .+\n)+"
+    r"    \)\n"
+)
+
+
+def strip_grpc_version_gate(content: str) -> str:
+    """Remove grpcio-tools version gate from generated *_grpc.py stubs."""
+    updated = _GRPC_VERSION_GATE_RE.sub("", content)
+    if updated != content:
+        updated = updated.replace("import warnings\n", "")
+    return updated
+
+
+def strip_grpc_version_gate_from_tree(rootdir: str) -> None:
+    for dirpath, _, fnames in os.walk(rootdir):
+        for fname in fnames:
+            if not fname.endswith("_grpc.py"):
+                continue
+
+            path = os.path.join(dirpath, fname)
+            with open(path, "r+t") as f:
+                content = f.read()
+                updated = strip_grpc_version_gate(content)
+                if updated == content:
+                    continue
+                f.seek(0)
+                f.write(updated)
+                f.truncate()
 
 
 def files_filter(dir, items: List[str]) -> List[str]:
@@ -56,11 +95,18 @@ def fix_file_contents(rootdir, protobuf_version: str):
 
                 # Add ignore style check
                 content = content.replace("# -*- coding: utf-8 -*-", "# -*- coding: utf-8 -*-\n" + flake_ignore_line)
+
+                if fname.endswith("_grpc.py"):
+                    content = strip_grpc_version_gate(content)
+
                 f.seek(0)
                 f.write(content)
+                f.truncate()
 
 
 def generate_protobuf(src_proto_dir: str, dst_dir, protobuf_version: str):
+    from grpc_tools import command
+
     shutil.rmtree(dst_dir, ignore_errors=True)
 
     shutil.copytree(src_proto_dir, dst_dir, ignore=files_filter)

tests/query/test_query_session_hints.py

@@ -91,17 +91,18 @@ def test_sync_pool_pessimizes_node_connection(self):
         pool = ConnectionPool.__new__(ConnectionPool)
         connection = mock.Mock()
         pool._store = mock.Mock()
-        pool._store.connections_by_node_id = {42: connection}
+        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.connections_by_node_id = {}
+        pool._store.get_connection_by_node_id.return_value = None
         pool._on_disconnected = mock.Mock()
 
         pool._pessimize_node(42)
@@ -115,11 +116,12 @@ async def test_async_pool_pessimizes_node_connection(self):
         connection = mock.Mock()
         disconnect = mock.AsyncMock()
         pool._store = mock.Mock()
-        pool._store.connections_by_node_id = {42: connection}
+        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/_grpc/v6/draft/protos/ydb_dynamic_config_pb2_grpc.py

@@ -1,6 +1,7 @@
 # Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
 """Client and server classes corresponding to protobuf-defined services."""
 import grpc
+<<<<<<< HEAD
 import warnings
 
 
@@ -22,3 +23,5 @@
         + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
         + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
     )
+=======
+>>>>>>> 989b70e (Fix Copilot review issues for attach hint PR)

ydb/_grpc/v6/draft/protos/ydb_federated_query_pb2_grpc.py

@@ -1,6 +1,7 @@
 # Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
 """Client and server classes corresponding to protobuf-defined services."""
 import grpc
+<<<<<<< HEAD
 import warnings
 
 
@@ -22,3 +23,5 @@
         + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
         + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
     )
+=======
+>>>>>>> 989b70e (Fix Copilot review issues for attach hint PR)

ydb/_grpc/v6/draft/protos/ydb_keyvalue_pb2_grpc.py

@@ -1,6 +1,7 @@
 # Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
 """Client and server classes corresponding to protobuf-defined services."""
 import grpc
+<<<<<<< HEAD
 import warnings
 
 
@@ -22,3 +23,5 @@
         + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
         + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
     )
+=======
+>>>>>>> 989b70e (Fix Copilot review issues for attach hint PR)

ydb/_grpc/v6/draft/protos/ydb_maintenance_pb2_grpc.py

@@ -1,6 +1,7 @@
 # Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
 """Client and server classes corresponding to protobuf-defined services."""
 import grpc
+<<<<<<< HEAD
 import warnings
 
 
@@ -22,3 +23,5 @@
         + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
         + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
     )
+=======
+>>>>>>> 989b70e (Fix Copilot review issues for attach hint PR)

ydb/_grpc/v6/draft/protos/ydb_object_storage_pb2_grpc.py

@@ -1,6 +1,7 @@
 # Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
 """Client and server classes corresponding to protobuf-defined services."""
 import grpc
+<<<<<<< HEAD
 import warnings
 
 
@@ -22,3 +23,5 @@
         + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
         + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
     )
+=======
+>>>>>>> 989b70e (Fix Copilot review issues for attach hint PR)

ydb/_grpc/v6/draft/ydb_dynamic_config_v1_pb2_grpc.py

@@ -1,10 +1,10 @@
 # Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
 """Client and server classes corresponding to protobuf-defined services."""
 import grpc
-import warnings
 
 from ydb._grpc.v6.draft.protos import ydb_dynamic_config_pb2 as draft_dot_protos_dot_ydb__dynamic__config__pb2
 
+<<<<<<< HEAD
 GRPC_GENERATED_VERSION = '1.72.1'
 GRPC_VERSION = grpc.__version__
 _version_not_supported = False
@@ -23,6 +23,8 @@
         + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
         + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
     )
+=======
+>>>>>>> 989b70e (Fix Copilot review issues for attach hint PR)
 
 
 class DynamicConfigServiceStub(object):

ydb/_grpc/v6/draft/ydb_federated_query_v1_pb2_grpc.py

@@ -1,10 +1,10 @@
 # Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
 """Client and server classes corresponding to protobuf-defined services."""
 import grpc
-import warnings
 
 from ydb._grpc.v6.draft.protos import ydb_federated_query_pb2 as draft_dot_protos_dot_ydb__federated__query__pb2
 
+<<<<<<< HEAD
 GRPC_GENERATED_VERSION = '1.72.1'
 GRPC_VERSION = grpc.__version__
 _version_not_supported = False
@@ -23,6 +23,8 @@
         + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
         + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
     )
+=======
+>>>>>>> 989b70e (Fix Copilot review issues for attach hint PR)
 
 
 class FederatedQueryServiceStub(object):