Support multiple rate strategy for one route. (long2ice#3)

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## 0.1
 
+### 0.1.3
+
+- Support multiple rate strategy for one route. (#3)
+
 ### 0.1.2
 
 - Use milliseconds instead of seconds as default unit of expiration.

README.md

@@ -72,7 +72,7 @@ async def default_identifier(request: Request):
     forwarded = request.headers.get("X-Forwarded-For")
     if forwarded:
         return forwarded.split(",")[0]
-    return request.client.host
+    return request.client.host + ":" + request.scope["path"]
 ```
 
 ### callback

examples/main.py

@@ -19,5 +19,16 @@ async def index():
     return {"msg": "Hello World"}
 
 
+@app.get(
+    "/multiple",
+    dependencies=[
+        Depends(RateLimiter(times=1, seconds=5)),
+        Depends(RateLimiter(times=2, seconds=15)),
+    ],
+)
+async def multiple():
+    return {"msg": "Hello World"}
+
+
 if __name__ == "__main__":
     uvicorn.run("main:app", debug=True, reload=True)

fastapi_limiter/__init__.py

@@ -12,7 +12,7 @@ async def default_identifier(request: Request):
     forwarded = request.headers.get("X-Forwarded-For")
     if forwarded:
         return forwarded.split(",")[0]
-    return request.client.host
+    return request.client.host + ":" + request.scope["path"]
 
 
 async def default_callback(request: Request, response: Response, pexpire: int):

fastapi_limiter/depends.py

@@ -26,12 +26,19 @@ def __init__(
     async def __call__(self, request: Request, response: Response):
         if not FastAPILimiter.redis:
             raise Exception("You must call FastAPILimiter.init in startup event of fastapi!")
+        index = 0
+        for route in request.app.routes:
+            if route.path == request.scope["path"]:
+                for idx, dependency in enumerate(route.dependencies):
+                    if self is dependency.dependency:
+                        index = idx
+                        break
         # moved here because constructor run before app startup
         identifier = self.identifier or FastAPILimiter.identifier
         callback = self.callback or FastAPILimiter.callback
         redis = FastAPILimiter.redis
         rate_key = await identifier(request)
-        key = FastAPILimiter.prefix + ":" + rate_key
+        key = f"{FastAPILimiter.prefix}:{rate_key}:{index}"
         tr = redis.multi_exec()
         tr.incrby(key, 1)
         tr.pttl(key)

pyproject.toml

@@ -12,7 +12,7 @@ packages = [
 ]
 readme = "README.md"
 repository = "https://github.com/long2ice/fastapi-limiter.git"
-version = "0.1.2"
+version = "0.1.3"
 
 [tool.poetry.dependencies]
 aioredis = "*"

tests/test_depends.py

@@ -9,7 +9,6 @@ def test_limiter():
     with TestClient(app) as client:
         response = client.get("/")
         assert response.status_code == 200
-        assert response.json() == {"msg": "Hello World"}
 
         client.get("/")
 
@@ -19,4 +18,23 @@ def test_limiter():
 
         response = client.get("/")
         assert response.status_code == 200
-        assert response.json() == {"msg": "Hello World"}
+
+
+def test_limiter_multiple():
+    with TestClient(app) as client:
+        response = client.get("/multiple")
+        assert response.status_code == 200
+
+        response = client.get("/multiple")
+        assert response.status_code == 429
+        sleep(5)
+
+        response = client.get("/multiple")
+        assert response.status_code == 200
+
+        response = client.get("/multiple")
+        assert response.status_code == 429
+        sleep(10)
+
+        response = client.get("/multiple")
+        assert response.status_code == 200