Refactor database connection handling in LocalOAuthProvider to use async context manager

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()