xiaoxia-saas/alembic/versions/009_remove_workspace_concept.py

"""Remove workspace concept - Projects now directly under User

Revision ID: 009
Revises: 008
Create Date: 2026-06-26

This migration:
1. Moves subscription/quota fields from workspaces to users table
2. Converts projects.workspace_id to projects.owner_user_id
3. Adds shared_users JSON field to projects table
4. Removes workspace_id from all tables that had it
5. Drops workspace-related tables: workspaces, workspace_members, workspace_invitations
"""

import sqlalchemy as sa
from sqlalchemy import text

from alembic import op

# revision identifiers
revision = "009"
down_revision = "008"
branch_labels = None
depends_on = None


def upgrade() -> None:
    conn = op.get_bind()

    # Step 1: Add subscription/quota fields to users table
    conn.execute(text("""
        ALTER TABLE users
        ADD COLUMN IF NOT EXISTS subscription_plan VARCHAR(20) NOT NULL DEFAULT 'free'
    """))
    conn.execute(text("""
        ALTER TABLE users
        ADD COLUMN IF NOT EXISTS subscription_status VARCHAR(20) NOT NULL DEFAULT 'active'
    """))
    conn.execute(text("""
        ALTER TABLE users
        ADD COLUMN IF NOT EXISTS subscription_expires_at TIMESTAMP
    """))
    conn.execute(text("""
        ALTER TABLE users
        ADD COLUMN IF NOT EXISTS max_projects FLOAT NOT NULL DEFAULT 3
    """))
    conn.execute(text("""
        ALTER TABLE users
        ADD COLUMN IF NOT EXISTS max_storage_gb FLOAT NOT NULL DEFAULT 10
    """))
    conn.execute(text("""
        ALTER TABLE users
        ADD COLUMN IF NOT EXISTS used_storage_gb FLOAT NOT NULL DEFAULT 0
    """))

    # Step 2: Copy subscription data from workspaces to users
    conn.execute(text("""
        UPDATE users SET
            subscription_plan = w.subscription_plan,
            subscription_status = w.subscription_status,
            subscription_expires_at = w.subscription_expires_at,
            max_projects = w.max_projects,
            max_storage_gb = w.max_storage_gb,
            used_storage_gb = w.used_storage_gb
        FROM workspaces w
        WHERE w.owner_user_id = users.id
    """))

    # Step 3: Add owner_user_id and shared_users to projects table
    conn.execute(text("""
        ALTER TABLE projects
        ADD COLUMN IF NOT EXISTS owner_user_id VARCHAR(32)
    """))
    conn.execute(text("""
        ALTER TABLE projects
        ADD COLUMN IF NOT EXISTS shared_users JSON
    """))

    # Step 4: Migrate workspace_id to owner_user_id (from workspace_members where role=owner)
    conn.execute(text("""
        UPDATE projects SET
            owner_user_id = wm.user_id
        FROM workspace_members wm
        WHERE wm.workspace_id = projects.workspace_id
        AND wm.role = 'owner'
    """))

    # Set shared_users to empty array for all projects
    conn.execute(text("""
        UPDATE projects SET shared_users = '[]'::json
        WHERE shared_users IS NULL
    """))

    # Step 5: Remove workspace_id from all tables
    tables_with_workspace_id = [
        "asset_libraries",
        "assets",
        "classification_jobs",
        "edit_plans",
        "edit_templates",
        "generation_tasks",
        "generated_videos",
        "ingest_jobs",
        "milestones",
        "project_titles",
        "tasks",
        "task_issues",
    ]

    for table in tables_with_workspace_id:
        conn.execute(text(f"""
            ALTER TABLE {table} DROP COLUMN IF EXISTS workspace_id
        """))

    # Step 6: Drop workspace-related tables
    conn.execute(text("""
        DROP TABLE IF EXISTS workspace_invitations
    """))
    conn.execute(text("""
        DROP TABLE IF EXISTS workspace_members
    """))
    conn.execute(text("""
        DROP TABLE IF EXISTS workspaces
    """))

    # Step 7: Drop workspace_id from projects table
    conn.execute(text("""
        ALTER TABLE projects DROP COLUMN IF EXISTS workspace_id
    """))


def downgrade() -> None:
    conn = op.get_bind()

    # Add back workspace tables (simplified - in real scenario would need full recreation)
    conn.execute(text("""
        CREATE TABLE IF NOT EXISTS workspaces (
            id VARCHAR(36) PRIMARY KEY,
            name VARCHAR(100) NOT NULL,
            owner_user_id VARCHAR(36) NOT NULL,
            subscription_plan VARCHAR(20) NOT NULL DEFAULT 'free',
            subscription_status VARCHAR(20) NOT NULL DEFAULT 'active',
            subscription_expires_at TIMESTAMP,
            max_projects FLOAT NOT NULL DEFAULT 3,
            max_storage_gb FLOAT NOT NULL DEFAULT 10,
            used_storage_gb FLOAT NOT NULL DEFAULT 0,
            created_at TIMESTAMP NOT NULL DEFAULT NOW()
        )
    """))

    conn.execute(text("""
        CREATE TABLE IF NOT EXISTS workspace_members (
            id VARCHAR(36) PRIMARY KEY,
            workspace_id VARCHAR(36) NOT NULL,
            user_id VARCHAR(36) NOT NULL,
            role VARCHAR(20) NOT NULL,
            invited_by VARCHAR(36),
            joined_at TIMESTAMP NOT NULL DEFAULT NOW(),
            UNIQUE(workspace_id, user_id)
        )
    """))

    conn.execute(text("""
        CREATE TABLE IF NOT EXISTS workspace_invitations (
            id VARCHAR(36) PRIMARY KEY,
            workspace_id VARCHAR(36) NOT NULL,
            inviter_user_id VARCHAR(36) NOT NULL,
            invitee_email VARCHAR(255) NOT NULL,
            role VARCHAR(20) NOT NULL,
            invitation_token VARCHAR(255) NOT NULL UNIQUE,
            status VARCHAR(20) NOT NULL DEFAULT 'pending',
            expires_at TIMESTAMP,
            accepted_at TIMESTAMP,
            created_at TIMESTAMP NOT NULL DEFAULT NOW()
        )
    """))

    # Add back workspace_id column to projects
    conn.execute(text("""
        ALTER TABLE projects ADD COLUMN workspace_id VARCHAR(32)
    """))

    # Add back workspace_id columns to other tables
    tables_with_workspace_id = [
        "asset_libraries",
        "assets",
        "classification_jobs",
        "edit_plans",
        "edit_templates",
        "generation_tasks",
        "generated_videos",
        "ingest_jobs",
        "milestones",
        "project_titles",
        "tasks",
        "task_issues",
    ]

    for table in tables_with_workspace_id:
        conn.execute(text(f"""
            ALTER TABLE {table} ADD COLUMN workspace_id VARCHAR(36)
        """))

    # Note: This downgrade is incomplete - projects.owner_user_id data would need to be
    # converted back to workspace_ids, which requires reconstructing workspace records.