aiagent/backend/app/core/database.py

"""
数据库配置
"""
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from app.core.config import settings

# 创建数据库引擎（MySQL）
engine = create_engine(
    settings.DATABASE_URL,
    pool_pre_ping=True,
    pool_size=10,
    max_overflow=20,
    echo=settings.SQL_ECHO  # SQL 日志独立开关，默认关闭
)

# 创建会话工厂
# expire_on_commit=False：commit 后仍可读取已加载标量，避免 FastAPI 在序列化 ExecutionResponse 时
# 因会话已关闭而再次触发懒加载，从而出现「仅 HTTP 报 DATABASE_ERROR、TestClient 正常」的现象。
SessionLocal = sessionmaker(
    autocommit=False, autoflush=False, bind=engine, expire_on_commit=False
)

# 创建基础模型类
Base = declarative_base()


def get_db():
    """获取数据库会话"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def init_db():
    """初始化数据库，创建所有表"""
    # 导入所有模型，确保它们被注册
    import app.models.user
    import app.models.workflow
    import app.models.agent
    import app.models.execution
    import app.models.model_config
    import app.models.workflow_template
    import app.models.permission
    import app.models.alert_rule
    import app.models.agent_llm_log
    import app.models.agent_vector_memory
    import app.models.agent_learning_pattern
    import app.models.agent_schedule
    import app.models.knowledge_base
    import app.models.notification
    import app.models.orchestration_template
    import app.models.plugin
    import app.models.goal
    import app.models.task
    import app.models.tool
    import app.models.data_source
    import app.models.execution_log
    import app.models.agent_execution_log
    import app.models.feedback_record
    import app.models.knowledge_entry
    import app.models.node_template
    import app.models.persistent_user_memory
    import app.models.shadow_comparison
    import app.models.user_behavior
    import app.models.user_feishu_open_id
    import app.models.user_fingerprint
    import app.models.workflow_version
    import app.models.audit_log
    import app.models.conversation_branch
    import app.models.push_subscription
    import app.models.fcm_token
    import app.models.workspace
    import app.models.scene_contract
    import app.models.team
    import app.models.chat_message
    import app.models.agent_session
    import app.models.billing
    Base.metadata.create_all(bind=engine)

    # v1.1.0: 手机验证字段迁移（safe ALTER）
    _run_safe_migrations()

    # v1.3: Auto-seed RBAC roles & permissions
    _seed_rbac()


def _run_safe_migrations():
    """安全迁移：对可能缺失的列执行 ALTER TABLE（幂等，失败不中断）。"""
    import logging
    from sqlalchemy import text
    logger = logging.getLogger(__name__)
    migrations = [
        # v1.0: 手机验证
        "ALTER TABLE users ADD COLUMN phone_verified TINYINT(1) DEFAULT 0",
        "ALTER TABLE users ADD COLUMN phone_verified_at DATETIME NULL",
        # v1.2: 订阅与用量
        "ALTER TABLE users ADD COLUMN subscription_tier VARCHAR(20) DEFAULT 'free'",
        "ALTER TABLE users ADD COLUMN subscription_expires_at DATETIME NULL",
        "ALTER TABLE users ADD COLUMN daily_usage_count INT DEFAULT 0",
        "ALTER TABLE users ADD COLUMN daily_usage_date DATE NULL",
        # v1.3: 多租户 workspace_id
        "ALTER TABLE agent_sessions ADD COLUMN workspace_id CHAR(36) NULL",
        "ALTER TABLE agent_sessions ADD INDEX idx_session_workspace (workspace_id)",
        "ALTER TABLE chat_messages ADD COLUMN workspace_id CHAR(36) NULL",
        "ALTER TABLE chat_messages ADD INDEX idx_msg_workspace (workspace_id)",
        "ALTER TABLE agent_execution_logs ADD COLUMN workspace_id CHAR(36) NULL",
        "ALTER TABLE agent_execution_logs ADD INDEX idx_exec_workspace (workspace_id)",
        "ALTER TABLE agent_llm_logs ADD COLUMN workspace_id CHAR(36) NULL",
        "ALTER TABLE agent_llm_logs ADD INDEX idx_llm_workspace (workspace_id)",
        "ALTER TABLE conversation_branches ADD COLUMN workspace_id CHAR(36) NULL",
        "ALTER TABLE conversation_branches ADD INDEX idx_branch_workspace (workspace_id)",
        "ALTER TABLE node_templates ADD COLUMN workspace_id CHAR(36) NULL",
        "ALTER TABLE node_templates ADD INDEX idx_nt_workspace (workspace_id)",
        "ALTER TABLE node_plugins ADD COLUMN workspace_id CHAR(36) NULL",
        "ALTER TABLE node_plugins ADD INDEX idx_plugin_workspace (workspace_id)",
        "ALTER TABLE knowledge_entries ADD COLUMN workspace_id CHAR(36) NULL",
        "ALTER TABLE knowledge_entries ADD INDEX idx_ke_workspace (workspace_id)",
        "ALTER TABLE alert_rules ADD COLUMN workspace_id CHAR(36) NULL",
        "ALTER TABLE alert_rules ADD INDEX idx_ar_workspace (workspace_id)",
        "ALTER TABLE alert_logs ADD COLUMN workspace_id CHAR(36) NULL",
        "ALTER TABLE alert_logs ADD INDEX idx_al_workspace (workspace_id)",
        "ALTER TABLE agent_ratings ADD COLUMN workspace_id CHAR(36) NULL",
        "ALTER TABLE agent_ratings ADD INDEX idx_arating_workspace (workspace_id)",
        "ALTER TABLE agent_favorites ADD COLUMN workspace_id CHAR(36) NULL",
        "ALTER TABLE agent_favorites ADD INDEX idx_afav_workspace (workspace_id)",
        "ALTER TABLE template_ratings ADD COLUMN workspace_id CHAR(36) NULL",
        "ALTER TABLE template_ratings ADD INDEX idx_trating_workspace (workspace_id)",
        "ALTER TABLE template_favorites ADD COLUMN workspace_id CHAR(36) NULL",
        "ALTER TABLE template_favorites ADD INDEX idx_tfav_workspace (workspace_id)",
    ]
    for raw_sql in migrations:
        try:
            with engine.begin() as conn:
                conn.execute(text(raw_sql))
            logger.info("迁移执行: %s", raw_sql)
        except Exception:
            # 列已存在或其他错误，跳过
            pass

    # 回填已有数据的 workspace_id
    _backfill_workspace_ids(logger)


def _backfill_workspace_ids(logger):
    """回填已有数据的 workspace_id（通过父表 FK 链推导）。"""
    from sqlalchemy import text

    backfills = [
        # (table, fk_column, parent_table) — 通过父表的 workspace_id 回填
        ("agent_sessions", "agent_id", "agents"),
        ("agent_execution_logs", "agent_id", "agents"),
        ("agent_llm_logs", "agent_id", "agents"),
        ("agent_ratings", "agent_id", "agents"),
        ("agent_favorites", "agent_id", "agents"),
        ("template_ratings", "template_id", "workflow_templates"),
        ("template_favorites", "template_id", "workflow_templates"),
    ]

    for table, fk_col, parent_table in backfills:
        try:
            with engine.begin() as conn:
                result = conn.execute(text(
                    f"UPDATE {table} t "
                    f"JOIN {parent_table} p ON t.{fk_col} = p.id "
                    f"SET t.workspace_id = p.workspace_id "
                    f"WHERE t.workspace_id IS NULL AND p.workspace_id IS NOT NULL"
                ))
                if result.rowcount:
                    logger.info("回填 %s.workspace_id: %s 行", table, result.rowcount)
        except Exception:
            pass  # 父表可能没有 workspace_id 列

    # chat_messages: 通过 agent_sessions 回填
    try:
        with engine.begin() as conn:
            result = conn.execute(text(
                "UPDATE chat_messages cm "
                "JOIN agent_sessions s ON cm.session_id = s.id "
                "SET cm.workspace_id = s.workspace_id "
                "WHERE cm.workspace_id IS NULL AND s.workspace_id IS NOT NULL"
            ))
            if result.rowcount:
                logger.info("回填 chat_messages.workspace_id: %s 行", result.rowcount)
    except Exception:
        pass

    # conversation_branches: 通过 agent_sessions 回填
    try:
        with engine.begin() as conn:
            result = conn.execute(text(
                "UPDATE conversation_branches cb "
                "JOIN agent_sessions s ON cb.parent_session_id = s.id "
                "SET cb.workspace_id = s.workspace_id "
                "WHERE cb.workspace_id IS NULL AND s.workspace_id IS NOT NULL"
            ))
            if result.rowcount:
                logger.info("回填 conversation_branches.workspace_id: %s 行", result.rowcount)
    except Exception:
        pass


def _seed_rbac():
    """Auto-seed system roles & permissions on startup (idempotent)."""
    import logging
    import uuid
    from sqlalchemy import text

    logger = logging.getLogger(__name__)

    # Check if RBAC tables exist (may fail on fresh DB before create_all runs)
    try:
        from app.models.permission import Role, Permission
    except Exception:
        logger.warning("RBAC model not available, skip seeding")
        return

    db = SessionLocal()
    try:
        # ── Permissions ──
        PERMISSIONS = [
            {"name": "工作流-创建", "code": "workflow:create", "resource": "workflow", "action": "create"},
            {"name": "工作流-查看", "code": "workflow:read", "resource": "workflow", "action": "read"},
            {"name": "工作流-更新", "code": "workflow:update", "resource": "workflow", "action": "update"},
            {"name": "工作流-删除", "code": "workflow:delete", "resource": "workflow", "action": "delete"},
            {"name": "工作流-执行", "code": "workflow:execute", "resource": "workflow", "action": "execute"},
            {"name": "工作流-分享", "code": "workflow:share", "resource": "workflow", "action": "share"},
            {"name": "Agent-创建", "code": "agent:create", "resource": "agent", "action": "create"},
            {"name": "Agent-查看", "code": "agent:read", "resource": "agent", "action": "read"},
            {"name": "Agent-更新", "code": "agent:update", "resource": "agent", "action": "update"},
            {"name": "Agent-删除", "code": "agent:delete", "resource": "agent", "action": "delete"},
            {"name": "Agent-执行", "code": "agent:execute", "resource": "agent", "action": "execute"},
            {"name": "Agent-部署", "code": "agent:deploy", "resource": "agent", "action": "deploy"},
            {"name": "执行-查看", "code": "execution:read", "resource": "execution", "action": "read"},
            {"name": "执行-取消", "code": "execution:cancel", "resource": "execution", "action": "cancel"},
            {"name": "数据源-创建", "code": "data_source:create", "resource": "data_source", "action": "create"},
            {"name": "数据源-查看", "code": "data_source:read", "resource": "data_source", "action": "read"},
            {"name": "数据源-更新", "code": "data_source:update", "resource": "data_source", "action": "update"},
            {"name": "数据源-删除", "code": "data_source:delete", "resource": "data_source", "action": "delete"},
            {"name": "模型配置-创建", "code": "model_config:create", "resource": "model_config", "action": "create"},
            {"name": "模型配置-查看", "code": "model_config:read", "resource": "model_config", "action": "read"},
            {"name": "模型配置-更新", "code": "model_config:update", "resource": "model_config", "action": "update"},
            {"name": "模型配置-删除", "code": "model_config:delete", "resource": "model_config", "action": "delete"},
            {"name": "权限-管理", "code": "permission:manage", "resource": "permission", "action": "manage"},
        ]

        permission_map = {}
        for perm_data in PERMISSIONS:
            existing = db.query(Permission).filter(Permission.code == perm_data["code"]).first()
            if existing:
                permission_map[perm_data["code"]] = existing
            else:
                perm = Permission(
                    id=str(uuid.uuid4()),
                    name=perm_data["name"],
                    code=perm_data["code"],
                    resource=perm_data["resource"],
                    action=perm_data["action"],
                )
                db.add(perm)
                permission_map[perm_data["code"]] = perm
        db.commit()

        # ── Roles ──
        SYSTEM_ROLES = [
            {"name": "admin", "description": "系统管理员，拥有所有权限", "is_system": True},
            {"name": "developer", "description": "开发者，可以创建和管理工作流、Agent", "is_system": True},
            {"name": "viewer", "description": "查看者，只能查看工作流和执行记录", "is_system": True},
            {"name": "operator", "description": "操作员，可以执行工作流，但不能修改", "is_system": True},
        ]

        role_map = {}
        for role_data in SYSTEM_ROLES:
            existing = db.query(Role).filter(Role.name == role_data["name"]).first()
            if existing:
                role_map[role_data["name"]] = existing
            else:
                role = Role(
                    id=str(uuid.uuid4()),
                    name=role_data["name"],
                    description=role_data["description"],
                    is_system=role_data["is_system"],
                )
                db.add(role)
                role_map[role_data["name"]] = role
        db.commit()

        # ── Role-Permission mappings ──
        ROLE_PERMISSIONS = {
            "admin": ["*"],
            "developer": [
                "workflow:create", "workflow:read", "workflow:update", "workflow:delete", "workflow:execute", "workflow:share",
                "agent:create", "agent:read", "agent:update", "agent:delete", "agent:execute", "agent:deploy",
                "execution:read", "execution:cancel",
                "data_source:create", "data_source:read", "data_source:update", "data_source:delete",
                "model_config:create", "model_config:read", "model_config:update", "model_config:delete",
            ],
            "viewer": ["workflow:read", "agent:read", "execution:read", "data_source:read", "model_config:read"],
            "operator": ["workflow:read", "workflow:execute", "agent:read", "agent:execute", "execution:read", "execution:cancel"],
        }

        for role_name, perm_codes in ROLE_PERMISSIONS.items():
            role = role_map.get(role_name)
            if not role:
                continue
            if perm_codes == ["*"]:
                role.permissions = list(permission_map.values())
            else:
                role.permissions = [permission_map[c] for c in perm_codes if c in permission_map]
        db.commit()

        logger.info("RBAC seeded: %d roles, %d permissions", len(role_map), len(permission_map))

    except Exception:
        db.rollback()
        logger.warning("RBAC seed skipped (tables may not exist yet)", exc_info=True)
    finally:
        db.close()