"""departments, per-department categories and budgets, approval workflows Existing data is kept: everything created with version 1 is moved into a department called "IT". Revision ID: 0002 Revises: 0001 """ import sqlalchemy as sa from alembic import op revision = "0002" down_revision = "0001" branch_labels = None depends_on = None def _drop_unique_on(table: str, columns: list[str]) -> None: insp = sa.inspect(op.get_bind()) for uc in insp.get_unique_constraints(table): if uc["column_names"] == columns: op.drop_constraint(uc["name"], table, type_="unique") def upgrade() -> None: op.create_table( "departments", sa.Column("id", sa.Integer(), primary_key=True), sa.Column("name", sa.String(length=120), nullable=False, unique=True), sa.Column("active", sa.Boolean(), nullable=False, server_default=sa.true()), sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()), ) op.create_table( "workflow_steps", sa.Column("id", sa.Integer(), primary_key=True), sa.Column("department_id", sa.Integer(), sa.ForeignKey("departments.id", ondelete="CASCADE"), nullable=False), sa.Column("position", sa.Integer(), nullable=False, server_default="0"), sa.Column("name", sa.String(length=120), nullable=False), sa.Column("approver_mode", sa.String(length=10), nullable=False, server_default="role"), sa.Column("approver_ids", sa.JSON(), nullable=True), ) op.create_index("ix_workflow_steps_department_id", "workflow_steps", ["department_id"]) op.add_column("users", sa.Column("department_id", sa.Integer(), nullable=True)) op.add_column("categories", sa.Column("department_id", sa.Integer(), nullable=True)) op.add_column("budget_years", sa.Column("department_id", sa.Integer(), nullable=True)) op.add_column("budget_years", sa.Column("current_step", sa.Integer(), nullable=True)) conn = op.get_bind() has_data = any(conn.execute(sa.text(f"SELECT 1 FROM {t} LIMIT 1")).first() for t in ("users", "categories", "budget_years")) if has_data: conn.execute(sa.text("INSERT INTO departments (name, active, created_at) VALUES ('IT', true, now())")) dept_id = conn.execute(sa.text("SELECT id FROM departments WHERE name = 'IT'")).scalar_one() for table in ("users", "categories", "budget_years"): conn.execute(sa.text(f"UPDATE {table} SET department_id = :d"), {"d": dept_id}) conn.execute(sa.text("UPDATE budget_years SET current_step = 0 WHERE status = 'in_review'")) op.alter_column("categories", "department_id", nullable=False) op.alter_column("budget_years", "department_id", nullable=False) op.create_foreign_key("fk_users_department", "users", "departments", ["department_id"], ["id"]) op.create_foreign_key("fk_categories_department", "categories", "departments", ["department_id"], ["id"], ondelete="CASCADE") op.create_foreign_key("fk_budget_years_department", "budget_years", "departments", ["department_id"], ["id"]) op.create_index("ix_categories_department_id", "categories", ["department_id"]) op.create_index("ix_budget_years_department_id", "budget_years", ["department_id"]) _drop_unique_on("categories", ["name_en"]) _drop_unique_on("budget_years", ["year"]) op.create_unique_constraint("uq_categories_department_name", "categories", ["department_id", "name_en"]) op.create_unique_constraint("uq_budget_years_department_year", "budget_years", ["department_id", "year"]) def downgrade() -> None: raise NotImplementedError("Restore a backup taken before the upgrade instead.")