jude-cehua/models.py

"""
数据库模型定义
根据提供的SQL文件创建对应的SQLAlchemy模型
"""

from datetime import datetime
from typing import Optional
from flask_sqlalchemy import SQLAlchemy
from sqlalchemy import text
from sqlalchemy import BigInteger, String, Text, DateTime, Boolean, Integer, Enum
from utils import get_beijing_time

# 创建数据库实例
db = SQLAlchemy()

class User(db.Model):
    """增强版用户表"""
    __tablename__ = 'users'
    
    # 基础信息
    user_id = db.Column(BigInteger, primary_key=True, autoincrement=True)
    username = db.Column(String(50), nullable=False, unique=True, comment='用户名')
    
    # 登录方式相关
    login_type = db.Column(Enum('dingtalk', 'wecom', 'email', 'phone'), nullable=False, comment='登录方式')
    dingtalk_id = db.Column(String(200), unique=True, comment='钉钉用户ID')
    wecom_id = db.Column(String(200), unique=True, comment='企业微信用户ID')
    
    # 联系信息
    email = db.Column(String(255), unique=True, comment='邮箱地址')
    phone = db.Column(String(20), unique=True, comment='手机号码')
    
    # 密码安全（邮箱/手机登录使用）
    password_hash = db.Column(String(255), comment='密码哈希值')
    salt = db.Column(String(100), comment='密码盐值')
    failed_login_attempts = db.Column(Integer, default=0, comment='失败登录次数')
    locked_until = db.Column(DateTime, nullable=True, comment='账户锁定到期时间')
    
    # 个人信息
    avatar_url = db.Column(String(500), comment='用户头像URL')
    nickname = db.Column(String(100), comment='用户昵称')
    
    # 权限与状态
    role = db.Column(Enum('admin', 'user'), default='user', comment='用户角色')
    status = db.Column(Enum('active', 'inactive', 'suspended'), default='active', comment='用户状态')
    is_deleted = db.Column(Boolean, default=False, comment='是否已删除')
    deleted_at = db.Column(DateTime, nullable=True, comment='删除时间')
    
    # 第三方集成信息
    dingtalk_unionid = db.Column(String(200), comment='钉钉UnionID')
    wecom_unionid = db.Column(String(200), comment='企业微信UnionID')
    openid = db.Column(String(200), comment='通用OpenID')
    
    # 时间戳
    created_at = db.Column(DateTime, default=get_beijing_time, comment='创建时间')
    updated_at = db.Column(DateTime, default=get_beijing_time, onupdate=get_beijing_time, comment='更新时间')
    last_login_at = db.Column(DateTime, nullable=True, comment='最后登录时间')
    last_login_ip = db.Column(String(45), comment='最后登录IP')
    
    # 关系
    conversations = db.relationship('Conversation', backref='user', lazy=True, cascade='all, delete-orphan')
    
    def __repr__(self):
        return f'<User {self.username}>'
    
    def to_dict(self):
        """转换为字典格式"""
        return {
            'user_id': self.user_id,
            'username': self.username,
            'login_type': self.login_type,
            'dingtalk_id': self.dingtalk_id,
            'wecom_id': self.wecom_id,
            'email': self.email,
            'phone': self.phone,
            'avatar_url': self.avatar_url,
            'nickname': self.nickname,
            'role': self.role,
            'status': self.status,
            'is_deleted': self.is_deleted,
            'deleted_at': self.deleted_at.isoformat() if self.deleted_at else None,
            'dingtalk_unionid': self.dingtalk_unionid,
            'wecom_unionid': self.wecom_unionid,
            'openid': self.openid,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'updated_at': self.updated_at.isoformat() if self.updated_at else None,
            'last_login_at': self.last_login_at.isoformat() if self.last_login_at else None,
            'last_login_ip': self.last_login_ip
        }

class Conversation(db.Model):
    """对话会话表"""
    __tablename__ = 'conversations'
    
    conversation_id = db.Column(BigInteger, primary_key=True, autoincrement=True)
    user_id = db.Column(BigInteger, db.ForeignKey('users.user_id'), nullable=False, comment='用户ID')
    title = db.Column(String(500), nullable=False, comment='对话标题')
    conversation_type = db.Column(Enum('chat', 'qa', 'assistant'), default='chat', comment='对话类型')
    session_id = db.Column(String(255), nullable=True, unique=True, comment='AI会话ID，用于上下文隔离')
    is_archived = db.Column(Boolean, default=False, comment='是否已归档')
    is_deleted = db.Column(Boolean, default=False, comment='是否已删除')
    deleted_at = db.Column(DateTime, nullable=True, comment='删除时间')
    message_count = db.Column(Integer, default=0, comment='消息数量')
    last_message_at = db.Column(DateTime, nullable=True, comment='最后消息时间')
    created_at = db.Column(DateTime, default=get_beijing_time, comment='创建时间')
    updated_at = db.Column(DateTime, default=get_beijing_time, onupdate=get_beijing_time, comment='更新时间')
    
    # 关系
    messages = db.relationship('Message', backref='conversation', lazy=True, cascade='all, delete-orphan')
    
    def __repr__(self):
        return f'<Conversation {self.conversation_id}: {self.title}>'
    
    def to_dict(self):
        """转换为字典格式"""
        return {
            'conversation_id': self.conversation_id,
            'user_id': self.user_id,
            'title': self.title,
            'conversation_type': self.conversation_type,
            'session_id': self.session_id,
            'is_archived': self.is_archived,
            'is_deleted': self.is_deleted,
            'deleted_at': self.deleted_at.isoformat() if self.deleted_at else None,
            'message_count': self.message_count,
            'last_message_at': self.last_message_at.isoformat() if self.last_message_at else None,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'updated_at': self.updated_at.isoformat() if self.updated_at else None
        }

