Merge pull request #4 from rowrowrowrow/support-milliseconds

Support milliseconds

CHANGELOG.md

@@ -7,3 +7,8 @@
 - Configuring the global default through the FastAPILimiter.init method.
 - Update status to 429 when too many requests.
 - Update default_callback params and add `Retry-After` response header.
+
+### 0.1.2
+
+- Use milliseconds instead of seconds as default unit of expiration.
+- Update default_callback, round milliseconds up to nearest second for `Retry-After` value.

fastapi_limiter/__init__.py

@@ -1,6 +1,7 @@
 from typing import Callable
 
 import aioredis
+from math import ceil
 from fastapi import HTTPException
 from starlette.requests import Request
 from starlette.status import HTTP_429_TOO_MANY_REQUESTS
@@ -13,13 +14,15 @@ async def default_identifier(request: Request):
     return request.client.host
 
 
-async def default_callback(request: Request, expire: int):
+async def default_callback(request: Request, pexpire: int):
     """
     default callback when too many requests
     :param request:
-    :param expire: The remaining seconds
+    :param pexpire: The remaining milliseconds
     :return:
     """
+    expire = ceil(pexpire / 1000)
+
     raise HTTPException(
         HTTP_429_TOO_MANY_REQUESTS, "Too Many Requests", headers={"Retry-After": str(expire)}
     )

fastapi_limiter/depends.py

@@ -11,14 +11,15 @@ class RateLimiter:
     def __init__(
         self,
         times: conint(ge=0) = 1,
+        milliseconds: conint(ge=-1) = 0,
         seconds: conint(ge=-1) = 0,
         minutes: conint(ge=-1) = 0,
         hours: conint(ge=-1) = 0,
         identifier: Optional[Callable] = None,
         callback: Optional[Callable] = None,
     ):
         self.times = times
-        self.seconds = seconds + 60 * minutes + 3600 * hours
+        self.milliseconds = milliseconds + 1000 * seconds + 60000 * minutes + 3600000 * hours
         self.identifier = identifier
         self.callback = callback
 
@@ -33,9 +34,9 @@ async def __call__(self, request: Request, response: Response):
         key = FastAPILimiter.prefix + ":" + rate_key
         p = redis.pipeline()
         p.incrby(key, 1)
-        p.ttl(key)
-        num, expire = await p.execute()
+        p.pttl(key)
+        num, pexpire = await p.execute()
         if num == 1:
-            await redis.expire(key, self.seconds)
+            await redis.pexpire(key, self.milliseconds)
         if num > self.times:
-            return await callback(request, expire)
+            return await callback(request, pexpire)

pyproject.toml

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