Spaces:
Sleeping
Sleeping
File size: 5,684 Bytes
b2be963 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.orm import Session
from pydantic import BaseModel
from typing import Optional
from datetime import datetime, timezone
from app.db.base import get_db
from app.models.notification import Notification
from app.api.auth import get_current_user
from app.core.security import get_current_user_sse
from app.models.user import User, UserRole
from app.core.notifications import notification_manager
from fastapi.responses import StreamingResponse
import asyncio
import json
router = APIRouter(prefix="/notifications", tags=["Notifications"])
def _check_admin(user_id: int, db: Session):
user = db.query(User).filter(User.id == user_id).first()
if not user or user.role != UserRole.ADMIN:
raise HTTPException(status_code=403, detail="Admin access required")
return user
# ── Schemas ───────────────────────────────────────────────────
class NotificationCreate(BaseModel):
type: str # checkout_shipping | checkout_payment | checkout_complete
title: str
message: Optional[str] = None
data: Optional[dict] = None
class NotificationOut(BaseModel):
id: int
type: str
title: str
message: Optional[str]
data: Optional[dict]
is_read: bool
created_at: datetime
class Config:
from_attributes = True
# ── POST /notifications — create (no auth, called from checkout) ──
@router.post("/")
async def create_notification(payload: NotificationCreate, db: Session = Depends(get_db)):
notif = Notification(
type=payload.type,
title=payload.title,
message=payload.message,
data=payload.data,
is_read=False,
created_at=datetime.now(timezone.utc),
)
db.add(notif)
db.commit()
db.refresh(notif)
# Broadcast to active SSE clients
await notification_manager.broadcast(
NotificationOut.model_validate(notif).model_dump()
)
return {
"isSuccess": True,
"value": {"id": notif.id},
"statusCode": 201,
}
# ── GET /notifications — list for admin ──
@router.get("/")
async def list_notifications(
skip: int = 0,
limit: int = 50,
unread_only: bool = False,
current_user_id: int = Depends(get_current_user),
db: Session = Depends(get_db),
):
_check_admin(current_user_id, db)
query = db.query(Notification)
if unread_only:
query = query.filter(Notification.is_read == False) # noqa: E712
query = query.order_by(Notification.created_at.desc())
total = query.count()
notifications = query.offset(skip).limit(limit).all()
return {
"isSuccess": True,
"value": {
"notifications": [
NotificationOut.model_validate(n).model_dump() for n in notifications
],
"total": total,
"unread_count": db.query(Notification).filter(Notification.is_read == False).count(), # noqa: E712
},
"statusCode": 200,
}
# ── PUT /notifications/{id}/read — mark single as read ──
@router.put("/{notification_id}/read")
async def mark_read(
notification_id: int,
current_user_id: int = Depends(get_current_user),
db: Session = Depends(get_db),
):
_check_admin(current_user_id, db)
notif = db.query(Notification).filter(Notification.id == notification_id).first()
if not notif:
raise HTTPException(status_code=404, detail="Notification not found")
notif.is_read = True
db.commit()
return {"isSuccess": True, "value": None, "statusCode": 200}
# ── PUT /notifications/read-all — mark all as read ──
@router.put("/read-all")
async def mark_all_read(
current_user_id: int = Depends(get_current_user),
db: Session = Depends(get_db),
):
_check_admin(current_user_id, db)
db.query(Notification).filter(Notification.is_read == False).update( # noqa: E712
{"is_read": True}
)
db.commit()
return {"isSuccess": True, "value": None, "statusCode": 200}
# ── GET /notifications/stream — SSE stream for admin ──
@router.get("/stream")
async def stream_notifications(
current_user_id: int = Depends(get_current_user_sse),
db: Session = Depends(get_db),
):
"""
SSE endpoint for real-time notifications.
Includes a heartbeat to keep the connection alive on Hugging Face / Vercel.
"""
_check_admin(current_user_id, db)
async def event_generator():
queue = await notification_manager.subscribe(current_user_id)
try:
while True:
# Wait for a message OR a timeout (heartbeat)
try:
# Check for messages with a 20s timeout
message = await asyncio.wait_for(queue.get(), timeout=20.0)
yield f"data: {message}\n\n"
except asyncio.TimeoutError:
# Send a heartbeat comment to keep the connection alive
yield ": heartbeat\n\n"
finally:
await notification_manager.unsubscribe(current_user_id, queue)
return StreamingResponse(
event_generator(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no", # Disable buffering for Nginx/Proxies
},
)
|