修改消息相关索引

This commit is contained in:
2026-06-29 09:09:58 +08:00
parent 627087797b
commit b21f35582e
4 changed files with 220 additions and 72 deletions
@@ -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")
+5 -1
View File
@@ -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 sqlalchemy.orm import Mapped, mapped_column
from app.models.base import Base, TimestampMixin from app.models.base import Base, TimestampMixin
@@ -16,3 +16,7 @@ class Notification(Base, TimestampMixin):
type: Mapped[str] = mapped_column(String(32), default="system") type: Mapped[str] = mapped_column(String(32), default="system")
is_read: Mapped[bool] = mapped_column(Boolean, default=False) is_read: Mapped[bool] = mapped_column(Boolean, default=False)
related_id: Mapped[str | None] = mapped_column(String(64), nullable=True) 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'),
)
@@ -1,4 +1,4 @@
from sqlalchemy import String, ForeignKey from sqlalchemy import String, ForeignKey, Index, UniqueConstraint
from sqlalchemy.orm import Mapped, mapped_column from sqlalchemy.orm import Mapped, mapped_column
from app.models.base import Base, TimestampMixin from app.models.base import Base, TimestampMixin
@@ -12,3 +12,8 @@ class NotificationRead(Base, TimestampMixin):
String(32), ForeignKey("notifications.id"), index=True String(32), ForeignKey("notifications.id"), index=True
) )
user_id: Mapped[str] = mapped_column(String(32), 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'),
)
+160 -70
View File
@@ -1,7 +1,7 @@
import logging import logging
from datetime import datetime, timezone, timedelta 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 sqlalchemy.ext.asyncio import AsyncSession
from app.models.notification import Notification from app.models.notification import Notification
@@ -76,66 +76,157 @@ async def create_notification(
async def get_notifications( async def get_notifications(
db: AsyncSession, user_id: str, page: int = 1, page_size: int = 20, is_read: bool = None db: AsyncSession, user_id: str, page: int = 1, page_size: int = 20, is_read: bool = None
) -> tuple[list[dict], int]: ) -> tuple[list[dict], int]:
"""Get notifications for a user (personal + broadcast), with per-user read state.""" """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 = ( Optimized query approach:
select( - Use NOT EXISTS for broadcast read status instead of LEFT JOIN + CASE
Notification, - Split conditions to enable index usage
user_is_read_expr, """
) # Personal notifications query (direct is_read column)
.outerjoin( personal_cond = Notification.user_id == user_id
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: if is_read is not None:
query = query.where(user_is_read_expr == is_read) personal_cond = personal_cond & (Notification.is_read == is_read)
count_query = count_query.where(user_is_read_expr == 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 # For broadcast notifications, user_is_read is True if exists in notification_reads
result = await db.execute(query.offset((page - 1) * page_size).limit(page_size)) broadcast_read_flag = (
rows = result.all() 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 = [] items = []
for notif, user_is_read in rows: for row in paginated:
items.append({ items.append({
"id": notif.id, "id": row.id,
"user_id": notif.user_id, "user_id": row.user_id,
"title": notif.title, "title": row.title,
"content": notif.content, "content": row.content,
"type": notif.type, "type": row.type,
"is_read": bool(user_is_read), "is_read": bool(row.user_is_read),
"related_id": notif.related_id, "related_id": row.related_id,
"created_at": notif.created_at, "created_at": row.created_at,
}) })
return items, total 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: async def get_unread_count(db: AsyncSession, user_id: str) -> int:
"""Get unread notification count for a user.""" """Get unread notification count for a user. Optimized with EXISTS."""
# Count personal unread # Count personal unread (uses composite index)
personal_count = ( personal_count = (
await db.execute( await db.execute(
select(func.count(Notification.id)).where( select(func.count(Notification.id)).where(
@@ -223,23 +314,22 @@ async def get_unread_count(db: AsyncSession, user_id: str) -> int:
) )
).scalar() or 0 ).scalar() or 0
# Count broadcast unread (not in notification_reads) # Count broadcast unread using NOT EXISTS (efficient, can use index on notification_reads)
read_broadcast_ids = ( broadcast_count = (
await db.execute( await db.execute(
select(NotificationRead.notification_id).where( select(func.count(Notification.id)).where(
NotificationRead.user_id == user_id, Notification.user_id.is_(None),
~(
select(NotificationRead.id)
.where(
(NotificationRead.notification_id == Notification.id)
& (NotificationRead.user_id == user_id)
)
.exists()
)
) )
) )
).all() ).scalar() or 0
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 return personal_count + broadcast_count