class Message(db.Model):
    """消息记录表"""
    __tablename__ = 'messages'
    
    message_id = db.Column(BigInteger, primary_key=True, autoincrement=True)
    conversation_id = db.Column(BigInteger, db.ForeignKey('conversations.conversation_id'), nullable=False, comment='会话ID')
    role = db.Column(Enum('user', 'assistant', 'system', 'thinking'), nullable=False, comment='消息角色')
    content = db.Column(Text, nullable=False, comment='消息内容')
    thinking_content = db.Column(Text, nullable=True, comment='AI思考过程内容')
    content_type = db.Column(Enum('text', 'image', 'file'), default='text', comment='内容类型')
    message_order = db.Column(Integer, nullable=False, comment='消息顺序(从1开始)')
    token_count = db.Column(Integer, default=0, comment='Token数量')
    is_deleted = db.Column(Boolean, default=False, comment='是否已删除')
    deleted_at = db.Column(DateTime, nullable=True, comment='删除时间')
    created_at = db.Column(DateTime, default=get_beijing_time, comment='创建时间')
    
    def __repr__(self):
        return f'<Message {self.message_id}: {self.role}>'
    
    def to_dict(self):
        """转换为字典格式"""
        return {
            'message_id': self.message_id,
            'conversation_id': self.conversation_id,
            'role': self.role,
            'content': self.content,
            'thinking_content': self.thinking_content,
            'content_type': self.content_type,
            'message_order': self.message_order,
            'token_count': self.token_count,
            'is_deleted': self.is_deleted,
            'deleted_at': self.deleted_at.isoformat() if self.deleted_at else None,
            'created_at': self.created_at.isoformat() if self.created_at else None
        }

class SystemConfig(db.Model):
    """系统配置表"""
    __tablename__ = 'system_config'
    
    config_id = db.Column(Integer, primary_key=True, autoincrement=True)
    config_key = db.Column(String(100), nullable=False, unique=True, comment='配置键')
    config_value = db.Column(Text, nullable=True, comment='配置值')
    config_type = db.Column(Enum('string', 'number', 'boolean', 'json'), default='string', comment='配置类型')
    description = db.Column(String(500), nullable=True, comment='配置描述')
    created_at = db.Column(DateTime, default=get_beijing_time)
    updated_at = db.Column(DateTime, default=get_beijing_time, onupdate=get_beijing_time)
    
    def __repr__(self):
        return f'<SystemConfig {self.config_key}>'
    
    def to_dict(self):
        """转换为字典格式"""
        return {
            'config_id': self.config_id,
            'config_key': self.config_key,
            'config_value': self.config_value,
            'config_type': self.config_type,
            'description': self.description,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'updated_at': self.updated_at.isoformat() if self.updated_at else None
        }

def init_database(app):
    """初始化数据库"""
    db.init_app(app)
    
    with app.app_context():
        # 创建所有表
        db.create_all()
        
        # 创建索引
        create_indexes()
        
        # 插入默认配置
        insert_default_configs()

