Fix the mysql retry middleware.

oauth2-authenticator/arxiv_oauth2/main.py

@@ -155,7 +155,7 @@ def create_app(*args, **kwargs) -> FastAPI:
 
     app.add_middleware(SessionMiddleware, secret_key=settings.SECRET_KEY)
 
-    app.add_middleware(MySQLRetryMiddleware, retry_attempts=3)
+    app.add_middleware(MySQLRetryMiddleware, engine=_classic_engine,  retry_attempts=3)
 
     app.include_router(auth_router)
 

oauth2-authenticator/arxiv_oauth2/mysql_retry.py

@@ -21,6 +21,7 @@
 #
 # MySQLdb.OperationalError: (2006, 'Server has gone away')
 #
+from typing import Callable
 
 import MySQLdb  # type: ignore
 import logging
@@ -29,21 +30,23 @@
 from fastapi import FastAPI, Request
 from sqlalchemy.engine import Engine
 import asyncio
+from starlette.types import Scope, Receive, Send
 
 
 class MySQLRetryMiddleware:
-    def __init__(self, app: FastAPI, retry_attempts: int = 2, retry_delay: int = 1):
+    def __init__(self, app: FastAPI, engine: Engine = None, retry_attempts: int = 2, retry_delay: int = 1):
         self.app = app
         self.retry_attempts = retry_attempts
         self.retry_delay = retry_delay
+        self.engine = engine
 
-    async def __call__(self, request: Request, call_next):
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         attempt = 0
         while attempt <= self.retry_attempts:
             attempt += 1
             try:
-                response = await call_next(request)
-                return response
+                await self.app(scope, receive, send)
+                return
             except MySQLdb.OperationalError:
                 logging.error(f"MySQL OperationalError detected (attempt {attempt}). Disposing engine.")
                 if attempt > self.retry_attempts:
@@ -54,6 +57,5 @@ async def __call__(self, request: Request, call_next):
                 if attempt > self.retry_attempts:
                     raise
 
-            engine: Engine = request.app.extra['arxiv_db_engine']
-            engine.dispose()
+            self.engine.dispose()
             await asyncio.sleep(self.retry_delay)