graph_rag_server/src/api/sdk/api_manage.py

"""
API key management endpoints
"""
from fastapi import FastAPI, HTTPException, Depends
from typing import List, Dict, Any, Optional
from src.utils.mysql import get_global_mysql_client
from datetime import datetime, timedelta
import secrets
import string
from src.common.result import Result

app = FastAPI(
    title="API Key Management",
    description="API for managing API keys",
    version="1.0.0"
)


def generate_api_key() -> str:
    """
    Generate a new API key in format sk-<random-string>
    """
    # Generate random string
    alphabet = string.ascii_letters + string.digits
    random_string = ''.join(secrets.choice(alphabet) for _ in range(32))
    return f"sk-{random_string}"


@app.post("/generate")
async def generate_api_key_endpoint(expiry_days: Optional[int] = None):
    """
    Generate a new API key
    """
    try:
        api_key = generate_api_key()
        mysql_client = get_global_mysql_client()
        
        # Calculate expiration date if provided
        expired_at = None
        if expiry_days:
            expired_at = datetime.now() + timedelta(days=expiry_days)
        
        # Insert API key into database
        query = """
        INSERT INTO api_keys (api_key, expired_at, is_active)
        VALUES (%s, %s, TRUE)
        """
        mysql_client.execute(query, [api_key, expired_at])
        
        data = {
            "api_key": api_key,
            "created_at": datetime.now().isoformat(),
            "expired_at": expired_at.isoformat() if expired_at else None
        }
        
        return Result.success(data=data, message="API key 生成成功")
        
    except Exception as e:
        return Result.error(code=500, message=f"生成 API key 失败: {str(e)}")


@app.get("/list")
async def list_api_keys():
    """
    List all API keys
    """
    try:
        mysql_client = get_global_mysql_client()
        query = """
        SELECT id, api_key, created_at, expired_at, is_active
        FROM api_keys
        ORDER BY created_at DESC
        """
        result = mysql_client.fetch_all(query)
        
        return Result.success(data=result, message="获取 API keys 成功")
        
    except Exception as e:
        return Result.error(code=500, message=f"获取 API keys 失败: {str(e)}")


@app.put("/toggle/{api_key_id}")
async def toggle_api_key(api_key_id: int):
    """
    Toggle API key activation status
    """
    try:
        mysql_client = get_global_mysql_client()
        
        # Check if API key exists
        query = "SELECT is_active FROM api_keys WHERE id = %s"
        result = mysql_client.fetch_one(query, [api_key_id])
        
        if not result:
            return Result.error(code=404, message="API key 不存在")
        
        # Toggle status
        new_status = not result["is_active"]
        update_query = "UPDATE api_keys SET is_active = %s WHERE id = %s"
        mysql_client.execute(update_query, [new_status, api_key_id])
        
        data = {
            "api_key_id": api_key_id,
            "is_active": new_status
        }
        
        return Result.success(data=data, message="API key 状态更新成功")
        
    except Exception as e:
        return Result.error(code=500, message=f"更新 API key 状态失败: {str(e)}")


@app.delete("/delete/{api_key_id}")
async def delete_api_key(api_key_id: int):
    """
    Delete API key
    """
    try:
        mysql_client = get_global_mysql_client()
        
        # Check if API key exists
        query = "SELECT id FROM api_keys WHERE id = %s"
        result = mysql_client.fetch_one(query, [api_key_id])
        
        if not result:
            return Result.error(code=404, message="API key 不存在")
        
        # Delete API key
        delete_query = "DELETE FROM api_keys WHERE id = %s"
        mysql_client.execute(delete_query, [api_key_id])
        
        data = {
            "api_key_id": api_key_id
        }
        
        return Result.success(data=data, message="API key 删除成功")
        
    except Exception as e:
        return Result.error(code=500, message=f"删除 API key 失败: {str(e)}")