Skip to content

Commit

Permalink
Fix auth credentials
Browse files Browse the repository at this point in the history
  • Loading branch information
vgvoleg committed Dec 23, 2024
1 parent 0dfe8e7 commit 973dd4e
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 5 deletions.
5 changes: 4 additions & 1 deletion ydb/_topic_reader/topic_reader_asyncio.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
5 changes: 4 additions & 1 deletion ydb/_topic_writer/topic_writer_asyncio.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
14 changes: 11 additions & 3 deletions ydb/aio/credentials.py
Original file line number Diff line number Diff line change
@@ -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):
Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit 973dd4e

Please sign in to comment.