Spaces:
Sleeping
Sleeping
File size: 2,004 Bytes
3a8b459 8807ee2 3a8b459 |
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 |
import logging
import os
from dotenv import load_dotenv
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import QueuePool
load_dotenv()
logger = logging.getLogger(__name__)
DATABASE_URL = os.getenv("DATABASE_URL")
if not DATABASE_URL:
# Fallback for local dev if not set
logger.warning("DATABASE_URL not set, using SQLite fallback")
DATABASE_URL = "sqlite:///./test.db"
# Configure engine with connection pooling and SSL for Neon
engine_args = {
"pool_pre_ping": True, # Verify connections before using
"pool_recycle": 300, # Recycle connections every 5 minutes
}
# Only add pooling for non-SQLite databases
if not DATABASE_URL.startswith("sqlite"):
engine_args.update({
"poolclass": QueuePool,
"pool_size": 5,
"max_overflow": 10,
})
# SSL is handled via the connection string for Neon (sslmode=require)
logger.info("Using PostgreSQL with connection pooling")
else:
# SQLite needs check_same_thread=False for FastAPI
engine_args["connect_args"] = {"check_same_thread": False}
logger.info("Using SQLite fallback database")
engine = create_engine(DATABASE_URL, **engine_args)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()
def get_db():
"""
Database session dependency for FastAPI.
Yields a database session and ensures cleanup after request.
"""
db = SessionLocal()
try:
yield db
finally:
db.close()
def check_db_connection() -> bool:
"""
Check if database connection is working.
Returns True if connection successful, 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: {e}")
return False
|