From c6ef4ec01b3a6c089b56c3f8e0a0a8963bf61b5a Mon Sep 17 00:00:00 2001 From: Timofey Koolin Date: Thu, 21 Sep 2023 11:28:18 +0300 Subject: [PATCH 1/3] type hints to pool --- ydb/pool.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ydb/pool.py b/ydb/pool.py index 1f33cf8a..484ab1e9 100644 --- a/ydb/pool.py +++ b/ydb/pool.py @@ -9,7 +9,7 @@ from . import connection as connection_impl, issues, resolver, _utilities, tracing from abc import abstractmethod -from .connection import Connection +from .connection import Connection, EndpointKey logger = logging.getLogger(__name__) @@ -123,7 +123,7 @@ def subscribe(self): return subscription @tracing.with_trace() - def get(self, preferred_endpoint=None) -> Connection: + def get(self, preferred_endpoint: Optional[EndpointKey] = None) -> Connection: with self.lock: if preferred_endpoint is not None and preferred_endpoint.node_id in self.connections_by_node_id: return self.connections_by_node_id[preferred_endpoint.node_id] From c4a8d5f6331e7e86091a4141a77626a7663db1a7 Mon Sep 17 00:00:00 2001 From: Timofey Koolin Date: Thu, 21 Sep 2023 11:29:51 +0300 Subject: [PATCH 2/3] Update pool.py --- ydb/aio/pool.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/ydb/aio/pool.py b/ydb/aio/pool.py index fd4e8016..c637a7ca 100644 --- a/ydb/aio/pool.py +++ b/ydb/aio/pool.py @@ -1,11 +1,12 @@ import asyncio import logging import random +import typing from ydb import issues from ydb.pool import ConnectionsCache as _ConnectionsCache, IConnectionPool -from .connection import Connection +from .connection import Connection, EndpointKey from . import resolver @@ -21,7 +22,7 @@ def __init__(self, use_all_nodes: bool = False): self._fast_fail_error = None - async def get(self, preferred_endpoint=None, fast_fail=False, wait_timeout=10): + async def get(self, preferred_endpoint: typing.Optional[EndpointKey] = None, fast_fail=False, wait_timeout=10): if fast_fail: await asyncio.wait_for(self._fast_fail_event.wait(), timeout=wait_timeout) From d6e917661b5494af3f92ee65a9722bbc16c18ea3 Mon Sep 17 00:00:00 2001 From: Timofey Koolin Date: Thu, 21 Sep 2023 11:30:20 +0300 Subject: [PATCH 3/3] Update pool.py --- ydb/pool.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ydb/pool.py b/ydb/pool.py index 484ab1e9..1e75950e 100644 --- a/ydb/pool.py +++ b/ydb/pool.py @@ -5,6 +5,7 @@ from concurrent import futures import collections import random +import typing from . import connection as connection_impl, issues, resolver, _utilities, tracing from abc import abstractmethod @@ -123,7 +124,7 @@ def subscribe(self): return subscription @tracing.with_trace() - def get(self, preferred_endpoint: Optional[EndpointKey] = None) -> Connection: + def get(self, preferred_endpoint: typing.Optional[EndpointKey] = None) -> Connection: with self.lock: if preferred_endpoint is not None and preferred_endpoint.node_id in self.connections_by_node_id: return self.connections_by_node_id[preferred_endpoint.node_id]