Throttle model refactor to support ninja v1.2.0

Author: eadwinCode

ninja_extra/conf/settings.py

@@ -41,7 +41,12 @@ class Config:
     THROTTLE_RATES: Dict[str, Optional[str]] = Field(
         {"user": "1000/day", "anon": "100/day"}
     )
-    THROTTLE_CLASSES: List[Any] = []
+    THROTTLE_CLASSES: List[Any] = Field(
+        [
+            "ninja_extra.throttling.AnonRateThrottle",
+            "ninja_extra.throttling.UserRateThrottle",
+        ]
+    )
     NUM_PROXIES: Optional[int] = None
     INJECTOR_MODULES: List[Any] = []
     ORDERING_CLASS: Any = Field(

ninja_extra/constants.py

@@ -15,6 +15,7 @@
 TRACE = "TRACE"
 ROUTE_METHODS = [POST, PUT, PATCH, DELETE, GET, HEAD, OPTIONS, TRACE]
 THROTTLED_FUNCTION = "__throttled_endpoint__"
+THROTTLED_OBJECTS = "__throttled_objects__"
 ROUTE_FUNCTION = "__route_function__"
 
 ROUTE_CONTEXT_VAR: contextvars.ContextVar[t.Optional["RouteContext"]] = (

ninja_extra/controllers/base.py

@@ -26,12 +26,13 @@
 from django.urls import path as django_path
 from injector import inject, is_decorated_with_inject
 from ninja import NinjaAPI, Router
-from ninja.constants import NOT_SET
+from ninja.constants import NOT_SET, NOT_SET_TYPE
 from ninja.security.base import AuthBase
 from ninja.signature import is_async
+from ninja.throttling import BaseThrottle
 from ninja.utils import normalize_path
 
-from ninja_extra.constants import ROUTE_FUNCTION, THROTTLED_FUNCTION
+from ninja_extra.constants import ROUTE_FUNCTION, THROTTLED_FUNCTION, THROTTLED_OBJECTS
 from ninja_extra.exceptions import APIException, NotFound, PermissionDenied, bad_request
 from ninja_extra.helper import get_function_name
 from ninja_extra.operation import Operation, PathView
@@ -51,7 +52,6 @@
 
 if TYPE_CHECKING:  # pragma: no cover
     from ninja_extra import NinjaExtraAPI
-    from ninja_extra.throttling import BaseThrottle
 
     from .route.context import RouteContext
 
@@ -126,10 +126,10 @@ def some_method_name(self):
     throttling_classes: List[Type["BaseThrottle"]] = []
     throttling_init_kwargs: Optional[Dict[Any, Any]] = None
 
-    Ok = Ok
-    Id = Id
-    Detail = Detail
-    bad_request = bad_request
+    Ok = Ok  # TODO: remove soonest
+    Id = Id  # TODO: remove soonest
+    Detail = Detail  # TODO: remove soonest
+    bad_request = bad_request  # TODO: remove soonest
 
     @classmethod
     def get_api_controller(cls) -> "APIController":
@@ -294,6 +294,7 @@ def __init__(
         prefix: str,
         *,
         auth: Any = NOT_SET,
+        throttle: Union[BaseThrottle, List[BaseThrottle], NOT_SET_TYPE] = NOT_SET,
         tags: Union[Optional[List[str]], str] = None,
         permissions: Optional["PermissionType"] = None,
         auto_import: bool = True,
@@ -303,6 +304,7 @@ def __init__(
         self.auth: Optional[AuthBase] = auth
 
         self.tags = tags  # type: ignore
+        self.throttle = throttle
 
         self.auto_import: bool = auto_import  # set to false and it would be ignored when api.auto_discover is called
         # `controller_class` target class that the APIController wraps
@@ -348,8 +350,6 @@ def tags(self, value: Union[str, List[str], None]) -> None:
         self._tags = tag
 
     def __call__(self, cls: ControllerClassType) -> ControllerClassType:
-        from ninja_extra.throttling import throttle
-
         self.auto_import = getattr(cls, "auto_import", self.auto_import)
         if not issubclass(cls, ControllerBase):
             # We force the cls to inherit from `ControllerBase` by creating another type.
@@ -360,8 +360,15 @@ def __call__(self, cls: ControllerClassType) -> ControllerClassType:
         assert isinstance(
             cls.throttling_classes, (list, tuple)
         ), f"Controller[{cls.__name__}].throttling_class must be a list or tuple"
-        has_throttling_classes = len(cls.throttling_classes) > 0
-        throttling_init_kwargs = cls.throttling_init_kwargs or {}
+
+        throttling_objects: Union[BaseThrottle, List[BaseThrottle], NOT_SET_TYPE]
+        if cls.throttling_classes:
+            throttling_init_kwargs = cls.throttling_init_kwargs or {}
+            throttling_objects = [
+                item(**throttling_init_kwargs) for item in cls.throttling_classes
+            ]
+        else:
+            throttling_objects = self.throttle
 
         if not self.tags:
             tag = str(cls.__name__).lower().replace("controller", "")
@@ -386,10 +393,13 @@ def __call__(self, cls: ControllerClassType) -> ControllerClassType:
 
         for _, v in self._controller_class_route_functions.items():
             throttled_endpoint = v.as_view.__dict__.get(THROTTLED_FUNCTION)
-            if not throttled_endpoint and has_throttling_classes:
-                v.route.view_func = throttle(
-                    *cls.throttling_classes, **throttling_init_kwargs
-                )(v.route.view_func)
+
+            if throttled_endpoint or throttling_objects is not NOT_SET:
+                v.route.route_params.throttle = v.as_view.__dict__.get(
+                    THROTTLED_OBJECTS, lambda: throttling_objects
+                )()
+                setattr(v.route.view_func, THROTTLED_FUNCTION, True)
+
             self._add_operation_from_route_function(v)
 
         if not is_decorated_with_inject(cls.__init__):
@@ -463,6 +473,7 @@ def add_api_operation(
         view_func: Callable,
         *,
         auth: Any = NOT_SET,
+        throttle: Union[BaseThrottle, List[BaseThrottle], NOT_SET_TYPE] = NOT_SET,
         response: Any = NOT_SET,
         operation_id: Optional[str] = None,
         summary: Optional[str] = None,
@@ -504,13 +515,14 @@ def add_api_operation(
             url_name=url_name,
             include_in_schema=include_in_schema,
             openapi_extra=openapi_extra,
+            throttle=throttle,
         )
         return operation
 
 
 @overload
 def api_controller(
-    prefix_or_class: Type[T],
+    prefix_or_class: Union[ControllerClassType, Type[T]],
 ) -> Union[Type[ControllerBase], Type[T]]:  # pragma: no cover
     ...
 
@@ -522,12 +534,14 @@ def api_controller(
     tags: Union[Optional[List[str]], str] = None,
     permissions: Optional["PermissionType"] = None,
     auto_import: bool = True,
-) -> Callable[[Type[T]], Union[Type[ControllerBase], Type[T]]]:  # pragma: no cover
+) -> Callable[
+    [Union[Type, Type[T]]], Union[Type[ControllerBase], Type[T]]
+]:  # pragma: no cover
     ...
 
 
 def api_controller(
-    prefix_or_class: Union[str, ControllerClassType] = "",
+    prefix_or_class: Union[str, Union[ControllerClassType, Type]] = "",
     auth: Any = NOT_SET,
     tags: Union[Optional[List[str]], str] = None,
     permissions: Optional["PermissionType"] = None,

ninja_extra/controllers/route/__init__.py

@@ -1,7 +1,8 @@
 import typing as t
 
-from ninja.constants import NOT_SET
+from ninja.constants import NOT_SET, NOT_SET_TYPE
 from ninja.signature import is_async
+from ninja.throttling import BaseThrottle
 from ninja.types import TCallable
 
 from ninja_extra.constants import (
@@ -52,6 +53,7 @@ def __init__(
         path: str,
         methods: t.List[str],
         auth: t.Any = NOT_SET,
+        throttle: t.Union[BaseThrottle, t.List[BaseThrottle], NOT_SET_TYPE] = NOT_SET,
         response: t.Union[t.Any, t.List[t.Any]] = NOT_SET,
         operation_id: t.Optional[str] = None,
         summary: t.Optional[str] = None,
@@ -97,6 +99,7 @@ def __init__(
             path=path,
             methods=methods,
             auth=auth,
+            throttle=throttle,
             response=_response,
             operation_id=operation_id,
             summary=summary,
@@ -124,6 +127,7 @@ def _create_route_function(
         path: str,
         methods: t.List[str],
         auth: t.Any = NOT_SET,
+        throttle: t.Union[BaseThrottle, t.List[BaseThrottle], NOT_SET_TYPE] = NOT_SET,
         response: t.Union[t.Any, t.List[t.Any]] = NOT_SET,
         operation_id: t.Optional[str] = None,
         summary: t.Optional[str] = None,
@@ -164,6 +168,7 @@ def _create_route_function(
             include_in_schema=include_in_schema,
             permissions=permissions,
             openapi_extra=openapi_extra,
+            throttle=throttle,
         )
         route_function_class = RouteFunction
         if route_obj.is_async:
@@ -178,6 +183,7 @@ def get(
         path: str = "",
         *,
         auth: t.Any = NOT_SET,
+        throttle: t.Union[BaseThrottle, t.List[BaseThrottle], NOT_SET_TYPE] = NOT_SET,
         response: t.Union[t.Any, t.List[t.Any]] = NOT_SET,
         operation_id: t.Optional[str] = None,
         summary: t.Optional[str] = None,
@@ -242,6 +248,7 @@ def decorator(view_func: TCallable) -> TCallable:
                 include_in_schema=include_in_schema,
                 permissions=permissions,
                 openapi_extra=openapi_extra,
+                throttle=throttle,
             )
 
         return decorator
@@ -252,6 +259,7 @@ def post(
         path: str = "",
         *,
         auth: t.Any = NOT_SET,
+        throttle: t.Union[BaseThrottle, t.List[BaseThrottle], NOT_SET_TYPE] = NOT_SET,
         response: t.Union[t.Any, t.List[t.Any]] = NOT_SET,
         operation_id: t.Optional[str] = None,
         summary: t.Optional[str] = None,
@@ -316,6 +324,7 @@ def decorator(view_func: TCallable) -> TCallable:
                 include_in_schema=include_in_schema,
                 permissions=permissions,
                 openapi_extra=openapi_extra,
+                throttle=throttle,
             )
 
         return decorator
@@ -326,6 +335,7 @@ def delete(
         path: str = "",
         *,
         auth: t.Any = NOT_SET,
+        throttle: t.Union[BaseThrottle, t.List[BaseThrottle], NOT_SET_TYPE] = NOT_SET,
         response: t.Union[t.Any, t.List[t.Any]] = NOT_SET,
         operation_id: t.Optional[str] = None,
         summary: t.Optional[str] = None,
@@ -390,6 +400,7 @@ def decorator(view_func: TCallable) -> TCallable:
                 include_in_schema=include_in_schema,
                 permissions=permissions,
                 openapi_extra=openapi_extra,
+                throttle=throttle,
             )
 
         return decorator
@@ -400,6 +411,7 @@ def patch(
         path: str = "",
         *,
         auth: t.Any = NOT_SET,
+        throttle: t.Union[BaseThrottle, t.List[BaseThrottle], NOT_SET_TYPE] = NOT_SET,
         response: t.Union[t.Any, t.List[t.Any]] = NOT_SET,
         operation_id: t.Optional[str] = None,
         summary: t.Optional[str] = None,
@@ -465,6 +477,7 @@ def decorator(view_func: TCallable) -> TCallable:
                 include_in_schema=include_in_schema,
                 permissions=permissions,
                 openapi_extra=openapi_extra,
+                throttle=throttle,
             )
 
         return decorator
@@ -475,6 +488,7 @@ def put(
         path: str = "",
         *,
         auth: t.Any = NOT_SET,
+        throttle: t.Union[BaseThrottle, t.List[BaseThrottle], NOT_SET_TYPE] = NOT_SET,
         response: t.Union[t.Any, t.List[t.Any]] = NOT_SET,
         operation_id: t.Optional[str] = None,
         summary: t.Optional[str] = None,
@@ -540,6 +554,7 @@ def decorator(view_func: TCallable) -> TCallable:
                 include_in_schema=include_in_schema,
                 permissions=permissions,
                 openapi_extra=openapi_extra,
+                throttle=throttle,
             )
 
         return decorator
@@ -551,6 +566,7 @@ def generic(
         *,
         methods: t.List[str],
         auth: t.Any = NOT_SET,
+        throttle: t.Union[BaseThrottle, t.List[BaseThrottle], NOT_SET_TYPE] = NOT_SET,
         response: t.Union[t.Any, t.List[t.Any]] = NOT_SET,
         operation_id: t.Optional[str] = None,
         summary: t.Optional[str] = None,
@@ -617,6 +633,7 @@ def decorator(view_func: TCallable) -> TCallable:
                 include_in_schema=include_in_schema,
                 permissions=permissions,
                 openapi_extra=openapi_extra,
+                throttle=throttle,
             )
 
         return decorator

ninja_extra/operation.py

@@ -19,7 +19,7 @@
 from django.http import HttpRequest
 from django.http.response import HttpResponse, HttpResponseBase
 from django.utils.encoding import force_str
-from ninja.constants import NOT_SET
+from ninja.constants import NOT_SET, NOT_SET_TYPE
 from ninja.errors import AuthenticationError
 from ninja.operation import (
     AsyncOperation as NinjaAsyncOperation,
@@ -31,12 +31,13 @@
     PathView as NinjaPathView,
 )
 from ninja.signature import is_async
+from ninja.throttling import BaseThrottle
 from ninja.types import TCallable
 from ninja.utils import check_csrf
 
 from ninja_extra.compatible import asynccontextmanager
 from ninja_extra.constants import ROUTE_CONTEXT_VAR
-from ninja_extra.exceptions import APIException
+from ninja_extra.exceptions import APIException, Throttled
 from ninja_extra.helper import get_function_name
 from ninja_extra.logger import request_logger
 
@@ -204,6 +205,22 @@ def run(self, request: HttpRequest, **kw: Any) -> HttpResponseBase:
                 e.args = (msg,) + e.args[1:]
             return self.api.on_exception(request, e)
 
+    def _check_throttles(self, request: HttpRequest) -> Optional[HttpResponse]:
+        throttle_durations = []
+        for throttle in self.throttle_objects:
+            if not throttle.allow_request(request):
+                throttle_durations.append(throttle.wait())
+
+        if throttle_durations:
+            # Filter out `None` values which may happen in case of config / rate
+            durations = [
+                duration for duration in throttle_durations if duration is not None
+            ]
+
+            duration = max(durations, default=None)
+            raise Throttled(wait=duration)
+        return None
+
 
 class AsyncOperation(Operation, NinjaAsyncOperation):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
@@ -218,15 +235,21 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
 
     async def _run_checks(self, request: HttpRequest) -> Optional[HttpResponse]:  # type: ignore
         """Runs security checks for each operation"""
+        # csrf:
+        if self.api.csrf:
+            error = check_csrf(request, self.view_func)
+            if error:
+                return error
+
         # auth:
         if self.auth_callbacks:
-            error = await self._run_authentication(request)
+            error = await self._run_authentication(request)  # type: ignore[assignment]
             if error:
                 return error
 
-        # csrf:
-        if self.api.csrf:
-            error = check_csrf(request, self.view_func)
+        # Throttling:
+        if self.throttle_objects:
+            error = self._check_throttles(request)  # type: ignore
             if error:
                 return error
 
@@ -295,7 +318,7 @@ async def run(self, request: HttpRequest, **kw: Any) -> HttpResponseBase:  # typ
                 result = await self.view_func(request, **values)
                 _processed_results = await self._result_to_response(
                     request, result, temporal_response
-                )  # type: ignore
+                )
                 return cast(HttpResponseBase, _processed_results)
         except Exception as e:
             return self.api.on_exception(request, e)
@@ -317,6 +340,7 @@ def add_operation(
         view_func: Callable,
         *,
         auth: Optional[Union[Sequence[Callable], Callable, object]] = NOT_SET,
+        throttle: Union[BaseThrottle, List[BaseThrottle], NOT_SET_TYPE] = NOT_SET,
         response: Any = NOT_SET,
         operation_id: Optional[str] = None,
         summary: Optional[str] = None,
@@ -352,6 +376,7 @@ def add_operation(
             include_in_schema=include_in_schema,
             url_name=url_name,
             openapi_extra=openapi_extra,
+            throttle=throttle,
         )
 
         self.operations.append(operation)