diff --git a/src/quota/quota_exceed_error.py b/src/quota/quota_exceed_error.py new file mode 100644 index 00000000..3adb0a10 --- /dev/null +++ b/src/quota/quota_exceed_error.py @@ -0,0 +1,38 @@ +"""Any exception that can occur when user does not have enough tokens available.""" + +# pylint: disable=line-too-long) + + +class QuotaExceedError(Exception): + """Any exception that can occur when user does not have enough tokens available.""" + + def __init__( + self, subject_id: str, subject_type: str, available: int, needed: int = 0 + ) -> None: + """Construct exception object.""" + message: str = "" + + if needed == 0 and available <= 0: + match subject_type: + case "u": + message = f"User {subject_id} has no available tokens" + case "c": + message = "Cluster has no available tokens" + case _: + message = f"Unknown subject {subject_id} has no available tokens" + else: + match subject_type: + case "u": + message = f"User {subject_id} has {available} tokens, but {needed} tokens are needed" # noqa: E501 + case "c": + message = f"Cluster has {available} tokens, but {needed} tokens are needed" + case _: + message = f"Unknown subject {subject_id} has {available} tokens, but {needed} tokens are needed" # noqa: E501 + + # call the base class constructor with the parameters it needs + super().__init__(message) + + # custom attributes + self.subject_id = subject_id + self.available = available + self.needed = needed diff --git a/src/quota/quota_limiter.py b/src/quota/quota_limiter.py new file mode 100644 index 00000000..d2731741 --- /dev/null +++ b/src/quota/quota_limiter.py @@ -0,0 +1,49 @@ +"""Abstract class that is parent for all quota limiter implementations.""" + +import logging +from abc import ABC, abstractmethod + + +logger = logging.getLogger(__name__) + + +class QuotaLimiter(ABC): + """Abstract class that is parent for all quota limiter implementations.""" + + @abstractmethod + def available_quota(self, subject_id: str) -> int: + """Retrieve available quota for given user.""" + + @abstractmethod + def revoke_quota(self) -> None: + """Revoke quota for given user.""" + + @abstractmethod + def increase_quota(self) -> None: + """Increase quota for given user.""" + + @abstractmethod + def ensure_available_quota(self, subject_id: str = "") -> None: + """Ensure that there's available quota left.""" + + @abstractmethod + def consume_tokens( + self, input_tokens: int, output_tokens: int, subject_id: str = "" + ) -> None: + """Consume tokens by given user.""" + + @abstractmethod + def __init__(self) -> None: + """Initialize connection config.""" + + @abstractmethod + def _initialize_tables(self) -> None: + """Initialize tables and indexes.""" + + # pylint: disable=W0201 + def connect(self) -> None: + """Initialize connection to database.""" + + def connected(self) -> bool: + """Check if connection to cache is alive.""" + return True diff --git a/src/quota/quota_limiter_factory.py b/src/quota/quota_limiter_factory.py new file mode 100644 index 00000000..92f743f1 --- /dev/null +++ b/src/quota/quota_limiter_factory.py @@ -0,0 +1,32 @@ +"""Quota limiter factory class.""" + +import logging + +from models.config import QuotaHandlersConfiguration + +from quota.quota_limiter import QuotaLimiter + +logger = logging.getLogger(__name__) + + +# pylint: disable=too-few-public-methods + + +class QuotaLimiterFactory: + """Quota limiter factory class.""" + + @staticmethod + def quota_limiters(config: QuotaHandlersConfiguration) -> list[QuotaLimiter]: + """Create instances of quota limiters based on loaded configuration. + + Returns: + List of instances of 'QuotaLimiter', + """ + limiters: list[QuotaLimiter] = [] + + limiters_config = config.limiters + if limiters_config is None: + logger.warning("Quota limiters are not specified in configuration") + return limiters + + return limiters