Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 16 additions & 1 deletion src/quota/quota_exceed_error.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,22 @@ class QuotaExceedError(Exception):
def __init__(
self, subject_id: str, subject_type: str, available: int, needed: int = 0
) -> None:
"""Construct exception object."""
"""Construct exception object.

Initialize the QuotaExceedError with the subject identity and token counts.

Parameters:
subject_id (str): Identifier of the subject (user id or cluster id).
subject_type (str): Subject kind: "u" for user, "c" for cluster,
any other value treated as unknown.
available (int): Number of tokens currently available to the subject.
needed (int): Number of tokens required; defaults to 0.

Attributes:
subject_id (str): Copied from the `subject_id` parameter.
available (int): Copied from the `available` parameter.
needed (int): Copied from the `needed` parameter.
"""
message: str = ""

if needed == 0 and available <= 0:
Expand Down
80 changes: 72 additions & 8 deletions src/quota/quota_limiter.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,15 +51,37 @@ class QuotaLimiter(ABC):

@abstractmethod
def available_quota(self, subject_id: str) -> int:
"""Retrieve available quota for given user."""
"""Retrieve available quota for given user.

Get the remaining quota for the specified subject.

Parameters:
subject_id (str): Identifier of the subject (user or service) whose quota to retrieve.

Returns:
available_quota (int): Number of quota units currently available for the subject.
"""

@abstractmethod
def revoke_quota(self) -> None:
"""Revoke quota for given user."""
"""Revoke quota for given user.

Revoke the quota for the limiter's target subject by setting its available quota to zero.

This operation removes or disables any remaining allowance so
subsequent checks will report no available quota.
"""

@abstractmethod
def increase_quota(self) -> None:
"""Increase quota for given user."""
"""Increase quota for given user.

Increase the available quota for the limiter's subject according to its
configured increase policy.

Updates persistent storage to add the configured quota increment to the
subject's stored available quota.
"""

@abstractmethod
def ensure_available_quota(self, subject_id: str = "") -> None:
Expand All @@ -69,23 +91,59 @@ def ensure_available_quota(self, subject_id: str = "") -> None:
def consume_tokens(
self, input_tokens: int, output_tokens: int, subject_id: str = ""
) -> None:
"""Consume tokens by given user."""
"""Consume tokens by given user.

Consume the specified input and output tokens from a subject's available quota.

Parameters:
input_tokens (int): Number of input tokens to deduct from the subject's quota.
output_tokens (int): Number of output tokens to deduct from the subject's quota.
subject_id (str): Identifier of the subject (user or service) whose
quota will be reduced. If omitted, applies to the default subject.
"""

@abstractmethod
def __init__(self) -> None:
"""Initialize connection configuration(s)."""
"""Initialize connection configuration(s).

Create a QuotaLimiter instance and initialize database connection configuration attributes.

Attributes:
sqlite_connection_config (Optional[SQLiteDatabaseConfiguration]):
SQLite connection configuration or `None` when not configured.
postgres_connection_config
(Optional[PostgreSQLDatabaseConfiguration]): PostgreSQL connection
configuration or `None` when not configured.
"""
self.sqlite_connection_config: Optional[SQLiteDatabaseConfiguration] = None
self.postgres_connection_config: Optional[PostgreSQLDatabaseConfiguration] = (
None
)

@abstractmethod
def _initialize_tables(self) -> None:
"""Initialize tables and indexes."""
"""Initialize tables and indexes.

Create any database tables and indexes required by the quota limiter implementation.

Implementations must ensure the database schema and indexes needed for
storing and querying quota state exist; calling this method when the
schema already exists should be safe (idempotent). Raise an exception
on irrecoverable initialization failures.
"""

# pylint: disable=W0201
def connect(self) -> None:
"""Initialize connection to database."""
"""Initialize connection to database.

Establish the configured database connection, initialize required
tables, and enable autocommit.

If a PostgreSQL or SQLite configuration is present, a connection to
that backend will be created, then _initialize_tables() will be called
to prepare storage. If table initialization fails, the connection is
closed and the original exception is propagated.
"""
logger.info("Initializing connection to quota limiter database")
if self.postgres_connection_config is not None:
self.connection = connect_pg(self.postgres_connection_config)
Expand All @@ -102,7 +160,13 @@ def connect(self) -> None:
self.connection.autocommit = True

def connected(self) -> bool:
"""Check if connection to quota limiter database is alive."""
"""Check if connection to quota limiter database is alive.

Determine whether the storage connection is alive.

Returns:
`true` if the connection is alive, `false` otherwise.
"""
if self.connection is None:
logger.warning("Not connected, need to reconnect later")
return False
Expand Down
27 changes: 25 additions & 2 deletions src/quota/quota_limiter_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,14 @@ class QuotaLimiterFactory:
def quota_limiters(config: QuotaHandlersConfiguration) -> list[QuotaLimiter]:
"""Create instances of quota limiters based on loaded configuration.

Parameters:
config (QuotaHandlersConfiguration): Configuration containing
storage settings and limiter definitions.

Returns:
List of instances of 'QuotaLimiter',
list[QuotaLimiter]: List of initialized quota limiter instances.
Returns an empty list if storage configuration or limiter
definitions are missing.
"""
limiters: list[QuotaLimiter] = []

Expand Down Expand Up @@ -56,7 +62,24 @@ def create_limiter(
initial_quota: int,
increase_by: int,
) -> QuotaLimiter:
"""Create selected quota limiter."""
"""Create selected quota limiter.

Instantiate a quota limiter instance for the given limiter type.

Parameters:
configuration (QuotaHandlersConfiguration): Configuration used to
initialize the limiter.
limiter_type (str): Identifier of the limiter to create; expected values are
`constants.USER_QUOTA_LIMITER` or `constants.CLUSTER_QUOTA_LIMITER`.
initial_quota (int): Starting quota value assigned to the limiter.
increase_by (int): Amount by which the quota increases when replenished.

Returns:
QuotaLimiter: A configured quota limiter instance of the requested type.

Raises:
ValueError: If `limiter_type` is not a recognized limiter identifier.
"""
match limiter_type:
case constants.USER_QUOTA_LIMITER:
return UserQuotaLimiter(configuration, initial_quota, increase_by)
Expand Down
Loading
Loading