diff --git a/README.md b/README.md index 60dbf0e..c3526b3 100644 --- a/README.md +++ b/README.md @@ -18,7 +18,19 @@ The data can be stored using any async SQLAlchemy supported database: - External databases: PostgreSQL, MySQL, MariaDB, CockroachDB - Internal database: SQLite -By default, an internal SQLite database is used. +If none of the below SPOOLMAN_DB_* environment variables are set, a SQLite database located in the user directory will be created and used. + +Database configuration: +| Variable | Description | +|---------------------------|------------------------------------------------------------------------------------------------------------------------------| +| SPOOLMAN_DB_TYPE | Type of database, any of: "postgres", "mysql", "sqlite", "cockroachdb" | +| SPOOLMAN_DB_HOST | Database hostname | +| SPOOLMAN_DB_PORT | Database port | +| SPOOLMAN_DB_NAME | Database name | +| SPOOLMAN_DB_USERNAME | Database username | +| SPOOLMAN_DB_PASSWORD_FILE | Path of file which contains the database password. This is more secure than using SPOOLMAN_DB_PASSWORD. | +| SPOOLMAN_DB_PASSWORD | Database password | +| SPOOLMAN_DB_QUERY | Query parameters for the database connection, e.g. set to "unix_socket=/path/to/mysql.sock" to connect using a MySQL socket. | ## Data diff --git a/pyproject.toml b/pyproject.toml index 5fe939f..189af35 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,7 +12,7 @@ dependencies = {file = ["requirements.txt"]} [tool.ruff] select = ["ALL"] -ignore = ["A003", "D101", "D104", "D407", "S104", "TRY201", "TRY003", "EM101", "EM102"] +ignore = ["A003", "D101", "D104", "D406", "D407", "S104", "TRY201", "TRY003", "EM101", "EM102"] line-length = 120 target-version = "py39" diff --git a/requirements.txt b/requirements.txt index a75ec07..7f3325a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,9 @@ uvicorn[standard]==0.21.1 fastapi==0.95.0 -SQLAlchemy[asyncio]==2.0.8 +SQLAlchemy[asyncio,aiomysql,postgresql_asyncpg,aiosqlite]==2.0.8 pydantic==1.10.7 +platformdirs==3.2.0 -aiosqlite==0.18.0 +# CockroachDB support +sqlalchemy-cockroachdb==2.0.0 +asyncpg==0.27.0 diff --git a/spoolman/database/database.py b/spoolman/database/database.py index 262822c..b8e59c9 100644 --- a/spoolman/database/database.py +++ b/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: diff --git a/spoolman/env.py b/spoolman/env.py new file mode 100644 index 0000000..66b3d37 --- /dev/null +++ b/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") diff --git a/spoolman/main.py b/spoolman/main.py index be61cbd..c3e6e4d 100644 --- a/spoolman/main.py +++ b/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__":