ai-agent-admin/backend-fastapi/core/database_manager/api.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
数据库管理API（异步版本）
"""
import logging
from typing import List

from fastapi import APIRouter, Depends, HTTPException, Query
from sqlalchemy.ext.asyncio import AsyncSession

from app.database import get_db
from core.database_manager.schema import (
    ColumnInfo,
    ConstraintInfo,
    DatabaseConfig,
    DatabaseCreateIn,
    DatabaseInfo,
    DatabaseOperationOut,
    DatabaseRenameIn,
    DataOperationOut,
    DeleteDataIn,
    ExecuteDDLIn,
    ExecuteDDLOut,
    ExecuteSQLIn,
    ExecuteSQLOut,
    IndexInfo,
    InsertDataIn,
    QueryDataIn,
    QueryDataOut,
    SchemaCreateIn,
    SchemaInfo,
    SchemaOperationOut,
    SchemaRenameIn,
    TableInfo,
    TableStructure,
    UpdateDataIn,
    ViewInfo,
    ViewStructure,
)
from core.database_manager.service import (
    AsyncDatabaseManagerService,
    DatabaseManagerForbidden,
    DatabaseManagerValidationError,
)

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/database_manager", tags=["数据库管理"])


async def _get_service(db_name: str, db: AsyncSession) -> AsyncDatabaseManagerService:
    try:
        return await AsyncDatabaseManagerService.create(db_name, db)
    except ValueError as exc:
        raise HTTPException(status_code=404, detail=str(exc)) from exc


def _handle_manager_error(exc: Exception) -> None:
    if isinstance(exc, DatabaseManagerForbidden):
        raise HTTPException(status_code=403, detail=str(exc)) from exc
    if isinstance(exc, DatabaseManagerValidationError):
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    raise exc


# ============ 数据库配置 ============
@router.get("/configs", response_model=List[DatabaseConfig], summary="获取数据库配置列表")
async def get_database_configs(db: AsyncSession = Depends(get_db)):
    """获取所有配置的数据库连接信息"""
    return await AsyncDatabaseManagerService.get_database_configs_async(db)


@router.post("/{db_name}/test", summary="测试数据库连接")
async def test_database_connection(db_name: str, db: AsyncSession = Depends(get_db)):
    """测试指定数据库的连接"""
    service = await _get_service(db_name, db)
    return await service.test_connection()


# ============ 数据库管理 ============
@router.get("/{db_name}/databases", response_model=List[DatabaseInfo], summary="获取数据库列表")
async def get_databases(db_name: str, db: AsyncSession = Depends(get_db)):
    """获取所有数据库"""
    service = await _get_service(db_name, db)
    return await service.get_databases()


@router.post("/{db_name}/databases", response_model=DatabaseOperationOut, summary="创建数据库")
async def create_database(
    db_name: str,
    data: DatabaseCreateIn,
    db: AsyncSession = Depends(get_db),
):
    """创建新数据库"""
    try:
        service = await _get_service(db_name, db)
        success = await service.create_database(
            name=data.name,
            owner=data.owner,
            encoding=data.encoding,
            template=data.template,
            charset=data.charset,
            collation=data.collation,
        )
        return DatabaseOperationOut(
            success=success,
            message="数据库创建成功" if success else "数据库创建失败",
        )
    except HTTPException:
        raise
    except (DatabaseManagerForbidden, DatabaseManagerValidationError) as e:
        _handle_manager_error(e)
    except Exception as e:
        logger.error(f"Failed to create database: {e}")
        return DatabaseOperationOut(success=False, message=str(e))


@router.patch(
    "/{db_name}/databases/{database_name}",
    response_model=DatabaseOperationOut,
    summary="重命名数据库",
)
async def rename_database(
    db_name: str,
    database_name: str,
    data: DatabaseRenameIn,
    db: AsyncSession = Depends(get_db),
):
    """重命名数据库（PostgreSQL / SQL Server）"""
    try:
        service = await _get_service(db_name, db)
        success = await service.rename_database(database_name, data.new_name)
        return DatabaseOperationOut(
            success=success,
            message="数据库重命名成功" if success else "数据库重命名失败",
        )
    except HTTPException:
        raise
    except (DatabaseManagerForbidden, DatabaseManagerValidationError) as e:
        _handle_manager_error(e)
    except Exception as e:
        logger.error(f"Failed to rename database: {e}")
        return DatabaseOperationOut(success=False, message=str(e))


@router.delete("/{db_name}/databases/{database_name}", response_model=DatabaseOperationOut, summary="删除数据库")
async def drop_database(
    db_name: str,
    database_name: str,
    db: AsyncSession = Depends(get_db),
):
    """删除数据库"""
    try:
        service = await _get_service(db_name, db)
        success = await service.drop_database(database_name)
        return DatabaseOperationOut(
            success=success,
            message="数据库删除成功" if success else "数据库删除失败",
        )
    except HTTPException:
        raise
    except (DatabaseManagerForbidden, DatabaseManagerValidationError) as e:
        _handle_manager_error(e)
    except Exception as e:
        logger.error(f"Failed to drop database: {e}")
        return DatabaseOperationOut(success=False, message=str(e))


# ============ Schema管理（PostgreSQL） ============
@router.get("/{db_name}/schemas", response_model=List[SchemaInfo], summary="获取Schema列表")
async def get_schemas(
    db_name: str,
    database: str = None,
    db: AsyncSession = Depends(get_db),
):
    """获取所有Schema（仅PostgreSQL）"""
    service = await _get_service(db_name, db)
    return await service.get_schemas(database)


@router.post("/{db_name}/schemas", response_model=SchemaOperationOut, summary="创建Schema")
async def create_schema(
    db_name: str,
    data: SchemaCreateIn,
    db: AsyncSession = Depends(get_db),
):
    """创建 Schema（PostgreSQL / SQL Server）"""
    try:
        service = await _get_service(db_name, db)
        success = await service.create_schema(data.name, data.database, data.owner)
        return SchemaOperationOut(
            success=success,
            message="Schema 创建成功" if success else "Schema 创建失败",
        )
    except HTTPException:
        raise
    except (DatabaseManagerForbidden, DatabaseManagerValidationError) as e:
        _handle_manager_error(e)
    except Exception as e:
        logger.error(f"Failed to create schema: {e}")
        return SchemaOperationOut(success=False, message=str(e))


@router.delete(
    "/{db_name}/schemas/{schema_name}",
    response_model=SchemaOperationOut,
    summary="删除Schema",
)
async def drop_schema(
    db_name: str,
    schema_name: str,
    database: str = Query(None, description="数据库名"),
    cascade: bool = Query(True, description="是否 CASCADE"),
    db: AsyncSession = Depends(get_db),
):
    """删除 Schema（PostgreSQL / SQL Server）"""
    try:
        service = await _get_service(db_name, db)
        success = await service.drop_schema(schema_name, database, cascade)
        return SchemaOperationOut(
            success=success,
            message="Schema 删除成功" if success else "Schema 删除失败",
        )
    except HTTPException:
        raise
    except (DatabaseManagerForbidden, DatabaseManagerValidationError) as e:
        _handle_manager_error(e)
    except Exception as e:
        logger.error(f"Failed to drop schema: {e}")
        return SchemaOperationOut(success=False, message=str(e))


@router.patch(
    "/{db_name}/schemas/{schema_name}",
    response_model=SchemaOperationOut,
    summary="重命名Schema",
)
async def rename_schema(
    db_name: str,
    schema_name: str,
    data: SchemaRenameIn,
    db: AsyncSession = Depends(get_db),
):
    """重命名 Schema（仅 PostgreSQL）"""
    try:
        service = await _get_service(db_name, db)
        success = await service.rename_schema(schema_name, data.new_name, data.database)
        return SchemaOperationOut(
            success=success,
            message="Schema 重命名成功" if success else "Schema 重命名失败",
        )
    except HTTPException:
        raise
    except (DatabaseManagerForbidden, DatabaseManagerValidationError) as e:
        _handle_manager_error(e)
    except Exception as e:
        logger.error(f"Failed to rename schema: {e}")
        return SchemaOperationOut(success=False, message=str(e))


# ============ 表管理 ============
@router.get("/{db_name}/tables", response_model=List[TableInfo], summary="获取表列表")
async def get_tables(
    db_name: str,
    database: str = Query(None, description="数据库名"),
    schema_name: str = Query(None, description="Schema名"),
    db: AsyncSession = Depends(get_db),
):
    """获取指定schema/database的所有表"""
    service = await _get_service(db_name, db)
    return await service.get_tables(database=database, schema_name=schema_name)


@router.get("/{db_name}/tables/{table_name}/structure", response_model=TableStructure, summary="获取表结构")
async def get_table_structure(
    db_name: str,
    table_name: str,
    database: str = Query(None, description="数据库名"),
    schema_name: str = Query(None, description="Schema名"),
    db: AsyncSession = Depends(get_db),
):
    """获取表的详细结构"""
    service = await _get_service(db_name, db)
    try:
        return await service.get_table_structure(
            table_name, database=database, schema_name=schema_name
        )
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e)) from e


@router.get("/{db_name}/tables/{table_name}/ddl", summary="获取表DDL")
async def get_table_ddl(
    db_name: str,
    table_name: str,
    database: str = Query(None, description="数据库名"),
    schema_name: str = Query(None, description="Schema名"),
    db: AsyncSession = Depends(get_db),
):
    """获取表的DDL语句"""
    service = await _get_service(db_name, db)
    ddl = await service.get_table_ddl(table_name, schema_name, database)
    return {"ddl": ddl}


@router.get("/{db_name}/tables/{table_name}/columns", response_model=List[ColumnInfo], summary="获取表字段")
async def get_table_columns(
    db_name: str,
    table_name: str,
    database: str = Query(None, description="数据库名"),
    schema_name: str = Query(None, description="Schema名"),
    db: AsyncSession = Depends(get_db),
):
    """获取表的字段信息"""
    service = await _get_service(db_name, db)
    return await service.get_table_columns(table_name, schema_name, database)


@router.get("/{db_name}/tables/{table_name}/indexes", response_model=List[IndexInfo], summary="获取表索引")
async def get_table_indexes(
    db_name: str,
    table_name: str,
    database: str = Query(None, description="数据库名"),
    schema_name: str = Query(None, description="Schema名"),
    db: AsyncSession = Depends(get_db),
):
    """获取表的索引信息"""
    service = await _get_service(db_name, db)
    return await service.get_table_indexes(table_name, schema_name, database)


@router.get("/{db_name}/tables/{table_name}/constraints", response_model=List[ConstraintInfo], summary="获取表约束")
async def get_table_constraints(
    db_name: str,
    table_name: str,
    database: str = Query(None, description="数据库名"),
    schema_name: str = Query(None, description="Schema名"),
    db: AsyncSession = Depends(get_db),
):
    """获取表的约束信息"""
    service = await _get_service(db_name, db)
    return await service.get_table_constraints(table_name, schema_name, database)


# ============ 视图管理 ============
@router.get("/{db_name}/views", response_model=List[ViewInfo], summary="获取视图列表")
async def get_views(
    db_name: str,
    database: str = Query(None, description="数据库名"),
    schema_name: str = Query(None, description="Schema名"),
    db: AsyncSession = Depends(get_db),
):
    """获取指定schema/database的所有视图"""
    service = await _get_service(db_name, db)
    return await service.get_views(database=database, schema_name=schema_name)


@router.get("/{db_name}/views/{view_name}/structure", response_model=ViewStructure, summary="获取视图结构")
async def get_view_structure(
    db_name: str,
    view_name: str,
    schema_name: str = Query(None, description="Schema名"),
    db: AsyncSession = Depends(get_db),
):
    """获取视图的详细结构"""
    service = await _get_service(db_name, db)
    try:
        return await service.get_view_structure(view_name, schema_name)
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e)) from e


@router.get("/{db_name}/views/{view_name}/definition", summary="获取视图定义")
async def get_view_definition(
    db_name: str,
    view_name: str,
    schema_name: str = Query(None, description="Schema名"),
    db: AsyncSession = Depends(get_db),
):
    """获取视图的定义SQL"""
    service = await _get_service(db_name, db)
    definition = await service.get_view_definition(view_name, schema_name)
    return {"definition": definition}


@router.get("/{db_name}/views/{view_name}/dependencies", response_model=List[str], summary="获取视图依赖")
async def get_view_dependencies(
    db_name: str,
    view_name: str,
    schema_name: str = Query(None, description="Schema名"),
    db: AsyncSession = Depends(get_db),
):
    """获取视图依赖的表列表"""
    service = await _get_service(db_name, db)
    return await service.get_view_dependencies(view_name, schema_name)


# ============ 数据查询 ============
@router.post("/{db_name}/query", response_model=QueryDataOut, summary="查询表数据")
async def query_data(
    db_name: str,
    data: QueryDataIn,
    db: AsyncSession = Depends(get_db),
):
    """分页查询表数据"""
    service = await _get_service(db_name, db)
    return await service.query_data(
        table_name=data.table_name,
        schema_name=data.schema_name,
        database=data.database,
        page=data.page,
        page_size=data.page_size,
        where=data.where,
        order_by=data.order_by,
    )


# ============ SQL执行 ============
@router.post("/{db_name}/execute", response_model=ExecuteSQLOut, summary="执行SQL")
async def execute_sql(
    db_name: str,
    data: ExecuteSQLIn,
    db: AsyncSession = Depends(get_db),
):
    """执行自定义SQL语句"""
    service = await _get_service(db_name, db)
    return await service.execute_sql(data.sql, data.is_query)


# ============ 数据操作 ============
@router.post("/{db_name}/data/insert", response_model=DataOperationOut, summary="插入数据")
async def insert_data(
    db_name: str,
    data: InsertDataIn,
    db: AsyncSession = Depends(get_db),
):
    """向表中插入数据"""
    service = await _get_service(db_name, db)
    return await service.insert_data(data.table_name, data.data, data.schema_name)


@router.post("/{db_name}/data/update", response_model=DataOperationOut, summary="更新数据")
async def update_data(
    db_name: str,
    data: UpdateDataIn,
    db: AsyncSession = Depends(get_db),
):
    """更新表中的数据"""
    service = await _get_service(db_name, db)
    return await service.update_data(
        data.table_name, data.data, data.where, data.schema_name
    )


@router.post("/{db_name}/data/delete", response_model=DataOperationOut, summary="删除数据")
async def delete_data(
    db_name: str,
    data: DeleteDataIn,
    db: AsyncSession = Depends(get_db),
):
    """删除表中的数据"""
    service = await _get_service(db_name, db)
    return await service.delete_data(data.table_name, data.where, data.schema_name)


# ============ DDL操作 ============
@router.post("/{db_name}/execute/ddl", response_model=ExecuteDDLOut, summary="执行DDL语句")
async def execute_ddl(
    db_name: str,
    data: ExecuteDDLIn,
    db: AsyncSession = Depends(get_db),
):
    """执行DDL语句（CREATE TABLE, ALTER TABLE, DROP TABLE等）"""
    try:
        service = await _get_service(db_name, db)
        result = await service.execute_ddl(data.sql, data.database, data.schema_name)
        return ExecuteDDLOut(**result)
    except HTTPException:
        raise
    except (DatabaseManagerForbidden, DatabaseManagerValidationError) as e:
        _handle_manager_error(e)