feat: Add quota check decorator and implement tenant-level API rate limiting
- Add quota check decorator module quota_stub.py, providing community edition stub implementation - Add quota check decorators to multiple controllers - Implement tenant-level API call rate limiting - Remove redundant plan fields from tenant_model.py - Optimize user permission check logic with added error handling
This commit is contained in:
@@ -96,6 +96,38 @@ def require_api_key(
|
||||
resource_id=api_key_obj.resource_id,
|
||||
)
|
||||
|
||||
# ── Tenant 级别限速(来自套餐配额 api_ops_rate_limit)──────────
|
||||
try:
|
||||
from app.models.workspace_model import Workspace
|
||||
from premium.platform_admin.package_plan_service import TenantSubscriptionService
|
||||
|
||||
workspace = db.query(Workspace).filter(
|
||||
Workspace.id == api_key_obj.workspace_id
|
||||
).first()
|
||||
if workspace:
|
||||
quota = TenantSubscriptionService(db).get_effective_quota(workspace.tenant_id)
|
||||
tenant_qps_limit = quota.get("api_ops_rate_limit") if quota else None
|
||||
if tenant_qps_limit:
|
||||
rate_limiter = RateLimiterService()
|
||||
tenant_ok, tenant_info = await rate_limiter.check_tenant_rate_limit(
|
||||
workspace.tenant_id, tenant_qps_limit
|
||||
)
|
||||
if not tenant_ok:
|
||||
raise RateLimitException(
|
||||
"租户 API 调用速率超限",
|
||||
BizCode.API_KEY_QPS_LIMIT_EXCEEDED,
|
||||
rate_headers={
|
||||
"X-RateLimit-Tenant-Limit": str(tenant_info["limit"]),
|
||||
"X-RateLimit-Tenant-Remaining": str(tenant_info["remaining"]),
|
||||
"X-RateLimit-Tenant-Reset": str(tenant_info["reset"]),
|
||||
}
|
||||
)
|
||||
except RateLimitException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.warning(f"Tenant 限速检查异常,跳过: {e}")
|
||||
# ─────────────────────────────────────────────────────────────
|
||||
|
||||
rate_limiter = RateLimiterService()
|
||||
is_allowed, error_msg, rate_headers = await rate_limiter.check_all_limits(api_key_obj)
|
||||
if not is_allowed:
|
||||
|
||||
44
api/app/core/quota_stub.py
Normal file
44
api/app/core/quota_stub.py
Normal file
@@ -0,0 +1,44 @@
|
||||
"""
|
||||
配额检查 stub - 社区版使用,所有检查直接放行。
|
||||
企业版通过 premium.platform_admin.quota_decorator 提供真实实现。
|
||||
"""
|
||||
from functools import wraps
|
||||
from typing import Callable
|
||||
|
||||
|
||||
def _noop_decorator(func: Callable) -> Callable:
|
||||
"""空装饰器,直接放行"""
|
||||
return func
|
||||
|
||||
|
||||
def _noop_check(*args, **kwargs):
|
||||
"""空检查函数,直接放行"""
|
||||
pass
|
||||
|
||||
|
||||
try:
|
||||
from premium.platform_admin.quota_decorator import (
|
||||
check_workspace_quota,
|
||||
check_skill_quota,
|
||||
check_app_quota,
|
||||
check_knowledge_capacity_quota,
|
||||
check_memory_engine_quota,
|
||||
check_end_user_quota,
|
||||
check_ontology_project_quota,
|
||||
check_model_quota,
|
||||
check_model_activation_quota,
|
||||
get_quota_usage,
|
||||
_check_quota,
|
||||
)
|
||||
except ModuleNotFoundError:
|
||||
check_workspace_quota = _noop_decorator
|
||||
check_skill_quota = _noop_decorator
|
||||
check_app_quota = _noop_decorator
|
||||
check_knowledge_capacity_quota = _noop_decorator
|
||||
check_memory_engine_quota = _noop_decorator
|
||||
check_end_user_quota = _noop_decorator
|
||||
check_ontology_project_quota = _noop_decorator
|
||||
check_model_quota = _noop_decorator
|
||||
check_model_activation_quota = _noop_decorator
|
||||
get_quota_usage = lambda db, tenant_id: {}
|
||||
_check_quota = _noop_check
|
||||
Reference in New Issue
Block a user