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, email TEXT 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] = [] async with conn.execute( """ SELECT h.id, h.name, h.slug 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(id=int(row[0]), name=row[1], slug=row[2])) 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] = [] async with conn.execute( """ SELECT u.id, u.display_name, m.role 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(id=int(row[0]), display_name=row[1], role=row[2]) ) return members async def create_invitation( conn, household_id: int, email: str, invited_by_user_id: int, token: str, expires_at: str ) -> bool: try: await conn.execute( """ INSERT INTO HouseholdInvitation (household_id, email, invited_by_user_id, token, expires_at, status) VALUES (?, ?, ?, ?, ?, ?) """, (household_id, email, 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: async with conn.execute( "SELECT id, household_id, status FROM HouseholdInvitation WHERE token = ?", (token,), ) as c: row = await c.fetchone() if not row: return None return HouseholdInvitation(id=int(row[0]), household_id=int(row[1]), status=row[2]) 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: async with conn.execute( "SELECT id, name, slug FROM Household WHERE id = ?", (household_id,), ) as c: row = await c.fetchone() if not row: return None return Household(id=int(row[0]), name=row[1], slug=row[2])