from __future__ import annotations
from typing import Union
from dynaconf import Dynaconf
from pydantic import Field, ValidationError, field_validator
from pydantic_settings import BaseSettings
import sqlalchemy as sa
import sqlalchemy.orm as so
## Supported databases
valid_db_types: list[str] = ["sqlite", "postgres", "mssql"]
## Load database settings from environment
DYNACONF_DB_SETTINGS = Dynaconf(
environments=True,
## If you aren't using [dev] and [prod] envs,
# uncomment line below and add a [db] section to your settings.local.toml
# env="db",
envvar_prefix="DB",
settings_files=["settings.toml", ".secrets.toml"]
)
class DBSettings(BaseSettings):
"""Store database configuration.
Params:
type (str): [required] The type of database backend, i.e. `sqlite`, `postgres`, `mysql`, etc.
drivername (str): [required] The SQLAlchemy drivername string. Examples:
`['sqlite+pysqlite', 'postgresql+psycopg2', 'mysql+pymysql', 'mssql+pyodbc']`
username (str): The database user's username.
**Note**: Some OSes will use the logged-in user's $USERNAME, even when this value is `None`. If you have
unexpected results, try changing this to `user`. YOU MUST UPDATE THIS EVERYWHERE, in any function that
references `DBSettings.username`.
password (str): The database user's password.
host (str): The host address/IP/FQDN of the database server.
port (int): The port where the database is listening on the remote server.
database (str): The database name (or path, if using `SQLite`). Example:
(assumes sqlite) `./path/to/app.sqlite`
echo (bool): When `True`, the SQLAlchemy `Engine` will print its output to the console.
"""
type: str = Field(default=DYNACONF_DB_SETTINGS.DB_TYPE, env="DB_TYPE")
drivername: str = Field(
default=DYNACONF_DB_SETTINGS.DB_DRIVERNAME, env="DB_DRIVERNAME"
)
## If DBSettings is using your Windows/Unix username, try changing 'user' -> 'username' or vice-versa.
username: str | None = Field(
default=DYNACONF_DB_SETTINGS.DB_USERNAME, env="DB_USERNAME"
)
password: str | None = Field(
default=DYNACONF_DB_SETTINGS.DB_PASSWORD, env="DB_PASSWORD", repr=False
)
host: str | None = Field(default=DYNACONF_DB_SETTINGS.DB_HOST, env="DB_HOST")
port: Union[str, int, None] = Field(
default=DYNACONF_DB_SETTINGS.DB_PORT, env="DB_PORT"
)
database: str = Field(default=DYNACONF_DB_SETTINGS.DB_DATABASE, env="DB_DATABASE")
echo: bool = Field(default=DYNACONF_DB_SETTINGS.DB_ECHO, env="DB_ECHO")
@field_validator("port")
def validate_db_port(cls, v) -> int:
if v is None or v == "":
return None
elif isinstance(v, int):
return v
elif isinstance(v, str):
return int(v)
else:
raise ValidationError
def get_db_uri(self) -> sa.URL:
"""Build SQLAlchemy database URI.
Returns:
(sqlalchemy.URL): A formatted SQLAlchemy `URL` class instance.
"""
try:
_uri: sa.URL = sa.URL.create(
drivername=self.drivername,
username=self.username,
password=self.password,
host=self.host,
port=self.port,
database=self.database,
)
return _uri
except Exception as exc:
msg = Exception(
f"Unhandled exception getting SQLAlchemy database URL. Details: {exc}"
)
raise msg
def get_engine(self) -> sa.Engine:
"""Build a SQLAlchemy `Engine` object.
Returns:
(sqlalchemy.Engine): An initialized `Engine` object.
"""
assert self.get_db_uri() is not None, ValueError("db_uri is not None")
assert isinstance(self.get_db_uri(), sa.URL), TypeError(
f"db_uri must be of type sqlalchemy.URL. Got type: ({type(self.get_db_uri())})"
)
try:
engine: sa.Engine = sa.create_engine(
url=self.get_db_uri().render_as_string(hide_password=False),
echo=self.echo,
)
return engine
except Exception as exc:
msg = Exception(
f"Unhandled exception getting database engine. Details: {exc}"
)
raise msg
def get_session_pool(self) -> so.sessionmaker[so.Session]:
"""Build a SQLAlchemy `Session` pool.
Usage:
Create a variable, like `session_pool = DBSettings.get_session_pool()`. Then, call the session
pool as a context manager like: `with sesion_pool() as session: ...`
Returns:
(sqlalchemy.orm.sessionmaker[sqlalchemy.orm.Session]): An initialized pool of database sessions.
"""
engine: sa.Engine = self.get_engine()
assert engine is not None, ValueError("engine cannot be None")
assert isinstance(engine, sa.Engine), TypeError(
f"engine must be of type sqlalchemy.Engine. Got type: ({type(engine)})"
)
session_pool: so.sessionmaker[so.Session] = so.sessionmaker(bind=engine)
return session_pool
db_settings: DBSettings = DBSettings()