Skip to content

LCORE-1438: Organize imports in quota handler #1323

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tisnik merged 1 commit into from
Mar 15, 2026
Merged

LCORE-1438: Organize imports in quota handler #1323

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
2 changes: 1 addition & 1 deletion src/quota/cluster_quota_limiter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
"""Simple cluster quota limiter where quota is fixed for the whole cluster."""

from models.config import QuotaHandlersConfiguration
from log import get_logger
from models.config import QuotaHandlersConfiguration
from quota.revokable_quota_limiter import RevokableQuotaLimiter

logger = get_logger(__name__)
Expand Down
1 change: 1 addition & 0 deletions src/quota/connect_pg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
"""PostgreSQL connection handler."""

from typing import Any

import psycopg2

from log import get_logger
Expand Down
7 changes: 3 additions & 4 deletions src/quota/quota_limiter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,16 +30,15 @@
- reset quota to 10,000,000 tokens each month
"""

import datetime
import sqlite3
from abc import ABC, abstractmethod

from typing import Optional

import datetime
import sqlite3
import psycopg2

from log import get_logger
from models.config import SQLiteDatabaseConfiguration, PostgreSQLDatabaseConfiguration
from models.config import PostgreSQLDatabaseConfiguration, SQLiteDatabaseConfiguration
from quota.connect_pg import connect_pg
from quota.connect_sqlite import connect_sqlite

Expand Down
5 changes: 2 additions & 3 deletions src/quota/quota_limiter_factory.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,11 @@
"""Quota limiter factory class."""

from log import get_logger
import constants
from log import get_logger
from models.config import QuotaHandlersConfiguration

from quota.user_quota_limiter import UserQuotaLimiter
from quota.cluster_quota_limiter import ClusterQuotaLimiter
from quota.quota_limiter import QuotaLimiter
from quota.user_quota_limiter import UserQuotaLimiter

logger = get_logger(__name__)

Expand Down
14 changes: 7 additions & 7 deletions src/quota/revokable_quota_limiter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,24 +1,24 @@
"""Simple quota limiter where quota can be revoked."""

from datetime import datetime, UTC
from datetime import UTC, datetime

from models.config import QuotaHandlersConfiguration
from log import get_logger
from utils.connection_decorator import connection
from models.config import QuotaHandlersConfiguration
from quota.quota_exceed_error import QuotaExceedError
from quota.quota_limiter import QuotaLimiter
from quota.sql import (
CREATE_QUOTA_TABLE_PG,
CREATE_QUOTA_TABLE_SQLITE,
UPDATE_AVAILABLE_QUOTA_PG,
UPDATE_AVAILABLE_QUOTA_SQLITE,
INIT_QUOTA_PG,
INIT_QUOTA_SQLITE,
Comment on lines +12 to +13

@coderabbitai coderabbitai Bot Mar 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

Make quota initialization idempotent.

available_quota() follows a read-then-initialize flow. With INIT_QUOTA_PG / INIT_QUOTA_SQLITE defined as plain INSERTs in src/quota/sql.py:60-68, two concurrent first requests for the same subject can both see "missing" and one will fail on the duplicate row. Please make the init SQL conflict-tolerant, or catch duplicate-key violations before returning initial_quota.

Possible fix 
-INIT_QUOTA_PG = """
-    INSERT INTO quota_limits (id, subject, quota_limit, available, revoked_at)
-    VALUES (%s, %s, %s, %s, %s)
-    """
+INIT_QUOTA_PG = """
+    INSERT INTO quota_limits (id, subject, quota_limit, available, revoked_at)
+    VALUES (%s, %s, %s, %s, %s)
+    ON CONFLICT DO NOTHING
+    """

-INIT_QUOTA_SQLITE = """
-    INSERT INTO quota_limits (id, subject, quota_limit, available, revoked_at)
-    VALUES (?, ?, ?, ?, ?)
-    """
+INIT_QUOTA_SQLITE = """
+    INSERT OR IGNORE INTO quota_limits (id, subject, quota_limit, available, revoked_at)
+    VALUES (?, ?, ?, ?, ?)
+    """
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@src/quota/revokable_quota_limiter.py` around lines 12 - 13, available_quota()
races on initial row creation because INIT_QUOTA_PG and INIT_QUOTA_SQLITE are
plain INSERTs; make initialization idempotent by changing INIT_QUOTA_PG to use a
conflict-tolerant INSERT (e.g. INSERT ... ON CONFLICT DO NOTHING) and
INIT_QUOTA_SQLITE to use SQLite's conflict form (e.g. INSERT OR IGNORE), or
alternatively wrap the INSERT in available_quota() with a duplicate-key handler
that swallows the unique-constraint error and proceeds to return initial_quota;
update the SQL constants INIT_QUOTA_PG and INIT_QUOTA_SQLITE (or add the
duplicate-key catch in available_quota()) so concurrent first requests cannot
cause a duplicate-row failure.

SELECT_QUOTA_PG,
SELECT_QUOTA_SQLITE,
SET_AVAILABLE_QUOTA_PG,
SET_AVAILABLE_QUOTA_SQLITE,
INIT_QUOTA_PG,
INIT_QUOTA_SQLITE,
UPDATE_AVAILABLE_QUOTA_PG,
UPDATE_AVAILABLE_QUOTA_SQLITE,
)
from utils.connection_decorator import connection

logger = get_logger(__name__)

Expand Down
16 changes: 7 additions & 9 deletions src/quota/token_usage_history.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,25 +6,23 @@
"""

import sqlite3
from datetime import datetime, UTC
from datetime import UTC, datetime
from typing import Any, Optional

import psycopg2

from log import get_logger

from models.config import (
PostgreSQLDatabaseConfiguration,
QuotaHandlersConfiguration,
SQLiteDatabaseConfiguration,
)
from quota.connect_pg import connect_pg
from quota.connect_sqlite import connect_sqlite
from quota.sql import (
CREATE_TOKEN_USAGE_TABLE,
CONSUME_TOKENS_FOR_USER_PG,
CONSUME_TOKENS_FOR_USER_SQLITE,
)

from models.config import (
QuotaHandlersConfiguration,
SQLiteDatabaseConfiguration,
PostgreSQLDatabaseConfiguration,
CREATE_TOKEN_USAGE_TABLE,
)
from utils.connection_decorator import connection

Expand Down
2 changes: 1 addition & 1 deletion src/quota/user_quota_limiter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
"""Simple user quota limiter where each user has a fixed quota."""

from models.config import QuotaHandlersConfiguration
from log import get_logger
from models.config import QuotaHandlersConfiguration
from quota.revokable_quota_limiter import RevokableQuotaLimiter

logger = get_logger(__name__)
Expand Down
Loading