munch-ease-backend/scripts/migration_to_households.py

from __future__ import annotations

import asyncio
from typing import Optional

import aiosqlite

import db
from settings import settings


async def column_exists(conn: aiosqlite.Connection, table: str, column: str) -> bool:
    async with conn.execute(f"PRAGMA table_info({table});") as cursor:
        async for row in cursor:
            if row[1] == column:
                return True
    return False


async def table_exists(conn: aiosqlite.Connection, table: str) -> bool:
    async with conn.execute(
        "SELECT name FROM sqlite_master WHERE type='table' AND name=?;", (table,)
    ) as c:
        row = await c.fetchone()
        return row is not None


async def add_column_if_missing(conn: aiosqlite.Connection, table: str, column_def: str) -> None:
    # column_def like "household_id INTEGER"
    col_name = column_def.split()[0]
    if not await column_exists(conn, table, col_name):
        await conn.execute(f"ALTER TABLE {table} ADD COLUMN {column_def};")


async def ensure_default_household(conn: aiosqlite.Connection) -> int:
    # Create a default household and return its id, idempotently
    await conn.execute(
        """
        INSERT INTO Household (name, slug)
        VALUES ('My Household', 'default')
        ON CONFLICT(slug) DO NOTHING
        """
    )
    async with conn.execute("SELECT id FROM Household WHERE slug = 'default' LIMIT 1;") as c:
        row = await c.fetchone()
        assert row is not None
        return int(row[0])


async def backfill_table_household_id(
    conn: aiosqlite.Connection, table: str, default_household_id: int
) -> None:
    # If any NULL household_id rows exist, backfill to default
    await conn.execute(
        f"""
        UPDATE {table}
        SET household_id = ?
        WHERE household_id IS NULL
        """,
        (default_household_id,),
    )


async def run_migration(conn: Optional[aiosqlite.Connection] = None):
    owned = False
    if conn is None:
        conn = await db.connect(settings.database_path)
        owned = True
    try:
        # Ensure new domain tables exist
        from users import repository as users_db
        from households import repository as households_db

        await users_db.create(conn)
        await households_db.create(conn)

        # Add household_id columns to tenant tables
        for table in [
            "Recipe",
            "Ingredient",
            "Meal",
            "MealParticipant",
            "MealRecipe",
            "ShoppingList",
            "ShoppingListItem",
        ]:
            await add_column_if_missing(conn, table, "household_id INTEGER")

        # Backfill default household
        default_hid = await ensure_default_household(conn)
        for table in [
            "Recipe",
            "Ingredient",
            "Meal",
            "MealParticipant",
            "MealRecipe",
            "ShoppingList",
            "ShoppingListItem",
        ]:
            await backfill_table_household_id(conn, table, default_hid)

        # Create indices on household_id for efficient scoping
        for table in [
            "Recipe",
            "Ingredient",
            "Meal",
            "MealParticipant",
            "MealRecipe",
            "ShoppingList",
            "ShoppingListItem",
        ]:
            await conn.execute(
                f"CREATE INDEX IF NOT EXISTS idx_{table.lower()}_household_id ON {table}(household_id);"
            )

        # Port persons -> users (idempotent) if legacy table exists
        if await table_exists(conn, "Person"):
            async with conn.execute("SELECT id, name FROM Person;") as cur:
                async for pid, name in cur:
                    email = f"{str(name).lower()}@example.com"
                    display_name = str(name)
                    await conn.execute(
                        "INSERT OR IGNORE INTO User (id, email, display_name) VALUES (?, ?, ?);",
                        (int(pid), email, display_name),
                    )

        # Ensure all users are members of default household (idempotent)
        await conn.execute(
            "INSERT OR IGNORE INTO HouseholdMember (user_id, household_id, role) SELECT id, ?, 'admin' FROM User;",
            (default_hid,),
        )

        await conn.commit()
    finally:
        if owned:
            await conn.close()


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