Added proper support for all databases

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