Added proper support for all databases

spoolman/database/database.py

@@ -3,17 +3,18 @@
 from collections.abc import AsyncGenerator
 from typing import Optional
 
+from sqlalchemy import URL
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
 
 from spoolman.database.models import Base
 
 
 class Database:
-    connection_url: str
-    engine: AsyncEngine
-    session_maker: async_sessionmaker[AsyncSession]
+    connection_url: URL
+    engine: Optional[AsyncEngine]
+    session_maker: Optional[async_sessionmaker[AsyncSession]]
 
-    def __init__(self: "Database", connection_url: str) -> None:
+    def __init__(self: "Database", connection_url: URL) -> None:
         """Construct the Database wrapper and set config parameters."""
         self.connection_url = connection_url
 
@@ -24,6 +25,8 @@ class Database:
 
     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)
 
@@ -31,8 +34,12 @@ class Database:
 __db: Optional[Database] = None
 
 
-async def setup_db(connection_url: str) -> None:
-    """Connect the singleton DB object."""
+async def setup_db(connection_url: URL) -> None:
+    """Connect to the database.
+
+    Args:
+        connection_url: The URL to connect to the database.
+    """
     global __db  # noqa: PLW0603
     __db = Database(connection_url)
     __db.connect()
@@ -41,7 +48,7 @@ async def setup_db(connection_url: str) -> None:
 
 async def get_db_session() -> AsyncGenerator[AsyncSession, None]:
     """Get a DB session to be used with FastAPI's dependency system."""
-    if __db is None:
+    if __db is None or __db.session_maker is None:
         raise RuntimeError("DB is not setup.")
     async with __db.session_maker() as session:
         try:

spoolman/env.py

@@ -0,0 +1,155 @@
+"""Utilities for grabbing config from environment variables."""
+
+import os
+from enum import Enum
+from pathlib import Path
+from typing import Optional
+from urllib import parse
+
+
+class DatabaseType(Enum):
+    """The database type."""
+
+    POSTGRES = "postgres"
+    MYSQL = "mysql"
+    SQLITE = "sqlite"
+    COCKROACHDB = "cockroachdb"
+
+    def to_drivername(self: "DatabaseType") -> str:
+        """Get the drivername for the database type.
+
+        Returns:
+            str: The drivername.
+        """
+        if self is DatabaseType.POSTGRES:
+            return "postgresql+asyncpg"
+        if self is DatabaseType.MYSQL:
+            return "mysql+aiomysql"
+        if self is DatabaseType.SQLITE:
+            return "sqlite+aiosqlite"
+        if self is DatabaseType.COCKROACHDB:
+            return "cockroachdb+asyncpg"
+        raise ValueError(f"Unknown database type '{self}'.")
+
+
+def get_database_type() -> Optional[DatabaseType]:
+    """Get the database type from environment variables.
+
+    Returns None if no environment variable was set for the database type.
+
+    Returns:
+        Optional[DatabaseType]: The database type.
+    """
+    database_type = os.getenv("SPOOLMAN_DB_TYPE")
+    if database_type is None:
+        return None
+    if database_type == "postgres":
+        return DatabaseType.POSTGRES
+    if database_type == "mysql":
+        return DatabaseType.MYSQL
+    if database_type == "sqlite":
+        return DatabaseType.SQLITE
+    if database_type == "cockroachdb":
+        return DatabaseType.COCKROACHDB
+    raise ValueError(f"Failed to parse SPOOLMAN_DB_TYPE variable: Unknown database type '{database_type}'.")
+
+
+def get_host() -> Optional[str]:
+    """Get the DB host from environment variables.
+
+    Returns None if no environment variable was set for the host.
+
+    Returns:
+        Optional[str]: The host.
+    """
+    return os.getenv("SPOOLMAN_DB_HOST")
+
+
+def get_port() -> Optional[int]:
+    """Get the DB port from environment variables.
+
+    Returns None if no environment variable was set for the port.
+
+    Returns:
+        Optional[str]: The port.
+    """
+    port = os.getenv("SPOOLMAN_DB_PORT")
+    if port is None:
+        return None
+    try:
+        return int(port)
+    except ValueError as exc:
+        raise ValueError(f"Failed to parse SPOOLMAN_DB_PORT variable: {str(exc)}") from exc
+
+
+def get_database() -> Optional[str]:
+    """Get the DB name from environment variables.
+
+    Returns None if no environment variable was set for the name.
+
+    Returns:
+        Optional[str]: The name.
+    """
+    return os.getenv("SPOOLMAN_DB_NAME")
+
+
+def get_query() -> Optional[dict[str, str]]:
+    """Get the DB query from environment variables.
+
+    Returns None if no environment variable was set for the query.
+
+    Returns:
+        Optional[dict]: The query.
+    """
+    query = os.getenv("SPOOLMAN_DB_QUERY")
+    if query is None:
+        return None
+    try:
+        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
+
+
+def get_username() -> Optional[str]:
+    """Get the DB username from environment variables.
+
+    Returns None if no environment variable was set for the username.
+
+    Returns:
+        Optional[str]: The username.
+    """
+    return os.getenv("SPOOLMAN_DB_USERNAME")
+
+
+def get_password() -> Optional[str]:
+    """Get the DB password from environment variables.
+
+    Returns None if no environment variables were set for the password.
+
+    Raises:
+        ValueError: If it failed to read the password from a password file.
+
+    Returns:
+        Optional[str]: The password.
+    """
+    # First attempt: grab password from a file. This is the most secure way of storing passwords.
+    file_path = os.getenv("SPOOLMAN_DB_PASSWORD_FILE")
+    if file_path is not None:
+        file = Path(file_path)
+        if not file.exists() or not file.is_file():
+            raise ValueError(
+                "Failed to parse SPOOLMAN_DB_PASSWORD_FILE variable: "
+                f'Database password file "{file_path}" does not exist.',
+            )
+        try:
+            with file.open(encoding="utf-8") as f:
+                return f.read()
+        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)}.',
+            ) from exc
+
+    # Second attempt: grab directly from an environment variable.
+    return os.getenv("SPOOLMAN_DB_PASSWORD")

spoolman/main.py

@@ -1,20 +1,77 @@
 """Main entrypoint to the server."""
 
+import logging
+from pathlib import Path
 
 import uvicorn
 from fastapi import FastAPI
+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
 
+# Get the data directory
+data_dir = Path(user_data_dir("spoolman"))
+data_dir.mkdir(parents=True, exist_ok=True)
+
+# Setup file logger
+logging.basicConfig(
+    filename=data_dir.joinpath("spoolman.log"),
+    filemode="w",
+    level=logging.INFO,
+    format="%(asctime)s:%(levelname)s:%(message)s",
+    datefmt="%Y-%m-%d %I:%M:%S%p",
+)
+
+# Setup console logger
+formatter = logging.Formatter("%(name)-12s: %(levelname)-8s %(message)s")
+console = logging.StreamHandler()
+console.setLevel(logging.INFO)
+console.setFormatter(formatter)
+logging.getLogger("").addHandler(console)
+
+# Get logger instance
+logger = logging.getLogger(__name__)
+
+# Setup FastAPI
 app = FastAPI(debug=True)
 app.mount("/api/v1", v1_app)
 
 
+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("sqlite+aiosqlite:///./sql_app.db")
+    await setup_db(get_connection_url())
 
 
 if __name__ == "__main__":