102 lines
3.0 KiB
Python
102 lines
3.0 KiB
Python
"""SQLAlchemy database setup."""
|
|
|
|
import logging
|
|
from collections.abc import AsyncGenerator
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
from platformdirs import user_data_dir
|
|
from sqlalchemy import URL
|
|
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
|
|
|
|
from spoolman import env
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
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
|
|
|
|
data_dir = Path(user_data_dir("spoolman"))
|
|
data_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
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,
|
|
)
|
|
|
|
|
|
class Database:
|
|
connection_url: URL
|
|
engine: Optional[AsyncEngine]
|
|
session_maker: Optional[async_sessionmaker[AsyncSession]]
|
|
|
|
def __init__(self: "Database", connection_url: URL) -> None:
|
|
"""Construct the Database wrapper and set config parameters."""
|
|
self.connection_url = connection_url
|
|
|
|
def connect(self: "Database") -> None:
|
|
"""Connect to the database."""
|
|
if env.get_logging_level() == logging.DEBUG:
|
|
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
|
|
|
|
connect_args = {}
|
|
if self.connection_url.drivername == "sqlite+aiosqlite":
|
|
connect_args["timeout"] = 60
|
|
|
|
self.engine = create_async_engine(
|
|
self.connection_url,
|
|
connect_args=connect_args,
|
|
pool_pre_ping=True,
|
|
)
|
|
self.session_maker = async_sessionmaker(self.engine, autocommit=False, autoflush=True, expire_on_commit=False)
|
|
|
|
|
|
__db: Optional[Database] = None
|
|
|
|
|
|
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()
|
|
|
|
|
|
async def get_db_session() -> AsyncGenerator[AsyncSession, None]:
|
|
"""Get a DB session to be used with FastAPI's dependency system."""
|
|
if __db is None or __db.session_maker is None:
|
|
raise RuntimeError("DB is not setup.")
|
|
async with __db.session_maker() as session:
|
|
try:
|
|
yield session
|
|
await session.commit()
|
|
except Exception as exc:
|
|
await session.rollback()
|
|
raise exc
|
|
finally:
|
|
await session.close()
|