Skip to content

Commit d34e5cb

Browse files
committed
LCORE-1092: better docstring for quota scheduler
1 parent 2fbbe74 commit d34e5cb

1 file changed

Lines changed: 108 additions & 10 deletions

File tree

src/runners/quota_scheduler.py

Lines changed: 108 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,20 @@
2828

2929

3030
def quota_scheduler(config: QuotaHandlersConfiguration) -> bool:
31-
"""Quota scheduler task."""
31+
"""
32+
Run the quota scheduler loop that applies configured quota limiters periodically.
33+
34+
Parameters:
35+
config (QuotaHandlersConfiguration): Configuration containing storage
36+
settings (sqlite or postgres), a list of limiter configurations, and
37+
scheduler.period in seconds. If configuration is invalid or no
38+
storage/limiters are configured, the scheduler will not start.
39+
40+
Returns:
41+
bool: `True` if the scheduler started (unreachable in normal execution
42+
because the function enters an infinite loop), `False` if validation
43+
failed or a database connection could not be established.
44+
"""
3245
if config is None:
3346
logger.warning("Quota limiters are not configured, skipping")
3447
return False
@@ -74,14 +87,29 @@ def quota_scheduler(config: QuotaHandlersConfiguration) -> bool:
7487

7588

7689
def get_increase_quota_statement(config: QuotaHandlersConfiguration) -> str:
77-
"""Get the SQL statement to increase quota."""
90+
"""
91+
Select the SQL statement used to increase stored quota according to the database backend.
92+
93+
Parameters:
94+
config (QuotaHandlersConfiguration): Configuration that indicates which
95+
storage backend (SQLite or PostgreSQL) is in use.
96+
97+
Returns:
98+
str: SQL statement to perform a quota increase appropriate for the configured backend.
99+
"""
78100
if config.sqlite is not None:
79101
return INCREASE_QUOTA_STATEMENT_SQLITE
80102
return INCREASE_QUOTA_STATEMENT_PG
81103

82104

