Merge pull request #8 from r3tr0ananas/main

Add SlowAPI

Author: THEGOLDENPRO

api/__init__.py

@@ -1 +1 @@
-__version__ = "1.3.2"
+__version__ = "1.3.5"

api/errors.py

@@ -1,5 +1,16 @@
+from fastapi import Request
+from slowapi.errors import RateLimitExceeded
+from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 
+__all__ = (
+    "APIException",
+    "CategoryNotFound",
+    "BookNotFound",
+    "RateLimited",
+    "rate_limit_handler"
+)
+
 class APIException(Exception):
     def __init__(self, msg) -> None:
         self.msg = msg
@@ -34,4 +45,35 @@ class BookNotFound(BaseModel):
                 }
             ]
         }
-    }
+    }
+
+class RateLimited(BaseModel):
+    error: str
+    message: str
+
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {
+                    "error": "RateLimited",
+                    "message": "Rate Limit exceeded: 3 per 1 second"
+                }
+            ]
+        }
+    }
+
+
+def rate_limit_handler(request: Request, exc: RateLimitExceeded):
+    response = JSONResponse(
+        status_code = 429,
+        content = {
+            "error": "RateLimited", 
+            "message": f"Rate limit exceeded: {exc.detail} (Follow the rates: https://github.com/THEGOLDENPRO/aghpb_api/wiki#rate-limiting)"
+        }
+    )
+
+    response = request.app.state.limiter._inject_headers(
+        response, request.state.view_rate_limit
+    )
+
+    return response

api/main.py

@@ -1,4 +1,3 @@
-from __future__ import annotations
 from typing import TYPE_CHECKING, List # DON'T YOU DARE PUT THIS UNDER TYPE_CHECKING!!! I'm warning you!
 
 if TYPE_CHECKING:
@@ -9,9 +8,13 @@
 from . import errors, __version__
 from .anime_girls import AGHPB, CategoryNotFound, Book, BookDict
 
-from fastapi import FastAPI, Query
+from fastapi import FastAPI, Query, Request
 from fastapi.responses import FileResponse, JSONResponse, RedirectResponse
 
+from slowapi import Limiter
+from slowapi.util import get_remote_address
+from slowapi.errors import RateLimitExceeded
+
 ROOT_PATH = (lambda x: x if x is not None else "")(os.environ.get("ROOT_PATH")) # Like: /aghpb/v1
 
 TAGS_METADATA = [
@@ -27,13 +30,24 @@
 ]
 
 DESCRIPTION = """
-Behold the **anime girls holding programming books** API. ✴️
+<div align="center">
+
+  <img src="https://raw.githubusercontent.com/THEGOLDENPRO/aghpb_api/main/assets/logo.png" alt="Logo" width="180">
+
+  Behold the **anime girls holding programming books** API. ✴️
+
+  This is a ✨ feature rich 🌟 [open source](https://github.com/THEGOLDENPRO/aghpb_api) API I made for the anime girls holding programming books [github repo](https://github.com/cat-milk/Anime-Girls-Holding-Programming-Books) because I was bored.
+
+  🐞 Report bugs [over here](https://github.com/THEGOLDENPRO/aghpb_api/issues).
 
-This is a ✨ feature rich 🌟 [open source](https://github.com/THEGOLDENPRO/aghpb_api) API I made for the anime girls holding programming books [github repo](https://github.com/cat-milk/Anime-Girls-Holding-Programming-Books) because I was bored.
+</div>
 
-🐞 Report bugs [over here](https://github.com/THEGOLDENPRO/aghpb_api/issues).
+<br>
+
+Rate limiting applies to the ``/random`` and ``/get`` endpoints. Check out the rate limits [over here](https://github.com/THEGOLDENPRO/aghpb_api/wiki#rate-limiting).
 """
 
+limiter = Limiter(key_func=get_remote_address, headers_enabled=True)
 app = FastAPI(
     title = "AGHPB API",
     description = DESCRIPTION,
@@ -46,6 +60,9 @@
 
     root_path = ROOT_PATH
 )
+app.state.limiter = limiter
+app.add_exception_handler(RateLimitExceeded, errors.rate_limit_handler)
+
 @app.get(
     "/",
     name = "Takes you to these docs.",
@@ -74,10 +91,15 @@ async def root():
         404: {
             "model": errors.CategoryNotFound, 
             "description": "The category was not Found."
+        },
+        429: {
+            "model": errors.RateLimited,
+            "description": "Rate limit exceeded!"
         }
     },
 )
-async def random(category: str = None) -> FileResponse:
+@limiter.limit("3/second")
+async def random(request: Request, category: str = None) -> FileResponse:
     """Returns a random book."""
     if category is None:
         category = aghpb.random_category()
@@ -138,7 +160,6 @@ async def search(
         book[1].to_dict() for book in books
     ]
 
-
 @app.get(
     "/get/id/{search_id}",
     name = "Allows you to get a book by search id.",
@@ -155,10 +176,15 @@ async def search(
         404: {
             "model": errors.BookNotFound, 
             "description": "The book was not Found."
+        },
+        429: {
+            "model": errors.RateLimited,
+            "description": "Rate Limit exceeded"
         }
     },
 )
-async def get_id(search_id: str) -> FileResponse:
+@limiter.limit("3/second")
+async def get_id(request: Request, search_id: str) -> FileResponse:
     """Returns the book found."""
     for book in aghpb.books:
 

requirements.txt

@@ -2,4 +2,5 @@ thefuzz
 fastapi[all]
 devgoldyutils
 typing-extensions
-uvicorn[standard]
+uvicorn[standard]
+slowapi