Added Alembic for migrations

spoolman/database/database.py

@@ -2,13 +2,48 @@
 
 import logging
 from collections.abc import AsyncGenerator
+from pathlib import Path
 from typing import Optional
 
+from platformdirs import user_data_dir
 from sqlalchemy import URL
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
 
-from spoolman.database.models import Base
-from spoolman.env import get_logging_level
+from spoolman import env
+
+logger = logging.getLogger(__name__)
+
+
+def get_connection_url() -> URL:
+    """Construct the connection URL for the database based on environment variables."""
+    db_type = env.get_database_type()
+    host = env.get_host()
+    port = env.get_port()
+    database = env.get_database()
+    query = env.get_query()
+    username = env.get_username()
+    password = env.get_password()
+
+    if db_type is None:
+        db_type = env.DatabaseType.SQLITE
+
+        data_dir = Path(user_data_dir("spoolman"))
+        data_dir.mkdir(parents=True, exist_ok=True)
+
+        database = str(data_dir.joinpath("spoolman.db"))
+        logger.info('No database type specified, using a default SQLite database located at "%s"', database)
+    else:
+        logger.info('Connecting to database of type "%s" at "%s:%s"', db_type, host, port)
+
+    return URL.create(
+        drivername=db_type.to_drivername(),
+        host=host,
+        port=port,
+        database=database,
+        query=query or {},
+        username=username,
+        password=password,
+    )
 
 
 class Database:
@@ -22,7 +57,7 @@ class Database:
 
     def connect(self: "Database") -> None:
         """Connect to the database."""
-        if get_logging_level() == logging.DEBUG:
+        if env.get_logging_level() == logging.DEBUG:
             logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
 
         connect_args = {}
@@ -39,18 +74,11 @@ class Database:
         )
         self.session_maker = async_sessionmaker(self.engine, autocommit=False, autoflush=True)
 
-    async def create_tables(self: "Database") -> None:
-        """Create tables for all defined models."""
-        if self.engine is None:
-            raise RuntimeError("DB is not connected.")
-        async with self.engine.begin() as conn:
-            await conn.run_sync(Base.metadata.create_all)
-
 
 __db: Optional[Database] = None
 
 
-async def setup_db(connection_url: URL) -> None:
+def setup_db(connection_url: URL) -> None:
     """Connect to the database.
 
     Args:
@@ -59,7 +87,6 @@ async def setup_db(connection_url: URL) -> None:
     global __db  # noqa: PLW0603
     __db = Database(connection_url)
     __db.connect()
-    await __db.create_tables()
 
 
 async def get_db_session() -> AsyncGenerator[AsyncSession, None]:

spoolman/env.py

@@ -80,7 +80,7 @@ def get_port() -> Optional[int]:
     try:
         return int(port)
     except ValueError as exc:
-        raise ValueError(f"Failed to parse SPOOLMAN_DB_PORT variable: {str(exc)}") from exc
+        raise ValueError(f"Failed to parse SPOOLMAN_DB_PORT variable: {exc!s}") from exc
 
 
 def get_database() -> Optional[str]:
@@ -109,7 +109,7 @@ def get_query() -> Optional[dict[str, str]]:
         parsed_dict = parse.parse_qs(query, strict_parsing=True)
         return {key: value[0] for key, value in parsed_dict.items()}
     except ValueError as exc:
-        raise ValueError(f"Failed to parse SPOOLMAN_DB_QUERY variable: {str(exc)}") from exc
+        raise ValueError(f"Failed to parse SPOOLMAN_DB_QUERY variable: {exc!s}") from exc
 
 
 def get_username() -> Optional[str]:
@@ -149,7 +149,7 @@ def get_password() -> Optional[str]:
         except OSError as exc:
             raise ValueError(
                 "Failed to parse SPOOLMAN_DB_PASSWORD_FILE variable: "
-                f'Failed to read password from file "{file_path}": {str(exc)}.',
+                f'Failed to read password from file "{file_path}": {exc!s}.',
             ) from exc
 
     # Second attempt: grab directly from an environment variable.

spoolman/main.py

@@ -1,6 +1,7 @@
 """Main entrypoint to the server."""
 
 import logging
+import subprocess
 from pathlib import Path
 
 import uvicorn
@@ -8,11 +9,10 @@ from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 from platformdirs import user_data_dir
-from sqlalchemy import URL
 
 from spoolman import env
 from spoolman.api.v1.router import app as v1_app
-from spoolman.database.database import setup_db
+from spoolman.database.database import get_connection_url, setup_db
 
 # Get the data directory
 data_dir = Path(user_data_dir("spoolman"))
@@ -59,38 +59,20 @@ if env.is_debug_mode():
     )
 
 
-def get_connection_url() -> URL:
-    """Construct the connection URL for the database based on environment variables."""
-    db_type = env.get_database_type()
-    host = env.get_host()
-    port = env.get_port()
-    database = env.get_database()
-    query = env.get_query()
-    username = env.get_username()
-    password = env.get_password()
-
-    if db_type is None:
-        db_type = env.DatabaseType.SQLITE
-        database = str(data_dir.joinpath("spoolman.db"))
-        logger.info('No database type specified, using a default SQLite database located at "%s"', database)
-    else:
-        logger.info('Connecting to database of type "%s" at "%s:%s"', db_type, host, port)
-
-    return URL.create(
-        drivername=db_type.to_drivername(),
-        host=host,
-        port=port,
-        database=database,
-        query=query or {},
-        username=username,
-        password=password,
-    )
-
-
 @app.on_event("startup")
 async def startup() -> None:
     """Run the service's startup sequence."""
-    await setup_db(get_connection_url())
+    logger.info("Setting up database...")
+    setup_db(get_connection_url())
+
+    logger.info("Performing migrations...")
+    # Run alembic in a subprocess.
+    # There is some issue with the uvicorn worker that causes the process to hang when running alembic directly.
+    # See: https://github.com/sqlalchemy/alembic/discussions/1155
+    project_root = Path(__file__).parent.parent
+    subprocess.run(["alembic", "upgrade", "head"], check=True, cwd=project_root)  # noqa: S603, S607
+
+    logger.info("Startup complete.")
 
 
 if __name__ == "__main__":