munch-ease-backend/scripts/migration_to_households.py

136 lines
4.3 KiB
Python

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 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 and create memberships in default household
# Only perform if users table currently empty
async with conn.execute("SELECT COUNT(1) FROM User;") as c:
row = await c.fetchone()
user_count = int(row[0]) if row else 0
if user_count == 0:
async with conn.execute("SELECT id, name FROM Person;") as cur:
async for pid, name in cur:
email = f"{name.lower()}@example.com"
display_name = name
# Insert user
await conn.execute(
"INSERT INTO User (id, email, display_name) VALUES (?, ?, ?)\n ON CONFLICT(id) DO NOTHING;",
(pid, email, display_name),
)
# Create membership
await conn.execute(
"INSERT OR IGNORE INTO HouseholdMember (user_id, household_id, role) VALUES (?, ?, ?);",
(pid, default_hid, "admin"),
)
await conn.commit()
finally:
if owned:
await conn.close()
if __name__ == "__main__":
asyncio.run(run_migration())