0bdb83d2f3
Each call into opencloud_token_store is wrapped in its own asyncio.run() (app/app_config.py, main.py's token popover), so every call gets a fresh event loop. The asyncpg connection pool is bound to the loop it was created in, so the next asyncio.run() could try to reuse a connection from an already-closed loop and fail — unrelated to actual DB availability, but the broad except Exception in resolve_opencloud_access_token() silently swallowed it and fell through to "token not found", masking the real cause. - add_token/get_token/remove_token/list_usernames now call engine.dispose() in a finally block so each asyncio.run() starts with a clean pool. - resolve_opencloud_access_token() now surfaces the real exception via st.error() before falling back to the .env token (kept only for local dev without a DB). - Added tests/unit/test_opencloud_token_store.py covering that dispose() is called after every store operation, including when commit() raises.
108 lines
4.4 KiB
Python
108 lines
4.4 KiB
Python
from __future__ import annotations
|
||
|
||
from cryptography.fernet import Fernet, InvalidToken
|
||
from sqlalchemy import select
|
||
|
||
from app.core.config import settings
|
||
from app.db.base import SessionLocal, engine
|
||
from app.db.models.opencloud_token import OpenCloudToken
|
||
|
||
|
||
class OpenCloudTokenStoreError(RuntimeError):
|
||
"""Ошибка при работе с таблицей токенов OpenCloud (шифрование/БД)."""
|
||
|
||
|
||
# engine.dispose() в finally каждой функции ниже — обязателен. Приложение
|
||
# синхронное (Streamlit), каждый вызов этого модуля оборачивается в
|
||
# отдельный asyncio.run() (см. app/app_config.py, main.py), который создаёт
|
||
# новый event loop и закрывает его по завершении. Пул соединений asyncpg
|
||
# привязан к loop'у, в котором был создан: без dispose() следующий
|
||
# asyncio.run() (уже в новом loop'е) попытается забрать из пула соединение,
|
||
# открытое в уже закрытом loop'е предыдущего вызова, и упадёт — эта ошибка
|
||
# не имеет отношения к реальной доступности БД, но выглядит как она.
|
||
|
||
|
||
def _get_fernet() -> Fernet:
|
||
if not settings.FERNET_KEY:
|
||
raise OpenCloudTokenStoreError(
|
||
"Не задан OPENBAO__SETTINGS__FERNET_KEY — шифрование токенов "
|
||
"OpenCloud недоступно."
|
||
)
|
||
try:
|
||
return Fernet(settings.FERNET_KEY.encode())
|
||
except (ValueError, TypeError) as exc:
|
||
raise OpenCloudTokenStoreError(
|
||
f"OPENBAO__SETTINGS__FERNET_KEY некорректен: {exc}"
|
||
) from exc
|
||
|
||
|
||
async def add_token(username: str, token: str, added_by: str | None = None) -> None:
|
||
"""Сохраняет (или перезаписывает) App Token заказчика, зашифровав его."""
|
||
if not username:
|
||
raise OpenCloudTokenStoreError("Не задан username для сохранения токена.")
|
||
if not token:
|
||
raise OpenCloudTokenStoreError("Не задан токен для сохранения.")
|
||
|
||
encrypted = _get_fernet().encrypt(token.encode()).decode()
|
||
|
||
try:
|
||
async with SessionLocal() as session:
|
||
obj = await session.get(OpenCloudToken, username)
|
||
if obj is not None:
|
||
obj.encrypted_token = encrypted
|
||
if added_by:
|
||
obj.added_by = added_by
|
||
else:
|
||
session.add(
|
||
OpenCloudToken(
|
||
username=username, encrypted_token=encrypted, added_by=added_by
|
||
)
|
||
)
|
||
await session.commit()
|
||
finally:
|
||
await engine.dispose()
|
||
|
||
|
||
async def get_token(username: str) -> str | None:
|
||
"""username -> расшифрованный App Token, либо None, если записи нет."""
|
||
if not username:
|
||
return None
|
||
try:
|
||
async with SessionLocal() as session:
|
||
obj = await session.get(OpenCloudToken, username)
|
||
finally:
|
||
await engine.dispose()
|
||
if obj is None:
|
||
return None
|
||
try:
|
||
return _get_fernet().decrypt(obj.encrypted_token.encode()).decode()
|
||
except InvalidToken as exc:
|
||
raise OpenCloudTokenStoreError(
|
||
f"Не удалось расшифровать токен для '{username}' — "
|
||
"проверьте OPENBAO__SETTINGS__FERNET_KEY."
|
||
) from exc
|
||
|
||
|
||
async def remove_token(username: str) -> bool:
|
||
"""Удаляет запись пользователя. Возвращает True, если запись была найдена."""
|
||
try:
|
||
async with SessionLocal() as session:
|
||
obj = await session.get(OpenCloudToken, username)
|
||
if obj is None:
|
||
return False
|
||
await session.delete(obj)
|
||
await session.commit()
|
||
return True
|
||
finally:
|
||
await engine.dispose()
|
||
|
||
|
||
async def list_usernames() -> list[str]:
|
||
"""Список username, для которых сохранён токен (без расшифровки)."""
|
||
try:
|
||
async with SessionLocal() as session:
|
||
result = await session.execute(select(OpenCloudToken.username))
|
||
return list(result.scalars().all())
|
||
finally:
|
||
await engine.dispose()
|