89 lines
3.4 KiB
Python
89 lines
3.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""迁移脚本：添加邮箱、真实姓名和邮箱验证码表。"""
|
|
import logging
|
|
from sqlalchemy import text
|
|
from insurance.db.compat import db
|
|
|
|
|
|
def migrate():
|
|
"""添加用户邮箱字段和验证码表。"""
|
|
try:
|
|
from sqlalchemy import inspect
|
|
|
|
inspector = inspect(db.engine)
|
|
user_columns = [col["name"] for col in inspector.get_columns("wecom_user_mapping")]
|
|
dialect = db.engine.dialect.name
|
|
|
|
additions = [
|
|
("email", "VARCHAR(255)"),
|
|
("real_name", "VARCHAR(128)"),
|
|
("email_verified", "VARCHAR(16) DEFAULT 'false'"),
|
|
]
|
|
for name, column_type in additions:
|
|
if name not in user_columns:
|
|
db.session.execute(text(
|
|
f"ALTER TABLE wecom_user_mapping ADD COLUMN {name} {column_type}"
|
|
))
|
|
logging.info("成功添加 wecom_user_mapping.%s 字段", name)
|
|
|
|
index_names = {idx["name"] for idx in inspector.get_indexes("wecom_user_mapping")}
|
|
if "idx_wecom_user_mapping_email_unique" not in index_names:
|
|
duplicate_email = db.session.execute(text("""
|
|
SELECT email
|
|
FROM wecom_user_mapping
|
|
WHERE email IS NOT NULL AND email <> ''
|
|
GROUP BY email
|
|
HAVING COUNT(*) > 1
|
|
LIMIT 1
|
|
""")).fetchone()
|
|
if duplicate_email:
|
|
logging.warning(
|
|
"wecom_user_mapping 存在重复邮箱 %s，跳过邮箱唯一索引创建",
|
|
duplicate_email[0],
|
|
)
|
|
else:
|
|
db.session.execute(text("""
|
|
CREATE UNIQUE INDEX IF NOT EXISTS idx_wecom_user_mapping_email_unique
|
|
ON wecom_user_mapping (email)
|
|
WHERE email IS NOT NULL AND email <> ''
|
|
"""))
|
|
logging.info("成功创建 wecom_user_mapping.email 唯一索引")
|
|
|
|
if not inspector.has_table("email_verification_codes"):
|
|
id_type = "INTEGER PRIMARY KEY AUTOINCREMENT" if dialect == "sqlite" else "SERIAL PRIMARY KEY"
|
|
db.session.execute(text(f"""
|
|
CREATE TABLE email_verification_codes (
|
|
id {id_type},
|
|
email VARCHAR(255) NOT NULL,
|
|
purpose VARCHAR(32) NOT NULL,
|
|
code_hash VARCHAR(128) NOT NULL,
|
|
expires_at TIMESTAMP NOT NULL,
|
|
consumed_at TIMESTAMP NULL,
|
|
attempts INTEGER DEFAULT 0,
|
|
ip VARCHAR(64),
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
"""))
|
|
db.session.execute(text(
|
|
"CREATE INDEX idx_email_verification_email_purpose "
|
|
"ON email_verification_codes (email, purpose)"
|
|
))
|
|
logging.info("成功创建 email_verification_codes 表")
|
|
|
|
db.session.commit()
|
|
except Exception as e:
|
|
logging.error("添加邮箱验证码能力失败: %s", e)
|
|
db.session.rollback()
|
|
raise
|
|
|
|
|
|
def downgrade():
|
|
"""回滚邮箱验证码表。"""
|
|
try:
|
|
db.session.execute(text("DROP TABLE IF EXISTS email_verification_codes"))
|
|
db.session.commit()
|
|
except Exception as e:
|
|
logging.error("删除 email_verification_codes 表失败: %s", e)
|
|
db.session.rollback()
|
|
raise
|