Spaces:
Sleeping
Sleeping
File size: 3,852 Bytes
b64de39 | 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 | """
Database Configuration and Session Management
"""
import logging
from contextlib import contextmanager
from sqlalchemy import create_engine, event, exc
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.pool import QueuePool, NullPool
from sqlalchemy.engine import Engine
from app.config import settings
logger = logging.getLogger(__name__)
# Determine pool class based on database type
# SQLite doesn't support connection pooling, use NullPool
if settings.DATABASE_URL.startswith("sqlite"):
poolclass = NullPool
connect_args = {"check_same_thread": False}
pool_pre_ping = False
else:
poolclass = QueuePool
connect_args = {}
pool_pre_ping = True
# Build engine kwargs; NullPool does not accept pool_size/max_overflow/pool_timeout
_engine_kwargs: dict = dict(
poolclass=poolclass,
pool_pre_ping=pool_pre_ping,
echo=settings.DEBUG,
connect_args=connect_args,
)
if poolclass is not NullPool:
_engine_kwargs.update(
pool_size=10,
max_overflow=20,
pool_recycle=3600,
pool_timeout=30,
)
engine = create_engine(settings.DATABASE_URL, **_engine_kwargs)
# Session factory
SessionLocal = sessionmaker(
autocommit=False,
autoflush=False,
bind=engine,
expire_on_commit=False
)
Base = declarative_base()
@event.listens_for(Engine, "connect")
def set_sqlite_pragma(dbapi_conn, connection_record):
"""
Set SQLite pragmas for better performance and foreign key support
"""
if settings.DATABASE_URL.startswith("sqlite"):
cursor = dbapi_conn.cursor()
cursor.execute("PRAGMA foreign_keys=ON")
cursor.execute("PRAGMA journal_mode=WAL")
cursor.execute("PRAGMA synchronous=NORMAL")
cursor.close()
@event.listens_for(Engine, "checkout")
def receive_checkout(dbapi_conn, connection_record, connection_proxy):
"""
Log connection checkout events
"""
if settings.DEBUG:
logger.debug("Connection checked out from pool")
def get_db():
"""
Database session dependency for FastAPI
Yields:
Session: SQLAlchemy database session
Example:
```python
@app.get("/items")
def get_items(db: Session = Depends(get_db)):
return db.query(Item).all()
```
"""
db = SessionLocal()
try:
yield db
except exc.SQLAlchemyError as e:
logger.error(f"Database error: {str(e)}")
db.rollback()
raise
finally:
db.close()
@contextmanager
def get_db_context():
"""
Context manager for database sessions (for use outside FastAPI)
Yields:
Session: SQLAlchemy database session
Example:
```python
with get_db_context() as db:
items = db.query(Item).all()
```
"""
db = SessionLocal()
try:
yield db
db.commit()
except exc.SQLAlchemyError as e:
logger.error(f"Database error: {str(e)}")
db.rollback()
raise
finally:
db.close()
def init_db():
"""
Initialize database by creating all tables
"""
try:
Base.metadata.create_all(bind=engine)
logger.info("Database tables created successfully")
except Exception as e:
logger.error(f"Error creating database tables: {str(e)}")
raise
def check_db_connection() -> bool:
"""
Check if database connection is healthy
Returns:
bool: True if connection is healthy, False otherwise
"""
try:
from sqlalchemy import text
with engine.connect() as conn:
conn.execute(text("SELECT 1"))
return True
except Exception as e:
logger.error(f"Database connection check failed: {str(e)}")
return False
|