Maas-Agent/apps/application/views/rate_limit_view.py

# coding=utf-8
"""
应用 API 限流视图
提供限流配置的 CRUD API
"""
from django.utils.translation import gettext as _
from rest_framework.request import Request
from rest_framework.views import APIView

from common.auth import TokenAuth
from common.exception.app_exception import AppApiException
from common.result import result
from application.models.rate_limit import RateLimit, RateLimitType
from application.models.application import Application


class RateLimitView(APIView):
    """应用 API 限流管理"""
    authentication_classes = [TokenAuth]

    class Get(APIView):
        """获取应用限流配置"""
        authentication_classes = [TokenAuth]

        def get(self, request: Request, application_id: str):
            try:
                application = Application.objects.get(id=application_id)
            except Application.DoesNotExist:
                raise AppApiException(404, _('应用不存在'))

            try:
                rate_limit = RateLimit.objects.get(application=application)
                data = {
                    'id': str(rate_limit.id),
                    'application_id': str(rate_limit.application_id),
                    'is_enabled': rate_limit.is_enabled,
                    'rate_type': rate_limit.rate_type,
                    'max_requests': rate_limit.max_requests,
                    'burst_size': rate_limit.burst_size,
                    'window_seconds': rate_limit.window_seconds,
                    'create_time': rate_limit.create_time.isoformat() if rate_limit.create_time else None,
                    'update_time': rate_limit.update_time.isoformat() if rate_limit.update_time else None,
                }
            except RateLimit.DoesNotExist:
                data = {
                    'application_id': application_id,
                    'is_enabled': False,
                    'rate_type': 'QPM',
                    'max_requests': 60,
                    'burst_size': 10,
                    'window_seconds': 60,
                }

            return result.success(data)

    class Update(APIView):
        """更新应用限流配置"""
        authentication_classes = [TokenAuth]

        def put(self, request: Request, application_id: str):
            try:
                application = Application.objects.get(id=application_id)
            except Application.DoesNotExist:
                raise AppApiException(404, _('应用不存在'))

            is_enabled = request.data.get('is_enabled', False)
            rate_type = request.data.get('rate_type', 'QPM')
            max_requests = request.data.get('max_requests', 60)
            burst_size = request.data.get('burst_size', 10)
            window_seconds = request.data.get('window_seconds', 60)

            # 验证参数
            if rate_type not in [t[0] for t in RateLimitType.choices]:
                raise AppApiException(400, _('无效的限流类型'))

            if max_requests <= 0:
                raise AppApiException(400, _('最大请求数必须大于0'))

            if window_seconds <= 0:
                raise AppApiException(400, _('时间窗口必须大于0'))

            rate_limit, created = RateLimit.objects.update_or_create(
                application=application,
                defaults={
                    'is_enabled': is_enabled,
                    'rate_type': rate_type,
                    'max_requests': max_requests,
                    'burst_size': burst_size,
                    'window_seconds': window_seconds,
                }
            )

            return result.success({
                'id': str(rate_limit.id),
                'application_id': str(rate_limit.application_id),
                'is_enabled': rate_limit.is_enabled,
                'rate_type': rate_limit.rate_type,
                'max_requests': rate_limit.max_requests,
                'burst_size': rate_limit.burst_size,
                'window_seconds': rate_limit.window_seconds,
                'create_time': rate_limit.create_time.isoformat() if rate_limit.create_time else None,
                'update_time': rate_limit.update_time.isoformat() if rate_limit.update_time else None,
            })

    class Reset(APIView):
        """重置应用限流计数"""
        authentication_classes = [TokenAuth]

        def post(self, request: Request, application_id: str):
            try:
                application = Application.objects.get(id=application_id)
            except Application.DoesNotExist:
                raise AppApiException(404, _('应用不存在'))

            # 这里可以重置内存中的限流计数
            # 实际实现需要访问 RateLimitMiddleware 的实例
            return result.success({'message': '限流计数已重置'})