105 lines
2.8 KiB
Python
105 lines
2.8 KiB
Python
"""Add multi_tenancy
|
|
|
|
Revision ID: ff9052a97835
|
|
Revises: ff8052a97834
|
|
Create Date: 2026-03-05 10:00:00.000000
|
|
|
|
"""
|
|
|
|
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 = "ff9052a97835"
|
|
down_revision: Union[str, None] = "ff8052a97834"
|
|
branch_labels: Union[str, Sequence[str], None] = None
|
|
depends_on: Union[str, Sequence[str], None] = None
|
|
|
|
|
|
def upgrade() -> None:
|
|
# 1. Create tenants table
|
|
op.execute(
|
|
"""
|
|
CREATE TABLE tenants (
|
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
name VARCHAR(255) NOT NULL,
|
|
slug VARCHAR(100) NOT NULL UNIQUE,
|
|
is_active BOOLEAN NOT NULL DEFAULT TRUE,
|
|
created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
|
updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
|
|
);
|
|
"""
|
|
)
|
|
|
|
|
|
# 3. Add tenant_id columns (nullable first)
|
|
tables = [
|
|
"users",
|
|
"projects",
|
|
"drive_folders",
|
|
"drive_files",
|
|
"drive_activities",
|
|
"chatbot_documents",
|
|
"notifications",
|
|
"user_files",
|
|
]
|
|
for table in tables:
|
|
op.execute(
|
|
f"ALTER TABLE {table} ADD COLUMN tenant_id UUID REFERENCES tenants(id);"
|
|
)
|
|
|
|
# 4. Backfill existing data
|
|
for table in tables:
|
|
op.execute(
|
|
f"UPDATE {table} SET tenant_id = '00000000-0000-0000-0000-000000000001';"
|
|
)
|
|
|
|
# 5. Set NOT NULL (except users)
|
|
not_null_tables = [
|
|
"projects",
|
|
"drive_folders",
|
|
"drive_files",
|
|
"drive_activities",
|
|
"chatbot_documents",
|
|
"notifications",
|
|
"user_files",
|
|
]
|
|
for table in not_null_tables:
|
|
op.execute(f"ALTER TABLE {table} ALTER COLUMN tenant_id SET NOT NULL;")
|
|
|
|
# 6. Create indexes
|
|
for table in tables:
|
|
op.execute(f"CREATE INDEX ix_{table}_tenant_id ON {table}(tenant_id);")
|
|
|
|
# 7. Constraint unique (tenant_id, email)
|
|
op.execute("ALTER TABLE users DROP CONSTRAINT IF EXISTS users_email_key;")
|
|
op.execute("ALTER TABLE users DROP CONSTRAINT IF EXISTS uq_users_email;")
|
|
op.execute(
|
|
"ALTER TABLE users ADD CONSTRAINT uq_user_tenant_email UNIQUE (tenant_id, email);"
|
|
)
|
|
|
|
|
|
def downgrade() -> None:
|
|
op.execute("ALTER TABLE users DROP CONSTRAINT IF EXISTS uq_user_tenant_email;")
|
|
op.execute("ALTER TABLE users ADD CONSTRAINT users_email_key UNIQUE (email);")
|
|
|
|
tables = [
|
|
"users",
|
|
"projects",
|
|
"drive_folders",
|
|
"drive_files",
|
|
"drive_activities",
|
|
"chatbot_documents",
|
|
"notifications",
|
|
"user_files",
|
|
]
|
|
for table in tables:
|
|
op.execute(f"DROP INDEX IF EXISTS ix_{table}_tenant_id;")
|
|
op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS tenant_id;")
|
|
|
|
op.execute("DROP TABLE IF EXISTS tenants;")
|