Added proper support for all databases

This commit is contained in:
Donkie
2023-04-06 16:19:33 +02:00
parent 77ffab61c5
commit 1a148c0991
6 changed files with 246 additions and 12 deletions

View File

@@ -18,7 +18,19 @@ The data can be stored using any async SQLAlchemy supported database:
- External databases: PostgreSQL, MySQL, MariaDB, CockroachDB - External databases: PostgreSQL, MySQL, MariaDB, CockroachDB
- Internal database: SQLite - 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 ## Data

View File

@@ -12,7 +12,7 @@ dependencies = {file = ["requirements.txt"]}
[tool.ruff] [tool.ruff]
select = ["ALL"] 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 line-length = 120
target-version = "py39" target-version = "py39"

View File

@@ -1,6 +1,9 @@
uvicorn[standard]==0.21.1 uvicorn[standard]==0.21.1
fastapi==0.95.0 fastapi==0.95.0
SQLAlchemy[asyncio]==2.0.8 SQLAlchemy[asyncio,aiomysql,postgresql_asyncpg,aiosqlite]==2.0.8
pydantic==1.10.7 pydantic==1.10.7
platformdirs==3.2.0
aiosqlite==0.18.0 # CockroachDB support
sqlalchemy-cockroachdb==2.0.0
asyncpg==0.27.0

View File

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

155
spoolman/env.py Normal file
View File

@@ -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")

View File

@@ -1,20 +1,77 @@
"""Main entrypoint to the server.""" """Main entrypoint to the server."""
import logging
from pathlib import Path
import uvicorn import uvicorn
from fastapi import FastAPI 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.api.v1.router import app as v1_app
from spoolman.database.database import setup_db 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 = FastAPI(debug=True)
app.mount("/api/v1", v1_app) 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") @app.on_event("startup")
async def startup() -> None: async def startup() -> None:
"""Run the service's startup sequence.""" """Run the service's startup sequence."""
await setup_db("sqlite+aiosqlite:///./sql_app.db") await setup_db(get_connection_url())
if __name__ == "__main__": if __name__ == "__main__":