30 lines
1.3 KiB
Python
30 lines
1.3 KiB
Python
|
"""Gate business traffic during deployment and database recovery, including sockets."""
|
||
|
|
import asyncio
|
||
|
|
from pathlib import Path
|
||
|
|
from starlette.types import ASGIApp, Receive, Scope, Send
|
||
|
|
from starlette.responses import JSONResponse
|
||
|
|
|
||
|
|
MAINTENANCE_FILE = Path('/var/lib/replypier-updater/maintenance')
|
||
|
|
|
||
|
|
|
||
|
|
class SystemMaintenanceMiddleware:
|
||
|
|
def __init__(self, app: ASGIApp) -> None:
|
||
|
|
self.app = app
|
||
|
|
|
||
|
|
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||
|
|
path = scope.get('path', '')
|
||
|
|
exempt = path == '/health' or (path == '/api/v1/system-updates' and scope.get('method') == 'GET')
|
||
|
|
if scope['type'] in {'http', 'websocket'} and MAINTENANCE_FILE.exists() and not exempt:
|
||
|
|
if scope['type'] == 'websocket':
|
||
|
|
await send({'type': 'websocket.close', 'code': 1013})
|
||
|
|
else:
|
||
|
|
response = JSONResponse({'detail': '系统更新中，接待将在更新完成后恢复。'},
|
||
|
|
status_code=503, headers={'Retry-After': '5'})
|
||
|
|
await response(scope, receive, send)
|
||
|
|
return
|
||
|
|
await self.app(scope, receive, send)
|
||
|
|
|
||
|
|
|
||
|
|
async def wait_for_deployment() -> None:
|
||
|
|
while MAINTENANCE_FILE.exists():
|
||
|
|
await asyncio.sleep(1)
|