limits 4.0.1__py3-none-any.whl → 4.1__py3-none-any.whl

limits/storage/mongodb.py

@@ -3,27 +3,26 @@ from __future__ import annotations
 import datetime
 import time
 from abc import ABC, abstractmethod
-from typing import cast
 
 from deprecated.sphinx import versionadded, versionchanged
 
 from limits.typing import (
-    Dict,
-    List,
     MongoClient,
     MongoCollection,
     MongoDatabase,
     Optional,
-    Tuple,
     Type,
     Union,
+    cast,
 )
 
 from ..util import get_dependency
-from .base import MovingWindowSupport, Storage
+from .base import MovingWindowSupport, SlidingWindowCounterSupport, Storage
 
 
-class MongoDBStorageBase(Storage, MovingWindowSupport, ABC):
+class MongoDBStorageBase(
+    Storage, MovingWindowSupport, SlidingWindowCounterSupport, ABC
+):
     """
     Rate limit storage with MongoDB as backend.
 
@@ -48,7 +47,8 @@ class MongoDBStorageBase(Storage, MovingWindowSupport, ABC):
          collections.
         :param counter_collection_name: The collection name to use for individual counters
          used in fixed window strategies
-        :param window_collection_name: The collection name to use for moving window storage
+        :param window_collection_name: The collection name to use for sliding & moving window
+         storage
         :param wrap_exceptions: Whether to wrap storage exceptions in
          :exc:`limits.errors.StorageError` before raising it.
         :param options: all remaining keyword arguments are passed to the
@@ -98,7 +98,7 @@ class MongoDBStorageBase(Storage, MovingWindowSupport, ABC):
     @property
     def base_exceptions(
         self,
-    ) -> Union[Type[Exception], Tuple[Type[Exception], ...]]:  # pragma: no cover
+    ) -> Union[Type[Exception], tuple[Type[Exception], ...]]:  # pragma: no cover
         return self.lib_errors.PyMongoError  # type: ignore
 
     def __initialize_database(self) -> None:
@@ -203,7 +203,7 @@ class MongoDBStorageBase(Storage, MovingWindowSupport, ABC):
         except:  # noqa: E722
             return False
 
-    def get_moving_window(self, key: str, limit: int, expiry: int) -> Tuple[float, int]:
+    def get_moving_window(self, key: str, limit: int, expiry: int) -> tuple[float, int]:
         """
         returns the starting point and the number of entries in the moving
         window
@@ -257,9 +257,9 @@ class MongoDBStorageBase(Storage, MovingWindowSupport, ABC):
 
         timestamp = time.time()
         try:
