aiagent/backend/app/api/data_sources.py

"""
数据源管理API
"""
from fastapi import APIRouter, Depends, Query
from sqlalchemy.orm import Session
from pydantic import BaseModel
from typing import List, Optional, Dict, Any
from datetime import datetime
import logging
from app.core.database import get_db
from app.models.data_source import DataSource
from app.api.auth import get_current_user
from app.api.deps import require_workspace_admin, WorkspaceContext
from app.models.user import User
from app.core.exceptions import NotFoundError, ValidationError
from app.services.data_source_connector import DataSourceConnector

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/v1/data-sources", tags=["data-sources"])


class DataSourceCreate(BaseModel):
    """数据源创建模型"""
    name: str
    type: str
    description: Optional[str] = None
    config: Dict[str, Any]


class DataSourceUpdate(BaseModel):
    """数据源更新模型"""
    name: Optional[str] = None
    description: Optional[str] = None
    config: Optional[Dict[str, Any]] = None
    status: Optional[str] = None


class DataSourceResponse(BaseModel):
    """数据源响应模型"""
    id: str
    name: str
    type: str
    description: Optional[str]
    config: Dict[str, Any]
    status: str
    user_id: str
    last_connected_at: Optional[datetime]
    created_at: datetime
    updated_at: datetime

    class Config:
        from_attributes = True


@router.get("", response_model=List[DataSourceResponse])
async def get_data_sources(
    skip: int = 0,
    limit: int = 100,
    type: Optional[str] = None,
    status: Optional[str] = None,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """获取数据源列表"""
    query = db.query(DataSource).filter(
        DataSource.user_id == current_user.id
    )
    
    if type:
        query = query.filter(DataSource.type == type)
    
    if status:
        query = query.filter(DataSource.status == status)
    
    data_sources = query.order_by(DataSource.created_at.desc()).offset(skip).limit(limit).all()
    return data_sources


@router.post("", response_model=DataSourceResponse, status_code=201)
async def create_data_source(
    data_source_data: DataSourceCreate,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """创建数据源"""
    # 验证数据源类型
    valid_types = ['mysql', 'postgresql', 'mongodb', 'redis', 'csv', 'json', 'api', 's3']
    if data_source_data.type not in valid_types:
        raise ValidationError(f"不支持的数据源类型: {data_source_data.type}")
    
    # 测试连接（可选，如果配置了连接信息）
    if data_source_data.type in ['mysql', 'postgresql', 'mongodb', 'redis']:
        try:
            connector = DataSourceConnector(data_source_data.type, data_source_data.config)
            connector.test_connection()
        except Exception as e:
            logger.warning(f"数据源连接测试失败: {str(e)}")
            # 不阻止创建，但标记状态为error
    
    data_source = DataSource(
        name=data_source_data.name,
        type=data_source_data.type,
        description=data_source_data.description,
        config=data_source_data.config,
        user_id=current_user.id,
        status="active"
    )
    db.add(data_source)
    db.commit()
    db.refresh(data_source)
    return data_source


@router.get("/{data_source_id}", response_model=DataSourceResponse)
async def get_data_source(
    data_source_id: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """获取数据源详情"""
    data_source = db.query(DataSource).filter(
        DataSource.id == data_source_id,
        DataSource.user_id == current_user.id
    ).first()
    
    if not data_source:
        raise NotFoundError("数据源", data_source_id)
    
    return data_source


@router.put("/{data_source_id}", response_model=DataSourceResponse)
async def update_data_source(
    data_source_id: str,
    data_source_data: DataSourceUpdate,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """更新数据源"""
    data_source = db.query(DataSource).filter(
        DataSource.id == data_source_id,
        DataSource.user_id == current_user.id
    ).first()
    
    if not data_source:
        raise NotFoundError("数据源", data_source_id)
    
    if data_source_data.name is not None:
        data_source.name = data_source_data.name
    if data_source_data.description is not None:
        data_source.description = data_source_data.description
    if data_source_data.config is not None:
        data_source.config = data_source_data.config
    if data_source_data.status is not None:
        data_source.status = data_source_data.status
    
    db.commit()
    db.refresh(data_source)
    return data_source


@router.delete("/{data_source_id}", status_code=204)
async def delete_data_source(
    data_source_id: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
    _ws_admin: WorkspaceContext = Depends(require_workspace_admin),
):
    """删除数据源（仅工作区管理员和平台管理员可操作）"""
    data_source = db.query(DataSource).filter(
        DataSource.id == data_source_id,
        DataSource.user_id == current_user.id
    ).first()
    
    if not data_source:
        raise NotFoundError("数据源", data_source_id)
    
    db.delete(data_source)
    db.commit()
    return None


@router.post("/{data_source_id}/test", status_code=200)
async def test_data_source_connection(
    data_source_id: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """测试数据源连接"""
    data_source = db.query(DataSource).filter(
        DataSource.id == data_source_id,
        DataSource.user_id == current_user.id
    ).first()
    
    if not data_source:
        raise NotFoundError("数据源", data_source_id)
    
    try:
        connector = DataSourceConnector(data_source.type, data_source.config)
        result = connector.test_connection()
        
        # 更新最后连接时间
        from datetime import datetime
        data_source.last_connected_at = datetime.utcnow()
        data_source.status = "active"
        db.commit()
        
        return {
            "success": True,
            "message": "连接成功",
            "details": result
        }
    except Exception as e:
        # 更新状态为error
        data_source.status = "error"
        db.commit()
        
        return {
            "success": False,
            "message": f"连接失败: {str(e)}",
            "error": str(e)
        }


@router.post("/{data_source_id}/query", status_code=200)
async def query_data_source(
    data_source_id: str,
    query: Dict[str, Any],
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """查询数据源"""
    data_source = db.query(DataSource).filter(
        DataSource.id == data_source_id,
        DataSource.user_id == current_user.id
    ).first()
    
    if not data_source:
        raise NotFoundError("数据源", data_source_id)
    
    try:
        connector = DataSourceConnector(data_source.type, data_source.config)
        result = connector.query(query)
        
        # 更新最后连接时间
        from datetime import datetime
        data_source.last_connected_at = datetime.utcnow()
        data_source.status = "active"
        db.commit()
        
        return {
            "success": True,
            "data": result
        }
    except Exception as e:
        logger.error(f"查询数据源失败: {str(e)}")
        data_source.status = "error"
        db.commit()
        
        return {
            "success": False,
            "error": str(e)
        }