/ spoolman / main.py
main.py
  1  """Main entrypoint to the server."""
  2  
  3  import logging
  4  import subprocess
  5  from logging.handlers import TimedRotatingFileHandler
  6  from pathlib import Path
  7  
  8  import uvicorn
  9  from fastapi import FastAPI
 10  from fastapi.middleware.cors import CORSMiddleware
 11  from fastapi.middleware.gzip import GZipMiddleware
 12  from fastapi.responses import PlainTextResponse, RedirectResponse, Response
 13  from prometheus_client import generate_latest
 14  from scheduler.asyncio.scheduler import Scheduler
 15  
 16  from spoolman import env, externaldb
 17  from spoolman.api.v1.router import app as v1_app
 18  from spoolman.client import SinglePageApplication
 19  from spoolman.database import database
 20  from spoolman.prometheus.metrics import registry
 21  
 22  # Define a console logger
 23  console_handler = logging.StreamHandler()
 24  console_handler.setFormatter(logging.Formatter("%(name)-26s %(levelname)-8s %(message)s"))
 25  
 26  # Setup the spoolman logger, which all spoolman modules will use
 27  log_level = env.get_logging_level()
 28  root_logger = logging.getLogger()
 29  root_logger.setLevel(log_level)
 30  root_logger.addHandler(console_handler)
 31  
 32  # Fix uvicorn logging
 33  logging.getLogger("uvicorn").setLevel(log_level)
 34  if logging.getLogger("uvicorn").handlers:
 35      logging.getLogger("uvicorn").removeHandler(logging.getLogger("uvicorn").handlers[0])
 36  logging.getLogger("uvicorn").addHandler(console_handler)
 37  
 38  logging.getLogger("uvicorn.error").setLevel(log_level)
 39  
 40  access_handlers = logging.getLogger("uvicorn.access").handlers
 41  if access_handlers:
 42      logging.getLogger("uvicorn.access").setLevel(log_level)
 43      logging.getLogger("uvicorn.access").removeHandler(access_handlers[0])
 44      logging.getLogger("uvicorn.access").addHandler(console_handler)
 45  
 46  # Get logger instance for this module
 47  logger = logging.getLogger(__name__)
 48  
 49  
 50  # Setup FastAPI
 51  app = FastAPI(
 52      debug=env.is_debug_mode(),
 53      title="Spoolman",
 54      version=env.get_version(),
 55  )
 56  app.add_middleware(GZipMiddleware)
 57  app.mount(env.get_base_path() + "/api/v1", v1_app)
 58  
 59  
 60  # WA for prometheus /metrics bind with SinglePageApp at root
 61  @app.get(
 62      env.get_base_path() + "/metrics",
 63      response_class=PlainTextResponse,
 64      name="Get metrics for prometheus",
 65      description=(
 66          "Get app metrics for prometheusIf enabled SPOOLMAN_METRICS_ENABLED returned metrics by Spools and Filaments"
 67      ),
 68  )
 69  def get_metrics() -> bytes:
 70      """Return prometheus metrics."""
 71      return generate_latest(registry)
 72  
 73  
 74  base_path = env.get_base_path()
 75  if base_path != "":
 76      logger.info("Base path is: %s", base_path)
 77  
 78      # If base path is set, add a redirect from non-slash suffix to slash
 79      # suffix. Otherwise it won't work.
 80      @app.get(base_path)
 81      def root_redirect() -> Response:
 82          """Redirect to base path."""
 83          return RedirectResponse(base_path + "/")
 84  
 85  
 86  # Return a dynamic js config file
 87  # This is so that the client side can access the base path variable.
 88  @app.get(env.get_base_path() + "/config.js")
 89  def get_configjs() -> Response:
 90      """Return a dynamic js config file."""
 91      if '"' in base_path:
 92          raise ValueError("Base path contains quotes, which are not allowed.")
 93  
 94      return Response(
 95          content=f"""
 96  window.SPOOLMAN_BASE_PATH = "{base_path}";
 97  """,
 98          media_type="text/javascript",
 99      )
