From b21f35582e550b4d1a7ba7c93153f288f4bfdaa5 Mon Sep 17 00:00:00 2001 From: wwwwwwwww <526125649@qq.com> Date: Mon, 29 Jun 2026 09:09:58 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E6=B6=88=E6=81=AF=E7=9B=B8?= =?UTF-8?q?=E5=85=B3=E7=B4=A2=E5=BC=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../n123456789ab_add_notification_indexes.py | 49 ++++ video-gen-api/app/models/notification.py | 6 +- video-gen-api/app/models/notification_read.py | 7 +- video-gen-api/app/services/notification.py | 230 ++++++++++++------ 4 files changed, 220 insertions(+), 72 deletions(-) create mode 100644 video-gen-api/alembic/versions/n123456789ab_add_notification_indexes.py diff --git a/video-gen-api/alembic/versions/n123456789ab_add_notification_indexes.py b/video-gen-api/alembic/versions/n123456789ab_add_notification_indexes.py new file mode 100644 index 00000000..eecbe7e3 --- /dev/null +++ b/video-gen-api/alembic/versions/n123456789ab_add_notification_indexes.py @@ -0,0 +1,49 @@ +"""add notification indexes for query optimization + +Revision ID: n123456789ab +Revises: c72a6f69e641 +Create Date: 2026-06-29 12:00:00.000000 +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = "n123456789ab" +down_revision: Union[str, None] = "c72a6f69e641" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # Composite index for personal notifications: (user_id, is_read, created_at) + op.create_index( + "idx_notif_user_isread_created", + "notifications", + ["user_id", "is_read", "created_at"], + unique=False, + ) + + # Composite index for notification_reads lookups + op.create_index( + "idx_notif_read_user_notif", + "notification_reads", + ["user_id", "notification_id"], + unique=False, + ) + + # Unique constraint to prevent duplicate read records + # Check if constraint already exists first (in case of partial migration) + op.create_unique_constraint( + "uq_notif_read_user", + "notification_reads", + ["notification_id", "user_id"], + ) + + +def downgrade() -> None: + op.drop_constraint("uq_notif_read_user", "notification_reads", type_="unique") + op.drop_index("idx_notif_read_user_notif", table_name="notification_reads") + op.drop_index("idx_notif_user_isread_created", table_name="notifications") diff --git a/video-gen-api/app/models/notification.py b/video-gen-api/app/models/notification.py index 660d8476..9b084705 100644 --- a/video-gen-api/app/models/notification.py +++ b/video-gen-api/app/models/notification.py @@ -1,4 +1,4 @@ -from sqlalchemy import Boolean, ForeignKey, String, Text +from sqlalchemy import Boolean, ForeignKey, String, Text, Index from sqlalchemy.orm import Mapped, mapped_column from app.models.base import Base, TimestampMixin @@ -16,3 +16,7 @@ class Notification(Base, TimestampMixin): type: Mapped[str] = mapped_column(String(32), default="system") is_read: Mapped[bool] = mapped_column(Boolean, default=False) related_id: Mapped[str | None] = mapped_column(String(64), nullable=True) + + __table_args__ = ( + Index('idx_notif_user_isread_created', 'user_id', 'is_read', 'created_at'), + ) diff --git a/video-gen-api/app/models/notification_read.py b/video-gen-api/app/models/notification_read.py index 9b0cbe79..743d7402 100644 --- a/video-gen-api/app/models/notification_read.py +++ b/video-gen-api/app/models/notification_read.py @@ -1,4 +1,4 @@ -from sqlalchemy import String, ForeignKey +from sqlalchemy import String, ForeignKey, Index, UniqueConstraint from sqlalchemy.orm import Mapped, mapped_column from app.models.base import Base, TimestampMixin @@ -12,3 +12,8 @@ class NotificationRead(Base, TimestampMixin): String(32), ForeignKey("notifications.id"), index=True ) user_id: Mapped[str] = mapped_column(String(32), index=True) + + __table_args__ = ( + UniqueConstraint('notification_id', 'user_id', name='uq_notif_read_user'), + Index('idx_notif_read_user_notif', 'user_id', 'notification_id'), + ) diff --git a/video-gen-api/app/services/notification.py b/video-gen-api/app/services/notification.py index 36c82cee..0c2ba19a 100644 --- a/video-gen-api/app/services/notification.py +++ b/video-gen-api/app/services/notification.py @@ -1,7 +1,7 @@ import logging from datetime import datetime, timezone, timedelta -from sqlalchemy import select, func, update, case, or_ +from sqlalchemy import select, func, update from sqlalchemy.ext.asyncio import AsyncSession from app.models.notification import Notification @@ -76,66 +76,157 @@ async def create_notification( 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") + """Get notifications for a user (personal + broadcast), with per-user read state. - 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), - ) - ) - ) - + Optimized query approach: + - Use NOT EXISTS for broadcast read status instead of LEFT JOIN + CASE + - Split conditions to enable index usage + """ + # Personal notifications query (direct is_read column) + personal_cond = Notification.user_id == user_id 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) + personal_cond = personal_cond & (Notification.is_read == is_read) + + personal_query = ( + select( + Notification.id, + Notification.user_id, + Notification.title, + Notification.content, + Notification.type, + Notification.is_read.label("user_is_read"), + Notification.related_id, + Notification.created_at, + Notification.updated_at, + ) + .where(personal_cond) + ) - query = query.order_by(Notification.created_at.desc()) + # Broadcast notifications query (use NOT EXISTS for unread check) + broadcast_base = Notification.user_id.is_(None) + + if is_read is True: + # Read broadcasts: EXISTS in notification_reads + broadcast_cond = broadcast_base & ( + select(NotificationRead.id) + .where( + (NotificationRead.notification_id == Notification.id) + & (NotificationRead.user_id == user_id) + ) + .exists() + ) + elif is_read is False: + # Unread broadcasts: NOT EXISTS in notification_reads + broadcast_cond = broadcast_base & ~( + select(NotificationRead.id) + .where( + (NotificationRead.notification_id == Notification.id) + & (NotificationRead.user_id == user_id) + ) + .exists() + ) + else: + broadcast_cond = broadcast_base - 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() + # For broadcast notifications, user_is_read is True if exists in notification_reads + broadcast_read_flag = ( + select(NotificationRead.id) + .where( + (NotificationRead.notification_id == Notification.id) + & (NotificationRead.user_id == user_id) + ) + .exists() + .label("user_is_read") + ) + + broadcast_query = ( + select( + Notification.id, + Notification.user_id, + Notification.title, + Notification.content, + Notification.type, + broadcast_read_flag, + Notification.related_id, + Notification.created_at, + Notification.updated_at, + ) + .where(broadcast_cond) + ) + + # Count queries - optimized + personal_count_query = select(func.count(Notification.id)).where(personal_cond) + + if is_read is True: + broadcast_count_cond = broadcast_base & ( + select(NotificationRead.id) + .where( + (NotificationRead.notification_id == Notification.id) + & (NotificationRead.user_id == user_id) + ) + .exists() + ) + elif is_read is False: + broadcast_count_cond = broadcast_base & ~( + select(NotificationRead.id) + .where( + (NotificationRead.notification_id == Notification.id) + & (NotificationRead.user_id == user_id) + ) + .exists() + ) + else: + broadcast_count_cond = broadcast_base + + broadcast_count_query = select(func.count(Notification.id)).where(broadcast_count_cond) + + # Get counts + personal_count = (await db.execute(personal_count_query)).scalar() or 0 + broadcast_count = (await db.execute(broadcast_count_query)).scalar() or 0 + total = personal_count + broadcast_count + + # Get paginated results: fetch enough from both sides to merge and sort + # Fetch up to page_size from each, then merge, sort, and slice in memory + # This is efficient because page sizes are small (20-50) + offset = (page - 1) * page_size + fetch_limit = offset + page_size + + personal_items = [] + if personal_count > 0: + p_result = await db.execute( + personal_query + .order_by(Notification.created_at.desc()) + .limit(fetch_limit) + ) + personal_items = p_result.all() + + broadcast_items = [] + if broadcast_count > 0: + b_result = await db.execute( + broadcast_query + .order_by(Notification.created_at.desc()) + .limit(fetch_limit) + ) + broadcast_items = b_result.all() + + # Merge and sort by created_at descending + all_items = list(personal_items) + list(broadcast_items) + all_items.sort(key=lambda x: x.created_at, reverse=True) + + # Apply pagination + paginated = all_items[offset:offset + page_size] items = [] - for notif, user_is_read in rows: + for row in paginated: 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, + "id": row.id, + "user_id": row.user_id, + "title": row.title, + "content": row.content, + "type": row.type, + "is_read": bool(row.user_is_read), + "related_id": row.related_id, + "created_at": row.created_at, }) return items, total @@ -212,8 +303,8 @@ async def mark_all_read(db: AsyncSession, user_id: str) -> None: async def get_unread_count(db: AsyncSession, user_id: str) -> int: - """Get unread notification count for a user.""" - # Count personal unread + """Get unread notification count for a user. Optimized with EXISTS.""" + # Count personal unread (uses composite index) personal_count = ( await db.execute( select(func.count(Notification.id)).where( @@ -223,23 +314,22 @@ async def get_unread_count(db: AsyncSession, user_id: str) -> int: ) ).scalar() or 0 - # Count broadcast unread (not in notification_reads) - read_broadcast_ids = ( + # Count broadcast unread using NOT EXISTS (efficient, can use index on notification_reads) + broadcast_count = ( await db.execute( - select(NotificationRead.notification_id).where( - NotificationRead.user_id == user_id, + select(func.count(Notification.id)).where( + Notification.user_id.is_(None), + ~( + select(NotificationRead.id) + .where( + (NotificationRead.notification_id == Notification.id) + & (NotificationRead.user_id == user_id) + ) + .exists() + ) ) ) - ).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 + ).scalar() or 0 return personal_count + broadcast_count