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