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.distribution import Distribution, DistributionStatus, DistributionTargetType from app.models.document import Document from app.models.document_version import DocumentVersion from app.models.study import Study from app.models.study_member import StudyMember from app.schemas.notification import NotificationItem 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): role_target_exists = exists( select(StudyMember.id).where( StudyMember.study_id == Document.trial_id, StudyMember.user_id == user_id, StudyMember.is_active.is_(True), StudyMember.role_in_study == Distribution.target_id, ) ) target_matches = or_( and_( Distribution.target_type == DistributionTargetType.USER, Distribution.target_id == str(user_id), ), and_( Distribution.target_type == DistributionTargetType.ROLE, role_target_exists, ), ) return ( select( Distribution, Document, DocumentVersion, Study, DesktopNotificationDelivery, ) .join(Document, Distribution.document_id == Document.id) .join(DocumentVersion, Distribution.version_id == DocumentVersion.id) .join(Study, Document.trial_id == Study.id) .outerjoin( DesktopNotificationDelivery, and_( DesktopNotificationDelivery.distribution_id == Distribution.id, DesktopNotificationDelivery.user_id == user_id, ), ) .where( Distribution.status == DistributionStatus.ACTIVE, Distribution.created_at >= enabled_at, target_matches, or_( DesktopNotificationDelivery.id.is_(None), and_( DesktopNotificationDelivery.delivered_at.is_(None), or_( DesktopNotificationDelivery.claimed_at.is_(None), DesktopNotificationDelivery.claimed_at < lease_cutoff, ), ), ), ) .order_by(Distribution.created_at.asc()) .with_for_update(of=Distribution, skip_locked=True) ) async def claim_notifications( db: AsyncSession, user_id: uuid.UUID, limit: int, ) -> tuple[uuid.UUID | None, datetime | None, list[NotificationItem]]: subscription = await get_subscription(db, user_id) if not subscription or not subscription.enabled or not subscription.enabled_at: return None, None, [] 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[NotificationItem] = [] for distribution, document, version, study, delivery in rows: if delivery is None: delivery = DesktopNotificationDelivery( user_id=user_id, distribution_id=distribution.id, ) db.add(delivery) delivery.claim_token = token delivery.claimed_at = now items.append( NotificationItem( id=distribution.id, document_id=document.id, version_id=version.id, document_title=document.title, document_no=document.doc_no, version_no=version.version_no, change_summary=version.change_summary, effective_at=version.effective_at, created_at=distribution.created_at, study_id=study.id, study_name=study.name, delivered_at=delivery.delivered_at, read_at=delivery.read_at, ) ) 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.distribution_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, distribution_id: uuid.UUID, ) -> None: delivery = ( await db.execute( select(DesktopNotificationDelivery).where( DesktopNotificationDelivery.user_id == user_id, DesktopNotificationDelivery.distribution_id == distribution_id, ) ) ).scalar_one_or_none() if delivery is None: delivery = DesktopNotificationDelivery( user_id=user_id, distribution_id=distribution_id, ) db.add(delivery) delivery.read_at = datetime.now(timezone.utc) await db.commit()