278 lines
8.3 KiB
Python
278 lines
8.3 KiB
Python
import logging
|
|
from datetime import datetime, timezone, timedelta
|
|
|
|
from sqlalchemy import select, func, update, case, or_
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.models.notification import Notification
|
|
from app.models.notification_read import NotificationRead
|
|
from app.utils.id_gen import generate_id
|
|
|
|
logger = logging.getLogger("videogen")
|
|
CST = timezone(timedelta(hours=8))
|
|
|
|
|
|
def _to_local_str(dt):
|
|
if dt is None:
|
|
return None
|
|
if dt.tzinfo and dt.tzinfo.utcoffset(None) == timedelta(0):
|
|
dt = dt.astimezone(CST)
|
|
return dt.replace(tzinfo=None).isoformat()
|
|
|
|
|
|
async def create_notification(
|
|
db: AsyncSession,
|
|
user_id: str | None,
|
|
title: str,
|
|
content: str,
|
|
notif_type: str = "system",
|
|
related_id: str | None = None,
|
|
push_ws: bool = True,
|
|
) -> Notification:
|
|
"""Create a notification. user_id=None means broadcast."""
|
|
notif = Notification(
|
|
id=generate_id(),
|
|
user_id=user_id,
|
|
title=title,
|
|
content=content,
|
|
type=notif_type,
|
|
related_id=related_id,
|
|
)
|
|
db.add(notif)
|
|
await db.flush()
|
|
|
|
# Push via WebSocket if available
|
|
if push_ws:
|
|
try:
|
|
from app.utils.redis import get_redis
|
|
|
|
redis = get_redis()
|
|
if redis:
|
|
import json
|
|
|
|
payload = json.dumps(
|
|
{
|
|
"event": "notification",
|
|
"data": {
|
|
"id": notif.id,
|
|
"title": notif.title,
|
|
"content": notif.content,
|
|
"type": notif.type,
|
|
"user_id": notif.user_id,
|
|
"created_at": _to_local_str(notif.created_at),
|
|
},
|
|
}
|
|
)
|
|
if notif.user_id:
|
|
await redis.publish(f"user:{notif.user_id}:notifications", payload)
|
|
else:
|
|
await redis.publish("broadcast:notifications", payload)
|
|
except Exception:
|
|
pass
|
|
|
|
return notif
|
|
|
|
|
|
async def get_notifications(
|
|
db: AsyncSession, user_id: str, page: int = 1, page_size: int = 20, is_read: bool = None
|
|
) -> tuple[list[dict], int]:
|
|
"""Get notifications for a user (personal + broadcast), with per-user read state."""
|
|
user_is_read_expr = case(
|
|
(Notification.user_id.is_(None), NotificationRead.id.is_not(None)),
|
|
else_=Notification.is_read,
|
|
).label("user_is_read")
|
|
|
|
query = (
|
|
select(
|
|
Notification,
|
|
user_is_read_expr,
|
|
)
|
|
.outerjoin(
|
|
NotificationRead,
|
|
(Notification.id == NotificationRead.notification_id)
|
|
& (NotificationRead.user_id == user_id),
|
|
)
|
|
.where(
|
|
or_(
|
|
Notification.user_id == user_id,
|
|
Notification.user_id.is_(None),
|
|
)
|
|
)
|
|
)
|
|
|
|
count_query = (
|
|
select(func.count(Notification.id))
|
|
.outerjoin(
|
|
NotificationRead,
|
|
(Notification.id == NotificationRead.notification_id)
|
|
& (NotificationRead.user_id == user_id),
|
|
)
|
|
.where(
|
|
or_(
|
|
Notification.user_id == user_id,
|
|
Notification.user_id.is_(None),
|
|
)
|
|
)
|
|
)
|
|
|
|
if is_read is not None:
|
|
query = query.where(user_is_read_expr == is_read)
|
|
count_query = count_query.where(user_is_read_expr == is_read)
|
|
|
|
query = query.order_by(Notification.created_at.desc())
|
|
|
|
total = (await db.execute(count_query)).scalar() or 0
|
|
result = await db.execute(query.offset((page - 1) * page_size).limit(page_size))
|
|
rows = result.all()
|
|
|
|
items = []
|
|
for notif, user_is_read in rows:
|
|
items.append({
|
|
"id": notif.id,
|
|
"user_id": notif.user_id,
|
|
"title": notif.title,
|
|
"content": notif.content,
|
|
"type": notif.type,
|
|
"is_read": bool(user_is_read),
|
|
"related_id": notif.related_id,
|
|
"created_at": notif.created_at,
|
|
})
|
|
|
|
return items, total
|
|
|
|
|
|
async def mark_read(db: AsyncSession, notification_id: str, user_id: str) -> None:
|
|
"""Mark a notification as read for a specific user."""
|
|
result = await db.execute(
|
|
select(Notification).where(Notification.id == notification_id).limit(1)
|
|
)
|
|
notif = result.scalar_one_or_none()
|
|
if not notif:
|
|
return
|
|
|
|
if notif.user_id is not None:
|
|
# Personal notification: update the row directly
|
|
if notif.user_id == user_id:
|
|
notif.is_read = True
|
|
await db.flush()
|
|
else:
|
|
# Broadcast notification: use notification_reads table
|
|
existing = await db.execute(
|
|
select(NotificationRead).where(
|
|
NotificationRead.notification_id == notification_id,
|
|
NotificationRead.user_id == user_id,
|
|
)
|
|
.limit(1)
|
|
)
|
|
if not existing.scalar_one_or_none():
|
|
db.add(NotificationRead(
|
|
id=generate_id(),
|
|
notification_id=notification_id,
|
|
user_id=user_id,
|
|
))
|
|
await db.flush()
|
|
|
|
|
|
async def mark_all_read(db: AsyncSession, user_id: str) -> None:
|
|
"""Mark all notifications as read for a user."""
|
|
# Mark personal notifications
|
|
await db.execute(
|
|
update(Notification)
|
|
.where(
|
|
Notification.user_id == user_id,
|
|
Notification.is_read == False,
|
|
)
|
|
.values(is_read=True)
|
|
)
|
|
|
|
# Mark all broadcast notifications
|
|
broadcasts = await db.execute(
|
|
select(Notification.id)
|
|
.where(Notification.user_id.is_(None))
|
|
)
|
|
broadcast_ids = [row[0] for row in broadcasts.all()]
|
|
|
|
if broadcast_ids:
|
|
existing_reads = await db.execute(
|
|
select(NotificationRead.notification_id).where(
|
|
NotificationRead.notification_id.in_(broadcast_ids),
|
|
NotificationRead.user_id == user_id,
|
|
)
|
|
)
|
|
already_read = {row[0] for row in existing_reads.all()}
|
|
|
|
for bid in broadcast_ids:
|
|
if bid not in already_read:
|
|
db.add(NotificationRead(
|
|
id=generate_id(),
|
|
notification_id=bid,
|
|
user_id=user_id,
|
|
))
|
|
await db.flush()
|
|
|
|
|
|
async def get_unread_count(db: AsyncSession, user_id: str) -> int:
|
|
"""Get unread notification count for a user."""
|
|
# Count personal unread
|
|
personal_count = (
|
|
await db.execute(
|
|
select(func.count(Notification.id)).where(
|
|
Notification.user_id == user_id,
|
|
Notification.is_read == False,
|
|
)
|
|
)
|
|
).scalar() or 0
|
|
|
|
# Count broadcast unread (not in notification_reads)
|
|
read_broadcast_ids = (
|
|
await db.execute(
|
|
select(NotificationRead.notification_id).where(
|
|
NotificationRead.user_id == user_id,
|
|
)
|
|
)
|
|
).all()
|
|
read_ids = {row[0] for row in read_broadcast_ids}
|
|
|
|
broadcast_query = select(func.count(Notification.id)).where(
|
|
Notification.user_id.is_(None),
|
|
)
|
|
if read_ids:
|
|
broadcast_query = broadcast_query.where(~Notification.id.in_(read_ids))
|
|
|
|
broadcast_count = (await db.execute(broadcast_query)).scalar() or 0
|
|
|
|
return personal_count + broadcast_count
|
|
|
|
|
|
async def get_notification_read_users(
|
|
db: AsyncSession, notification_id: str, page: int = 1, page_size: int = 50
|
|
) -> tuple[list[dict], int]:
|
|
"""Get list of users who have read a specific notification."""
|
|
from app.models.user import User
|
|
|
|
query = (
|
|
select(NotificationRead, User.username)
|
|
.join(User, NotificationRead.user_id == User.id)
|
|
.where(NotificationRead.notification_id == notification_id)
|
|
.order_by(NotificationRead.created_at.desc())
|
|
)
|
|
|
|
count_query = (
|
|
select(func.count(NotificationRead.id))
|
|
.where(NotificationRead.notification_id == notification_id)
|
|
)
|
|
|
|
total = (await db.execute(count_query)).scalar() or 0
|
|
result = await db.execute(query.offset((page - 1) * page_size).limit(page_size))
|
|
rows = result.all()
|
|
|
|
items = []
|
|
for read_record, username in rows:
|
|
items.append({
|
|
"user_id": read_record.user_id,
|
|
"username": username,
|
|
"read_at": read_record.created_at,
|
|
})
|
|
|
|
return items, total
|