137 lines
4.3 KiB
Python
137 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())
|