LCORE-739: quota limiter stub code

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

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

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