44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
import os
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
from sqlalchemy.orm import DeclarativeBase
|
|
|
|
|
|
DB_PATH = os.environ.get("LOTUSGEN_DB_PATH", os.path.join(os.path.dirname(__file__), "..", "lotusgen.db"))
|
|
DB_URL = f"sqlite+aiosqlite:///{os.path.abspath(DB_PATH)}"
|
|
|
|
engine = create_async_engine(
|
|
DB_URL,
|
|
echo=False,
|
|
connect_args={"check_same_thread": False},
|
|
)
|
|
|
|
AsyncSessionLocal = async_sessionmaker(
|
|
bind=engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False,
|
|
)
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
pass
|
|
|
|
|
|
async def init_db() -> None:
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
def migrate_slugs(connection):
|
|
proj_cols = [r[1] for r in connection.exec_driver_sql("PRAGMA table_info(projects)").fetchall()]
|
|
if "slug" not in proj_cols:
|
|
connection.exec_driver_sql("ALTER TABLE projects ADD COLUMN slug VARCHAR(255) DEFAULT ''")
|
|
tab_cols = [r[1] for r in connection.exec_driver_sql("PRAGMA table_info(tables)").fetchall()]
|
|
if "slug" not in tab_cols:
|
|
connection.exec_driver_sql("ALTER TABLE tables ADD COLUMN slug VARCHAR(255) DEFAULT ''")
|
|
|
|
await conn.run_sync(migrate_slugs)
|
|
|
|
|
|
async def get_db():
|
|
async with AsyncSessionLocal() as session:
|
|
yield session
|