"""执行 auth.users status CHECK 约束迁移"""
import os
import psycopg2
from dotenv import load_dotenv

load_dotenv()
dsn = os.environ["APP_DB_DSN"]
conn = psycopg2.connect(dsn)
conn.autocommit = True
cur = conn.cursor()

print(f"连接: {dsn.split('@')[1] if '@' in dsn else dsn}")

# 删除旧约束
cur.execute("ALTER TABLE auth.users DROP CONSTRAINT IF EXISTS users_status_check")
print("已删除旧约束 users_status_check")

# 添加新约束（含 'new'）
cur.execute("""
    ALTER TABLE auth.users ADD CONSTRAINT users_status_check
    CHECK (status IN ('new', 'pending', 'approved', 'rejected', 'disabled'))
""")
print("已添加新约束（含 'new'）")

# 验证
cur.execute("""
    SELECT conname, pg_get_constraintdef(oid)
    FROM pg_constraint
    WHERE conrelid = 'auth.users'::regclass AND contype = 'c'
""")
for row in cur.fetchall():
    print(f"验证: {row[0]}: {row[1]}")

conn.close()
print("完成")