75 lines
2.1 KiB
Python
75 lines
2.1 KiB
Python
from sqlalchemy import create_engine, text
|
|
from sqlalchemy.orm import sessionmaker, Session
|
|
from contextlib import contextmanager
|
|
from typing import Generator
|
|
import os
|
|
|
|
|
|
import logging
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Get database URL from environment variable with fallback
|
|
DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///omegacard.db")
|
|
|
|
# Create engine with proper configuration
|
|
engine = create_engine(
|
|
DATABASE_URL,
|
|
pool_pre_ping=True, # Enable connection health checks
|
|
pool_size=5, # Set reasonable pool size
|
|
max_overflow=10 # Allow some overflow connections
|
|
)
|
|
|
|
# Create session factory
|
|
SessionLocal = sessionmaker(
|
|
bind=engine,
|
|
autocommit=False,
|
|
autoflush=False
|
|
)
|
|
|
|
@contextmanager
|
|
def get_db_session() -> Generator[Session, None, None]:
|
|
"""Context manager for database sessions"""
|
|
session = SessionLocal()
|
|
try:
|
|
yield session
|
|
except Exception as e:
|
|
logger.error(f"Database session error: {str(e)}")
|
|
session.rollback()
|
|
raise
|
|
finally:
|
|
session.close()
|
|
|
|
def get_db() -> Generator[Session, None, None]:
|
|
"""Dependency for FastAPI to get database sessions"""
|
|
with get_db_session() as session:
|
|
yield session
|
|
|
|
def init_db() -> None:
|
|
"""Initialize database tables"""
|
|
from .models import Base
|
|
try:
|
|
Base.metadata.create_all(bind=engine)
|
|
logger.info("Database tables created successfully")
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize database: {str(e)}")
|
|
raise
|
|
|
|
def check_db_connection() -> bool:
|
|
"""Check if database connection is working"""
|
|
try:
|
|
with get_db_session() as session:
|
|
session.execute(text("SELECT 1"))
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Database connection check failed: {str(e)}")
|
|
return False
|
|
|
|
def destroy_db() -> None:
|
|
"""Destroy all database tables"""
|
|
from .models import Base
|
|
try:
|
|
Base.metadata.drop_all(bind=engine)
|
|
logger.info("Database tables dropped successfully")
|
|
except Exception as e:
|
|
logger.error(f"Failed to destroy database: {str(e)}")
|
|
raise |