Phase 1: initial codes

stockmind/db/database.py

@@ -0,0 +1,39 @@
+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()