Files
iotDashboard/services/device_manager/database.py

84 lines
1.8 KiB
Python

"""
Database session management for FastAPI with SQLAlchemy.
Uses dependency injection pattern for database sessions.
"""
from collections.abc import Generator
from contextlib import contextmanager
from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker
from config import config
# Create engine with connection pooling
engine = create_engine(
config.DATABASE_URL,
pool_pre_ping=True,
pool_size=5,
max_overflow=10,
echo=False,
)
SessionLocal = sessionmaker(
autocommit=False,
autoflush=False,
bind=engine,
)
def get_db() -> Generator[Session]:
"""
FastAPI dependency that provides a database session.
Usage in endpoints:
@app.post("/devices")
async def create_device(db: Session = Depends(get_db)):
device = Device(...)
db.add(device)
db.commit()
return device
The session is automatically closed after the request completes.
"""
db = SessionLocal()
try:
yield db
finally:
db.close()
@contextmanager
def get_db_context():
"""
Context manager for database sessions outside of FastAPI endpoints.
Usage:
with get_db_context() as db:
device = db.query(Device).first()
"""
db = SessionLocal()
try:
yield db
db.commit()
except Exception:
db.rollback()
raise
finally:
db.close()
def check_db_connection() -> bool:
"""
Check if database connection is working.
Returns:
True if connection successful, False otherwise
"""
try:
with engine.connect() as conn:
conn.execute("SELECT 1")
return True
except Exception as e:
print(f"Database connection failed: {e}")
return False