many2many-select.py

import asyncio
from sqlalchemy import create_engine, Column, Integer, String, ForeignKey
from sqlalchemy.orm import declarative_base, relationship, selectin_polymorphic
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import Mapped
from sqlalchemy.future import select
from sqlalchemy.ext.asyncio import (
    AsyncSession,
    create_async_engine,
    async_sessionmaker,
    async_scoped_session,
)

DATABASE_URL = "sqlite+aiosqlite:///:memory:"

Base = declarative_base()


class GroupUserLink(Base):
    __tablename__ = "group_user_link"
    group_id: int = Column(Integer, ForeignKey("groups.id"), primary_key=True)
    user_id: int = Column(Integer, ForeignKey("users.id"), primary_key=True)


class User(Base):
    __tablename__ = "users"
    id: int = Column(Integer, primary_key=True)
    name: str = Column(String(50))
    groups: Mapped[list["Group"]] = relationship(
        "Group",
        secondary="group_user_link",
        # lazy="selectin",
        back_populates="users",
    )

    def __str__(self):
        return f"Usr:\"{self.name}\""

    def __repr__(self):
        return f"<Usr:\"{self.name}\">"


class Group(Base):
    __tablename__ = "groups"
    id = Column(Integer, primary_key=True)
    title = Column(String(100))
    users: Mapped[list[User]] = relationship(
        "User",
        secondary="group_user_link",
        # lazy="selectin",
    )

    def __str__(self):
        return f"Grp:\"{self.title}\""

    def __repr__(self):
        return f"<Grp:\"{self.title}\">"

    async def get_users(self, session):
        query = (
            select(User).join(GroupUserLink).where(GroupUserLink.group_id == self.id)
        )
        result = await session.execute(query)
        users = result.scalars().all()
        return users


engine = create_async_engine(url=DATABASE_URL, echo=False, future=True)
session_factory = async_sessionmaker(
    bind=engine,
    autoflush=False,
    autocommit=False,
    expire_on_commit=False,
)


async def create_tables() -> None:
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)


async def create_group_and_users(session: AsyncSession):
    user1 = User(name="User 1")
    user2 = User(name="User 2")
    user3 = User(name="User 3")
    session.add(user1)
    session.add(user2)
    session.add(user3)

    group1 = Group(title="Group 1")
    session.add(group1)
    group1.users.append(user1)
    group1.users.append(user2)

    group2 = Group(title="Group 2")
    session.add(group2)
    group2.users.append(user3)

    await session.commit()
    await session.refresh(group1)
    await session.refresh(group2)
    return group1, group2


async def get_group_by_id(group_id: int, session):
    result = await session.execute(select(Group).where(Group.id == group_id))
    group = result.scalars().first()
    return group


async def get_all_groups(session):
    result = await session.execute(select(Group))
    groups = result.scalars().all()
    return groups


async def get_users_for_group(group_id: int, session):
    query = select(User).join(GroupUserLink).where(GroupUserLink.group_id == group_id)
    result = await session.execute(query)
    users = result.scalars().all()
    return users


def get_groups_for_user(user_id: int, session):
    query = select(Group).join(GroupUserLink).where(GroupUserLink.user_id == user_id)
    result = session.execute(query)
    groups = result.scalars().all()
    return groups


async def main():
    await create_tables()
    async with session_factory() as session:
        group1, group2 = await create_group_and_users(session)
        print(f"Created groups: {group1}, {group2}")
        all_groups = await get_all_groups(session)
        for grp in all_groups:
            retrieved_group: Group = await get_group_by_id(grp.id, session)
            # users = await retrieved_group.users.all() # PSEUDO CODE!!!
            users = await get_users_for_group(grp.id, session)
            print(f"Retrieved group: {retrieved_group}, Users: {users}")


if __name__ == "__main__":
    asyncio.run(main())