"""migrate existing users and teams to accounts Revision ID: 018 Revises: 017 Create Date: 2026-02-07 This is the most critical migration. It creates a _team_account_mapping table for deterministic cross-migration lookups, then migrates all users to accounts. Three paths: A) Teams with users → Account with deterministic owner B) Teams with zero users → Account with owner_id=NULL, subscription status='orphaned' C) Users without a team → Personal account, user is owner """ import secrets import string from typing import Sequence, Union from alembic import op import sqlalchemy as sa from sqlalchemy.dialects.postgresql import UUID # revision identifiers, used by Alembic. revision: str = '018' down_revision: Union[str, None] = '017' branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None # Characters for display codes — exclude confusing chars DISPLAY_CODE_CHARS = string.ascii_uppercase + string.digits DISPLAY_CODE_CHARS = DISPLAY_CODE_CHARS.replace('0', '').replace('O', '').replace('I', '').replace('1', '').replace('L', '') def _generate_display_code(existing_codes: set) -> str: """Generate a unique 8-character display code.""" for _ in range(100): code = ''.join(secrets.choice(DISPLAY_CODE_CHARS) for _ in range(8)) if code not in existing_codes: existing_codes.add(code) return code raise RuntimeError("Failed to generate unique display code after 100 attempts") def upgrade() -> None: conn = op.get_bind() # Create mapping table for deterministic cross-migration lookups op.create_table( '_team_account_mapping', sa.Column('team_id', UUID(as_uuid=True), nullable=False), sa.Column('account_id', UUID(as_uuid=True), nullable=False), sa.Column('owner_user_id', UUID(as_uuid=True), nullable=True), sa.PrimaryKeyConstraint('team_id'), ) existing_codes: set = set() # --- Path A & B: Process all teams --- teams = conn.execute(sa.text("SELECT id, name FROM teams")).fetchall() for team in teams: team_id = team[0] team_name = team[1] display_code = _generate_display_code(existing_codes) # Find deterministic owner: team admin first, then earliest user owner_row = conn.execute(sa.text(""" SELECT id FROM users WHERE team_id = :tid ORDER BY is_team_admin DESC, created_at ASC, id ASC LIMIT 1 """), {"tid": team_id}).fetchone() owner_user_id = owner_row[0] if owner_row else None # Create account conn.execute(sa.text(""" INSERT INTO accounts (id, name, display_code, owner_id, created_at, updated_at) VALUES (gen_random_uuid(), :name, :code, :owner_id, NOW(), NOW()) """), {"name": team_name, "code": display_code, "owner_id": owner_user_id}) # Get the account we just created account_row = conn.execute(sa.text( "SELECT id FROM accounts WHERE display_code = :code" ), {"code": display_code}).fetchone() account_id = account_row[0] # Insert mapping conn.execute(sa.text(""" INSERT INTO _team_account_mapping (team_id, account_id, owner_user_id) VALUES (:tid, :aid, :uid) """), {"tid": team_id, "aid": account_id, "uid": owner_user_id}) if owner_user_id is not None: # Path A: Team with users # Create active subscription conn.execute(sa.text(""" INSERT INTO subscriptions (id, account_id, plan, status, created_at, updated_at) VALUES (gen_random_uuid(), :aid, 'free', 'active', NOW(), NOW()) """), {"aid": account_id}) # Update all users in this team # Team admins become owners, others keep their role conn.execute(sa.text(""" UPDATE users SET account_id = :aid, account_role = CASE WHEN is_team_admin = true THEN 'owner' ELSE role END WHERE team_id = :tid """), {"aid": account_id, "tid": team_id}) else: # Path B: Team with zero users (orphan) conn.execute(sa.text(""" INSERT INTO subscriptions (id, account_id, plan, status, created_at, updated_at) VALUES (gen_random_uuid(), :aid, 'free', 'orphaned', NOW(), NOW()) """), {"aid": account_id}) # --- Path C: Users without a team --- teamless_users = conn.execute(sa.text( "SELECT id, name FROM users WHERE team_id IS NULL AND account_id IS NULL" )).fetchall() for user in teamless_users: user_id = user[0] user_name = user[1] display_code = _generate_display_code(existing_codes) # Create personal account (owner_id set to NULL initially) conn.execute(sa.text(""" INSERT INTO accounts (id, name, display_code, owner_id, created_at, updated_at) VALUES (gen_random_uuid(), :name, :code, NULL, NOW(), NOW()) """), {"name": f"{user_name}'s Account", "code": display_code}) account_row = conn.execute(sa.text( "SELECT id FROM accounts WHERE display_code = :code" ), {"code": display_code}).fetchone() account_id = account_row[0] # Update user conn.execute(sa.text(""" UPDATE users SET account_id = :aid, account_role = 'owner' WHERE id = :uid """), {"aid": account_id, "uid": user_id}) # Set owner conn.execute(sa.text(""" UPDATE accounts SET owner_id = :uid WHERE id = :aid """), {"uid": user_id, "aid": account_id}) # Create free subscription conn.execute(sa.text(""" INSERT INTO subscriptions (id, account_id, plan, status, created_at, updated_at) VALUES (gen_random_uuid(), :aid, 'free', 'active', NOW(), NOW()) """), {"aid": account_id}) # --- Validation --- orphaned_users = conn.execute(sa.text( "SELECT COUNT(*) FROM users WHERE account_id IS NULL" )).scalar() if orphaned_users > 0: raise RuntimeError( f"Migration 018 failed validation: {orphaned_users} users still have NULL account_id" ) team_count = conn.execute(sa.text("SELECT COUNT(*) FROM teams")).scalar() mapping_count = conn.execute(sa.text("SELECT COUNT(*) FROM _team_account_mapping")).scalar() if mapping_count != team_count: raise RuntimeError( f"Migration 018 failed: mapping count ({mapping_count}) != team count ({team_count})" ) def downgrade() -> None: conn = op.get_bind() # Clear account data from users conn.execute(sa.text("UPDATE users SET account_id = NULL, account_role = NULL")) # Delete all subscriptions and accounts created by this migration conn.execute(sa.text("DELETE FROM subscriptions")) conn.execute(sa.text("DELETE FROM accounts")) # Drop mapping table op.drop_table('_team_account_mapping')