66 lines
2.9 KiB
Python
66 lines
2.9 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
from collections.abc import Coroutine
|
||
from typing import Any, TypeVar
|
||
|
||
from sqlalchemy.ext.asyncio import AsyncAttrs, async_sessionmaker, create_async_engine
|
||
from sqlalchemy.orm import DeclarativeBase
|
||
from sqlalchemy.pool import NullPool
|
||
|
||
from app.core.config import settings
|
||
|
||
T = TypeVar("T")
|
||
|
||
connect_args: dict = {}
|
||
if settings.DATABASE_SCHEMA:
|
||
connect_args["server_settings"] = {"search_path": settings.DATABASE_SCHEMA}
|
||
if "sqlite" in settings.DATABASE_URL:
|
||
connect_args["check_same_thread"] = False
|
||
|
||
# ПОЧЕМУ NullPool И run_db.
|
||
# Streamlit синхронный: каждый вызов async-кода из UI идёт через отдельный
|
||
# asyncio.run(), который создаёт новый event loop и закрывает его по
|
||
# завершении. Соединения asyncpg привязаны к loop'у, в котором открыты:
|
||
# обычный пул отдал бы следующему asyncio.run() соединение из уже закрытого
|
||
# loop'а, и запрос упал бы с ошибкой, похожей на недоступность БД. В SFM-parser
|
||
# (../sfm/app/opencloud_token_store.py) это лечили engine.dispose() в finally
|
||
# каждой функции — легко забыть в новой функции.
|
||
# Здесь системно: NullPool — соединение открывается на сессию и закрывается
|
||
# вместе с ней, пула между loop'ами нет. При наших объёмах (единицы
|
||
# пользователей, десятки запросов в день) цена — лишний коннект на сессию —
|
||
# несущественна. Если нагрузка вырастет, первым делом смотреть сюда
|
||
# (PgBouncer перед Postgres или общий фоновый loop вместо asyncio.run).
|
||
engine = create_async_engine(
|
||
settings.DATABASE_URL,
|
||
poolclass=NullPool,
|
||
connect_args=connect_args,
|
||
)
|
||
|
||
SessionLocal = async_sessionmaker(
|
||
bind=engine,
|
||
autoflush=False,
|
||
expire_on_commit=False,
|
||
)
|
||
|
||
|
||
class Base(AsyncAttrs, DeclarativeBase):
|
||
pass
|
||
|
||
|
||
def run_db(coro: Coroutine[Any, Any, T]) -> T:
|
||
"""Единственная точка входа из синхронного кода (UI, джоба) в async-код БД.
|
||
|
||
Каждый вызов — новый event loop (asyncio.run). Безопасно только вместе с
|
||
NullPool (см. комментарий выше). Внутри уже работающего loop'а вызывать
|
||
нельзя — там нужно просто `await`.
|
||
"""
|
||
try:
|
||
asyncio.get_running_loop()
|
||
except RuntimeError:
|
||
return asyncio.run(coro)
|
||
coro.close()
|
||
raise RuntimeError(
|
||
"run_db() вызван внутри работающего event loop — используйте await."
|
||
)
|