Spaces:
Running
Running
File size: 2,842 Bytes
69be42f dc3879e 69be42f dc3879e 69be42f dc3879e 69be42f dc3879e 69be42f |
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 |
"""Database connection and session management.
[Task]: T010
[From]: specs/001-user-auth/plan.md
[Task]: T004
[From]: specs/004-ai-chatbot/plan.md
"""
from sqlmodel import create_engine, Session
from typing import Generator
from core.config import get_settings
settings = get_settings()
# Create database engine with connection pooling
# Optimized for conversation/message table queries in Phase III
# SQLite doesn't support connection pooling, so we conditionally apply parameters
is_sqlite = settings.database_url.startswith("sqlite:")
is_postgresql = settings.database_url.startswith("postgresql:") or settings.database_url.startswith("postgres://")
if is_sqlite:
# SQLite configuration (no pooling)
engine = create_engine(
settings.database_url,
echo=settings.environment == "development", # Log SQL in development
connect_args={"check_same_thread": False} # Allow multithreaded access
)
elif is_postgresql:
# PostgreSQL configuration with connection pooling
engine = create_engine(
settings.database_url,
echo=settings.environment == "development", # Log SQL in development
pool_pre_ping=True, # Verify connections before using
pool_size=10, # Number of connections to maintain
max_overflow=20, # Additional connections beyond pool_size
pool_recycle=3600, # Recycle connections after 1 hour (prevents stale connections)
pool_timeout=30, # Timeout for getting connection from pool
connect_args={
"connect_timeout": 10, # Connection timeout
}
)
else:
# Default configuration for other databases
engine = create_engine(
settings.database_url,
echo=settings.environment == "development",
pool_pre_ping=True
)
def get_session() -> Generator[Session, None, None]:
"""Get database session.
Yields:
Session: SQLModel database session
Example:
```python
@app.get("/users")
def read_users(session: Session = Depends(get_session)):
users = session.exec(select(User)).all()
return users
```
"""
with Session(engine) as session:
yield session
# Alias for compatibility with chat.py
get_db = get_session
def init_db():
"""Initialize database tables.
Creates all tables defined in SQLModel models.
Should be called on application startup.
"""
from sqlmodel import SQLModel
import models.user # Import models to register them with SQLModel
import models.task # Import task model
# Phase III: Import conversation and message models
try:
import models.conversation
import models.message
except ImportError:
# Models not yet created (Phase 2 pending)
pass
SQLModel.metadata.create_all(engine)
|