将多个SQLAlchemy查询优化为一个查询

Question

我正在尝试写一个函数，目的是为特定用户返回他们参与的每个聊天室中在线用户和离线用户的数量。

下面列出的方法可以实现我想要的功能，但它会向Postgres数据库发送多条查询，我希望能够通过一条查询来完成。

class Users(Base):
    id: Mapped[uuid:UUID] = mapped_column(
        UUID(as_uuid=True), unique=True, default=uuid.uuid4, autoincrement=True
    )
    username: Mapped[str] = mapped_column(String, primary_key=True, unique=True)
    password: Mapped[str] = mapped_column(String, nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, default=datetime.now
    )

    sessions: Mapped[List["Sessions"]] = relationship(
        back_populates="user", lazy="selectin"
    )

    messages: Mapped[List["Messages"]] = relationship(
        back_populates="user", lazy="selectin"
    )

    created_chatrooms: Mapped[List["Chatrooms"]] = relationship(
        back_populates="created_by", lazy="selectin"
    )

    chatrooms: Mapped[List["Chatrooms"]] = relationship(
        secondary=user_chatroom_table, back_populates="users"
    )
    online: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)

    async def get_chatrooms_and_user_counts(self, db: AsyncSession):
        chatroom_ids = await self.get_user_chatrooms(db=db)
        chatrooms_info = []
        for chatroom_id in chatroom_ids:
            online_count_stmt = (
                select(func.count())
                .select_from(
                    user_chatroom_table.join(
                        Users, user_chatroom_table.c.user_id == Users.id
                    )
                )
                .where(
                    user_chatroom_table.c.chatroom_id == chatroom_id,
                    Users.online == True,
                )
            )
            offline_count_stmt = (
                select(func.count())
                .select_from(
                    user_chatroom_table.join(
                        Users, user_chatroom_table.c.user_id == Users.id
                    )
                )
                .where(
                    user_chatroom_table.c.chatroom_id == chatroom_id,
                    Users.online == False,
                )
            )

            online_count_result = await db.execute(online_count_stmt)
            offline_count_result = await db.execute(offline_count_stmt)

            online_count = online_count_result.scalar()
            offline_count = offline_count_result.scalar()

            chatrooms_info.append(
                {
                    "chatroom_id": str(chatroom_id),
                    "online": online_count,
                    "offline": offline_count,
                }
            )
        return chatrooms_info

    async def get_user_chatrooms(self, db: AsyncSession) -> List["Chatrooms"]:
        """
        Asynchronously get a list of chatrooms that the given user is a part of.

        :param user_id: UUID of the user to check.
        :param db: SQLAlchemy AsyncSession.
        :return: List of Chatrooms the user is a part of.
        """
        # Create a query that selects chatrooms joined by the given user_id
        stmt = (
            select(Chatrooms.id)
            .join(
                user_chatroom_table, Chatrooms.id == user_chatroom_table.c.chatroom_id
            )
            .where(user_chatroom_table.c.user_id == self.id)
        )

        result = await db.execute(stmt)
        chatrooms = result.scalars().all()

        return chatrooms


class Chatrooms(Base):
    id: Mapped[uuid:UUID] = mapped_column(
        UUID(as_uuid=True),
        primary_key=True,
        unique=True,
        default=uuid.uuid4,
    )

    name: Mapped[str] = mapped_column(unique=True, nullable=False)

    created_by: Mapped["Users"] = relationship(
        back_populates="created_chatrooms", lazy="selectin"
    )

    user_id: Mapped[uuid:UUID] = mapped_column(
        UUID(as_uuid=True), ForeignKey("users.id")
    )
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, default=datetime.now
    )
    category: Mapped[str] = mapped_column(nullable=False)

    messages: Mapped[List["Messages"]] = relationship(
        back_populates="chatroom", lazy="selectin"
    )
    users: Mapped[List[Users]] = relationship(
        secondary=user_chatroom_table, back_populates="chatrooms"
    )

Answer 1

这个看起来是有效的，而且和我最近做的例子非常相似，只是多了一个条件，用来获取某个用户的聊天房间。虽然这个代码可能还可以进一步简化，但我觉得这是最灵活的版本。比如说，如果需要的话，你也可以在同一个查询中获取整个聊天房间的对象。

async def get_chatroom_counts_for_user(session, current_user_id):
    # Count up online users per chatroom.
    online_subq = select(
        user_chatroom_table.c.chatroom_id,
        func.count(user_chatroom_table.c.user_id).label('online_count')
    ).join(Users,
        user_chatroom_table.c.user_id == Users.id
    ).where(
        Users.online == True
    ).group_by(
        user_chatroom_table.c.chatroom_id
    ).subquery()

    # Count up online users per chatroom.
    offline_subq = select(
        user_chatroom_table.c.chatroom_id,
        func.count(user_chatroom_table.c.user_id).label('offline_count')
    ).join(Users,
        user_chatroom_table.c.user_id == Users.id
    ).where(
        Users.online == False
    ).group_by(
        user_chatroom_table.c.chatroom_id
    ).subquery()

    # All the chatrooms this user belongs to.
    current_user_chatroom_subq = select(
        user_chatroom_table.c.chatroom_id
    ).where(
        user_chatroom_table.c.user_id == current_user_id
    )
    # Now get chatrooms with those counts.
    # We use outerjoin because some chatrooms have no matching
    # online and/or offline count and we want those chatrooms
    # to still be included.
    # When those chatrooms are included that missing count would
    # normally be NULL/None but we use coalesce() to tell
    # the database to convert that value to 0.
    q = select(
        Chatrooms.id,
        func.coalesce(online_subq.c.online_count, 0).label('online_count'),
        func.coalesce(offline_subq.c.offline_count, 0).label('offline_count')
    ).outerjoin(
        online_subq, Chatrooms.id == online_subq.c.chatroom_id
    ).outerjoin(
        offline_subq, Chatrooms.id == offline_subq.c.chatroom_id
    ).where(
        Chatrooms.id.in_(current_user_chatroom_subq))
    return [{
        "chatroom_id": chatroom_id,
        "online_count": online_count,
        "offline_count": offline_count
     } for chatroom_id, online_count, offline_count in (await session.execute(q))]

我使用的第二个表格长这样：

user_chatroom_table = Table('users_chatrooms', Base.metadata,
    Column("id", Integer, primary_key=True),
    Column("chatroom_id", Integer, ForeignKey('chatrooms.id')),
    Column("user_id", Integer, ForeignKey('users.id')),
    UniqueConstraint('chatroom_id', 'user_id'),
)