100  
101  
102  # Mount the client side app
103  app.mount(base_path, app=SinglePageApplication(directory="client/dist", base_path=env.get_base_path()))
104  
105  
106  def add_cors_middleware() -> None:
107      """Add CORS middleware to the FastAPI app based on environment settings."""
108      origins = []
109      if env.is_debug_mode():
110          logger.warning("Running in debug mode, allowing all origins.")
111          origins = ["*"]
112      elif env.is_cors_defined():
113          cors_origins = env.get_cors_origin()
114          if cors_origins:
115              logger.info("CORS origins defined: %s", cors_origins)
116              origins = cors_origins
117          else:
118              logger.warning("CORS origins are not defined, no CORS will be applied.")
119  
120      if not origins:
121          return
122  
123      app.add_middleware(
124          CORSMiddleware,
125          allow_origins=origins,
126          allow_credentials=True,
127          allow_methods=["*"],
128          allow_headers=["*"],
129          expose_headers=["X-Total-Count"],
130      )
131  
132  
133  add_cors_middleware()
134  
135  
136  def add_file_logging() -> None:
137      """Add file logging to the root logger."""
138      # Define a file logger with log rotation
139      log_file = env.get_logs_dir().joinpath("spoolman.log")
140      file_handler = TimedRotatingFileHandler(log_file, when="midnight", backupCount=5)
141      file_handler.setFormatter(logging.Formatter("%(asctime)s:%(levelname)s:%(message)s", "%Y-%m-%d %H:%M:%S"))
142      root_logger.addHandler(file_handler)
143  
144      logging.getLogger("uvicorn").addHandler(file_handler)
145      access_handlers = logging.getLogger("uvicorn.access").handlers
146      if access_handlers:
147          logging.getLogger("uvicorn.access").addHandler(file_handler)
148  
149  
150  @app.on_event("startup")
151  async def startup() -> None:
152      """Run the service's startup sequence."""
153      # Check that the data directory is writable
154      env.check_write_permissions()
155  
156      # Don't add file logging until we have verified that the data directory is writable
157      add_file_logging()
158  
159      logger.info(
160          "Starting Spoolman v%s (commit: %s) (built: %s)",
161          app.version,
162          env.get_commit_hash(),
163          env.get_build_date(),
164      )
165  
166      logger.info("Using data directory: %s", env.get_data_dir().resolve())
167      logger.info("Using logs directory: %s", env.get_logs_dir().resolve())
168      logger.info("Using backups directory: %s", env.get_backups_dir().resolve())
169  
170      logger.info("Setting up database...")
171      database.setup_db(database.get_connection_url())
172  
173      logger.info("Performing migrations...")
174      # Run alembic in a subprocess.
175      # There is some issue with the uvicorn worker that causes the process to hang when running alembic directly.
176      # See: https://github.com/sqlalchemy/alembic/discussions/1155
177      project_root = Path(__file__).parent.parent
178      subprocess.run(["alembic", "upgrade", "head"], check=True, cwd=project_root)  # noqa: ASYNC221, S607
179  
180      # Setup scheduler
181      schedule = Scheduler()
182      database.schedule_tasks(schedule)
183      externaldb.schedule_tasks(schedule)
184  
185      logger.info("Startup complete.")
186  
187      if env.is_docker() and not env.is_data_dir_mounted():
188          logger.warning("!!!! WARNING !!!!")
189          logger.warning("!!!! WARNING !!!!")
190          logger.warning("The data directory is not mounted.")
191          logger.warning(
192              'Spoolman stores its database in the container directory "%s". '
193              "If this directory isn't mounted to the host OS, the database will be lost when the container is stopped.",
194              env.get_data_dir(),
195          )
196          logger.warning(
197              "Please carefully read the docker part of the README.md file, "
198              "and ensure your docker-compose file matches the example.",
199          )
200          logger.warning("!!!! WARNING !!!!")
201          logger.warning("!!!! WARNING !!!!")
202  
203  
204  if __name__ == "__main__":
205      uvicorn.run(app, host="0.0.0.0", port=8000)