diff --git a/api/src/api/v1/websocket.py b/api/src/api/v1/websocket.py index 5fe1bce..274ddff 100644 --- a/api/src/api/v1/websocket.py +++ b/api/src/api/v1/websocket.py @@ -11,6 +11,7 @@ from sqlalchemy.exc import IntegrityError from fastapi import APIRouter, WebSocket, WebSocketDisconnect, status from sqlalchemy.ext.asyncio import AsyncSession +from src.repository.user_repository import UserRepository from src.services.budget_line_service import BudgetLineService from src.api.v1.deps import get_user_by_token from src.db.models.app_user import AppUser @@ -25,10 +26,13 @@ from src.services.user_service import UserService @asynccontextmanager -async def get_db_session(): +async def get_db_session(user_id: int | None = None): """Контекстный менеджер для получения сессии базы данных.""" db = SessionLocal() try: + if user_id is not None: + user_repo = UserRepository(db) + await user_repo.set_app_user_id(user_id=user_id) yield db finally: await db.close() @@ -603,8 +607,6 @@ async def login(websocket: WebSocket, **kwargs) -> int: return user.id - - def resolve_cell_key(event_data: dict) -> dict | None: data = event_data.get("data") @@ -738,7 +740,7 @@ async def process_websocket( case _: pass - async with get_db_session() as db: + async with get_db_session(user_id=user_id) as db: processor = processor_cls(db) data["result"] = await processor.process( event_data=data,