airflow_scheduler/.venv/Lib/site-packages/limits/aio/storage/mongodb.py

from __future__ import annotations

import asyncio
import datetime
import time

from deprecated.sphinx import versionadded, versionchanged

from limits.aio.storage.base import (
    MovingWindowSupport,
    SlidingWindowCounterSupport,
    Storage,
)
from limits.typing import (
    ParamSpec,
    TypeVar,
    cast,
)
from limits.util import get_dependency

P = ParamSpec("P")
R = TypeVar("R")


@versionadded(version="2.1")
@versionchanged(
    version="3.14.0",
    reason="Added option to select custom collection names for windows & counters",
)
class MongoDBStorage(Storage, MovingWindowSupport, SlidingWindowCounterSupport):
    """
    Rate limit storage with MongoDB as backend.

    Depends on :pypi:`motor`
    """

    STORAGE_SCHEME = ["async+mongodb", "async+mongodb+srv"]
    """
    The storage scheme for MongoDB for use in an async context
    """

    DEPENDENCIES = ["motor.motor_asyncio", "pymongo"]

    def __init__(
        self,
        uri: str,
        database_name: str = "limits",
        counter_collection_name: str = "counters",
        window_collection_name: str = "windows",
        wrap_exceptions: bool = False,
        **options: float | str | bool,
    ) -> None:
        """
        :param uri: uri of the form ``async+mongodb://[user:password]@host:port?...``,
         This uri is passed directly to :class:`~motor.motor_asyncio.AsyncIOMotorClient`
        :param database_name: The database to use for storing the rate limit
         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 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 constructor of :class:`~motor.motor_asyncio.AsyncIOMotorClient`
        :raise ConfigurationError: when the :pypi:`motor` or :pypi:`pymongo` are
         not available
        """

        uri = uri.replace("async+mongodb", "mongodb", 1)

        super().__init__(uri, wrap_exceptions=wrap_exceptions, **options)

        self.dependency = self.dependencies["motor.motor_asyncio"]
        self.proxy_dependency = self.dependencies["pymongo"]
        self.lib_errors, _ = get_dependency("pymongo.errors")

        self.storage = self.dependency.module.AsyncIOMotorClient(uri, **options)
        # TODO: Fix this hack. It was noticed when running a benchmark
        # with FastAPI - however - doesn't appear in unit tests or in an isolated
        # use. Reference: https://jira.mongodb.org/browse/MOTOR-822
        self.storage.get_io_loop = asyncio.get_running_loop

        self.__database_name = database_name
        self.__collection_mapping = {
            "counters": counter_collection_name,
            "windows": window_collection_name,
        }
        self.__indices_created = False

    @property
    def base_exceptions(
        self,
    ) -> type[Exception] | tuple[type[Exception], ...]:  # pragma: no cover
        return self.lib_errors.PyMongoError  # type: ignore

    @property
    def database(self):  # type: ignore
        return self.storage.get_database(self.__database_name)

    async def create_indices(self) -> None:
        if not self.__indices_created:
            await asyncio.gather(
                self.database[self.__collection_mapping["counters"]].create_index(
                    "expireAt", expireAfterSeconds=0
                ),
                self.database[self.__collection_mapping["windows"]].create_index(
                    "expireAt", expireAfterSeconds=0
                ),
            )
        self.__indices_created = True

    async def reset(self) -> int | None:
        """
        Delete all rate limit keys in the rate limit collections (counters, windows)
        """
        num_keys = sum(
            await asyncio.gather(
                self.database[self.__collection_mapping["counters"]].count_documents(
                    {}
                ),
                self.database[self.__collection_mapping["windows"]].count_documents({}),
            )
        )
        await asyncio.gather(
            self.database[self.__collection_mapping["counters"]].drop(),
            self.database[self.__collection_mapping["windows"]].drop(),
        )

        return cast(int, num_keys)

    async def clear(self, key: str) -> None:
        """
        :param key: the key to clear rate limits for
        """
        await asyncio.gather(
            self.database[self.__collection_mapping["counters"]].find_one_and_delete(
                {"_id": key}
            ),
            self.database[self.__collection_mapping["windows"]].find_one_and_delete(
                {"_id": key}
            ),
        )

    async def get_expiry(self, key: str) -> float:
        """
        :param key: the key to get the expiry for
        """
        counter = await self.database[self.__collection_mapping["counters"]].find_one(
            {"_id": key}
        )
        return (
            (counter["expireAt"] if counter else datetime.datetime.now())
            .replace(tzinfo=datetime.timezone.utc)
            .timestamp()
        )

    async def get(self, key: str) -> int:
        """
        :param key: the key to get the counter value for
        """
        counter = await self.database[self.__collection_mapping["counters"]].find_one(
            {
                "_id": key,
                "expireAt": {"$gte": datetime.datetime.now(datetime.timezone.utc)},
            },
            projection=["count"],
        )

        return counter and counter["count"] or 0

    async def incr(
        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 elastic_expiry: whether to keep extending the rate limit
         window every hit.
        :param amount: the number to increment by
        """
        await self.create_indices()

        expiration = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(
            seconds=expiry
        )

        response = await self.database[
            self.__collection_mapping["counters"]
        ].find_one_and_update(
            {"_id": key},
            [
                {
                    "$set": {
                        "count": {
                            "$cond": {
                                "if": {"$lt": ["$expireAt", "$$NOW"]},
                                "then": amount,
                                "else": {"$add": ["$count", amount]},
                            }
                        },
                        "expireAt": {
                            "$cond": {
                                "if": {"$lt": ["$expireAt", "$$NOW"]},
                                "then": expiration,
                                "else": (expiration if elastic_expiry else "$expireAt"),
                            }
                        },
                    }
                },
            ],
            upsert=True,
            projection=["count"],
            return_document=self.proxy_dependency.module.ReturnDocument.AFTER,
        )

        return int(response["count"])

    async def check(self) -> bool:
        """
        Check if storage is healthy by calling
        :meth:`motor.motor_asyncio.AsyncIOMotorClient.server_info`
        """
        try:
            await self.storage.server_info()

            return True
        except:  # noqa: E722
            return False

    async 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 str key: rate limit key
        :param int expiry: expiry of entry
        :return: (start of window, number of acquired entries)
        """
        timestamp = time.time()
        if result := (
            await self.database[self.__collection_mapping["windows"]]
            .aggregate(
                [
                    {"$match": {"_id": key}},
                    {
                        "$project": {
                            "entries": {
                                "$filter": {
                                    "input": "$entries",
                                    "as": "entry",
                                    "cond": {"$gte": ["$$entry", timestamp - expiry]},
                                }
                            }
                        }
                    },
                    {"$unwind": "$entries"},
                    {
                        "$group": {
                            "_id": "$_id",
                            "min": {"$min": "$entries"},
                            "count": {"$sum": 1},
                        }
                    },
                ]
            )
            .to_list(length=1)
        ):
            return result[0]["min"], result[0]["count"]
        return timestamp, 0

    async 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 of entries to acquire
        """
        await self.create_indices()

        if amount > limit:
            return False

        timestamp = time.time()
        try:
            updates: dict[
                str,
                dict[str, datetime.datetime | dict[str, list[float] | int]],
            ] = {
                "$push": {
                    "entries": {
                        "$each": [timestamp] * amount,
                        "$position": 0,
                        "$slice": limit,
                    }
                },
                "$set": {
                    "expireAt": (
                        datetime.datetime.now(datetime.timezone.utc)
                        + datetime.timedelta(seconds=expiry)
                    )
                },
            }

            await self.database[self.__collection_mapping["windows"]].update_one(
                {
                    "_id": key,
                    f"entries.{limit - amount}": {"$not": {"$gte": timestamp - expiry}},
                },
                updates,
                upsert=True,
            )

            return True
        except self.proxy_dependency.module.errors.DuplicateKeyError:
            return False

    async def acquire_sliding_window_entry(
        self, key: str, limit: int, expiry: int, amount: int = 1
    ) -> bool:
        await self.create_indices()
        expiry_ms = expiry * 1000
        result = await self.database[
            self.__collection_mapping["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.proxy_dependency.module.ReturnDocument.AFTER,
            upsert=True,
        )

        return cast(bool, result["_acquired"])

    async def get_sliding_window(
        self, key: str, expiry: int
    ) -> tuple[int, float, int, float]:
        expiry_ms = expiry * 1000
        if result := await self.database[
            self.__collection_mapping["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.proxy_dependency.module.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 __del__(self) -> None:
        self.storage and self.storage.close()