Merge pull request #540 from ydb-platform/fix_aio_creds

Fix auth credentials

Author: vgvoleg

tests/aio/test_credentials.py

@@ -36,9 +36,10 @@ async def test_yandex_service_account_credentials():
         tests.auth.test_credentials.PRIVATE_KEY,
         server.get_endpoint(),
     )
-    t = (await credentials.auth_metadata())[0][1]
+    t = await credentials.get_auth_token()
     assert t == "test_token"
     assert credentials.get_expire_time() <= 42
+
     server.stop()
 
 

ydb/_topic_reader/topic_reader_asyncio.py

@@ -516,7 +516,10 @@ async def _read_messages_loop(self):
     async def _update_token_loop(self):
         while True:
             await asyncio.sleep(self._update_token_interval)
-            await self._update_token(token=self._get_token_function())
+            token = self._get_token_function()
+            if asyncio.iscoroutine(token):
+                token = await token
+            await self._update_token(token=token)
 
     async def _update_token(self, token: str):
         await self._update_token_event.wait()

ydb/_topic_writer/topic_writer_asyncio.py

@@ -686,7 +686,10 @@ def write(self, messages: List[InternalMessage]):
     async def _update_token_loop(self):
         while True:
             await asyncio.sleep(self._update_token_interval)
-            await self._update_token(token=self._get_token_function())
+            token = self._get_token_function()
+            if asyncio.iscoroutine(token):
+                token = await token
+            await self._update_token(token=token)
 
     async def _update_token(self, token: str):
         await self._update_token_event.wait()

ydb/aio/credentials.py

@@ -1,11 +1,13 @@
-import time
-
 import abc
 import asyncio
 import logging
-from ydb import issues, credentials
+import time
+
+from ydb import credentials
+from ydb import issues
 
 logger = logging.getLogger(__name__)
+YDB_AUTH_TICKET_HEADER = "x-ydb-auth-ticket"
 
 
 class _OneToManyValue(object):
@@ -64,6 +66,12 @@ def __init__(self):
     async def _make_token_request(self):
         pass
 
+    async def get_auth_token(self) -> str:
+        for header, token in await self.auth_metadata():
+            if header == YDB_AUTH_TICKET_HEADER:
+                return token
+        return ""
+
     async def _refresh(self):
         current_time = time.time()
         self._log_refresh_start(current_time)