mealie/alembic/versions/2024-07-12-16.16.29_feecc8ffb956_add_households.py
2024-09-16 13:52:12 +00:00

322 lines
13 KiB
Python

"""add households
Revision ID: feecc8ffb956
Revises: 32d69327997b
Create Date: 2024-07-12 16:16:29.973929
"""
from datetime import datetime, timezone
from textwrap import dedent
from typing import Any
from uuid import uuid4
import sqlalchemy as sa
from slugify import slugify
from sqlalchemy import orm
import mealie.db.migration_types
from alembic import op
from mealie.core.config import get_app_settings
# revision identifiers, used by Alembic.
revision = "feecc8ffb956"
down_revision = "32d69327997b"
branch_labels: str | tuple[str, ...] | None = None
depends_on: str | tuple[str, ...] | None = None
settings = get_app_settings()
def is_postgres():
return op.get_context().dialect.name == "postgresql"
def generate_id() -> str:
"""See GUID.convert_value_to_guid"""
val = uuid4()
if is_postgres():
return str(val)
else:
return f"{val.int:032x}"
def dedupe_cookbook_slugs():
bind = op.get_bind()
session = orm.Session(bind=bind)
with session:
sql = sa.text(
dedent(
"""
SELECT slug, group_id, COUNT(*)
FROM cookbooks
GROUP BY slug, group_id
HAVING COUNT(*) > 1
"""
)
)
rows = session.execute(sql).fetchall()
for slug, group_id, _ in rows:
sql = sa.text(
dedent(
"""
SELECT id
FROM cookbooks
WHERE slug = :slug AND group_id = :group_id
ORDER BY id
"""
)
)
cookbook_ids = session.execute(sql, {"slug": slug, "group_id": group_id}).fetchall()
for i, (cookbook_id,) in enumerate(cookbook_ids):
if i == 0:
continue
sql = sa.text(
dedent(
"""
UPDATE cookbooks
SET slug = :slug || '-' || :i
WHERE id = :id
"""
)
)
session.execute(sql, {"slug": slug, "i": i, "id": cookbook_id})
def create_household(session: orm.Session, group_id: str) -> str:
# create/insert household
household_id = generate_id()
timestamp = datetime.now(timezone.utc).isoformat()
household_data = {
"id": household_id,
"name": settings.DEFAULT_HOUSEHOLD,
"slug": slugify(settings.DEFAULT_HOUSEHOLD),
"group_id": group_id,
"created_at": timestamp,
"update_at": timestamp,
}
columns = ", ".join(household_data.keys())
placeholders = ", ".join(f":{key}" for key in household_data.keys())
sql_statement = f"INSERT INTO households ({columns}) VALUES ({placeholders})"
session.execute(sa.text(sql_statement), household_data)
# fetch group preferences so we can copy them over to household preferences
migrated_field_defaults = {
"private_group": True, # this is renamed later
"first_day_of_week": 0,
"recipe_public": True,
"recipe_show_nutrition": False,
"recipe_show_assets": False,
"recipe_landscape_view": False,
"recipe_disable_comments": False,
"recipe_disable_amount": True,
}
sql_statement = (
f"SELECT {', '.join(migrated_field_defaults.keys())} FROM group_preferences WHERE group_id = :group_id"
)
group_preferences = session.execute(sa.text(sql_statement), {"group_id": group_id}).fetchone()
# build preferences data
if group_preferences:
preferences_data: dict[str, Any] = {}
for i, (field, default_value) in enumerate(migrated_field_defaults.items()):
value = group_preferences[i]
preferences_data[field] = value if value is not None else default_value
else:
preferences_data = migrated_field_defaults
preferences_data["id"] = generate_id()
preferences_data["household_id"] = household_id
preferences_data["created_at"] = timestamp
preferences_data["update_at"] = timestamp
preferences_data["private_household"] = preferences_data.pop("private_group")
# insert preferences data
columns = ", ".join(preferences_data.keys())
placeholders = ", ".join(f":{key}" for key in preferences_data.keys())
sql_statement = f"INSERT INTO household_preferences ({columns}) VALUES ({placeholders})"
session.execute(sa.text(sql_statement), preferences_data)
return household_id
def create_households_for_groups() -> dict[str, str]:
bind = op.get_bind()
session = orm.Session(bind=bind)
group_id_household_id_map: dict[str, str] = {}
with session:
rows = session.execute(sa.text("SELECT id FROM groups")).fetchall()
for row in rows:
group_id = row[0]
group_id_household_id_map[group_id] = create_household(session, group_id)
return group_id_household_id_map
def _do_assignment(session: orm.Session, table: str, group_id: str, household_id: str):
sql = sa.text(
dedent(
f"""
UPDATE {table}
SET household_id = :household_id
WHERE group_id = :group_id
""",
)
)
session.execute(sql, {"group_id": group_id, "household_id": household_id})
def assign_households(group_id_household_id_map: dict[str, str]):
tables = [
"cookbooks",
"group_events_notifiers",
"group_meal_plan_rules",
"invite_tokens",
"recipe_actions",
"users",
"webhook_urls",
]
bind = op.get_bind()
session = orm.Session(bind=bind)
with session:
for table in tables:
for group_id, household_id in group_id_household_id_map.items():
_do_assignment(session, table, group_id, household_id)
def populate_household_data():
group_id_household_id_map = create_households_for_groups()
assign_households(group_id_household_id_map)
def upgrade():
dedupe_cookbook_slugs()
# ### commands auto generated by Alembic - please adjust! ###
op.create_table(
"households",
sa.Column("id", mealie.db.migration_types.GUID(), nullable=False),
sa.Column("name", sa.String(), nullable=False),
sa.Column("slug", sa.String(), nullable=True),
sa.Column("group_id", mealie.db.migration_types.GUID(), nullable=False),
sa.Column("created_at", sa.DateTime(), nullable=True),
sa.Column("update_at", sa.DateTime(), nullable=True),
sa.ForeignKeyConstraint(
["group_id"],
["groups.id"],
),
sa.PrimaryKeyConstraint("id"),
sa.UniqueConstraint("group_id", "name", name="household_name_group_id_key"),
sa.UniqueConstraint("group_id", "slug", name="household_slug_group_id_key"),
)
op.create_index(op.f("ix_households_created_at"), "households", ["created_at"], unique=False)
op.create_index(op.f("ix_households_group_id"), "households", ["group_id"], unique=False)
op.create_index(op.f("ix_households_name"), "households", ["name"], unique=False)
op.create_index(op.f("ix_households_slug"), "households", ["slug"], unique=False)
op.create_table(
"household_preferences",
sa.Column("id", mealie.db.migration_types.GUID(), nullable=False),
sa.Column("household_id", mealie.db.migration_types.GUID(), nullable=False),
sa.Column("private_household", sa.Boolean(), nullable=True),
sa.Column("first_day_of_week", sa.Integer(), nullable=True),
sa.Column("recipe_public", sa.Boolean(), nullable=True),
sa.Column("recipe_show_nutrition", sa.Boolean(), nullable=True),
sa.Column("recipe_show_assets", sa.Boolean(), nullable=True),
sa.Column("recipe_landscape_view", sa.Boolean(), nullable=True),
sa.Column("recipe_disable_comments", sa.Boolean(), nullable=True),
sa.Column("recipe_disable_amount", sa.Boolean(), nullable=True),
sa.Column("created_at", sa.DateTime(), nullable=True),
sa.Column("update_at", sa.DateTime(), nullable=True),
sa.ForeignKeyConstraint(
["household_id"],
["households.id"],
),
sa.PrimaryKeyConstraint("id"),
)
op.create_index(op.f("ix_household_preferences_created_at"), "household_preferences", ["created_at"], unique=False)
op.create_index(
op.f("ix_household_preferences_household_id"), "household_preferences", ["household_id"], unique=False
)
with op.batch_alter_table("cookbooks") as batch_op:
batch_op.add_column(sa.Column("household_id", mealie.db.migration_types.GUID(), nullable=True))
batch_op.create_index(op.f("ix_cookbooks_household_id"), ["household_id"], unique=False)
batch_op.create_foreign_key("fk_cookbooks_household_id", "households", ["household_id"], ["id"])
# not directly related to households, but important for frontend routes
batch_op.create_unique_constraint("cookbook_slug_group_id_key", ["slug", "group_id"])
with op.batch_alter_table("group_events_notifiers") as batch_op:
batch_op.add_column(sa.Column("household_id", mealie.db.migration_types.GUID(), nullable=True))
batch_op.create_index(op.f("ix_group_events_notifiers_household_id"), ["household_id"], unique=False)
batch_op.create_foreign_key("fk_group_events_notifiers_household_id", "households", ["household_id"], ["id"])
with op.batch_alter_table("group_meal_plan_rules") as batch_op:
batch_op.add_column(sa.Column("household_id", mealie.db.migration_types.GUID(), nullable=True))
batch_op.create_index(op.f("ix_group_meal_plan_rules_household_id"), ["household_id"], unique=False)
batch_op.create_foreign_key("fk_group_meal_plan_rules_household_id", "households", ["household_id"], ["id"])
with op.batch_alter_table("invite_tokens") as batch_op:
batch_op.add_column(sa.Column("household_id", mealie.db.migration_types.GUID(), nullable=True))
batch_op.create_index(op.f("ix_invite_tokens_household_id"), ["household_id"], unique=False)
batch_op.create_foreign_key("fk_invite_tokens_household_id", "households", ["household_id"], ["id"])
with op.batch_alter_table("recipe_actions") as batch_op:
batch_op.add_column(sa.Column("household_id", mealie.db.migration_types.GUID(), nullable=True))
batch_op.create_index(op.f("ix_recipe_actions_household_id"), ["household_id"], unique=False)
batch_op.create_foreign_key("fk_recipe_actions_household_id", "households", ["household_id"], ["id"])
with op.batch_alter_table("users") as batch_op:
batch_op.add_column(sa.Column("household_id", mealie.db.migration_types.GUID(), nullable=True))
batch_op.create_index(op.f("ix_users_household_id"), ["household_id"], unique=False)
batch_op.create_foreign_key("fk_users_household_id", "households", ["household_id"], ["id"])
with op.batch_alter_table("webhook_urls") as batch_op:
batch_op.add_column(sa.Column("household_id", mealie.db.migration_types.GUID(), nullable=True))
batch_op.create_index(op.f("ix_webhook_urls_household_id"), ["household_id"], unique=False)
batch_op.create_foreign_key("fk_webhook_urls_household_id", "households", ["household_id"], ["id"])
# ### end Alembic commands ###
populate_household_data()
def downgrade():
# ### commands auto generated by Alembic - please adjust! ###
op.drop_constraint(None, "webhook_urls", type_="foreignkey")
op.drop_index(op.f("ix_webhook_urls_household_id"), table_name="webhook_urls")
op.drop_column("webhook_urls", "household_id")
op.drop_constraint(None, "users", type_="foreignkey")
op.drop_index(op.f("ix_users_household_id"), table_name="users")
op.drop_column("users", "household_id")
op.drop_constraint(None, "recipe_actions", type_="foreignkey")
op.drop_index(op.f("ix_recipe_actions_household_id"), table_name="recipe_actions")
op.drop_column("recipe_actions", "household_id")
op.drop_constraint(None, "invite_tokens", type_="foreignkey")
op.drop_index(op.f("ix_invite_tokens_household_id"), table_name="invite_tokens")
op.drop_column("invite_tokens", "household_id")
op.drop_constraint(None, "group_meal_plan_rules", type_="foreignkey")
op.drop_index(op.f("ix_group_meal_plan_rules_household_id"), table_name="group_meal_plan_rules")
op.drop_column("group_meal_plan_rules", "household_id")
op.drop_constraint(None, "group_events_notifiers", type_="foreignkey")
op.drop_index(op.f("ix_group_events_notifiers_household_id"), table_name="group_events_notifiers")
op.drop_column("group_events_notifiers", "household_id")
op.drop_constraint(None, "cookbooks", type_="foreignkey")
op.drop_index(op.f("ix_cookbooks_household_id"), table_name="cookbooks")
op.drop_column("cookbooks", "household_id")
op.drop_constraint("cookbook_slug_group_id_key", "cookbooks", type_="unique")
op.drop_index(op.f("ix_household_preferences_household_id"), table_name="household_preferences")
op.drop_index(op.f("ix_household_preferences_created_at"), table_name="household_preferences")
op.drop_table("household_preferences")
op.drop_index(op.f("ix_households_slug"), table_name="households")
op.drop_index(op.f("ix_households_name"), table_name="households")
op.drop_index(op.f("ix_households_group_id"), table_name="households")
op.drop_index(op.f("ix_households_created_at"), table_name="households")
op.drop_table("households")
# ### end Alembic commands ###