66 lines
2.5 KiB
Python
66 lines
2.5 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import logging
|
||
|
|
from typing import TYPE_CHECKING
|
||
|
|
|
||
|
|
from sqlalchemy import event
|
||
|
|
from sqlalchemy.pool import StaticPool
|
||
|
|
from sqlmodel import Session, SQLModel, create_engine
|
||
|
|
|
||
|
|
if TYPE_CHECKING:
|
||
|
|
from collections.abc import Generator
|
||
|
|
|
||
|
|
from sqlalchemy.engine import Connection
|
||
|
|
|
||
|
|
|
||
|
|
class Database:
|
||
|
|
def __init__(self, database_url: str | None = None, *, echo: bool = False, connect_args: dict | None = None) -> None:
|
||
|
|
self._database_url = database_url or "sqlite:///:memory:"
|
||
|
|
|
||
|
|
default_connect = {"check_same_thread": False, "timeout": 30} if self._database_url.startswith("sqlite") else {}
|
||
|
|
merged_connect = {**default_connect, **(connect_args or {})}
|
||
|
|
|
||
|
|
if self._database_url == "sqlite:///:memory:":
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
logger.warning("Using in-memory SQLite database; all data will be lost when the application stops.")
|
||
|
|
self._engine = create_engine(self._database_url, echo=echo, connect_args=merged_connect, poolclass=StaticPool)
|
||
|
|
else:
|
||
|
|
self._engine = create_engine(self._database_url, echo=echo, connect_args=merged_connect)
|
||
|
|
|
||
|
|
if self._database_url.startswith("sqlite"):
|
||
|
|
|
||
|
|
def _enable_sqlite_pragmas(dbapi_conn: Connection, _connection_record: object) -> None:
|
||
|
|
try:
|
||
|
|
cur = dbapi_conn.cursor()
|
||
|
|
cur.execute("PRAGMA journal_mode=WAL;")
|
||
|
|
cur.execute("PRAGMA synchronous=NORMAL;")
|
||
|
|
cur.execute("PRAGMA foreign_keys=ON;")
|
||
|
|
cur.execute("PRAGMA busy_timeout=30000;")
|
||
|
|
cur.close()
|
||
|
|
except Exception:
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
logger.exception("Failed to set sqlite pragmas on new connection: ")
|
||
|
|
|
||
|
|
event.listen(self._engine, "connect", _enable_sqlite_pragmas)
|
||
|
|
|
||
|
|
def init_db(self) -> None:
|
||
|
|
pass
|
||
|
|
|
||
|
|
def get_session(self) -> Generator[Session, None, None]:
|
||
|
|
session = Session(self._engine)
|
||
|
|
try:
|
||
|
|
yield session
|
||
|
|
session.commit()
|
||
|
|
except Exception:
|
||
|
|
session.rollback()
|
||
|
|
raise
|
||
|
|
finally:
|
||
|
|
session.close()
|
||
|
|
|
||
|
|
def dispose(self) -> None:
|
||
|
|
self._engine.dispose()
|
||
|
|
|
||
|
|
|
||
|
|
def create_database(database_url: str | None = None, *, echo: bool = False, connect_args: dict | None = None) -> Database:
|
||
|
|
return Database(database_url, echo=echo, connect_args=connect_args)
|