83105
def get_reset_quota_statement(config: QuotaHandlersConfiguration) -> str:
84-
"""Get the SQL statement to reset quota."""
106+
"""
107+
Return the SQL statement used to reset quota records for the configured database backend.
108+
109+
Returns:
110+
str: The SQLite reset SQL statement when `config.sqlite` is set,
111+
otherwise the PostgreSQL reset SQL statement.
112+
"""
85113
if config.sqlite is not None:
86114
return RESET_QUOTA_STATEMENT_SQLITE
87115
return RESET_QUOTA_STATEMENT_PG
@@ -93,7 +121,21 @@ def quota_revocation(
93121
increase_quota_statement: str,
94122
reset_quota_statement: str,
95123
) -> None:
96-
"""Quota revocation mechanism."""
124+
"""
125+
Apply configured quota updates for a quota limiter using the provided database connection.
126+
127+
Processes the given limiter: increases quota when `quota_increase` is set
128+
and resets initial quota when `initial_quota` is greater than zero, using
129+
the supplied SQL statements.
130+
131+
Parameters:
132+
quota_limiter (QuotaLimiterConfiguration): Limiter configuration to process.
133+
increase_quota_statement (str): SQL statement used to increment quota values.
134+
reset_quota_statement (str): SQL statement used to reset quota values.
135+
136+
Raises:
137+
ValueError: If the limiter's `type` or `period` is not set.
138+
"""
97139
logger.info(
98140
"Quota revocation mechanism for limiter '%s' of type '%s'",
99141
quota_limiter.name,
@@ -134,7 +176,22 @@ def increase_quota(
134176
increase_by: int,
135177
period: str,
136178
) -> None:
137-
"""Increase quota by specified amount."""
179+
"""
180+
Increase the stored quota for a subject by a specified amount for a given period.
181+
182+
Executes the provided SQL update statement on the given database connection
183+
to increment the quota value for the specified subject and period.
184+
185+
Parameters:
186+
connection (Any): Database connection object (Postgres or SQLite) to
187+
execute the statement on.
188+
update_statement (str): SQL update statement that accepts parameters
189+
(increase_by, subject_id, period).
190+
subject_id (str): Identifier for the subject whose quota is modified
191+
(e.g., "u" for user, "c" for cluster).
192+
increase_by (int): Amount to add to the subject's quota.
193+
period (str): Quota period identifier used to scope the update.
194+
"""
138195
logger.info(
139196
"Increasing quota for subject '%s' by %d when period %s is reached",
140197
subject_id,
@@ -165,7 +222,17 @@ def reset_quota(
165222
reset_to: int,
166223
period: str,
167224
) -> None:
168-
"""Reset quota to specified amount."""
225+
"""
226+
Set the stored quota for a subject to a specific value for the given period.
227+
228+
Parameters:
229+
connection (Any): Database connection object used to execute the update.
230+
update_statement (str): SQL statement that sets the quota value
231+
(expects parameters: new_value, subject_id, period).
232+
subject_id (str): Identifier for the quota subject (e.g., "u" for user, "c" for cluster).
233+
reset_to (int): Value to set the subject's quota to.
234+
period (str): Period identifier for which the quota is being set.
235+
"""
169236
logger.info(
170237
"Resetting quota for subject '%s' to %d when period %s is reached",
171238
subject_id,
@@ -190,7 +257,16 @@ def reset_quota(
190257

191258

192259
def get_subject_id(limiter_type: str) -> str:
193-
"""Get subject ID based on quota limiter type."""
260+
"""
261+
Map a quota limiter type to its subject identifier.
262+
263+
Parameters:
264+
limiter_type (str): Quota limiter type constant (e.g., user or cluster).
265+
266+
Returns:
267+
str: `"u"` for a user limiter, `"c"` for a cluster limiter, or `"?"` if the type
268+
is not recognized.
269+
"""
194270
match limiter_type:
195271
case constants.USER_QUOTA_LIMITER:
196272
return "u"
@@ -201,7 +277,17 @@ def get_subject_id(limiter_type: str) -> str:
201277

202278

203279
def connect(config: QuotaHandlersConfiguration) -> Any:
204-
"""Initialize connection to database."""
280+
"""
281+
Create and return a database connection for quota handlers based on the configured backend.
282+
283+
Parameters:
284+
config (QuotaHandlersConfiguration): Configuration containing
285+
`postgres` or `sqlite` connection settings.
286+
287+
Returns:
288+
A database connection suitable for quota operations, or `None` if
289+
neither Postgres nor SQLite is configured.
290+
"""
205291
logger.info("Initializing connection to quota limiter database")
206292
if config.postgres is not None:
207293
return connect_pg(config.postgres)
@@ -211,7 +297,13 @@ def connect(config: QuotaHandlersConfiguration) -> Any:
211297

212298

213299
def init_tables(connection: Any) -> None:
214-
"""Initialize tables used by quota limiter."""
300+
"""
301+
Create the quota table required by the quota limiter on the provided database connection.
302+
303+
Parameters:
304+
connection (Any): A DB-API compatible connection on which the quota
305+
table(s) will be created; changes are committed before returning.
306+
"""
215307
logger.info("Initializing tables for quota limiter")
216308
cursor = connection.cursor()
217309
cursor.execute(CREATE_QUOTA_TABLE)
@@ -220,7 +312,13 @@ def init_tables(connection: Any) -> None:
220312

221313

222314
def start_quota_scheduler(configuration: Configuration) -> None:
223-
"""Start user and cluster quota scheduler in separate thread."""
315+
"""
316+
Start the quota scheduler in a daemon thread using the provided configuration.
317+
318+
Parameters:
319+
configuration (Configuration): Global configuration whose `quota_handlers`
320+
attribute is passed to the scheduler thread.
321+
"""
224322
logger.info("Starting quota scheduler")
225323
thread = Thread(
226324
target=quota_scheduler,

0 commit comments

Comments
 (0)