def create_indexes():
    """创建数据库索引"""
    try:
        # 定义所有需要创建的索引
        indexes = [
            # 用户表索引
            ("idx_users_login_type_status", "CREATE INDEX idx_users_login_type_status ON users (login_type, status)"),
            ("idx_users_role_status", "CREATE INDEX idx_users_role_status ON users (role, status)"),
            ("idx_users_email", "CREATE INDEX idx_users_email ON users (email)"),
            ("idx_users_phone", "CREATE INDEX idx_users_phone ON users (phone)"),
            ("idx_users_dingtalk", "CREATE INDEX idx_users_dingtalk ON users (dingtalk_id)"),
            ("idx_users_wecom", "CREATE INDEX idx_users_wecom ON users (wecom_id)"),
            ("idx_users_created_at", "CREATE INDEX idx_users_created_at ON users (created_at)"),
            ("idx_users_last_login", "CREATE INDEX idx_users_last_login ON users (last_login_at)"),
            ("idx_users_is_deleted", "CREATE INDEX idx_users_is_deleted ON users (is_deleted)"),
            ("idx_users_status_deleted", "CREATE INDEX idx_users_status_deleted ON users (status, is_deleted)"),
            
            # 对话表索引
            ("idx_conversations_user_id", "CREATE INDEX idx_conversations_user_id ON conversations (user_id)"),
            ("idx_conversations_user_updated", "CREATE INDEX idx_conversations_user_updated ON conversations (user_id, updated_at)"),
            ("idx_conversations_user_archived", "CREATE INDEX idx_conversations_user_archived ON conversations (user_id, is_archived)"),
            ("idx_conversations_user_deleted", "CREATE INDEX idx_conversations_user_deleted ON conversations (user_id, is_deleted)"),
            ("idx_conversations_updated_at", "CREATE INDEX idx_conversations_updated_at ON conversations (updated_at)"),
            ("idx_conversations_last_message", "CREATE INDEX idx_conversations_last_message ON conversations (last_message_at)"),
            ("idx_conversations_type_status", "CREATE INDEX idx_conversations_type_status ON conversations (conversation_type, is_deleted)"),
            
            # 消息表索引
            ("idx_messages_conversation_id", "CREATE INDEX idx_messages_conversation_id ON messages (conversation_id)"),
            ("idx_messages_conversation_order", "CREATE INDEX idx_messages_conversation_order ON messages (conversation_id, message_order)"),
            ("idx_messages_conversation_deleted", "CREATE INDEX idx_messages_conversation_deleted ON messages (conversation_id, is_deleted)"),
            ("idx_messages_created_at", "CREATE INDEX idx_messages_created_at ON messages (created_at)"),
            ("idx_messages_role", "CREATE INDEX idx_messages_role ON messages (role)"),
            ("idx_messages_content_type", "CREATE INDEX idx_messages_content_type ON messages (content_type)")
        ]
        
        db.session.commit()
    except Exception as e:
        db.session.rollback()

def insert_default_configs():
    """插入默认系统配置"""
    try:
        # 检查是否已有配置
        if SystemConfig.query.first():
            return
        
        # 只插入必要的配置
        essential_configs = [
            ('system_name', 'AI聊天系统', 'string', '系统名称'),
            ('system_version', '2.0.0', 'string', '系统版本'),
            ('user_registration', 'true', 'boolean', '是否允许用户注册'),
            ('default_user_role', 'user', 'string', '默认用户角色')
        ]
        
        for config_key, config_value, config_type, description in essential_configs:
            config = SystemConfig(
                config_key=config_key,
                config_value=config_value,
                config_type=config_type,
                description=description
            )
            db.session.add(config)
        
        db.session.commit()
    except Exception as e:
        db.session.rollback()

def get_or_create_user(login_type: str, **kwargs) -> User:
    """获取或创建用户"""
    try:
        # 根据登录类型查找用户
        if login_type == 'dingtalk':
            dingtalk_id = kwargs.get('dingtalk_id')
            if dingtalk_id:
                user = User.query.filter_by(dingtalk_id=dingtalk_id, is_deleted=False).first()
                if user:
                    return user
        
        # 如果没找到，创建新用户
        # 确保username不重复且符合长度要求
        username = kwargs.get('username', f"{login_type}_user_{get_beijing_time().timestamp()}")
        
        # 确保用户名长度符合要求（至少3个字符）
        if len(username) < 3:
            username = f"{username}_{get_beijing_time().timestamp()}"
        
        # 检查username是否已存在，如果存在则生成新的
        counter = 1
        original_username = username
        while User.query.filter_by(username=username, is_deleted=False).first():
            username = f"{original_username}_{counter}"
            counter += 1
        
        # 从kwargs中移除username，避免重复参数
        if 'username' in kwargs:
            del kwargs['username']
        
        # 处理唯一约束冲突（email, phone等）
        # 如果email已存在，清空email
        if 'email' in kwargs and kwargs['email']:
            existing_email = User.query.filter_by(email=kwargs['email'], is_deleted=False).first()
            if existing_email:
                kwargs['email'] = ''
        
        # 如果phone已存在，清空phone
        if 'phone' in kwargs and kwargs['phone']:
            existing_phone = User.query.filter_by(phone=kwargs['phone'], is_deleted=False).first()
            if existing_phone:
                kwargs['phone'] = ''
        
        # 确保空字符串字段设置为None，避免数据库约束错误
        if 'email' in kwargs and kwargs['email'] == '':
            kwargs['email'] = None
        if 'phone' in kwargs and kwargs['phone'] == '':
            kwargs['phone'] = None
            
        user = User(
            login_type=login_type,
            username=username,
            **kwargs
        )
        
        db.session.add(user)
        db.session.commit()
        return user
        
    except Exception as e:
        db.session.rollback()
        raise

def update_user_login_info(user: User, login_ip: str = None):
    """更新用户登录信息"""
    try:
        user.last_login_at = get_beijing_time()
        if login_ip:
            user.last_login_ip = login_ip
        db.session.commit()
    except Exception as e:
        db.session.rollback()
        raise