From 7e9f488a5fec822acda0d2074e9a0e1f0bfb5587 Mon Sep 17 00:00:00 2001 From: Ryan Kegel Date: Wed, 28 Oct 2026 23:46:26 -0400 Subject: [PATCH] Refactor database connection handling in LocalOAuthProvider to use async context manager --- oauth.py | 39 ++++++++++++++++++++++----------------- 1 file changed, 22 insertions(+), 17 deletions(-) diff --git a/oauth.py b/oauth.py index a3e3127..a07fbf0 100644 --- a/oauth.py +++ b/oauth.py @@ -8,8 +8,9 @@ import html import json import secrets import time +from contextlib import asynccontextmanager from pathlib import Path -from typing import Any +from typing import Any, AsyncIterator from urllib.parse import urlencode import aiosqlite @@ -110,13 +111,17 @@ class LocalOAuthProvider: ) await db.commit() - async def _connect(self) -> aiosqlite.Connection: + @asynccontextmanager + async def _connect(self) -> AsyncIterator[aiosqlite.Connection]: db = await aiosqlite.connect(self.database_path) db.row_factory = aiosqlite.Row - return db + try: + yield db + finally: + await db.close() async def get_client(self, client_id: str) -> OAuthClientInformationFull | None: - async with await self._connect() as db: + async with self._connect() as db: async with db.execute( "SELECT metadata_json FROM oauth_clients WHERE client_id = ?", (client_id,) ) as cursor: @@ -129,7 +134,7 @@ class LocalOAuthProvider: if not client_info.client_id: raise ValueError("OAuth client registration did not include a client ID") data = client_info.model_dump_json() - async with await self._connect() as db: + async with self._connect() as db: async with db.execute( "SELECT COUNT(*) AS n FROM oauth_clients WHERE client_id != ?", (client_info.client_id,), @@ -158,7 +163,7 @@ class LocalOAuthProvider: raise AuthorizeError("invalid_target", "This authorization server only issues tokens for this MCP endpoint") flow_id = secrets.token_urlsafe(32) - async with await self._connect() as db: + async with self._connect() as db: now = int(time.time()) await db.execute("DELETE FROM oauth_flows WHERE expires_at < ?", (now,)) await db.execute("DELETE FROM oauth_codes WHERE expires_at < ?", (now,)) @@ -185,7 +190,7 @@ class LocalOAuthProvider: return f"{self.issuer_url}/login?{urlencode({'flow': flow_id})}" async def get_flow(self, flow_id: str) -> dict[str, Any] | None: - async with await self._connect() as db: + async with self._connect() as db: async with db.execute( "SELECT * FROM oauth_flows WHERE flow_id = ?", (flow_id,) ) as cursor: @@ -202,7 +207,7 @@ class LocalOAuthProvider: async def login_is_limited(self, client_ip: str) -> bool: cutoff = int(time.time()) - LOGIN_WINDOW_SECONDS - async with await self._connect() as db: + async with self._connect() as db: async with db.execute( "SELECT COUNT(*) AS n FROM oauth_login_failures WHERE client_ip = ? AND failed_at >= ?", (client_ip, cutoff), @@ -213,7 +218,7 @@ class LocalOAuthProvider: async def record_login_failure(self, client_ip: str) -> None: now = int(time.time()) cutoff = now - LOGIN_WINDOW_SECONDS - async with await self._connect() as db: + async with self._connect() as db: await db.execute("DELETE FROM oauth_login_failures WHERE failed_at < ?", (cutoff,)) await db.execute( "INSERT INTO oauth_login_failures (client_ip, failed_at) VALUES (?, ?)", @@ -222,7 +227,7 @@ class LocalOAuthProvider: await db.commit() async def clear_login_failures(self, client_ip: str) -> None: - async with await self._connect() as db: + async with self._connect() as db: await db.execute("DELETE FROM oauth_login_failures WHERE client_ip = ?", (client_ip,)) await db.commit() @@ -238,7 +243,7 @@ class LocalOAuthProvider: code_value = secrets.token_urlsafe(48) now = int(time.time()) - async with await self._connect() as db: + async with self._connect() as db: await db.execute( """INSERT INTO oauth_codes (code, client_id, redirect_uri, redirect_uri_provided_explicitly, @@ -267,7 +272,7 @@ class LocalOAuthProvider: async def load_authorization_code( self, client: OAuthClientInformationFull, authorization_code: str ) -> AuthorizationCode | None: - async with await self._connect() as db: + async with self._connect() as db: async with db.execute( "SELECT * FROM oauth_codes WHERE code = ? AND client_id = ?", (authorization_code, client.client_id), @@ -295,7 +300,7 @@ class LocalOAuthProvider: now = int(time.time()) access = secrets.token_urlsafe(48) refresh = secrets.token_urlsafe(48) - async with await self._connect() as db: + async with self._connect() as db: await db.execute("BEGIN IMMEDIATE") async with db.execute( "SELECT expires_at FROM oauth_codes WHERE code = ? AND client_id = ?", @@ -335,7 +340,7 @@ class LocalOAuthProvider: self, client: OAuthClientInformationFull, refresh_token: str ) -> RefreshToken | None: now = int(time.time()) - async with await self._connect() as db: + async with self._connect() as db: async with db.execute( """SELECT * FROM oauth_tokens WHERE refresh_token = ? AND client_id = ? AND refresh_expires_at >= ?""", @@ -365,7 +370,7 @@ class LocalOAuthProvider: raise TokenError("invalid_scope", "A refresh token cannot grant additional permissions") access = secrets.token_urlsafe(48) refresh = secrets.token_urlsafe(48) - async with await self._connect() as db: + async with self._connect() as db: await db.execute("BEGIN IMMEDIATE") async with db.execute( """SELECT * FROM oauth_tokens WHERE refresh_token = ? AND client_id = ? @@ -406,7 +411,7 @@ class LocalOAuthProvider: async def load_access_token(self, token: str) -> AccessToken | None: now = int(time.time()) - async with await self._connect() as db: + async with self._connect() as db: async with db.execute( "SELECT * FROM oauth_tokens WHERE access_token = ? AND access_expires_at >= ?", (token, now), @@ -425,7 +430,7 @@ class LocalOAuthProvider: async def revoke_token(self, token: AccessToken | RefreshToken) -> None: column = "access_token" if isinstance(token, AccessToken) else "refresh_token" - async with await self._connect() as db: + async with self._connect() as db: await db.execute(f"DELETE FROM oauth_tokens WHERE {column} = ?", (token.token,)) await db.commit()