diff --git a/src/quota/__init__.py b/src/quota/__init__.py new file mode 100644 index 00000000..b96d7cc6 --- /dev/null +++ b/src/quota/__init__.py @@ -0,0 +1 @@ +"""Quota management.""" diff --git a/src/quota/sql.py b/src/quota/sql.py new file mode 100644 index 00000000..3b8b64be --- /dev/null +++ b/src/quota/sql.py @@ -0,0 +1,45 @@ +"""SQL commands used by quota management package.""" + +CREATE_QUOTA_TABLE = """ + CREATE TABLE IF NOT EXISTS quota_limits ( + id text NOT NULL, + subject char(1) NOT NULL, + quota_limit int NOT NULL, + available int, + updated_at timestamp with time zone, + revoked_at timestamp with time zone, + PRIMARY KEY(id, subject) + ); + """ + + +INCREASE_QUOTA_STATEMENT_PG = """ + UPDATE quota_limits + SET available=available+%s, revoked_at=NOW() + WHERE subject=%s + AND revoked_at < NOW() - INTERVAL %s ; + """ + + +INCREASE_QUOTA_STATEMENT_SQLITE = """ + UPDATE quota_limits + SET available=available+?, revoked_at=datetime('now') + WHERE subject=? + AND revoked_at < datetime('now', ?); + """ + + +RESET_QUOTA_STATEMENT_PG = """ + UPDATE quota_limits + SET available=%s, revoked_at=NOW() + WHERE subject=%s + AND revoked_at < NOW() - INTERVAL %s ; + """ + + +RESET_QUOTA_STATEMENT_SQLITE = """ + UPDATE quota_limits + SET available=?, revoked_at=datetime('now') + WHERE subject=? + AND revoked_at < datetime('now', ?); + """ diff --git a/src/runners/quota_scheduler.py b/src/runners/quota_scheduler.py index ce42dcb6..4abe257e 100644 --- a/src/runners/quota_scheduler.py +++ b/src/runners/quota_scheduler.py @@ -17,52 +17,15 @@ SQLiteDatabaseConfiguration, ) -logger = get_logger(__name__) - +from quota.sql import ( + CREATE_QUOTA_TABLE, + INCREASE_QUOTA_STATEMENT_PG, + INCREASE_QUOTA_STATEMENT_SQLITE, + RESET_QUOTA_STATEMENT_PG, + RESET_QUOTA_STATEMENT_SQLITE, +) -CREATE_QUOTA_TABLE = """ - CREATE TABLE IF NOT EXISTS quota_limits ( - id text NOT NULL, - subject char(1) NOT NULL, - quota_limit int NOT NULL, - available int, - updated_at timestamp with time zone, - revoked_at timestamp with time zone, - PRIMARY KEY(id, subject) - ); - """ - - -INCREASE_QUOTA_STATEMENT_PG = """ - UPDATE quota_limits - SET available=available+%s, revoked_at=NOW() - WHERE subject=%s - AND revoked_at < NOW() - INTERVAL %s ; - """ - - -INCREASE_QUOTA_STATEMENT_SQLITE = """ - UPDATE quota_limits - SET available=available+?, revoked_at=datetime('now') - WHERE subject=? - AND revoked_at < datetime('now', ?); - """ - - -RESET_QUOTA_STATEMENT_PG = """ - UPDATE quota_limits - SET available=%s, revoked_at=NOW() - WHERE subject=%s - AND revoked_at < NOW() - INTERVAL %s ; - """ - - -RESET_QUOTA_STATEMENT_SQLITE = """ - UPDATE quota_limits - SET available=?, revoked_at=datetime('now') - WHERE subject=? - AND revoked_at < datetime('now', ?); - """ +logger = get_logger(__name__) def quota_scheduler(config: QuotaHandlersConfiguration) -> bool: