feat: implement plan-based API rate limits

backend/api/auth.py

@@ -19,7 +19,7 @@
 
 class ServiceAccountUser:
     """Mock ServiceAccount user"""
-    
+
     def __init__(self, service_account):
         self.userId = service_account.id
         self.id = service_account.id
@@ -63,14 +63,20 @@ def authenticate(self, request):
         if secret_id:
             found = False
             try:
-                secret = Secret.objects.get(id=secret_id)
+                # Pre-fetch environment, app, and organisation
+                secret = Secret.objects.select_related(
+                    "environment__app__organisation"
+                ).get(id=secret_id)
                 env = secret.environment
                 found = True
             except Secret.DoesNotExist:
                 pass
             if not found:
                 try:
-                    dyn_secret = DynamicSecret.objects.get(id=secret_id)
+                    # Pre-fetch environment, app, and organisation
+                    dyn_secret = DynamicSecret.objects.select_related(
+                        "environment__app__organisation"
+                    ).get(id=secret_id)
                     env = dyn_secret.environment
                     found = True
                 except DynamicSecret.DoesNotExist:
@@ -84,7 +90,10 @@ def authenticate(self, request):
             # Try resolving env from header
             if env_id:
                 try:
-                    env = Environment.objects.get(id=env_id)
+                    # Pre-fetch app and organisation
+                    env = Environment.objects.select_related("app__organisation").get(
+                        id=env_id
+                    )
                 except Environment.DoesNotExist:
                     raise exceptions.AuthenticationFailed("Environment not found")
 
@@ -99,7 +108,10 @@ def authenticate(self, request):
                         )
                     if not env_name:
                         raise exceptions.AuthenticationFailed("Missing env parameter")
-                    env = Environment.objects.get(app_id=app_id, name__iexact=env_name)
+                    # Pre-fetch app and organisation
+                    env = Environment.objects.select_related("app__organisation").get(
+                        app_id=app_id, name__iexact=env_name
+                    )
                 except Environment.DoesNotExist:
                     # Check if the app exists to give a more specific error
                     App = apps.get_model("api", "App")
@@ -134,14 +146,14 @@ def authenticate(self, request):
 
             try:
                 service_token = get_service_token(auth_token)
-                service_account = get_service_account_from_token(auth_token)          
-                
+                service_account = get_service_account_from_token(auth_token)
+
                 creator = getattr(service_token, "created_by", None)
                 if creator:
                     user = creator.user
                 else:
                     user = ServiceAccountUser(service_account)
-                    
+
                 auth["service_account"] = service_account
                 auth["service_account_token"] = service_token
 

backend/api/throttling.py

@@ -0,0 +1,51 @@
+from rest_framework.throttling import SimpleRateThrottle
+from django.conf import settings
+from api.models import Organisation
+
+
+class PlanBasedRateThrottle(SimpleRateThrottle):
+    """
+    Limits the rate of API calls based on the Organisation's plan.
+    Uses the pre-fetched organisation data from request.auth to avoid DB lookups.
+    """
+
+    scope = "plan_based"
+
+    def get_cache_key(self, request, view):
+        if not request.user.is_authenticated or not request.auth:
+            return None
+
+        # Identify the user or service account
+        ident = self.get_ident(request)
+        if request.auth.get("org_member"):
+            ident = f"user_{request.auth['org_member'].id}"
+        elif request.auth.get("service_account"):
+            ident = f"sa_{request.auth['service_account'].id}"
+
+        return self.cache_format % {"scope": self.scope, "ident": ident}
+
+    def allow_request(self, request, view):
+        """
+        Override allow_request to dynamically set the rate based on the request user's plan.
+        """
+        # Default fallback (reads from REST_FRAMEWORK['DEFAULT_THROTTLE_RATES']['plan_based'])
+        new_rate = self.get_rate()
+
+        if request.user.is_authenticated and request.auth:
+            env = request.auth.get("environment")
+            if env:
+                try:
+                    plan = env.app.organisation.plan
+                    new_rate = self.get_rate_for_plan(plan)
+                except AttributeError:
+                    pass
+
+        # Update the throttle configuration for this specific request
+        self.rate = new_rate
+        self.num_requests, self.duration = self.parse_rate(self.rate)
+
+        return super().allow_request(request, view)
+
+    @staticmethod
+    def get_rate_for_plan(plan):
+        return settings.PLAN_RATE_LIMITS.get(plan, settings.PLAN_RATE_LIMITS["DEFAULT"])

backend/backend/settings.py

@@ -241,13 +241,27 @@ def get_version():
     "USER_DETAILS_SERIALIZER": "api.serializers.CustomUserSerializer"
 }
 
+# Define rate limits per plan (keys match Organisation model constants: FR, PR, EN)
+PLAN_RATE_LIMITS = {
+    "FR": os.getenv("RATE_LIMIT_FREE", "120/min"),
+    "PR": os.getenv("RATE_LIMIT_PRO", "240/min"),
+    "EN": os.getenv("RATE_LIMIT_ENTERPRISE", "1000/min"),
+    "DEFAULT": os.getenv("RATE_LIMIT_DEFAULT", "120/min"),
+}
+
 REST_FRAMEWORK = {
     "DEFAULT_PERMISSION_CLASSES": [
         "rest_framework.permissions.IsAuthenticated",
     ],
     "DEFAULT_AUTHENTICATION_CLASSES": (
         "rest_framework.authentication.SessionAuthentication",
     ),
+    "DEFAULT_THROTTLE_CLASSES": [
+        "api.throttling.PlanBasedRateThrottle",
+    ],
+    "DEFAULT_THROTTLE_RATES": {
+        "plan_based": PLAN_RATE_LIMITS["DEFAULT"],
+    },
     "EXCEPTION_HANDLER": "backend.exceptions.custom_exception_handler",
     "DEFAULT_RENDERER_CLASSES": [
         "rest_framework.renderers.JSONRenderer",
@@ -294,6 +308,27 @@ def get_version():
     },
 }
 
+REDIS_HOST = os.getenv("REDIS_HOST")
+REDIS_PORT = os.getenv("REDIS_PORT")
+REDIS_PASSWORD = get_secret("REDIS_PASSWORD")
+REDIS_SSL = os.getenv("REDIS_SSL", "False").lower() == "true"
+REDIS_PROTOCOL = "rediss" if REDIS_SSL else "redis"
+REDIS_AUTH = f":{REDIS_PASSWORD}@" if REDIS_PASSWORD else ""
+
+CACHES = {
+    "default": {
+        "BACKEND": "django.core.cache.backends.redis.RedisCache",
+        "LOCATION": f"{REDIS_PROTOCOL}://{REDIS_AUTH}{REDIS_HOST}:{REDIS_PORT}/1",
+        "OPTIONS": (
+            {
+                "ssl_cert_reqs": None,
+            }
+            if REDIS_SSL
+            else {}
+        ),
+    }
+}
+
 DYNAMODB = {
     "TABLE": os.getenv("DYNAMODB_LOGS_TABLE"),
     "INDEX": os.getenv("DYNAMODB_LOGS_TIMESTAMP_INDEX"),