375 lines
17 KiB
Python
375 lines
17 KiB
Python
"""User system refactor - unified migration.
|
|
|
|
Revision ID: user_refactor_v1
|
|
Revises: a1b2c3d4e5f6
|
|
Create Date: 2026-03-27
|
|
"""
|
|
from typing import Sequence, Union
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.dialects import postgresql
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = 'user_refactor_v1'
|
|
down_revision: Union[str, None] = 'a1b2c3d4e5f6'
|
|
branch_labels: Union[str, Sequence[str], None] = None
|
|
depends_on: Union[str, Sequence[str], None] = None
|
|
|
|
|
|
def upgrade() -> None:
|
|
# ============================================
|
|
# 1. Create identity_providers table (no foreign key to allow soft coupling)
|
|
# ============================================
|
|
op.execute("""
|
|
CREATE TABLE IF NOT EXISTS identity_providers (
|
|
id UUID PRIMARY KEY,
|
|
provider_type VARCHAR(50) NOT NULL,
|
|
name VARCHAR(100) NOT NULL,
|
|
is_active BOOLEAN DEFAULT TRUE,
|
|
config JSON,
|
|
tenant_id UUID,
|
|
created_at TIMESTAMP WITH TIME ZONE DEFAULT now(),
|
|
updated_at TIMESTAMP WITH TIME ZONE DEFAULT now()
|
|
)
|
|
""")
|
|
|
|
# ============================================
|
|
# 2. Create sso_scan_sessions table (no foreign keys for soft coupling)
|
|
# ============================================
|
|
op.execute("""
|
|
CREATE TABLE IF NOT EXISTS sso_scan_sessions (
|
|
id UUID PRIMARY KEY,
|
|
status VARCHAR(50) DEFAULT 'pending',
|
|
provider_type VARCHAR(50),
|
|
error_msg TEXT,
|
|
tenant_id UUID,
|
|
user_id UUID,
|
|
access_token TEXT,
|
|
expires_at TIMESTAMP WITH TIME ZONE NOT NULL,
|
|
created_at TIMESTAMP WITH TIME ZONE DEFAULT now()
|
|
)
|
|
""")
|
|
|
|
# ============================================
|
|
# 3. Alter tenants - add SSO fields
|
|
# ============================================
|
|
op.execute("ALTER TABLE tenants ADD COLUMN IF NOT EXISTS sso_enabled BOOLEAN DEFAULT FALSE")
|
|
op.execute("ALTER TABLE tenants ADD COLUMN IF NOT EXISTS sso_domain VARCHAR(255)")
|
|
op.execute("CREATE UNIQUE INDEX IF NOT EXISTS ux_tenants_sso_domain ON tenants(sso_domain) WHERE sso_domain IS NOT NULL")
|
|
|
|
# ============================================
|
|
# 4. Alter org_departments (no foreign key - soft coupling via program)
|
|
# ============================================
|
|
op.execute("ALTER TABLE org_departments ADD COLUMN IF NOT EXISTS external_id VARCHAR(100)")
|
|
op.execute("ALTER TABLE org_departments ADD COLUMN IF NOT EXISTS provider_id UUID")
|
|
op.execute("CREATE INDEX IF NOT EXISTS ix_org_departments_external_id ON org_departments(external_id)")
|
|
# Note: provider_id is UUID without FK constraint - program should validate existence
|
|
|
|
# ============================================
|
|
# 5. Alter org_members (no foreign keys - soft coupling via program)
|
|
# ============================================
|
|
op.execute("""
|
|
DO $$
|
|
BEGIN
|
|
-- 5.1 Handle feishu_open_id to open_id
|
|
IF EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'org_members' AND column_name = 'feishu_open_id') THEN
|
|
IF EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'org_members' AND column_name = 'open_id') THEN
|
|
UPDATE org_members SET open_id = feishu_open_id WHERE open_id IS NULL;
|
|
ALTER TABLE org_members DROP COLUMN feishu_open_id;
|
|
ELSE
|
|
ALTER TABLE org_members RENAME COLUMN feishu_open_id TO open_id;
|
|
END IF;
|
|
ALTER TABLE org_members DROP CONSTRAINT IF EXISTS org_members_feishu_open_id_key;
|
|
END IF;
|
|
|
|
IF NOT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'org_members' AND column_name = 'open_id') THEN
|
|
ALTER TABLE org_members ADD COLUMN open_id VARCHAR(100);
|
|
END IF;
|
|
|
|
-- 5.2 Handle feishu_user_id to external_id
|
|
IF EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'org_members' AND column_name = 'feishu_user_id') THEN
|
|
IF EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'org_members' AND column_name = 'external_id') THEN
|
|
UPDATE org_members SET external_id = feishu_user_id WHERE external_id IS NULL;
|
|
ALTER TABLE org_members DROP COLUMN feishu_user_id;
|
|
ELSE
|
|
ALTER TABLE org_members RENAME COLUMN feishu_user_id TO external_id;
|
|
END IF;
|
|
END IF;
|
|
|
|
IF NOT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'org_members' AND column_name = 'external_id') THEN
|
|
ALTER TABLE org_members ADD COLUMN external_id VARCHAR(100);
|
|
END IF;
|
|
END $$;
|
|
""")
|
|
|
|
op.execute("ALTER TABLE org_members ADD COLUMN IF NOT EXISTS unionid VARCHAR(100)")
|
|
op.execute("ALTER TABLE org_members ADD COLUMN IF NOT EXISTS provider_id UUID")
|
|
op.execute("ALTER TABLE org_members ADD COLUMN IF NOT EXISTS user_id UUID")
|
|
|
|
op.execute("CREATE INDEX IF NOT EXISTS ix_org_members_open_id ON org_members(open_id)")
|
|
op.execute("CREATE INDEX IF NOT EXISTS ix_org_members_external_id ON org_members(external_id)")
|
|
op.execute("CREATE INDEX IF NOT EXISTS ix_org_members_unionid ON org_members(unionid)")
|
|
op.execute("CREATE INDEX IF NOT EXISTS ix_org_members_user_id ON org_members(user_id)")
|
|
|
|
|
|
# Note: provider_id and user_id are UUIDs without FK constraints - program should validate
|
|
|
|
# ============================================
|
|
# 5.1 Data migration - backfill org_members.provider_id (feishu)
|
|
# ============================================
|
|
op.execute("""
|
|
UPDATE org_members AS om
|
|
SET provider_id = ip.id
|
|
FROM (
|
|
SELECT DISTINCT ON (tenant_id) id, tenant_id
|
|
FROM identity_providers
|
|
WHERE provider_type = 'feishu'
|
|
ORDER BY tenant_id, created_at DESC NULLS LAST, id
|
|
) AS ip
|
|
WHERE om.tenant_id = ip.tenant_id
|
|
AND om.provider_id IS NULL
|
|
""")
|
|
|
|
# ============================================
|
|
# 6. Alter users - add new fields and constraints
|
|
# ============================================
|
|
op.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS primary_mobile VARCHAR(50)")
|
|
op.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS registration_source VARCHAR(50) DEFAULT 'web'")
|
|
op.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS external_id VARCHAR(255)")
|
|
op.execute("CREATE INDEX IF NOT EXISTS ix_users_primary_mobile ON users(primary_mobile)")
|
|
op.execute("CREATE INDEX IF NOT EXISTS ix_users_external_id ON users(external_id)")
|
|
|
|
# Add unique constraints (partial indexes - allow multiple NULL values)
|
|
op.execute("""
|
|
DO $$
|
|
BEGIN
|
|
IF NOT EXISTS (
|
|
SELECT 1 FROM pg_indexes WHERE indexname = 'ix_users_tenant_email_unique'
|
|
) THEN
|
|
CREATE UNIQUE INDEX ix_users_tenant_email_unique ON users(tenant_id, email) WHERE email IS NOT NULL;
|
|
END IF;
|
|
END $$
|
|
""")
|
|
|
|
# Remove deprecated user identity columns (open_id / union_id)
|
|
op.execute("""
|
|
DO $$
|
|
BEGIN
|
|
IF EXISTS (
|
|
SELECT 1 FROM information_schema.columns
|
|
WHERE table_name = 'users' AND column_name = 'feishu_open_id'
|
|
) THEN
|
|
ALTER TABLE users DROP CONSTRAINT IF EXISTS users_feishu_open_id_key;
|
|
DROP INDEX IF EXISTS ix_users_feishu_open_id;
|
|
ALTER TABLE users DROP COLUMN feishu_open_id;
|
|
END IF;
|
|
|
|
IF EXISTS (
|
|
SELECT 1 FROM information_schema.columns
|
|
WHERE table_name = 'users' AND column_name = 'feishu_union_id'
|
|
) THEN
|
|
ALTER TABLE users DROP COLUMN feishu_union_id;
|
|
END IF;
|
|
END $$
|
|
""")
|
|
op.execute("""
|
|
DO $$
|
|
BEGIN
|
|
IF NOT EXISTS (
|
|
SELECT 1 FROM pg_indexes WHERE indexname = 'ix_users_tenant_mobile_unique'
|
|
) THEN
|
|
CREATE UNIQUE INDEX ix_users_tenant_mobile_unique ON users(tenant_id, primary_mobile) WHERE primary_mobile IS NOT NULL;
|
|
END IF;
|
|
END $$
|
|
""")
|
|
|
|
# ============================================
|
|
# 7. Drop deprecated departments table
|
|
# ============================================
|
|
op.execute("ALTER TABLE users DROP COLUMN IF EXISTS department_id")
|
|
op.execute("DROP TABLE IF EXISTS departments")
|
|
|
|
# ============================================
|
|
# 8. Alter channel_config - extend app_secret length
|
|
# ============================================
|
|
# Note: This requires dropping and recreating the column due to PostgreSQL limitation
|
|
# Only do this if the column exists and is smaller than 512
|
|
op.execute("""
|
|
DO $$
|
|
BEGIN
|
|
IF EXISTS (
|
|
SELECT 1 FROM information_schema.columns
|
|
WHERE table_name = 'channel_config' AND column_name = 'app_secret'
|
|
AND character_maximum_length < 512
|
|
) THEN
|
|
ALTER TABLE channel_config ALTER COLUMN app_secret TYPE VARCHAR(512);
|
|
END IF;
|
|
END $$
|
|
""")
|
|
|
|
# Step 1: Get distinct tenant_ids from org_departments that haven't been migrated
|
|
connection = op.get_bind()
|
|
result = connection.execute(sa.text("""
|
|
SELECT DISTINCT od.tenant_id
|
|
FROM org_departments od
|
|
WHERE od.tenant_id IS NOT NULL
|
|
AND NOT EXISTS (
|
|
SELECT 1 FROM identity_providers ip
|
|
WHERE ip.tenant_id = od.tenant_id
|
|
)
|
|
"""))
|
|
|
|
tenant_ids = [row[0] for row in result.fetchall()]
|
|
|
|
for tenant_id in tenant_ids:
|
|
# Generate provider ID using PostgreSQL function
|
|
provider_id = connection.execute(sa.text("SELECT gen_random_uuid()")).scalar()
|
|
|
|
# Insert IdentityProvider (only if not exists)
|
|
connection.execute(
|
|
sa.text("""
|
|
INSERT INTO identity_providers (id, provider_type, name, is_active, config, tenant_id, created_at, updated_at)
|
|
VALUES (:provider_id, 'feishu', 'Feishu SSO', TRUE, :config, :tenant_id, NOW(), NOW())
|
|
"""),
|
|
{
|
|
"provider_id": provider_id,
|
|
"config": '{"app_id": "", "app_secret": ""}',
|
|
"tenant_id": tenant_id
|
|
}
|
|
)
|
|
|
|
# Step 2: Update org_departments - map feishu_id to external_id and link to provider
|
|
# Only update rows where external_id is NULL (hasn't been migrated)
|
|
connection.execute(sa.text("""
|
|
DO $$
|
|
BEGIN
|
|
IF EXISTS (
|
|
SELECT 1 FROM information_schema.columns
|
|
WHERE table_name = 'org_departments' AND column_name = 'feishu_id'
|
|
) THEN
|
|
UPDATE org_departments od
|
|
SET
|
|
external_id = od.feishu_id,
|
|
provider_id = ip.id
|
|
FROM identity_providers ip
|
|
WHERE od.tenant_id = ip.tenant_id
|
|
AND ip.provider_type = 'feishu'
|
|
AND od.tenant_id IS NOT NULL
|
|
AND od.external_id IS NULL;
|
|
END IF;
|
|
END $$
|
|
"""))
|
|
|
|
# Step 3: Drop feishu_id column after migration
|
|
op.execute("ALTER TABLE org_departments DROP COLUMN IF EXISTS feishu_id")
|
|
# Drop legacy foreign key constraints that were added unexpectedly
|
|
# Using IF EXISTS for safety across different environments
|
|
op.execute("ALTER TABLE org_departments DROP CONSTRAINT IF EXISTS fk_org_departments_provider")
|
|
op.execute("ALTER TABLE org_members DROP CONSTRAINT IF EXISTS fk_org_members_provider")
|
|
# Add status column to org_departments for soft deletion during sync
|
|
op.execute("ALTER TABLE org_departments ADD COLUMN IF NOT EXISTS status VARCHAR(20) DEFAULT 'active'")
|
|
|
|
def downgrade() -> None:
|
|
# ============================================
|
|
# 8. Revert channel_config
|
|
# ============================================
|
|
op.execute("""
|
|
DO $$
|
|
BEGIN
|
|
IF EXISTS (
|
|
SELECT 1 FROM information_schema.columns
|
|
WHERE table_name = 'channel_config' AND column_name = 'app_secret'
|
|
) THEN
|
|
ALTER TABLE channel_config ALTER COLUMN app_secret TYPE VARCHAR(255);
|
|
END IF;
|
|
END $$
|
|
""")
|
|
|
|
# ============================================
|
|
# 7. Recreate departments table
|
|
# ============================================
|
|
op.create_table('departments',
|
|
sa.Column('id', sa.UUID(), nullable=False),
|
|
sa.Column('name', sa.VARCHAR(length=200), nullable=False),
|
|
sa.Column('parent_id', sa.UUID(), nullable=True),
|
|
sa.Column('manager_id', sa.UUID(), nullable=True),
|
|
sa.Column('sort_order', sa.INTEGER(), nullable=True),
|
|
sa.Column('created_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.func.now(), nullable=False),
|
|
sa.ForeignKeyConstraint(['parent_id'], ['departments.id'], name='departments_parent_id_fkey'),
|
|
sa.ForeignKeyConstraint(['manager_id'], ['users.id'], name='departments_manager_id_fkey'),
|
|
sa.PrimaryKeyConstraint('id', name='departments_pkey')
|
|
)
|
|
|
|
# ============================================
|
|
# 6. Revert users constraints and columns
|
|
# ============================================
|
|
op.execute("DROP INDEX IF EXISTS ix_users_primary_mobile")
|
|
op.execute("DROP INDEX IF EXISTS ix_users_tenant_email_unique")
|
|
op.execute("DROP INDEX IF EXISTS ix_users_tenant_mobile_unique")
|
|
op.execute("ALTER TABLE users DROP COLUMN IF EXISTS primary_mobile")
|
|
op.execute("ALTER TABLE users DROP COLUMN IF EXISTS registration_source")
|
|
op.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS feishu_open_id VARCHAR(255)")
|
|
op.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS feishu_union_id VARCHAR(255)")
|
|
op.execute("CREATE INDEX IF NOT EXISTS ix_users_open_id ON users(open_id)")
|
|
op.execute("CREATE INDEX IF NOT EXISTS ix_users_union_id ON users(union_id)")
|
|
op.execute("CREATE UNIQUE INDEX IF NOT EXISTS ix_users_feishu_open_id ON users(feishu_open_id)")
|
|
|
|
# ============================================
|
|
# 5. Revert org_members
|
|
# ============================================
|
|
op.execute("DROP INDEX IF EXISTS ix_org_members_user_id")
|
|
op.execute("DROP INDEX IF EXISTS ix_org_members_unionid")
|
|
op.execute("DROP INDEX IF EXISTS ix_org_members_external_id")
|
|
op.execute("ALTER TABLE org_members DROP CONSTRAINT IF EXISTS fk_org_members_user")
|
|
op.execute("ALTER TABLE org_members DROP CONSTRAINT IF EXISTS fk_org_members_provider")
|
|
op.execute("ALTER TABLE org_members DROP COLUMN IF EXISTS user_id")
|
|
op.execute("ALTER TABLE org_members DROP COLUMN IF EXISTS provider_id")
|
|
op.execute("ALTER TABLE org_members DROP COLUMN IF EXISTS unionid")
|
|
op.execute("ALTER TABLE org_members DROP COLUMN IF EXISTS external_id")
|
|
|
|
# ============================================
|
|
# 4. Revert org_departments
|
|
# ============================================
|
|
op.execute("DROP INDEX IF EXISTS ix_org_departments_external_id")
|
|
op.execute("ALTER TABLE org_departments DROP CONSTRAINT IF EXISTS fk_org_departments_provider")
|
|
op.execute("ALTER TABLE org_departments DROP COLUMN IF EXISTS provider_id")
|
|
op.execute("ALTER TABLE org_departments DROP COLUMN IF EXISTS external_id")
|
|
|
|
# ============================================
|
|
# 3. Revert tenants
|
|
# ============================================
|
|
op.execute("DROP INDEX IF EXISTS ux_tenants_sso_domain")
|
|
op.execute("ALTER TABLE tenants DROP COLUMN IF EXISTS sso_domain")
|
|
op.execute("ALTER TABLE tenants DROP COLUMN IF EXISTS sso_enabled")
|
|
|
|
# ============================================
|
|
# 2. Drop sso_scan_sessions
|
|
# ============================================
|
|
op.drop_table('sso_scan_sessions')
|
|
|
|
# ============================================
|
|
# 1. Drop identity_providers
|
|
# ============================================
|
|
op.drop_table('identity_providers')
|
|
|
|
# Note: Downgrade is NOT idempotent - it resets data
|
|
# In production, you may want to skip this or make it optional
|
|
connection = op.get_bind()
|
|
|
|
# Add back feishu_id column
|
|
op.execute("ALTER TABLE org_departments ADD COLUMN IF NOT EXISTS feishu_id VARCHAR(100)")
|
|
|
|
# Restore feishu_id from external_id
|
|
connection.execute(sa.text("""
|
|
UPDATE org_departments
|
|
SET feishu_id = external_id
|
|
WHERE external_id IS NOT NULL
|
|
"""))
|
|
|
|
# Delete the identity providers created by this migration
|
|
connection.execute(sa.text("""
|
|
DELETE FROM identity_providers
|
|
WHERE provider_type = 'feishu'
|
|
AND config::text = '{"app_id": "", "app_secret": ""}'
|
|
""")) |