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
- 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

View File

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

View File

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

View File

@@ -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:

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."""
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__":