schedule-management-system/backend/app/routers/ws.py

from fastapi import APIRouter, HTTPException, Query, WebSocket, WebSocketDisconnect, status
from sqlalchemy import select

from backend.app.core.dependencies import get_current_user_by_token
from backend.app.db.session import SessionLocal
from backend.app.models import RolePermission
from backend.app.services.center_sync import get_sync_status
from backend.app.services.sync_ws import sync_broadcaster


router = APIRouter()


@router.websocket("/ws/sync")
async def sync_status_socket(websocket: WebSocket, token: str | None = Query(default=None)):
    if not token:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        return

    async with SessionLocal() as db:
        try:
            user = await get_current_user_by_token(token, db)
        except HTTPException:
            await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
            return

        result = await db.execute(
            select(RolePermission.permission_code).where(RolePermission.role_id == user.role_id)
        )
        permissions = {row[0] for row in result.fetchall()}
        if "users.view" not in permissions:
            await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
            return

        await sync_broadcaster.connect(websocket)
        await sync_broadcaster.send_status(websocket, await get_sync_status(db))

        try:
            while True:
                await websocket.receive_text()
        except WebSocketDisconnect:
            pass
        finally:
            await sync_broadcaster.disconnect(websocket)