"""Base server implementation."""
import logging
from abc import ABC, abstractmethod
from functools import cached_property
from typing import Callable, Optional
from starlette.applications import Starlette
from starlette.middleware.cors import CORSMiddleware
from starlette.responses import JSONResponse
from starlette.routing import Route
import uvicorn
from ..settings import ServerSettings
async def default_health_check() -> bool:
"""Default health check implementation."""
return True
class BaseServer(ABC):
"""Base server class."""
def __init__(
self,
settings: ServerSettings,
readiness_check: Optional[Callable[[], bool]] = None,
liveness_check: Optional[Callable[[], bool]] = None,
):
"""Initialize base server."""
self.settings = settings
self.readiness_check = readiness_check or default_health_check
self.liveness_check = liveness_check or default_health_check
self.logger = logging.getLogger(__name__)
@cached_property
@abstractmethod
def rest(self) -> Starlette:
"""Create REST server."""
pass
def add_health_check(self, app: Starlette) -> None:
"""Add health check endpoints."""
async def health(request):
"""Health check endpoint."""
return JSONResponse({"status": "healthy"})
async def readiness(request):
"""Readiness check endpoint."""
is_ready = await self.readiness_check()
status_code = 200 if is_ready else 503
return JSONResponse({"ready": is_ready}, status_code=status_code)
async def liveness(request):
"""Liveness check endpoint."""
is_alive = await self.liveness_check()
status_code = 200 if is_alive else 503
return JSONResponse({"alive": is_alive}, status_code=status_code)
# Add health check routes
app.routes.extend([
Route("/health", health, methods=["GET"]),
Route("/health/ready", readiness, methods=["GET"]),
Route("/health/live", liveness, methods=["GET"]),
])
def add_middleware(self, app: Starlette) -> None:
"""Add middleware to the application."""
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
async def start(self) -> None:
"""Start the server."""
self.logger.info(f"Starting server on {self.settings.HOST}:{self.settings.PORT}")
config = uvicorn.Config(
app=self.rest,
host=self.settings.HOST,
port=self.settings.PORT,
log_level="debug" if self.settings.DEBUG else "info",
)
server = uvicorn.Server(config)
await server.serve()
async def stop(self) -> None:
"""Stop the server."""
self.logger.info("Stopping server")
# Implement graceful shutdown logic here
