Spaces:
Sleeping
Sleeping
File size: 5,342 Bytes
8e1643b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 | """
Database configuration and session management.
"""
import os
from pathlib import Path
from contextlib import contextmanager
from typing import Generator
from sqlalchemy import create_engine, event
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.pool import StaticPool
from .models import Base
class DatabaseManager:
"""
Manages database connections and sessions.
Singleton pattern ensures only one database connection pool exists.
"""
_instance = None
_engine = None
_session_factory = None
def __new__(cls):
if cls._instance is None:
cls._instance = super().__new__(cls)
return cls._instance
def __init__(self, db_path: str = None, echo: bool = False):
"""
Initialize database manager.
Args:
db_path: Path to SQLite database file. If None, uses default from config.
echo: If True, SQL statements are logged (useful for debugging)
"""
if self._engine is None:
if db_path is None:
# Use default path from project root
project_root = Path(__file__).parent.parent.parent
db_dir = project_root / 'data'
db_dir.mkdir(exist_ok=True)
db_path = str(db_dir / 'pdf_visualizer.db')
# Create engine
self._engine = self._create_engine(db_path, echo)
# Enable foreign keys for SQLite
@event.listens_for(self._engine, "connect")
def set_sqlite_pragma(dbapi_conn, connection_record):
cursor = dbapi_conn.cursor()
cursor.execute("PRAGMA foreign_keys=ON")
cursor.close()
# Create session factory
self._session_factory = sessionmaker(bind=self._engine)
# Create all tables
self.create_tables()
@staticmethod
def _create_engine(db_path: str, echo: bool = False):
"""
Create SQLAlchemy engine.
Args:
db_path: Path to database file
echo: If True, log SQL statements
Returns:
SQLAlchemy engine
"""
# SQLite connection string
database_url = f"sqlite:///{db_path}"
# Create engine
# Use StaticPool for SQLite to avoid threading issues
engine = create_engine(
database_url,
echo=echo,
connect_args={'check_same_thread': False},
poolclass=StaticPool
)
return engine
def create_tables(self):
"""Create all tables in the database."""
Base.metadata.create_all(self._engine)
def drop_tables(self):
"""Drop all tables (use with caution!)."""
Base.metadata.drop_all(self._engine)
def get_session(self) -> Session:
"""
Get a new database session.
Returns:
SQLAlchemy session
Note:
Caller is responsible for closing the session.
"""
return self._session_factory()
@contextmanager
def session_scope(self) -> Generator[Session, None, None]:
"""
Provide a transactional scope around database operations.
Usage:
with db_manager.session_scope() as session:
session.add(obj)
# Changes are automatically committed on exit
# Or rolled back on exception
Yields:
SQLAlchemy session
"""
session = self.get_session()
try:
yield session
session.commit()
except Exception:
session.rollback()
raise
finally:
session.close()
def get_engine(self):
"""Get the SQLAlchemy engine."""
return self._engine
def __repr__(self):
return f"<DatabaseManager(engine={self._engine})>"
# Convenience function for getting a session
def get_db_session() -> Session:
"""
Get a new database session.
Returns:
SQLAlchemy session
Example:
session = get_db_session()
try:
snapshots = session.query(PDFSnapshot).all()
finally:
session.close()
"""
db_manager = DatabaseManager()
return db_manager.get_session()
# Convenience context manager
@contextmanager
def db_session() -> Generator[Session, None, None]:
"""
Context manager for database sessions.
Usage:
with db_session() as session:
snapshots = session.query(PDFSnapshot).all()
Yields:
SQLAlchemy session
"""
db_manager = DatabaseManager()
with db_manager.session_scope() as session:
yield session
if __name__ == "__main__":
# Test database setup
print("Testing database configuration...")
# Initialize database manager
db_manager = DatabaseManager(echo=True)
print(f"✅ Database manager created: {db_manager}")
# Check if tables were created
from sqlalchemy import inspect
inspector = inspect(db_manager.get_engine())
tables = inspector.get_table_names()
print(f"\n✅ Tables created: {tables}")
# Test session creation
with db_manager.session_scope() as session:
print(f"\n✅ Session created: {session}")
print("\n✅ All database configuration tests passed!")
|