156 lines
4.5 KiB
Python
156 lines
4.5 KiB
Python
"""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")
|