Skip to content

Commit 767e15c

Browse files
committed
Merge branch 'feat/compliance' into deploy/dev
2 parents ba7b222 + 58bf081 commit 767e15c

File tree

2 files changed

+17
-1
lines changed

2 files changed

+17
-1
lines changed

api/controllers/console/error.py

+6
Original file line numberDiff line numberDiff line change
@@ -101,3 +101,9 @@ class AccountInFreezeError(BaseHTTPException):
101101
"This email account has been deleted within the past 30 days"
102102
"and is temporarily unavailable for new account registration."
103103
)
104+
105+
106+
class CompilanceRateLimitError(BaseHTTPException):
107+
error_code = "compilance_rate_limit"
108+
description = "Rate limit exceeded for downloading compliance report."
109+
code = 429

api/services/billing_service.py

+11-1
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,16 @@
66
wait_fixed)
77

88
from extensions.ext_database import db
9+
from libs.helper import RateLimiter
910
from models.account import TenantAccountJoin, TenantAccountRole
1011

1112

1213
class BillingService:
1314
base_url = os.environ.get("BILLING_API_URL", "BILLING_API_URL")
1415
secret_key = os.environ.get("BILLING_API_SECRET_KEY", "BILLING_API_SECRET_KEY")
1516

17+
compliance_download_rate_limiter = RateLimiter("compliance_download_rate_limiter", 3, 60)
18+
1619
@classmethod
1720
def get_info(cls, tenant_id: str):
1821
params = {"tenant_id": tenant_id}
@@ -107,11 +110,18 @@ def get_compliance_download_link(
107110
ip: str,
108111
device_info: str,
109112
):
113+
limiter_key = f"{account_id}:{tenant_id}"
114+
if cls.compliance_download_rate_limiter.is_rate_limited(limiter_key):
115+
from controllers.console.error import CompilanceRateLimitError
116+
raise CompilanceRateLimitError()
117+
110118
json = {
111119
"doc_name": doc_name,
112120
"account_id": account_id,
113121
"tenant_id": tenant_id,
114122
"ip_address": ip,
115123
"device_info": device_info,
116124
}
117-
return cls._send_request("POST", "/compliance/download", json=json)
125+
res = cls._send_request("POST", "/compliance/download", json=json)
126+
cls.compliance_download_rate_limiter.increment(limiter_key)
127+
return res

0 commit comments

Comments
 (0)