munch-ease-backend/households/repository.py

from __future__ import annotations
from typing import List
from households.models import Household, HouseholdMember, HouseholdInvitation


async def create(conn):
    # Households table
    await conn.execute(
        """
        CREATE TABLE IF NOT EXISTS Household (
            id INTEGER PRIMARY KEY,
            name TEXT NOT NULL,
            slug TEXT NOT NULL UNIQUE
        );
        """
    )

    # Membership table
    await conn.execute(
        """
        CREATE TABLE IF NOT EXISTS HouseholdMember (
            user_id INTEGER NOT NULL,
            household_id INTEGER NOT NULL,
            role TEXT NOT NULL,
            PRIMARY KEY (user_id, household_id),
            FOREIGN KEY(user_id) REFERENCES User(id) ON DELETE CASCADE,
            FOREIGN KEY(household_id) REFERENCES Household(id) ON DELETE CASCADE
        );
        """
    )

    # Invitations table
    await conn.execute(
        """
        CREATE TABLE IF NOT EXISTS HouseholdInvitation (
            id INTEGER PRIMARY KEY,
            household_id INTEGER NOT NULL,
            invited_by_user_id INTEGER NOT NULL,
            token TEXT NOT NULL UNIQUE,
            expires_at DATETIME NOT NULL,
            status TEXT NOT NULL,
            FOREIGN KEY(household_id) REFERENCES Household(id) ON DELETE CASCADE,
            FOREIGN KEY(invited_by_user_id) REFERENCES User(id) ON DELETE SET NULL
        );
        """
    )

    # Indices
    await conn.execute("CREATE INDEX IF NOT EXISTS idx_household_slug ON Household(slug);")
    await conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_household_member_household ON HouseholdMember(household_id);"
    )


async def member_ids_in_household(conn, household_id: int, user_ids: list[int]) -> set[int]:
    """Return the subset of user_ids that are members of the given household.

    Uses HouseholdMember join User to ensure users exist.
    """
    if not user_ids:
        return set()
    placeholders = ",".join(["?"] * len(user_ids))
    query = f"""
        SELECT u.id
        FROM HouseholdMember hm
        JOIN User u ON u.id = hm.user_id
        WHERE hm.household_id = ? AND u.id IN ({placeholders})
    """
    valid: set[int] = set()
    async with conn.execute(query, (household_id, *sorted(user_ids))) as c:
        async for row in c:
            valid.add(int(row[0]))
    return valid


async def are_members(conn, household_id: int, user_ids: list[int]) -> bool:
    """True only if all provided user_ids are members of the household."""
    if not user_ids:
        return True
    found = await member_ids_in_household(conn, household_id, user_ids)
    return found == set(user_ids)


async def list_for_user(conn, user_id: int) -> List[Household]:
    results: List[Household] = []
    cols = ", ".join([f"h.{k}" for k in Household.KEYS])
    async with conn.execute(
        f"""
        SELECT {cols} FROM Household h
        JOIN HouseholdMember m ON m.household_id = h.id
        WHERE m.user_id = ?
        ORDER BY h.id
        """,
        (user_id,),
    ) as c:
        async for row in c:
            results.append(Household(**{k: v for k, v in zip(Household.KEYS, row)}))
    return results


async def create_for_user(conn, name: str, slug: str, user_id: int) -> Household | None:
    try:
        async with conn.execute(
            "INSERT INTO Household (name, slug) VALUES (?, ?)", (name, slug)
        ) as cur:
            lrid = cur.lastrowid
            if lrid is None:
                return None
            hid = int(lrid)
        await conn.execute(
            "INSERT INTO HouseholdMember (user_id, household_id, role) VALUES (?, ?, ?)",
            (user_id, hid, "admin"),
        )
        return Household(id=hid, name=name, slug=slug)
    except Exception:
        return None


async def list_members(conn, household_id: int) -> List[HouseholdMember]:
    members: List[HouseholdMember] = []
    cols = ", ".join([f"u.{k}" for k in ["id", "display_name"]] + ["m.role"])
    async with conn.execute(
        f"""
        SELECT {cols}
        FROM HouseholdMember m
        JOIN User u ON u.id = m.user_id
        WHERE m.household_id = ?
        ORDER BY lower(u.display_name), u.id
        """,
        (household_id,),
    ) as c:
        async for row in c:
            members.append(
                HouseholdMember(**{k: v for k, v in zip(HouseholdMember.KEYS, row)})
            )
    return members


async def create_invitation(
    conn, household_id: int, invited_by_user_id: int, token: str, expires_at: str
) -> bool:
    try:
        await conn.execute(
            """
            INSERT INTO HouseholdInvitation (household_id, invited_by_user_id, token, expires_at, status)
            VALUES (?, ?, ?, ?, ?)
            """,
            (household_id, invited_by_user_id, token, expires_at, "pending"),
        )
        return True
    except Exception:
        return False


async def get_invitation_by_token(conn, token: str) -> HouseholdInvitation | None:
    cols = ", ".join(HouseholdInvitation.KEYS)
    async with conn.execute(
        f"SELECT {cols} FROM HouseholdInvitation WHERE token = ?",
        (token,),
    ) as c:
        row = await c.fetchone()
    if not row:
        return None
    return HouseholdInvitation(**{k: v for k, v in zip(HouseholdInvitation.KEYS, row)})


async def accept_invitation(conn, invitation_id: int, user_id: int, household_id: int):
    await conn.execute(
        "INSERT OR IGNORE INTO HouseholdMember (user_id, household_id, role) VALUES (?, ?, ?)",
        (user_id, household_id, "member"),
    )
    await conn.execute(
        "UPDATE HouseholdInvitation SET status = 'accepted' WHERE id = ?",
        (invitation_id,),
    )


async def get_household_by_id(conn, household_id: int) -> Household | None:
    cols = ", ".join(Household.KEYS)
    async with conn.execute(
        f"SELECT {cols} FROM Household WHERE id = ?",
        (household_id,),
    ) as c:
        row = await c.fetchone()
    if not row:
        return None
    return Household(**{k: v for k, v in zip(Household.KEYS, row)})