Optimize multiple SQLAlchemy queries to one query

Question

I am trying to write a function for a given user to return a count of all online users and offline users for each chatroom the given user is a part of.

The methods listed below achieve what I want but it writes multiple queries to postgres and I want to be able to do it in one query.

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

This seems to be working and is pretty close to example I made recently except the additional where clause to get the chatrooms of a single user. This might be able to be consolidated more but I think this is the most flexible version. For example, you could get the whole chatroom object as well if needed in the same query.

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

The secondary table I used looks like this:

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