diff --git a/src/quota/cluster_quota_limiter.py b/src/quota/cluster_quota_limiter.py index 805791b9d..d2170f853 100644 --- a/src/quota/cluster_quota_limiter.py +++ b/src/quota/cluster_quota_limiter.py @@ -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__) diff --git a/src/quota/connect_pg.py b/src/quota/connect_pg.py index fbbf95109..e4ba2c9bf 100644 --- a/src/quota/connect_pg.py +++ b/src/quota/connect_pg.py @@ -1,6 +1,7 @@ """PostgreSQL connection handler.""" from typing import Any + import psycopg2 from log import get_logger diff --git a/src/quota/quota_limiter.py b/src/quota/quota_limiter.py index bc0d659f8..14543c590 100644 --- a/src/quota/quota_limiter.py +++ b/src/quota/quota_limiter.py @@ -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 diff --git a/src/quota/quota_limiter_factory.py b/src/quota/quota_limiter_factory.py index 9b8105c4d..3067f7493 100644 --- a/src/quota/quota_limiter_factory.py +++ b/src/quota/quota_limiter_factory.py @@ -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__) diff --git a/src/quota/revokable_quota_limiter.py b/src/quota/revokable_quota_limiter.py index 6a7819abe..9f5ff54cb 100644 --- a/src/quota/revokable_quota_limiter.py +++ b/src/quota/revokable_quota_limiter.py @@ -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, 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__) diff --git a/src/quota/token_usage_history.py b/src/quota/token_usage_history.py index c0fccc7a9..3ea14a452 100644 --- a/src/quota/token_usage_history.py +++ b/src/quota/token_usage_history.py @@ -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 diff --git a/src/quota/user_quota_limiter.py b/src/quota/user_quota_limiter.py index 13901fd05..eee9dd7c2 100644 --- a/src/quota/user_quota_limiter.py +++ b/src/quota/user_quota_limiter.py @@ -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__)