1  """Database session management."""
 2  from typing import Generator
 3  
 4  from sqlalchemy import create_engine
 5  from sqlalchemy.ext.declarative import declarative_base
 6  from sqlalchemy.orm import sessionmaker, Session
 7  
 8  from ..core.config import settings
 9  
10  # Create database engine
11  engine = create_engine(
12      settings.DATABASE_URI,
13      pool_pre_ping=True,
14      pool_size=20,
15      max_overflow=10,
16      pool_recycle=3600,
17  )
18  
19  # Create session factory
20  SessionLocal = sessionmaker(
21      autocommit=False,
22      autoflush=False,
23      bind=engine,
24      expire_on_commit=False,
25  )
26  
27  # Base class for models
28  Base = declarative_base()
29  
30  def get_db() -> Generator[Session, None, None]:
31      """Dependency function that yields database sessions.
32      
33      Yields:
34          Session: A database session
35      """
36      db = SessionLocal()
37      try:
38          yield db
39      finally:
40          db.close()
41  
42  def init_db() -> None:
43      """Initialize database tables."""
44      Base.metadata.create_all(bind=engine)
45  
46  def get_session() -> Session:
47      """Get a new database session."""
48      return SessionLocal()