from __future__ import annotations import uuid from datetime import datetime, timedelta, timezone from sqlalchemy import and_, exists, or_, select, update from sqlalchemy.ext.asyncio import AsyncSession from app.models.desktop_notification import ( DesktopNotificationDelivery, DesktopNotificationSubscription, ) from app.models.notification import Notification from app.models.study_member import StudyMember from app.models.user import User CLAIM_LEASE = timedelta(minutes=5) async def get_subscription(db: AsyncSession, user_id: uuid.UUID) -> DesktopNotificationSubscription | None: return await db.get(DesktopNotificationSubscription, user_id) async def set_subscription( db: AsyncSession, user_id: uuid.UUID, enabled: bool, ) -> DesktopNotificationSubscription: subscription = await get_subscription(db, user_id) now = datetime.now(timezone.utc) if subscription is None: subscription = DesktopNotificationSubscription( user_id=user_id, enabled=enabled, enabled_at=now if enabled else None, ) db.add(subscription) elif subscription.enabled != enabled: subscription.enabled = enabled subscription.enabled_at = now if enabled else None await db.commit() await db.refresh(subscription) return subscription def _eligible_query(user_id: uuid.UUID, enabled_at: datetime, lease_cutoff: datetime): active_membership = exists( select(StudyMember.id).where( StudyMember.study_id == Notification.study_id, StudyMember.user_id == user_id, StudyMember.is_active.is_(True), ) ) return ( select(Notification, DesktopNotificationDelivery) .outerjoin( DesktopNotificationDelivery, and_( DesktopNotificationDelivery.notification_id == Notification.id, DesktopNotificationDelivery.user_id == user_id, ), ) .where( Notification.recipient_id == user_id, Notification.created_at >= enabled_at, Notification.resolved_at.is_(None), Notification.read_at.is_(None), active_membership, or_( DesktopNotificationDelivery.id.is_(None), and_( DesktopNotificationDelivery.delivered_at.is_(None), or_( DesktopNotificationDelivery.claimed_at.is_(None), DesktopNotificationDelivery.claimed_at < lease_cutoff, ), ), ), ) .order_by(Notification.created_at.asc()) .with_for_update(of=Notification, skip_locked=True) ) async def claim_notifications( db: AsyncSession, user_id: uuid.UUID, limit: int, ) -> tuple[uuid.UUID | None, datetime | None, list[Notification]]: subscription = await get_subscription(db, user_id) if not subscription or not subscription.enabled or not subscription.enabled_at: return None, None, [] # Reconcile every active project before selecting desktop deliveries so permission # changes fail closed and time-based reminders do not depend on opening the web Feed. user = await db.get(User, user_id) if user is None or not user.is_active: return None, None, [] study_ids = (await db.scalars( select(StudyMember.study_id).where( StudyMember.user_id == user_id, StudyMember.is_active.is_(True), ) )).all() from app.services.project_reminder_service import sync_project_reminders for study_id in set(study_ids): await sync_project_reminders(db, study_id, user) now = datetime.now(timezone.utc) token = uuid.uuid4() rows = ( await db.execute( _eligible_query(user_id, subscription.enabled_at, now - CLAIM_LEASE) .limit(max(1, min(limit, 50))) ) ).all() items: list[Notification] = [] for notification, delivery in rows: if delivery is None: delivery = DesktopNotificationDelivery( user_id=user_id, notification_id=notification.id, ) db.add(delivery) delivery.claim_token = token delivery.claimed_at = now items.append(notification) await db.commit() if not items: return None, None, [] return token, now + CLAIM_LEASE, items async def acknowledge_notifications( db: AsyncSession, user_id: uuid.UUID, claim_token: uuid.UUID, delivered_ids: list[uuid.UUID], ) -> None: if delivered_ids: await db.execute( update(DesktopNotificationDelivery) .where( DesktopNotificationDelivery.user_id == user_id, DesktopNotificationDelivery.claim_token == claim_token, DesktopNotificationDelivery.notification_id.in_(delivered_ids), ) .values(delivered_at=datetime.now(timezone.utc)) ) await db.commit() async def mark_notification_read( db: AsyncSession, user_id: uuid.UUID, notification_id: uuid.UUID, ) -> None: item = await db.scalar(select(Notification).where( Notification.id == notification_id, Notification.recipient_id == user_id, Notification.resolved_at.is_(None), )) if item is None: return now = datetime.now(timezone.utc) item.read_at = item.read_at or now if not item.requires_action: item.resolved_at = item.resolved_at or now await db.commit()