54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
import os
|
|
from sqlalchemy import create_engine, event
|
|
from sqlalchemy.orm import sessionmaker, DeclarativeBase
|
|
from . import config, mode
|
|
|
|
|
|
def _sqlite_unicode_lower(conn, _):
|
|
# встроенный lower() в SQLite не знает кириллицу — подменяем, чтобы поиск не зависел от регистра
|
|
conn.create_function("lower", 1, lambda s: s.lower() if isinstance(s, str) else s)
|
|
|
|
|
|
def make_engine(url: str):
|
|
if url.startswith("sqlite"):
|
|
os.makedirs(config.DATA_DIR, exist_ok=True)
|
|
e = create_engine(url, connect_args={"check_same_thread": False})
|
|
event.listen(e, "connect", _sqlite_unicode_lower)
|
|
return e
|
|
return create_engine(url, pool_pre_ping=True)
|
|
|
|
|
|
real_engine = make_engine(config.DATABASE_URL) # боевая БД — с ней работают все, кто не в демо
|
|
engine = real_engine
|
|
SessionLocal = sessionmaker(bind=real_engine, autoflush=False, expire_on_commit=False) # всегда боевая БД
|
|
|
|
# демо-БД (персональный демо-режим): открывается только для запросов с демо-токеном
|
|
demo_engine = None
|
|
DemoSession = sessionmaker(autoflush=False, expire_on_commit=False)
|
|
|
|
|
|
def set_demo_engine(e):
|
|
global demo_engine
|
|
old, demo_engine = demo_engine, e
|
|
if e is not None:
|
|
DemoSession.configure(bind=e)
|
|
return old
|
|
|
|
|
|
def session_for(demo: bool):
|
|
s = DemoSession() if demo and demo_engine is not None else SessionLocal()
|
|
s.info["demo"] = bool(demo and demo_engine is not None)
|
|
return s
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
pass
|
|
|
|
|
|
def get_db():
|
|
db = session_for(mode.is_demo_request())
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close()
|