outman/dify
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add rate limiter to report download api

Browse Source
This commit is contained in:
GareArc 2025-02-05 22:59:43 -05:00
parent 7dfd6a95ab
commit 58bf08111a
2 changed files with 17 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
api/controllers/console/error.py
View File

@ -101,3 +101,9 @@ class AccountInFreezeError(BaseHTTPException):
"This email account has been deleted within the past 30 days" "This email account has been deleted within the past 30 days"
"and is temporarily unavailable for new account registration." "and is temporarily unavailable for new account registration."
) )
class CompilanceRateLimitError(BaseHTTPException):
error_code = "compilance_rate_limit"
description = "Rate limit exceeded for downloading compliance report."
code = 429

12
api/services/billing_service.py
View File

@ -6,6 +6,7 @@ from tenacity import (retry, retry_if_exception_type, stop_before_delay,
wait_fixed) wait_fixed)
from extensions.ext_database import db from extensions.ext_database import db
from libs.helper import RateLimiter
from models.account import TenantAccountJoin, TenantAccountRole from models.account import TenantAccountJoin, TenantAccountRole
@ -13,6 +14,8 @@ class BillingService:
base_url = os.environ.get("BILLING_API_URL", "BILLING_API_URL") base_url = os.environ.get("BILLING_API_URL", "BILLING_API_URL")
secret_key = os.environ.get("BILLING_API_SECRET_KEY", "BILLING_API_SECRET_KEY") secret_key = os.environ.get("BILLING_API_SECRET_KEY", "BILLING_API_SECRET_KEY")
compliance_download_rate_limiter = RateLimiter("compliance_download_rate_limiter", 3, 60)
@classmethod @classmethod
def get_info(cls, tenant_id: str): def get_info(cls, tenant_id: str):
params = {"tenant_id": tenant_id} params = {"tenant_id": tenant_id}
@ -107,6 +110,11 @@ class BillingService:
ip: str, ip: str,
device_info: str, device_info: str,
): ):
limiter_key = f"{account_id}:{tenant_id}"
if cls.compliance_download_rate_limiter.is_rate_limited(limiter_key):
from controllers.console.error import CompilanceRateLimitError
raise CompilanceRateLimitError()
json = { json = {
"doc_name": doc_name, "doc_name": doc_name,
"account_id": account_id, "account_id": account_id,
@ -114,4 +122,6 @@ class BillingService:
"ip_address": ip, "ip_address": ip,
"device_info": device_info, "device_info": device_info,
} }
return cls._send_request("POST", "/compliance/download", json=json) res = cls._send_request("POST", "/compliance/download", json=json)
cls.compliance_download_rate_limiter.increment(limiter_key)
return res