Update pool.py

ydb-platform · Sep 21, 2023 · c4a8d5f · c4a8d5f
1 parent c6ef4ec
commit c4a8d5f
Showing 1 changed file with 3 additions and 2 deletions.
diff --git 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)