-            updates: Dict[
+            updates: dict[
                 str,
-                Dict[str, Union[datetime.datetime, Dict[str, Union[List[float], int]]]],
+                dict[str, Union[datetime.datetime, dict[str, Union[list[float], int]]]],
             ] = {
                 "$push": {
                     "entries": {
@@ -289,6 +289,199 @@ class MongoDBStorageBase(Storage, MovingWindowSupport, ABC):
         except self.lib.errors.DuplicateKeyError:
             return False
 
+    def get_sliding_window(
+        self, key: str, expiry: int
+    ) -> tuple[int, float, int, float]:
+        expiry_ms = expiry * 1000
+        if result := self.windows.find_one_and_update(
+            {"_id": key},
+            [
+                {
+                    "$set": {
+                        "previousCount": {
+                            "$cond": {
+                                "if": {
+                                    "$lte": [
+                                        {"$subtract": ["$expiresAt", "$$NOW"]},
+                                        expiry_ms,
+                                    ]
+                                },
+                                "then": {"$ifNull": ["$currentCount", 0]},
+                                "else": {"$ifNull": ["$previousCount", 0]},
+                            }
+                        },
+                        "currentCount": {
+                            "$cond": {
+                                "if": {
+                                    "$lte": [
+                                        {"$subtract": ["$expiresAt", "$$NOW"]},
+                                        expiry_ms,
+                                    ]
+                                },
+                                "then": 0,
+                                "else": {"$ifNull": ["$currentCount", 0]},
+                            }
+                        },
+                        "expiresAt": {
+                            "$cond": {
+                                "if": {
+                                    "$lte": [
+                                        {"$subtract": ["$expiresAt", "$$NOW"]},
+                                        expiry_ms,
+                                    ]
+                                },
+                                "then": {
+                                    "$add": ["$expiresAt", expiry_ms],
+                                },
+                                "else": "$expiresAt",
+                            }
+                        },
+                    }
+                }
+            ],
+            return_document=self.lib.ReturnDocument.AFTER,
+            projection=["currentCount", "previousCount", "expiresAt"],
+        ):
+            expires_at = (
+                (result["expiresAt"].replace(tzinfo=datetime.timezone.utc).timestamp())
+                if result.get("expiresAt")
+                else time.time()
+            )
+            current_ttl = max(0, expires_at - time.time())
+            prev_ttl = max(0, current_ttl - expiry if result["previousCount"] else 0)
+
+            return (
+                result["previousCount"],
+                prev_ttl,
+                result["currentCount"],
+                current_ttl,
+            )
+        return 0, 0.0, 0, 0.0
+
+    def acquire_sliding_window_entry(
+        self, key: str, limit: int, expiry: int, amount: int = 1
+    ) -> bool:
+        expiry_ms = expiry * 1000
+        result = self.windows.find_one_and_update(
+            {"_id": key},
+            [
+                {
+                    "$set": {
+                        "previousCount": {
+                            "$cond": {
+                                "if": {
+                                    "$lte": [
+                                        {"$subtract": ["$expiresAt", "$$NOW"]},
+                                        expiry_ms,
+                                    ]
+                                },
+                                "then": {"$ifNull": ["$currentCount", 0]},
+                                "else": {"$ifNull": ["$previousCount", 0]},
+                            }
+                        },
+                    }
+                },
+                {
+                    "$set": {
+                        "currentCount": {
+                            "$cond": {
+                                "if": {
+                                    "$lte": [
+                                        {"$subtract": ["$expiresAt", "$$NOW"]},
+                                        expiry_ms,
+                                    ]
+                                },
+                                "then": 0,
+                                "else": {"$ifNull": ["$currentCount", 0]},
+                            }
+                        },
+                        "expiresAt": {
+                            "$cond": {
+                                "if": {
+                                    "$lte": [
+                                        {"$subtract": ["$expiresAt", "$$NOW"]},
+                                        expiry_ms,
+                                    ]
+                                },
+                                "then": {
+                                    "$cond": {
+                                        "if": {"$gt": ["$expiresAt", 0]},
+                                        "then": {"$add": ["$expiresAt", expiry_ms]},
+                                        "else": {"$add": ["$$NOW", 2 * expiry_ms]},
+                                    }
+                                },
+                                "else": "$expiresAt",
+                            }
+                        },
+                    }
+                },
+                {
+                    "$set": {
+                        "curWeightedCount": {
+                            "$floor": {
+                                "$add": [
+                                    {
+                                        "$multiply": [
+                                            "$previousCount",
+                                            {
+                                                "$divide": [
+                                                    {
+                                                        "$max": [
+                                                            0,
+                                                            {
+                                                                "$subtract": [
+                                                                    "$expiresAt",
+                                                                    {
+                                                                        "$add": [
+                                                                            "$$NOW",
+                                                                            expiry_ms,
+                                                                        ]
+                                                                    },
+                                                                ]
+                                                            },
+                                                        ]
+                                                    },
+                                                    expiry_ms,
+                                                ]
+                                            },
+                                        ]
+                                    },
+                                    "$currentCount",
+                                ]
+                            }
+                        }
+                    }
+                },
+                {
+                    "$set": {
+                        "currentCount": {
+                            "$cond": {
+                                "if": {
+                                    "$lte": [
+                                        {"$add": ["$curWeightedCount", amount]},
+                                        limit,
+                                    ]
+                                },
+                                "then": {"$add": ["$currentCount", amount]},
+                                "else": "$currentCount",
+                            }
+                        }
+                    }
+                },
+                {
+                    "$set": {
+                        "_acquired": {
+                            "$lte": [{"$add": ["$curWeightedCount", amount]}, limit]
+                        }
+                    }
+                },
+                {"$unset": ["curWeightedCount"]},
+            ],
+            return_document=self.lib.ReturnDocument.AFTER,
+            upsert=True,
+        )
+        return cast(bool, result["_acquired"])
+
 
 @versionadded(version="2.1")
 @versionchanged(

limits/storage/redis.py

@@ -1,20 +1,31 @@
 from __future__ import annotations
 
 import time
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 from packaging.version import Version
 
-from limits.typing import Optional, RedisClient, ScriptP, Tuple, Type, Union
+from limits.typing import Optional, RedisClient, Type, Union
 
 from ..util import get_package_data
-from .base import MovingWindowSupport, Storage
+from .base import MovingWindowSupport, SlidingWindowCounterSupport, Storage
 
 if TYPE_CHECKING:
     import redis
 
 
-class RedisInteractor:
+class RedisStorage(Storage, MovingWindowSupport, SlidingWindowCounterSupport):
+    """
+    Rate limit storage with redis as backend.
+
+    Depends on :pypi:`redis`.
+    """
+
+    STORAGE_SCHEME = ["redis", "rediss", "redis+unix"]
+    """The storage scheme for redis"""
+
+    DEPENDENCIES = {"redis": Version("3.0")}
+
     RES_DIR = "resources/redis/lua_scripts"
 
     SCRIPT_MOVING_WINDOW = get_package_data(f"{RES_DIR}/moving_window.lua")
@@ -24,123 +35,17 @@ class RedisInteractor:
     SCRIPT_CLEAR_KEYS = get_package_data(f"{RES_DIR}/clear_keys.lua")
     SCRIPT_INCR_EXPIRE = get_package_data(f"{RES_DIR}/incr_expire.lua")
 
-    lua_moving_window: ScriptP[Tuple[int, int]]
-    lua_acquire_window: ScriptP[bool]
-
-    PREFIX = "LIMITS"
-
-    def prefixed_key(self, key: str) -> str:
-        return f"{self.PREFIX}:{key}"
-
-    def get_moving_window(self, key: str, limit: int, expiry: int) -> Tuple[float, int]:
-        """
-        returns the starting point and the number of entries in the moving
-        window
-
-        :param key: rate limit key
-        :param expiry: expiry of entry
-        :return: (start of window, number of acquired entries)
-        """
-        key = self.prefixed_key(key)
-        timestamp = time.time()
-        if window := self.lua_moving_window([key], [timestamp - expiry, limit]):
-            return float(window[0]), window[1]
-
-        return timestamp, 0
-
-    def _incr(
-        self,
-        key: str,
-        expiry: int,
-        connection: RedisClient,
-        elastic_expiry: bool = False,
-        amount: int = 1,
-    ) -> int:
-        """
-        increments the counter for a given rate limit key
-
-        :param connection: Redis connection
-        :param key: the key to increment
-        :param expiry: amount in seconds for the key to expire in
-        :param amount: the number to increment by
-        """
-        key = self.prefixed_key(key)
-        value = connection.incrby(key, amount)
-
-        if elastic_expiry or value == amount:
-            connection.expire(key, expiry)
-
-        return value
-
-    def _get(self, key: str, connection: RedisClient) -> int:
-        """
-        :param connection: Redis connection
-        :param key: the key to get the counter value for
-        """
-
-        key = self.prefixed_key(key)
-        return int(connection.get(key) or 0)
-
-    def _clear(self, key: str, connection: RedisClient) -> None:
-        """
-        :param key: the key to clear rate limits for
-        :param connection: Redis connection
-        """
-        key = self.prefixed_key(key)
-        connection.delete(key)
-
-    def _acquire_entry(
-        self,
-        key: str,
-        limit: int,
-        expiry: int,
-        connection: RedisClient,
-        amount: int = 1,
-    ) -> bool:
-        """
-        :param key: rate limit key to acquire an entry in
-        :param limit: amount of entries allowed
-        :param expiry: expiry of the entry
-        :param connection: Redis connection
-        :param amount: the number of entries to acquire
-        """
-        key = self.prefixed_key(key)
-        timestamp = time.time()
-        acquired = self.lua_acquire_window([key], [timestamp, limit, expiry, amount])
-
-        return bool(acquired)
-
-    def _get_expiry(self, key: str, connection: RedisClient) -> float:
-        """
-        :param key: the key to get the expiry for
-        :param connection: Redis connection
-        """
-
-        key = self.prefixed_key(key)
-        return max(connection.ttl(key), 0) + time.time()
-
-    def _check(self, connection: RedisClient) -> bool:
-        """
-        :param connection: Redis connection
-        check if storage is healthy
-        """
-        try:
-            return connection.ping()
-        except:  # noqa
-            return False
-
-
-class RedisStorage(RedisInteractor, Storage, MovingWindowSupport):
-    """
-    Rate limit storage with redis as backend.
-
-    Depends on :pypi:`redis`.
-    """
+    SCRIPT_SLIDING_WINDOW = get_package_data(f"{RES_DIR}/sliding_window.lua")
+    SCRIPT_ACQUIRE_SLIDING_WINDOW = get_package_data(
+        f"{RES_DIR}/acquire_sliding_window.lua"
+    )
 
-    STORAGE_SCHEME = ["redis", "rediss", "redis+unix"]
-    """The storage scheme for redis"""
+    lua_moving_window: "redis.commands.core.Script"
+    lua_acquire_moving_window: "redis.commands.core.Script"
+    lua_sliding_window: "redis.commands.core.Script"
+    lua_acquire_sliding_window: "redis.commands.core.Script"
 
-    DEPENDENCIES = {"redis": Version("3.0")}
+    PREFIX = "LIMITS"
 
     def __init__(
         self,
@@ -179,73 +84,189 @@ class RedisStorage(RedisInteractor, Storage, MovingWindowSupport):
     @property
     def base_exceptions(
         self,
-    ) -> Union[Type[Exception], Tuple[Type[Exception], ...]]:  # pragma: no cover
+    ) -> Union[Type[Exception], tuple[Type[Exception], ...]]:  # pragma: no cover
         return self.dependency.RedisError  # type: ignore[no-any-return]
 
     def initialize_storage(self, _uri: str) -> None:
-        self.lua_moving_window = self.storage.register_script(self.SCRIPT_MOVING_WINDOW)
-        self.lua_acquire_window = self.storage.register_script(
+        self.lua_moving_window = self.get_connection().register_script(
+            self.SCRIPT_MOVING_WINDOW
+        )
+        self.lua_acquire_moving_window = self.get_connection().register_script(
             self.SCRIPT_ACQUIRE_MOVING_WINDOW
         )
-        self.lua_clear_keys = self.storage.register_script(self.SCRIPT_CLEAR_KEYS)
-        self.lua_incr_expire = self.storage.register_script(
-            RedisStorage.SCRIPT_INCR_EXPIRE
+        self.lua_clear_keys = self.get_connection().register_script(
+            self.SCRIPT_CLEAR_KEYS
+        )
+        self.lua_incr_expire = self.get_connection().register_script(
+            self.SCRIPT_INCR_EXPIRE
+        )
+        self.lua_sliding_window = self.get_connection().register_script(
+            self.SCRIPT_SLIDING_WINDOW
         )
+        self.lua_acquire_sliding_window = self.get_connection().register_script(
+            self.SCRIPT_ACQUIRE_SLIDING_WINDOW
+        )
+
+    def get_connection(self, readonly: bool = False) -> RedisClient:
+        return cast(RedisClient, self.storage)
+
+    def _current_window_key(self, key: str) -> str:
+        """
+        Return the current window's storage key (Sliding window strategy)
+
+        Contrary to other strategies that have one key per rate limit item,
+        this strategy has two keys per rate limit item than must be on the same machine.
+        To keep the current key and the previous key on the same Redis cluster node,
+        curly braces are added.
+
+        Eg: "{constructed_key}"
+        """
+        return f"{{{key}}}"
+
+    def _previous_window_key(self, key: str) -> str:
+        """
+        Return the previous window's storage key (Sliding window strategy).
+
+        Curvy braces are added on the common pattern with the current window's key,
+        so the current and the previous key are stored on the same Redis cluster node.
+
+        Eg: "{constructed_key}/-1"
+        """
+        return f"{self._current_window_key(key)}/-1"
+
+    def prefixed_key(self, key: str) -> str:
+        return f"{self.PREFIX}:{key}"
+
+    def get_moving_window(self, key: str, limit: int, expiry: int) -> tuple[float, int]:
+        """
+        returns the starting point and the number of entries in the moving
+        window
+
+        :param key: rate limit key
+        :param expiry: expiry of entry
+        :return: (start of window, number of acquired entries)
+        """
+        key = self.prefixed_key(key)
+        timestamp = time.time()
+        if window := self.lua_moving_window([key], [timestamp - expiry, limit]):
+            return float(window[0]), window[1]
+
+        return timestamp, 0
+
+    def get_sliding_window(
+        self, key: str, expiry: int
+    ) -> tuple[int, float, int, float]:
+        previous_key = self.prefixed_key(self._previous_window_key(key))
+        current_key = self.prefixed_key(self._current_window_key(key))
+        if window := self.lua_sliding_window([previous_key, current_key], [expiry]):
+            return (
+                int(window[0] or 0),
+                max(0, float(window[1] or 0)) / 1000,
+                int(window[2] or 0),
+                max(0, float(window[3] or 0)) / 1000,
+            )
+        return 0, 0.0, 0, 0.0
 
     def incr(
-        self, key: str, expiry: int, elastic_expiry: bool = False, amount: int = 1
+        self,
+        key: str,
+        expiry: int,
+        elastic_expiry: bool = False,
+        amount: int = 1,
     ) -> int:
         """
         increments the counter for a given rate limit key
 
+
         :param key: the key to increment
         :param expiry: amount in seconds for the key to expire in
         :param amount: the number to increment by
         """
-
+        key = self.prefixed_key(key)
         if elastic_expiry:
-            return super()._incr(key, expiry, self.storage, elastic_expiry, amount)
+            value = self.get_connection().incrby(key, amount)
+            self.get_connection().expire(key, expiry)
+            return value
         else:
-            key = self.prefixed_key(key)
             return int(self.lua_incr_expire([key], [expiry, amount]))
 
     def get(self, key: str) -> int:
         """
+
         :param key: the key to get the counter value for
         """
 
-        return super()._get(key, self.storage)
+        key = self.prefixed_key(key)
+        return int(self.get_connection(True).get(key) or 0)
 
     def clear(self, key: str) -> None:
         """
         :param key: the key to clear rate limits for
         """
+        key = self.prefixed_key(key)
+        self.get_connection().delete(key)
 
-        return super()._clear(key, self.storage)
-
-    def acquire_entry(self, key: str, limit: int, expiry: int, amount: int = 1) -> bool:
+    def acquire_entry(
+        self,
+        key: str,
+        limit: int,
+        expiry: int,
+        amount: int = 1,
+    ) -> bool:
         """
         :param key: rate limit key to acquire an entry in
         :param limit: amount of entries allowed
         :param expiry: expiry of the entry
-        :param amount: the number to increment by
+
+        :param amount: the number of entries to acquire
         """
+        key = self.prefixed_key(key)
+        timestamp = time.time()
+        acquired = self.lua_acquire_moving_window(
+            [key], [timestamp, limit, expiry, amount]
+        )
 
-        return super()._acquire_entry(key, limit, expiry, self.storage, amount)
+        return bool(acquired)
+
+    def acquire_sliding_window_entry(
+        self,
+        key: str,
+        limit: int,
+        expiry: int,
+        amount: int = 1,
+    ) -> bool:
+        """
+        Acquire an entry. Shift the current window to the previous window if it expired.
+
+        :param key: rate limit key to acquire an entry in
+        :param limit: amount of entries allowed
+        :param expiry: expiry of the entry
+        :param amount: the number of entries to acquire
+        """
+        previous_key = self.prefixed_key(self._previous_window_key(key))
+        current_key = self.prefixed_key(self._current_window_key(key))
+        acquired = self.lua_acquire_sliding_window(
+            [previous_key, current_key], [limit, expiry, amount]
+        )
+        return bool(acquired)
 
     def get_expiry(self, key: str) -> float:
         """
         :param key: the key to get the expiry for
+
         """
 
-        return super()._get_expiry(key, self.storage)
+        key = self.prefixed_key(key)
+        return max(self.get_connection(True).ttl(key), 0) + time.time()
 
     def check(self) -> bool:
         """
         check if storage is healthy
         """
-
-        return super()._check(self.storage)
+        try:
+            return self.get_connection().ping()
+        except:  # noqa
+            return False
 
     def reset(self) -> Optional[int]:
         """

limits/storage/redis_cluster.py

@@ -4,7 +4,7 @@ from deprecated.sphinx import versionchanged
 from packaging.version import Version
 
 from limits.storage.redis import RedisStorage
-from limits.typing import Dict, Optional, Union
+from limits.typing import Optional, Union
 
 
 @versionchanged(
@@ -34,7 +34,7 @@ class RedisClusterStorage(RedisStorage):
     STORAGE_SCHEME = ["redis+cluster"]
     """The storage scheme for redis cluster"""
 
-    DEFAULT_OPTIONS: Dict[str, Union[float, str, bool]] = {
+    DEFAULT_OPTIONS: dict[str, Union[float, str, bool]] = {
         "max_connections": 1000,
     }
     "Default options passed to the :class:`~redis.cluster.RedisCluster`"
@@ -60,7 +60,7 @@ class RedisClusterStorage(RedisStorage):
          available or if the redis cluster cannot be reached.
         """
         parsed = urllib.parse.urlparse(uri)
-        parsed_auth: Dict[str, Union[float, str, bool]] = {}
+        parsed_auth: dict[str, Union[float, str, bool]] = {}
 
         if parsed.username:
             parsed_auth["username"] = parsed.username