40 lines
939 B
Python
40 lines
939 B
Python
import os
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import sessionmaker, Session
|
|
from db.models import Base
|
|
|
|
_DB_PATH = Path(__file__).parent.parent.parent / "stockmind.db"
|
|
_ENGINE = None
|
|
_SessionLocal = None
|
|
|
|
|
|
def _get_engine():
|
|
global _ENGINE
|
|
if _ENGINE is None:
|
|
_ENGINE = create_engine(f"sqlite:///{_DB_PATH}", connect_args={"check_same_thread": False})
|
|
return _ENGINE
|
|
|
|
|
|
def init_db():
|
|
engine = _get_engine()
|
|
Base.metadata.create_all(engine)
|
|
global _SessionLocal
|
|
_SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False)
|
|
|
|
|
|
@contextmanager
|
|
def get_session() -> Session:
|
|
global _SessionLocal
|
|
if _SessionLocal is None:
|
|
init_db()
|
|
session = _SessionLocal()
|
|
try:
|
|
yield session
|
|
except Exception:
|
|
session.rollback()
|
|
raise
|
|
finally:
|
|
session.close()
|