argos/argos/server/main.py

87 lines
2.4 KiB
Python

import sys
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from pydantic import ValidationError
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from argos.logging import logger
from argos.server import models, queries, routes
from argos.server.settings import get_app_settings, read_yaml_config
def get_application() -> FastAPI:
settings = get_app_settings()
app = FastAPI()
config = read_config(app, settings)
# Settings is the pydantic settings object
# Config is the argos config object (built from yaml)
app.state.config = config
app.state.settings = settings
app.add_event_handler(
"startup",
create_start_app_handler(app),
)
app.add_event_handler(
"shutdown",
create_stop_app_handler(app),
)
app.include_router(routes.api, prefix="/api")
app.include_router(routes.views)
app.mount("/static", StaticFiles(directory="argos/server/static"), name="static")
return app
def create_start_app_handler(app):
async def read_config_and_populate_db():
setup_database(app)
db = await connect_to_db(app)
await queries.update_from_config(db, app.state.config)
return read_config_and_populate_db
async def connect_to_db(app):
app.state.db = app.state.SessionLocal()
return app.state.db
def create_stop_app_handler(app):
async def stop_app():
app.state.db.close()
return stop_app
def read_config(app, settings):
try:
config = read_yaml_config(settings.yaml_file)
app.state.config = config
return config
except ValidationError as e:
logger.error("Errors where found while reading configuration:")
for error in e.errors():
logger.error(f"{error['loc']} is {error['type']}")
sys.exit(1)
def setup_database(app):
settings = app.state.settings
# For sqlite, we need to add connect_args={"check_same_thread": False}
logger.debug(f"Using database URL {settings.database_url}")
if settings.database_url.startswith("sqlite:////tmp"):
logger.warning("Using sqlite in /tmp is not recommended for production")
engine = create_engine(settings.database_url)
app.state.SessionLocal = sessionmaker(
autocommit=False, autoflush=False, bind=engine
)
app.state.engine = engine
models.Base.metadata.create_all(bind=engine)
app = get_application()