74 lines
2.2 KiB
Python
74 lines
2.2 KiB
Python
"""
|
|
Shared SQLAlchemy engine / session + schema bootstrap
|
|
"""
|
|
from __future__ import annotations
|
|
import importlib, datetime
|
|
import os
|
|
|
|
from sqlalchemy import (
|
|
create_engine, event, Table, Column, Text, DateTime, Index, make_url
|
|
)
|
|
from sqlalchemy.engine import url
|
|
from sqlalchemy.orm import sessionmaker, registry
|
|
from backend.core.settings import (
|
|
SQLALCHEMY_DATABASE_URI, DB_POOL_SIZE, DB_ECHO, SKIP_SCHEMA_BOOTSTRAP
|
|
)
|
|
|
|
IS_PG = SQLALCHEMY_DATABASE_URI.startswith("postgresql")
|
|
|
|
parsed_url = make_url(SQLALCHEMY_DATABASE_URI) # string into URL object
|
|
|
|
# engine & session
|
|
connect_args = {"sslmode": "require"} if parsed_url.drivername.startswith("postgresql") else {}
|
|
engine = create_engine(
|
|
SQLALCHEMY_DATABASE_URI,
|
|
pool_size = DB_POOL_SIZE,
|
|
max_overflow = 20,
|
|
pool_timeout = 30,
|
|
echo = DB_ECHO,
|
|
future = True,
|
|
pool_pre_ping = True,
|
|
pool_recycle=3600,
|
|
connect_args = connect_args
|
|
)
|
|
|
|
# SQLite -> WAL for concurrency
|
|
if SQLALCHEMY_DATABASE_URI.startswith("sqlite:///"):
|
|
|
|
@event.listens_for(engine, "connect")
|
|
def _set_wal(dbapi_conn, _):
|
|
dbapi_conn.execute("PRAGMA journal_mode=WAL;")
|
|
|
|
SessionLocal = sessionmaker(bind=engine, autoflush=False,
|
|
expire_on_commit=False, future=True)
|
|
|
|
# metadata (tables from every module)
|
|
mapper_registry = registry()
|
|
metadata = mapper_registry.metadata
|
|
|
|
# download-cache table
|
|
download_cache = Table(
|
|
"download_cache", metadata,
|
|
Column("key", Text, primary_key=True),
|
|
Column("path", Text, nullable=False),
|
|
Column("ext", Text, nullable=False),
|
|
Column("created_at", DateTime, default=datetime.datetime.utcnow,
|
|
nullable=False, index=True),
|
|
)
|
|
Index("ix_download_cache_created", download_cache.c.created_at)
|
|
|
|
# auto-bootstrap all
|
|
def _bootstrap_schema() -> None:
|
|
"""Import modules then create."""
|
|
table_modules = (
|
|
"backend.core.db_xp",
|
|
"backend.web.db_extra",
|
|
"backend.core.formats",
|
|
)
|
|
for mod in table_modules:
|
|
importlib.import_module(mod)
|
|
|
|
metadata.create_all(engine)
|
|
if SKIP_SCHEMA_BOOTSTRAP != "1":
|
|
_bootstrap_schema()
|