diff --git a/src/quota/quota_exceed_error.py b/src/quota/quota_exceed_error.py index 2317b347f..297c1ff3b 100644 --- a/src/quota/quota_exceed_error.py +++ b/src/quota/quota_exceed_error.py @@ -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: diff --git a/src/quota/quota_limiter.py b/src/quota/quota_limiter.py index e7f3fa154..f4f9d0064 100644 --- a/src/quota/quota_limiter.py +++ b/src/quota/quota_limiter.py @@ -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: @@ -69,11 +91,30 @@ 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 @@ -81,11 +122,28 @@ def __init__(self) -> 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) @@ -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 diff --git a/src/quota/quota_limiter_factory.py b/src/quota/quota_limiter_factory.py index 02365d58b..9b8105c4d 100644 --- a/src/quota/quota_limiter_factory.py +++ b/src/quota/quota_limiter_factory.py @@ -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] = [] @@ -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) diff --git a/src/quota/revokable_quota_limiter.py b/src/quota/revokable_quota_limiter.py index 49340e2cd..9bb47ed3c 100644 --- a/src/quota/revokable_quota_limiter.py +++ b/src/quota/revokable_quota_limiter.py @@ -33,7 +33,20 @@ def __init__( increase_by: int, subject_type: str, ) -> None: - """Initialize quota limiter.""" + """Initialize quota limiter. + + Create a revokable quota limiter configured for a specific subject type. + + Parameters: + configuration (QuotaHandlersConfiguration): Configuration object + containing `sqlite` and `postgres` connection settings. + initial_quota (int): The starting quota value assigned when a + subject's quota is initialized or revoked. + increase_by (int): Number of quota units to add when increasing a subject's quota. + subject_type (str): Identifier for the kind of subject the limiter + applies to (e.g., user, customer); when set to "c" the limiter + treats subject IDs as empty strings. + """ self.subject_type = subject_type self.initial_quota = initial_quota self.increase_by = increase_by @@ -42,7 +55,18 @@ def __init__( @connection def available_quota(self, subject_id: str = "") -> int: - """Retrieve available quota for given subject.""" + """Retrieve available quota for given subject. + + Get the available quota for a subject. + + Parameters: + subject_id (str): Subject identifier. For limiters with + subject_type "c", this value is ignored and treated as an empty + string. + + Returns: + int: The available quota for the subject. Returns 0 if no backend is configured. + """ if self.subject_type == "c": subject_id = "" if self.sqlite_connection_config is not None: @@ -53,7 +77,21 @@ def available_quota(self, subject_id: str = "") -> int: return 0 def _read_available_quota(self, query_statement: str, subject_id: str) -> int: - """Read available quota from selected database.""" + """Read available quota from selected database. + + Fetches the available quota for a subject from the database. + + If no quota record exists for the given subject, initializes the quota + and returns the limiter's initial quota. + + Parameters: + query_statement (str): SQL statement used to select the quota. + subject_id (str): Identifier of the subject whose quota is requested. + + Returns: + int: The available quota for the subject; `initial_quota` if a new + record was initialized. + """ # it is not possible to use context manager there, because SQLite does # not support it cursor = self.connection.cursor() @@ -70,7 +108,15 @@ def _read_available_quota(self, query_statement: str, subject_id: str) -> int: @connection def revoke_quota(self, subject_id: str = "") -> None: - """Revoke quota for given subject.""" + """Revoke quota for given subject. + + Revoke a subject's quota and record the revocation timestamp in the configured backend. + + Parameters: + subject_id (str): Identifier of the subject whose quota will be + revoked. If the limiter's `subject_type` is `"c"`, this value + is ignored and treated as an empty string. + """ if self.subject_type == "c": subject_id = "" @@ -82,7 +128,17 @@ def revoke_quota(self, subject_id: str = "") -> None: return def _revoke_quota(self, set_statement: str, subject_id: str) -> None: - """Revoke quota in given database.""" + """Revoke quota in given database. + + Set the subject's available quota back to the configured initial quota + and record the revocation timestamp. + + Parameters: + set_statement (str): SQL statement that updates the available quota + and `revoked_at` for a subject. + subject_id (str): Identifier of the subject whose quota will be + revoked. + """ # timestamp to be used revoked_at = datetime.now() @@ -96,7 +152,16 @@ def _revoke_quota(self, set_statement: str, subject_id: str) -> None: @connection def increase_quota(self, subject_id: str = "") -> None: - """Increase quota for given subject.""" + """Increase quota for given subject. + + Increase the available quota for a subject by the limiter's configured increment. + + Parameters: + subject_id (str): Identifier of the subject whose quota will be + increased. When the limiter's `subject_type` is `"c"`, this value + is normalized to the empty string and treated as a + global/customer-level entry. + """ if self.subject_type == "c": subject_id = "" @@ -109,7 +174,19 @@ def increase_quota(self, subject_id: str = "") -> None: return def _increase_quota(self, set_statement: str, subject_id: str) -> None: - """Increase quota in given database.""" + """Increase quota in given database. + + Increase the stored quota for a subject by the configured increment and + record the update timestamp. + + Executes the provided SQL statement with parameters (increase amount, + update timestamp, subject_id, subject_type) and commits the + transaction. + + Parameters: + set_statement (str): SQL statement that increments the available quota for a subject. + subject_id (str): Identifier of the subject whose quota will be increased. + """ # timestamp to be used updated_at = datetime.now() @@ -121,7 +198,19 @@ def _increase_quota(self, set_statement: str, subject_id: str) -> None: self.connection.commit() def ensure_available_quota(self, subject_id: str = "") -> None: - """Ensure that there's avaiable quota left.""" + """Ensure that there's available quota left. + + Ensure the subject has available quota; raises if quota is exhausted. + + Parameters: + subject_id (str): Identifier of the subject to check. If this + limiter's `subject_type` is `"c"`, the value is ignored and + treated as an empty string. + + Raises: + QuotaExceedError: If the available quota for the subject is + less than or equal to zero. + """ if self.subject_type == "c": subject_id = "" available = self.available_quota(subject_id) @@ -139,7 +228,19 @@ def consume_tokens( output_tokens: int = 0, subject_id: str = "", ) -> None: - """Consume tokens by given subject.""" + """ + Consume tokens from a subject's available quota. + + Deducts the sum of `input_tokens` and `output_tokens` from the + subject's stored quota and persists the update to the configured + database backend. For subject type "c", the `subject_id` is normalized + to an empty string before performing the operation. + + Parameters: + input_tokens (int): Number of input tokens to consume. + output_tokens (int): Number of output tokens to consume. + subject_id (str): Identifier of the subject whose quota will be consumed. + """ if self.subject_type == "c": subject_id = "" logger.info( @@ -168,7 +269,22 @@ def _consume_tokens( output_tokens: int, subject_id: str, ) -> None: - """Consume tokens from selected database.""" + """Consume tokens from selected database. + + Deduct the sum of input and output tokens from the subject's available + quota and persist the update. + + Parameters: + update_statement (str): SQL statement used to apply the quota change. + input_tokens (int): Number of input tokens to consume. + output_tokens (int): Number of output tokens to consume. + subject_id (str): Identifier of the subject whose quota will be updated. + + Notes: + The function updates the quota by -(input_tokens + output_tokens) + and stamps the record with the current datetime, then commits the + change. + """ # timestamp to be used updated_at = datetime.now() @@ -183,7 +299,12 @@ def _consume_tokens( cursor.close() def _initialize_tables(self) -> None: - """Initialize tables used by quota limiter.""" + """Initialize tables used by quota limiter. + + Create quota-related tables in the configured database and commit the change. + + This ensures the database schema required by the quota limiter exists. + """ logger.info("Initializing tables for quota limiter") cursor = self.connection.cursor() if self.sqlite_connection_config is not None: @@ -194,7 +315,19 @@ def _initialize_tables(self) -> None: self.connection.commit() def _init_quota(self, subject_id: str = "") -> None: - """Initialize quota for given ID.""" + """Initialize quota for given ID. + + Create a quota record for the given subject and set its initial values. + + Inserts a quota row for `subject_id` with both available and total + quota set to the limiter's configured initial value and stamps the + revocation timestamp. The operation writes to whichever backend(s) are + configured (SQLite and/or PostgreSQL) and commits the transaction. + + Parameters: + subject_id (str): Identifier of the subject whose quota to + initialize. Defaults to empty string. + """ # timestamp to be used revoked_at = datetime.now()