修改消息相关索引

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")

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'),
+    )

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'),
+    )

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")
-    
-    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),
-            )
-        )
-    )
+    """Get notifications for a user (personal + broadcast), with per-user read state.
     
+    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)
     
-    query = query.order_by(Notification.created_at.desc())
+    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)
+    )
 
-    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()
+    # 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
+
+    # 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,
-            )
-        )
-    ).all()
-    read_ids = {row[0] for row in read_broadcast_ids}
-
-    broadcast_query = select(func.count(Notification.id)).where(
+            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)
                     )
-    if read_ids:
-        broadcast_query = broadcast_query.where(~Notification.id.in_(read_ids))
-
-    broadcast_count = (await db.execute(broadcast_query)).scalar() or 0
+                    .exists()
+                )
+            )
+        )
+    ).scalar() or 0
 
     return personal_count + broadcast_count