From 2db72d6e036e9a6f10a8c4ac3db8fa04c7992b2d Mon Sep 17 00:00:00 2001 From: Biruk Haileye Tabor Date: Tue, 28 Jul 2026 23:55:12 +0300 Subject: [PATCH 1/4] fix: engine too many connections Fix create_engine() connection leaks per request/call --- tests/browser/db_helpers.py | 10 ++--- tests/conftest.py | 10 +++-- utils/core/auth.py | 5 +-- utils/core/db.py | 81 +++++++++++++++++++++++++------------ utils/core/dependencies.py | 8 ++-- 5 files changed, 71 insertions(+), 43 deletions(-) diff --git a/tests/browser/db_helpers.py b/tests/browser/db_helpers.py index 2535cf8..7b3373e 100644 --- a/tests/browser/db_helpers.py +++ b/tests/browser/db_helpers.py @@ -8,10 +8,10 @@ from datetime import timedelta from dotenv import load_dotenv -from sqlmodel import Session, create_engine, select +from sqlmodel import Session, select from utils.core.auth import get_password_hash -from utils.core.db import create_default_roles, get_connection_url +from utils.core.db import create_default_roles, get_engine from utils.core.models import ( Account, Invitation, @@ -39,8 +39,7 @@ def browser_db_session(): saved = {key: os.environ.get(key) for key in env} os.environ.update(env) try: - engine = create_engine(get_connection_url()) - with Session(engine) as session: + with Session(get_engine()) as session: yield session finally: for key, value in saved.items(): @@ -168,8 +167,7 @@ def browser_csrf_db_session(): saved = {key: os.environ.get(key) for key in env} os.environ.update(env) try: - engine = create_engine(get_connection_url()) - with Session(engine) as session: + with Session(get_engine()) as session: yield session finally: for key, value in saved.items(): diff --git a/tests/conftest.py b/tests/conftest.py index f140a71..35b94f8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,11 +3,12 @@ from typing import Generator, cast pytest_plugins = ["tests.frontend.fixtures"] -from sqlmodel import create_engine, Session, select +from sqlmodel import Session, select from fastapi.testclient import TestClient from dotenv import load_dotenv from utils.core.db import ( get_connection_url, + get_engine, tear_down_db, set_up_db, create_default_roles, @@ -74,9 +75,12 @@ def engine(env_vars): Create a new SQLModel engine for the test database. Use PostgreSQL for testing to match production environment. """ - # Use PostgreSQL for testing to match production environment + # Use PostgreSQL for testing to match production environment. get_engine() + # returns the same cached pool the app itself uses via get_session(), so + # TestClient requests and direct test queries share one pool instead of + # opening a second one for the same database. ensure_database_exists(get_connection_url()) - engine = create_engine(get_connection_url()) + engine = get_engine() set_up_db(drop=True) yield engine diff --git a/utils/core/auth.py b/utils/core/auth.py index f631d81..5164993 100644 --- a/utils/core/auth.py +++ b/utils/core/auth.py @@ -13,7 +13,7 @@ from fastapi.templating import Jinja2Templates from fastapi import Cookie from starlette.responses import Response -from utils.core.db import create_engine, get_connection_url +from utils.core.db import get_engine from utils.core.models import ( AccountRecoveryToken, EmailVerificationToken, @@ -339,8 +339,7 @@ def send_reset_email_task(email: str) -> None: FastAPI background tasks should not reuse request-scoped resources from `yield` dependencies, because cleanup may run before the task executes. """ - engine = create_engine(get_connection_url()) - with Session(engine) as session: + with Session(get_engine()) as session: send_reset_email(email, session) diff --git a/utils/core/db.py b/utils/core/db.py index fac531e..f024a09 100644 --- a/utils/core/db.py +++ b/utils/core/db.py @@ -2,7 +2,7 @@ import logging from itertools import chain from typing import Union, Sequence -from sqlalchemy.engine import URL +from sqlalchemy.engine import Engine, URL from sqlmodel import create_engine, Session, SQLModel, select, text from utils.core.models import ( Account, @@ -28,17 +28,42 @@ # --- Database connection functions --- +_engine_cache: dict[str, Engine] = {} + + +def get_engine() -> Engine: + """Returns a process-wide cached engine for the current connection URL. + + Calling create_engine() per request/call leaks pooled connections until + Python's GC reclaims them (engines aren't freed by simple refcounting), + which can exhaust Postgres's max_connections under sustained load. Cache + by URL rather than as a single global so tests that swap DB_NAME within + the same process (e.g. browser vs. browser-csrf DBs) still get one + engine per database instead of sharing a mismatched pool. + """ + url = get_connection_url() + key = str(url) + engine = _engine_cache.get(key) + if engine is None: + engine = create_engine(url) + _engine_cache[key] = engine + return engine + + def ensure_database_exists(url: URL) -> None: dbname = url.database server_url = url.set(database="postgres") engine = create_engine(server_url, isolation_level="AUTOCOMMIT") - with engine.connect() as conn: - exists = conn.execute( - text("SELECT 1 FROM pg_database WHERE datname = :n"), - {"n": dbname}, - ).scalar() - if not exists: - conn.execute(text(f'CREATE DATABASE "{dbname}"')) + try: + with engine.connect() as conn: + exists = conn.execute( + text("SELECT 1 FROM pg_database WHERE datname = :n"), + {"n": dbname}, + ).scalar() + if not exists: + conn.execute(text(f'CREATE DATABASE "{dbname}"')) + finally: + engine.dispose() def get_connection_url() -> URL: @@ -259,19 +284,21 @@ def set_up_db(drop: bool = False) -> None: drop (bool): If True, drops all existing tables before creating new ones. """ engine = create_engine(get_connection_url()) - if drop: - SQLModel.metadata.drop_all(engine) - # Ensure the private schema exists before creating tables - with engine.connect() as conn: - conn.execute(text("CREATE SCHEMA IF NOT EXISTS private")) - conn.commit() - SQLModel.metadata.create_all(engine) - # Create default permissions and seed account emails - with Session(engine) as session: - create_permissions(session) - session.commit() - seed_account_emails(session) - engine.dispose() + try: + if drop: + SQLModel.metadata.drop_all(engine) + # Ensure the private schema exists before creating tables + with engine.connect() as conn: + conn.execute(text("CREATE SCHEMA IF NOT EXISTS private")) + conn.commit() + SQLModel.metadata.create_all(engine) + # Create default permissions and seed account emails + with Session(engine) as session: + create_permissions(session) + session.commit() + seed_account_emails(session) + finally: + engine.dispose() def tear_down_db() -> None: @@ -279,8 +306,10 @@ def tear_down_db() -> None: Tears down the database by dropping all tables and the private schema. """ engine = create_engine(get_connection_url()) - SQLModel.metadata.drop_all(engine) - with engine.connect() as conn: - conn.execute(text("DROP SCHEMA IF EXISTS private CASCADE")) - conn.commit() - engine.dispose() + try: + SQLModel.metadata.drop_all(engine) + with engine.connect() as conn: + conn.execute(text("DROP SCHEMA IF EXISTS private CASCADE")) + conn.commit() + finally: + engine.dispose() diff --git a/utils/core/dependencies.py b/utils/core/dependencies.py index 507a05d..7d86ce2 100644 --- a/utils/core/dependencies.py +++ b/utils/core/dependencies.py @@ -13,7 +13,7 @@ oauth2_scheme_cookie, verify_password, ) -from utils.core.db import create_engine, get_connection_url +from utils.core.db import get_engine from utils.core.models import ( User, Role, @@ -43,8 +43,7 @@ def get_session() -> Generator[Session, None, None]: Yields: Session: A SQLModel session object for database operations. """ - engine = create_engine(get_connection_url()) - with Session(engine) as session: + with Session(get_engine()) as session: yield session @@ -434,8 +433,7 @@ async def get_user_from_request(request: Request) -> Optional[User]: tokens = (access_token, refresh_token) # Get a database session - engine = create_engine(get_connection_url()) - with Session(engine) as session: + with Session(get_engine()) as session: user, new_access_token, new_refresh_token = get_user_from_tokens( tokens, session ) From bddc3d2b8b40214b8604a8a0799d267b0dcef0b8 Mon Sep 17 00:00:00 2001 From: Biruk Haileye Tabor Date: Wed, 29 Jul 2026 10:26:37 +0300 Subject: [PATCH 2/4] refactor: use lru cached engine Use lru cached engine --- utils/core/db.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/utils/core/db.py b/utils/core/db.py index f024a09..22299ba 100644 --- a/utils/core/db.py +++ b/utils/core/db.py @@ -1,5 +1,6 @@ import os import logging +from functools import lru_cache from itertools import chain from typing import Union, Sequence from sqlalchemy.engine import Engine, URL @@ -28,7 +29,13 @@ # --- Database connection functions --- -_engine_cache: dict[str, Engine] = {} +@lru_cache +def _cached_engine(url: URL) -> Engine: + # URL objects hash/compare on their real fields (including password), so + # caching on the URL itself is safe. str(URL) masks the password as + # "***" for safe logging — passing that to create_engine() instead would + # make every connection attempt authenticate with the literal "***". + return create_engine(url) def get_engine() -> Engine: @@ -41,13 +48,7 @@ def get_engine() -> Engine: the same process (e.g. browser vs. browser-csrf DBs) still get one engine per database instead of sharing a mismatched pool. """ - url = get_connection_url() - key = str(url) - engine = _engine_cache.get(key) - if engine is None: - engine = create_engine(url) - _engine_cache[key] = engine - return engine + return _cached_engine(get_connection_url()) def ensure_database_exists(url: URL) -> None: From fd5864621b0e5366566b9c002d75d81cdb7f47c5 Mon Sep 17 00:00:00 2001 From: Rafiuzzaman Khan Date: Wed, 29 Jul 2026 13:33:01 -0400 Subject: [PATCH 3/4] add pool_pre_ping and dispose() --- .env.example | 4 +++ main.py | 4 +-- tests/utils/test_db.py | 79 ++++++++++++++++++++++++++++++++++++++++++ utils/core/db.py | 44 ++++++++++++++--------- 4 files changed, 112 insertions(+), 19 deletions(-) diff --git a/.env.example b/.env.example index f58040a..204d6d6 100644 --- a/.env.example +++ b/.env.example @@ -7,6 +7,10 @@ BASE_URL=http://localhost:8000 # Database connection mode (0 = direct, 1 = pooled) USE_POOL=0 +# SQLAlchemy pool (per process; keep small when using PgBouncer) +# DB_POOL_SIZE=5 +# DB_MAX_OVERFLOW=5 + # Shared database settings DB_HOST=127.0.0.1 DB_SSLMODE=prefer diff --git a/main.py b/main.py index d2f7935..c89d72f 100644 --- a/main.py +++ b/main.py @@ -46,7 +46,7 @@ RateLimitError, ) from exceptions.exceptions import NeedsNewTokens -from utils.core.db import set_up_db +from utils.core.db import set_up_db, clear_engine_cache logger = logging.getLogger("uvicorn.error") logger.setLevel(logging.DEBUG) @@ -58,7 +58,7 @@ async def lifespan(app: FastAPI): load_dotenv() set_up_db() yield - # Optional shutdown logic + clear_engine_cache() # Initialize the FastAPI app diff --git a/tests/utils/test_db.py b/tests/utils/test_db.py index 4f01274..b2c281b 100644 --- a/tests/utils/test_db.py +++ b/tests/utils/test_db.py @@ -3,6 +3,8 @@ from sqlalchemy import Engine from utils.core.db import ( get_connection_url, + get_engine, + clear_engine_cache, assign_permissions_to_role, create_default_roles, create_permissions, @@ -146,6 +148,83 @@ def test_get_connection_url_missing_pool_vars(monkeypatch): get_connection_url() +# --- Engine cache --- + + +def _direct_db_env(monkeypatch, *, name: str = "testdb", password: str = "testpass"): + for var in ( + "USE_POOL", + "DB_HOST", + "DB_PORT", + "DB_NAME", + "DB_USER", + "DB_PASSWORD", + "DB_POOL_PORT", + "DB_POOL_NAME", + "DB_APPUSER", + "DB_APPUSER_PASSWORD", + "DB_POOL_SIZE", + "DB_MAX_OVERFLOW", + ): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("DB_HOST", "localhost") + monkeypatch.setenv("DB_PORT", "5432") + monkeypatch.setenv("DB_NAME", name) + monkeypatch.setenv("DB_USER", "testuser") + monkeypatch.setenv("DB_PASSWORD", password) + + +def test_get_engine_reuses_same_instance(monkeypatch): + clear_engine_cache() + _direct_db_env(monkeypatch) + assert get_engine() is get_engine() + clear_engine_cache() + + +def test_get_engine_different_urls_get_different_engines(monkeypatch): + clear_engine_cache() + _direct_db_env(monkeypatch, name="db_a") + engine_a = get_engine() + _direct_db_env(monkeypatch, name="db_b") + engine_b = get_engine() + assert engine_a is not engine_b + clear_engine_cache() + + +def test_get_engine_not_keyed_by_masked_str_password(monkeypatch): + """str(URL) masks passwords; cache must still separate credentials.""" + clear_engine_cache() + _direct_db_env(monkeypatch, password="secretA") + engine_a = get_engine() + assert "***" in str(get_connection_url()) + _direct_db_env(monkeypatch, password="secretB") + engine_b = get_engine() + assert engine_a is not engine_b + clear_engine_cache() + + +def test_clear_engine_cache_disposes_and_creates_new(monkeypatch): + clear_engine_cache() + _direct_db_env(monkeypatch) + first = get_engine() + clear_engine_cache() + second = get_engine() + assert first is not second + clear_engine_cache() + + +def test_get_engine_applies_pool_settings(monkeypatch): + clear_engine_cache() + _direct_db_env(monkeypatch) + monkeypatch.setenv("DB_POOL_SIZE", "3") + monkeypatch.setenv("DB_MAX_OVERFLOW", "2") + engine = get_engine() + assert engine.pool.size() == 3 + assert engine.pool._max_overflow == 2 + assert engine.pool._pre_ping is True + clear_engine_cache() + + # --- Permission and Role Tests --- diff --git a/utils/core/db.py b/utils/core/db.py index 22299ba..51fc56b 100644 --- a/utils/core/db.py +++ b/utils/core/db.py @@ -1,6 +1,5 @@ import os import logging -from functools import lru_cache from itertools import chain from typing import Union, Sequence from sqlalchemy.engine import Engine, URL @@ -29,26 +28,36 @@ # --- Database connection functions --- -@lru_cache -def _cached_engine(url: URL) -> Engine: - # URL objects hash/compare on their real fields (including password), so - # caching on the URL itself is safe. str(URL) masks the password as - # "***" for safe logging — passing that to create_engine() instead would - # make every connection attempt authenticate with the literal "***". - return create_engine(url) +# Cache by URL object (not str(url)): str(URL) masks passwords as "***", so +# different credentials would collide on one cache key. +_engines: dict[URL, Engine] = {} def get_engine() -> Engine: - """Returns a process-wide cached engine for the current connection URL. - - Calling create_engine() per request/call leaks pooled connections until - Python's GC reclaims them (engines aren't freed by simple refcounting), - which can exhaust Postgres's max_connections under sustained load. Cache - by URL rather than as a single global so tests that swap DB_NAME within - the same process (e.g. browser vs. browser-csrf DBs) still get one - engine per database instead of sharing a mismatched pool. + """Return a process-wide cached engine for the current connection URL. + + Calling create_engine() per request leaks pooled connections until GC + reclaims them. Cache one engine per URL so tests that swap DB_NAME still + get a separate pool per database. """ - return _cached_engine(get_connection_url()) + url = get_connection_url() + engine = _engines.get(url) + if engine is None: + engine = create_engine( + url, + pool_pre_ping=True, + pool_size=int(os.getenv("DB_POOL_SIZE", "5")), + max_overflow=int(os.getenv("DB_MAX_OVERFLOW", "5")), + ) + _engines[url] = engine + return engine + + +def clear_engine_cache() -> None: + """Dispose and drop all cached engines (shutdown / tests).""" + for engine in _engines.values(): + engine.dispose() + _engines.clear() def ensure_database_exists(url: URL) -> None: @@ -314,3 +323,4 @@ def tear_down_db() -> None: conn.commit() finally: engine.dispose() + clear_engine_cache() From 606321b388c22d861c5c6206cb7602a7c20e40fa Mon Sep 17 00:00:00 2001 From: chriscarrollsmith Date: Wed, 29 Jul 2026 14:27:19 -0400 Subject: [PATCH 4/4] fix: address review feedback on engine cache - Route rate_limit.py through the shared get_engine() cache instead of its own module-level singleton, so no second unmanaged pool survives clear_engine_cache() or goes stale when the connection URL changes. - Guard the engine cache with a threading.Lock: sync dependencies run in FastAPI's threadpool, so concurrent first requests could race in get_engine() and orphan an undisposed engine, and clear_engine_cache() could hit a dict-mutated-during-iteration error. - Remove the clear_engine_cache() side effect from tear_down_db(); test fixtures that relied on it now call it explicitly in their teardown. - Make engine-cache tests isolation-safe via an engine_cache fixture so cleanup runs even when an assertion fails mid-test. Co-Authored-By: Claude Fable 5 --- tests/browser/conftest.py | 3 +++ tests/conftest.py | 2 ++ tests/utils/test_db.py | 31 ++++++++++++++++--------------- utils/core/db.py | 36 ++++++++++++++++++++++-------------- utils/core/rate_limit.py | 25 ++++++++----------------- 5 files changed, 51 insertions(+), 46 deletions(-) diff --git a/tests/browser/conftest.py b/tests/browser/conftest.py index de5471c..f88cd56 100644 --- a/tests/browser/conftest.py +++ b/tests/browser/conftest.py @@ -8,6 +8,7 @@ from playwright.sync_api import Browser, Page from tests.browser.db_helpers import browser_db_env from utils.core.db import ( + clear_engine_cache, ensure_database_exists, get_connection_url, set_up_db, @@ -149,6 +150,7 @@ def live_server(browser_env): proc.wait(timeout=5) with _temporary_env(browser_env): tear_down_db() + clear_engine_cache() @pytest.fixture(scope="session") @@ -162,3 +164,4 @@ def live_server_csrf(browser_csrf_env): proc.wait(timeout=5) with _temporary_env(browser_csrf_env): tear_down_db() + clear_engine_cache() diff --git a/tests/conftest.py b/tests/conftest.py index 35b94f8..1067b15 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -7,6 +7,7 @@ from fastapi.testclient import TestClient from dotenv import load_dotenv from utils.core.db import ( + clear_engine_cache, get_connection_url, get_engine, tear_down_db, @@ -87,6 +88,7 @@ def engine(env_vars): # Clean up after tests tear_down_db() + clear_engine_cache() @pytest.fixture diff --git a/tests/utils/test_db.py b/tests/utils/test_db.py index b2c281b..dbccef8 100644 --- a/tests/utils/test_db.py +++ b/tests/utils/test_db.py @@ -151,6 +151,15 @@ def test_get_connection_url_missing_pool_vars(monkeypatch): # --- Engine cache --- +@pytest.fixture +def engine_cache(): + """Ensure engine-cache tests start clean and never leak cached engines, + even when an assertion fails mid-test.""" + clear_engine_cache() + yield + clear_engine_cache() + + def _direct_db_env(monkeypatch, *, name: str = "testdb", password: str = "testpass"): for var in ( "USE_POOL", @@ -174,55 +183,47 @@ def _direct_db_env(monkeypatch, *, name: str = "testdb", password: str = "testpa monkeypatch.setenv("DB_PASSWORD", password) -def test_get_engine_reuses_same_instance(monkeypatch): - clear_engine_cache() +def test_get_engine_reuses_same_instance(engine_cache, monkeypatch): _direct_db_env(monkeypatch) assert get_engine() is get_engine() - clear_engine_cache() -def test_get_engine_different_urls_get_different_engines(monkeypatch): - clear_engine_cache() +def test_get_engine_different_urls_get_different_engines(engine_cache, monkeypatch): _direct_db_env(monkeypatch, name="db_a") engine_a = get_engine() _direct_db_env(monkeypatch, name="db_b") engine_b = get_engine() assert engine_a is not engine_b - clear_engine_cache() -def test_get_engine_not_keyed_by_masked_str_password(monkeypatch): +def test_get_engine_not_keyed_by_masked_str_password(engine_cache, monkeypatch): """str(URL) masks passwords; cache must still separate credentials.""" - clear_engine_cache() _direct_db_env(monkeypatch, password="secretA") engine_a = get_engine() assert "***" in str(get_connection_url()) _direct_db_env(monkeypatch, password="secretB") engine_b = get_engine() assert engine_a is not engine_b - clear_engine_cache() -def test_clear_engine_cache_disposes_and_creates_new(monkeypatch): - clear_engine_cache() +def test_clear_engine_cache_disposes_and_creates_new(engine_cache, monkeypatch): _direct_db_env(monkeypatch) first = get_engine() clear_engine_cache() second = get_engine() assert first is not second - clear_engine_cache() -def test_get_engine_applies_pool_settings(monkeypatch): - clear_engine_cache() +def test_get_engine_applies_pool_settings(engine_cache, monkeypatch): _direct_db_env(monkeypatch) monkeypatch.setenv("DB_POOL_SIZE", "3") monkeypatch.setenv("DB_MAX_OVERFLOW", "2") engine = get_engine() assert engine.pool.size() == 3 + # No public accessors for these; private attrs are stable in practice but + # may need updating on a SQLAlchemy major upgrade. assert engine.pool._max_overflow == 2 assert engine.pool._pre_ping is True - clear_engine_cache() # --- Permission and Role Tests --- diff --git a/utils/core/db.py b/utils/core/db.py index 51fc56b..4448aec 100644 --- a/utils/core/db.py +++ b/utils/core/db.py @@ -1,5 +1,6 @@ import os import logging +import threading from itertools import chain from typing import Union, Sequence from sqlalchemy.engine import Engine, URL @@ -29,8 +30,11 @@ # Cache by URL object (not str(url)): str(URL) masks passwords as "***", so -# different credentials would collide on one cache key. +# different credentials would collide on one cache key. Guarded by a lock: +# sync dependencies run in FastAPI's threadpool, so concurrent first requests +# would otherwise race and orphan an undisposed engine. _engines: dict[URL, Engine] = {} +_engines_lock = threading.Lock() def get_engine() -> Engine: @@ -38,26 +42,31 @@ def get_engine() -> Engine: Calling create_engine() per request leaks pooled connections until GC reclaims them. Cache one engine per URL so tests that swap DB_NAME still - get a separate pool per database. + get a separate pool per database. Pool settings (DB_POOL_SIZE, + DB_MAX_OVERFLOW) are read only when the engine for a URL is first + created; later environment changes have no effect on a cached engine. """ url = get_connection_url() - engine = _engines.get(url) - if engine is None: - engine = create_engine( - url, - pool_pre_ping=True, - pool_size=int(os.getenv("DB_POOL_SIZE", "5")), - max_overflow=int(os.getenv("DB_MAX_OVERFLOW", "5")), - ) - _engines[url] = engine + with _engines_lock: + engine = _engines.get(url) + if engine is None: + engine = create_engine( + url, + pool_pre_ping=True, + pool_size=int(os.getenv("DB_POOL_SIZE", "5")), + max_overflow=int(os.getenv("DB_MAX_OVERFLOW", "5")), + ) + _engines[url] = engine return engine def clear_engine_cache() -> None: """Dispose and drop all cached engines (shutdown / tests).""" - for engine in _engines.values(): + with _engines_lock: + engines = list(_engines.values()) + _engines.clear() + for engine in engines: engine.dispose() - _engines.clear() def ensure_database_exists(url: URL) -> None: @@ -323,4 +332,3 @@ def tear_down_db() -> None: conn.commit() finally: engine.dispose() - clear_engine_cache() diff --git a/utils/core/rate_limit.py b/utils/core/rate_limit.py index 9c74505..acf4e65 100644 --- a/utils/core/rate_limit.py +++ b/utils/core/rate_limit.py @@ -10,23 +10,14 @@ from fastapi import Request, Form from pydantic import EmailStr from dotenv import load_dotenv -from sqlmodel import Session, col, create_engine, delete, select +from sqlmodel import Session, col, delete, select -from utils.core.db import get_connection_url +from utils.core.db import get_engine from utils.core.models import RateLimitAttempt logger = getLogger("uvicorn.error") load_dotenv() -_rate_limit_engine = None - - -def _get_rate_limit_engine(): - global _rate_limit_engine - if _rate_limit_engine is None: - _rate_limit_engine = create_engine(get_connection_url()) - return _rate_limit_engine - @runtime_checkable class RateLimiter(Protocol): @@ -235,7 +226,7 @@ def _recent_attempts(self, session: Session, key: str, now: datetime): def check(self, key: str) -> Tuple[bool, int]: now = datetime.now(UTC) - with Session(_get_rate_limit_engine()) as session: + with Session(get_engine()) as session: attempts = self._recent_attempts(session, key, now) if len(attempts) >= self.max_attempts: oldest = attempts[0].attempted_at @@ -250,7 +241,7 @@ def check(self, key: str) -> Tuple[bool, int]: return False, 0 def record(self, key: str) -> None: - with Session(_get_rate_limit_engine()) as session: + with Session(get_engine()) as session: session.add( RateLimitAttempt( scope=self.scope, key=key, attempted_at=datetime.now(UTC) @@ -260,12 +251,12 @@ def record(self, key: str) -> None: def remaining(self, key: str) -> int: now = datetime.now(UTC) - with Session(_get_rate_limit_engine()) as session: + with Session(get_engine()) as session: attempts = self._recent_attempts(session, key, now) return max(0, self.max_attempts - len(attempts)) def reset(self, key: str) -> None: - with Session(_get_rate_limit_engine()) as session: + with Session(get_engine()) as session: session.exec( delete(RateLimitAttempt).where( col(RateLimitAttempt.scope) == self.scope, @@ -276,7 +267,7 @@ def reset(self, key: str) -> None: def prune(self) -> None: cutoff = self._cutoff(datetime.now(UTC)) - with Session(_get_rate_limit_engine()) as session: + with Session(get_engine()) as session: session.exec( delete(RateLimitAttempt).where( col(RateLimitAttempt.scope) == self.scope, @@ -286,7 +277,7 @@ def prune(self) -> None: session.commit() def clear(self) -> None: - with Session(_get_rate_limit_engine()) as session: + with Session(get_engine()) as session: session.exec( delete(RateLimitAttempt).where( col(RateLimitAttempt.scope) == self.scope