"""Runtime configuration objects for Moltres."""
from __future__ import annotations
import os
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, cast
from moltres_core.config import EngineConfig, FetchFormat, EngineOptionValue
from sqlalchemy.engine import Engine
if TYPE_CHECKING:
from sqlalchemy.ext.asyncio import AsyncEngine
else:
AsyncEngine = None # type: ignore[assignment, misc]
__all__ = [
"EngineConfig",
"EngineOptionValue",
"FetchFormat",
"MoltresConfig",
"create_config",
"DEFAULT_CONFIG",
]
[docs]
@dataclass
class MoltresConfig:
"""Container for all runtime configuration knobs."""
engine: EngineConfig
default_schema: str | None = None
include_metadata: bool = False
allowed_paths: tuple[str, ...] | None = None
options: dict[str, object] = field(default_factory=dict)
def _load_env_config() -> dict[str, object]:
"""Load configuration from environment variables.
Returns:
Dictionary of configuration values from environment
"""
config: dict[str, object] = {}
# DSN is handled separately in create_config
if "MOLTRES_ECHO" in os.environ:
config["echo"] = os.environ["MOLTRES_ECHO"].lower() in ("true", "1", "yes", "on")
if "MOLTRES_FETCH_FORMAT" in os.environ:
config["fetch_format"] = os.environ["MOLTRES_FETCH_FORMAT"]
if "MOLTRES_DIALECT" in os.environ:
config["dialect"] = os.environ["MOLTRES_DIALECT"]
if "MOLTRES_POOL_SIZE" in os.environ:
try:
config["pool_size"] = int(os.environ["MOLTRES_POOL_SIZE"])
except ValueError as exc:
raise ValueError("MOLTRES_POOL_SIZE must be an integer") from exc
if "MOLTRES_MAX_OVERFLOW" in os.environ:
try:
config["max_overflow"] = int(os.environ["MOLTRES_MAX_OVERFLOW"])
except ValueError as exc:
raise ValueError("MOLTRES_MAX_OVERFLOW must be an integer") from exc
if "MOLTRES_POOL_TIMEOUT" in os.environ:
try:
config["pool_timeout"] = int(os.environ["MOLTRES_POOL_TIMEOUT"])
except ValueError as exc:
raise ValueError("MOLTRES_POOL_TIMEOUT must be an integer") from exc
if "MOLTRES_POOL_RECYCLE" in os.environ:
try:
config["pool_recycle"] = int(os.environ["MOLTRES_POOL_RECYCLE"])
except ValueError as exc:
raise ValueError("MOLTRES_POOL_RECYCLE must be an integer") from exc
if "MOLTRES_POOL_PRE_PING" in os.environ:
config["pool_pre_ping"] = os.environ["MOLTRES_POOL_PRE_PING"].lower() in (
"true",
"1",
"yes",
"on",
)
if "MOLTRES_QUERY_TIMEOUT" in os.environ:
try:
config["query_timeout"] = float(os.environ["MOLTRES_QUERY_TIMEOUT"])
except ValueError as exc:
raise ValueError("MOLTRES_QUERY_TIMEOUT must be a number") from exc
if "MOLTRES_ALLOWED_PATHS" in os.environ:
config["allowed_paths"] = tuple(
p.strip() for p in os.environ["MOLTRES_ALLOWED_PATHS"].split(os.pathsep) if p.strip()
)
return config
def validate_connection_string(dsn: str, *, is_async: bool = False) -> None:
"""Validate a database connection string."""
from .utils.exceptions import DatabaseConnectionError
if not dsn or not isinstance(dsn, str):
raise DatabaseConnectionError(
f"Connection string must be a non-empty string, got: {type(dsn).__name__}"
)
dsn_lower = dsn.lower()
if dsn_lower.startswith("sqlite"):
if is_async and "+aiosqlite" not in dsn_lower:
raise DatabaseConnectionError(
f"Async SQLite connection requires 'sqlite+aiosqlite://' prefix. Got: {dsn[:50]}...",
suggestion="Use 'sqlite+aiosqlite:///path/to/db.db' for async SQLite connections.",
)
elif dsn_lower.startswith("postgresql"):
if is_async and "+asyncpg" not in dsn_lower:
raise DatabaseConnectionError(
f"Async PostgreSQL connection requires 'postgresql+asyncpg://' prefix. Got: {dsn[:50]}...",
suggestion="Use 'postgresql+asyncpg://user:pass@host:port/dbname' for async PostgreSQL connections.",
)
elif dsn_lower.startswith("mysql"):
if is_async and "+aiomysql" not in dsn_lower:
raise DatabaseConnectionError(
f"Async MySQL connection requires 'mysql+aiomysql://' prefix. Got: {dsn[:50]}...",
suggestion="Use 'mysql+aiomysql://user:pass@host:port/dbname' for async MySQL connections.",
)
if "://" not in dsn:
raise DatabaseConnectionError(
f"Connection string must include '://' separator. Got: {dsn[:50]}...",
suggestion="Connection strings should follow the format: 'dialect://user:pass@host:port/dbname'",
)
[docs]
def create_config(
dsn: str | None = None,
engine: Engine | "AsyncEngine | None" = None,
session: object | None = None,
**kwargs: EngineOptionValue,
) -> MoltresConfig:
"""Convenience helper used by ``moltres.connect``.
Supports environment variables for configuration:
- MOLTRES_DSN: :class:`Database` connection string
- MOLTRES_ECHO: Enable SQLAlchemy echo mode (true/false)
- MOLTRES_FETCH_FORMAT: "records", "pandas", or "polars"
- MOLTRES_DIALECT: Override SQL dialect detection
- MOLTRES_POOL_SIZE: Connection pool size
- MOLTRES_MAX_OVERFLOW: Maximum pool overflow connections
- MOLTRES_POOL_TIMEOUT: Pool timeout in seconds
- MOLTRES_POOL_RECYCLE: Connection recycle time in seconds
- MOLTRES_POOL_PRE_PING: Enable connection health checks (true/false)
- MOLTRES_QUERY_TIMEOUT: Query execution timeout in seconds
- MOLTRES_ALLOWED_PATHS: OS-separated list of allowed filesystem roots for file I/O
Args:
dsn: :class:`Database` connection string (e.g., "sqlite:///example.db").
If None, will try MOLTRES_DSN environment variable.
Cannot be provided if engine or session is provided.
engine: SQLAlchemy Engine instance to use. If provided, dsn is ignored.
This gives users more flexibility to configure the engine themselves.
Cannot be provided if session is provided.
session: SQLAlchemy Session or AsyncSession instance to use. If provided,
dsn and engine are ignored. The session's bind (engine) will be used.
Cannot be provided if dsn or engine is provided.
**kwargs: Additional configuration options. Valid keys include:
- echo: Enable SQLAlchemy echo mode
- fetch_format: "records", "pandas", or "polars"
- dialect: Override SQL dialect detection
- pool_size: Connection pool size (ignored if engine is provided)
- max_overflow: Maximum pool overflow connections (ignored if engine is provided)
- pool_timeout: Pool timeout in seconds (ignored if engine is provided)
- pool_recycle: Connection recycle time in seconds (ignored if engine is provided)
- pool_pre_ping: Enable connection health checks (ignored if engine is provided)
- query_timeout: Query execution timeout in seconds
- allowed_paths: Tuple of directory roots permitted for file read/write paths
- future: Use SQLAlchemy 2.0 style (default: True)
- Other options are stored in config.options
Returns:
:class:`MoltresConfig`: MoltresConfig instance with parsed configuration
Raises:
ValueError: If neither dsn, engine, nor session is provided and MOLTRES_DSN is not set
ValueError: If multiple of dsn, engine, and session are provided
"""
# Check if session is provided in kwargs (for backward compatibility)
if session is None and "session" in kwargs:
session = kwargs.pop("session")
# Check if engine is provided in kwargs (for backward compatibility)
if engine is None and "engine" in kwargs:
engine_obj = kwargs.pop("engine")
if not isinstance(engine_obj, Engine):
raise TypeError("engine must be a SQLAlchemy Engine instance")
engine = engine_obj
# Get DSN from argument or environment variable (only if engine and session are not provided)
if engine is None and session is None:
if dsn is None:
dsn = os.environ.get("MOLTRES_DSN")
if dsn is None:
raise ValueError(
"Either 'dsn' or 'engine' must be provided as argument, or MOLTRES_DSN environment variable must be set"
)
# Normalize SQLite paths: convert backslashes to forward slashes for URLs
# SQLite URLs always use forward slashes, even on Windows
if dsn and (dsn.startswith("sqlite:///") or dsn.startswith("sqlite+aiosqlite:///")):
# Replace backslashes with forward slashes in the path part
dsn = dsn.replace("\\", "/")
if dsn is not None:
validate_connection_string(dsn, is_async=False)
else:
# If engine or session is provided, ignore dsn
if dsn is not None:
raise ValueError(
"Cannot provide 'dsn' with 'engine' or 'session'. "
"Provide either a connection string, an Engine instance, or a Session instance."
)
dsn = None
# Load configuration from environment variables if not provided in kwargs
env_config = _load_env_config()
# Merge: kwargs override env vars, env vars override defaults
merged_kwargs = {**env_config, **kwargs}
# If an engine or session object is provided, infer dialect name unless explicitly overridden
inferred_engine_dialect: str | None = None
if engine is not None:
inferred_engine_dialect = getattr(getattr(engine, "dialect", None), "name", None)
if inferred_engine_dialect and "+" in inferred_engine_dialect:
# Normalize driver variants similar to DSN parsing (e.g., "mysql+aiomysql")
inferred_engine_dialect = inferred_engine_dialect.split("+", 1)[0]
elif session is not None:
# Extract engine from session to infer dialect
if hasattr(session, "get_bind"):
bind = session.get_bind()
elif hasattr(session, "bind"):
bind = session.bind
else:
bind = None
if bind is not None:
inferred_engine_dialect = getattr(getattr(bind, "dialect", None), "name", None)
if inferred_engine_dialect and "+" in inferred_engine_dialect:
# Normalize driver variants similar to DSN parsing (e.g., "mysql+aiomysql")
inferred_engine_dialect = inferred_engine_dialect.split("+", 1)[0]
if "dialect" not in merged_kwargs and inferred_engine_dialect:
merged_kwargs["dialect"] = inferred_engine_dialect
engine_kwargs: dict[str, object] = {
k: merged_kwargs.pop(k)
for k in list(merged_kwargs)
if k in EngineConfig.__dataclass_fields__
}
# Construct EngineConfig with validated kwargs
# Using **kwargs with type checking would be ideal, but dataclass doesn't support it directly
# This approach extracts known fields and passes them safely
engine_config = EngineConfig(
dsn=dsn, engine=engine, session=session, **cast(Any, engine_kwargs)
)
allowed_paths = merged_kwargs.pop("allowed_paths", None)
if allowed_paths is not None and not isinstance(allowed_paths, tuple):
if isinstance(allowed_paths, (list, set)):
allowed_paths = tuple(str(p) for p in allowed_paths)
else:
allowed_paths = (str(allowed_paths),)
return MoltresConfig(
engine=engine_config,
options=merged_kwargs,
allowed_paths=cast("tuple[str, ...] | None", allowed_paths),
)
DEFAULT_CONFIG = create_config("sqlite:///:memory:")