Chapter 6: Storing Messages with SQLite

pipenv install sqlalchemy aiosqlite

# app/database.py

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, declarative_base

# Define the database URL
# For SQLite, this points to a file. For production, you'd use PostgreSQL, etc.
DATABASE_URL = "sqlite:///./chat.db"

# Create the SQLAlchemy engine
# connect_args are specific to SQLite and ensure only one thread can communicate
# if using synchronous operations (though we'll use async)
engine = create_engine(
    DATABASE_URL, connect_args={"check_same_thread": False}, echo=False # Set echo=True for SQL logs
)

# Create a SessionLocal class
# This will be used to create session instances for database interactions
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

# Base class for our ORM models
Base = declarative_base()

# Dependency to get a database session
def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

# app/models.py

from sqlalchemy import Column, Integer, String, DateTime, ForeignKey
from sqlalchemy.orm import relationship
from datetime import datetime

from .database import Base

class User(Base):
    __tablename__ = "users"

    id = Column(Integer, primary_key=True, index=True)
    username = Column(String, unique=True, index=True)
    hashed_password = Column(String)

    messages = relationship("Message", back_populates="owner")

class Message(Base):
    __tablename__ = "messages"

    id = Column(Integer, primary_key=True, index=True)
    content = Column(String, index=True)
    timestamp = Column(DateTime, default=datetime.utcnow)
    owner_id = Column(Integer, ForeignKey("users.id"))

    owner = relationship("User", back_populates="messages")

# app/main.py (TEMPORARY ADDITION for database creation)
# Add this BEFORE app = FastAPI() line

from .database import Base, engine # New import

# This will create the database tables defined in models.py
Base.metadata.create_all(bind=engine)

# (rest of your app/main.py content)

pipenv shell
python -c "from app.database import Base, engine; from app import models; Base.metadata.create_all(bind=engine)"

# app/schemas.py

from pydantic import BaseModel
from datetime import datetime

# Pydantic models for data validation and serialization

class UserCreate(BaseModel):
    username: str
    password: str

class UserResponse(BaseModel):
    id: int
    username: str

    class Config:
        from_attributes = True # Or orm_mode = True for Pydantic v1.x

class MessageCreate(BaseModel):
    content: str

class MessageResponse(BaseModel):
    id: int
    content: str
    timestamp: datetime
    owner_id: int
    owner: UserResponse # Include owner details

    class Config:
        from_attributes = True # Or orm_mode = True for Pydantic v1.x

# app/main.py (further updated)

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends, HTTPException, status
from fastapi.security import OAuth2PasswordRequestForm
from typing import Dict, List
from sqlalchemy.orm import Session # New import

from .auth import Hasher, create_access_token, get_current_user, ACCESS_TOKEN_EXPIRE_MINUTES, Token
from .connections import ConnectionManager
from .database import get_db # New import
from . import models, schemas # New import
from datetime import timedelta

app = FastAPI()

manager = ConnectionManager()

# --- User Registration Endpoint ---
@app.post("/register", response_model=schemas.UserResponse, status_code=status.HTTP_201_CREATED)
async def register_user(user: schemas.UserCreate, db: Session = Depends(get_db)):
    db_user = db.query(models.User).filter(models.User.username == user.username).first()
    if db_user:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Username already registered")

    hashed_password = Hasher.get_password_hash(user.password)
    db_user = models.User(username=user.username, hashed_password=hashed_password)
    db.add(db_user)
    db.commit()
    db.refresh(db_user)
    return db_user

@app.post("/token", response_model=Token)
async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends(), db: Session = Depends(get_db)):
    user = db.query(models.User).filter(models.User.username == form_data.username).first()
    if not user or not Hasher.verify_password(form_data.password, user.hashed_password):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect username or password",
            headers={"WWW-Authenticate": "Bearer"},
        )
    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    access_token = create_access_token(
        data={"sub": user.username}, expires_delta=access_token_expires
    )
    return {"access_token": access_token, "token_type": "bearer"}

# Update get_current_user to return the actual User model, not just username
from .auth import SECRET_KEY, ALGORITHM, ACCESS_TOKEN_EXPIRE_MINUTES, TokenData
async def get_current_user_db(token: str = Depends(oauth2_scheme), db: Session = Depends(get_db)):
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            raise credentials_exception
        token_data = TokenData(username=username)
    except JWTError:
        raise credentials_exception
    user = db.query(models.User).filter(models.User.username == token_data.username).first()
    if user is None:
        raise credentials_exception
    return user

@app.get("/users/me", response_model=schemas.UserResponse)
async def read_users_me(current_user: models.User = Depends(get_current_user_db)):
    return current_user

# --- Chat History Endpoint ---
@app.get("/messages", response_model=List[schemas.MessageResponse])
async def get_chat_history(db: Session = Depends(get_db), current_user: models.User = Depends(get_current_user_db), skip: int = 0, limit: int = 100):
    messages = db.query(models.Message).order_by(models.Message.timestamp.desc()).offset(skip).limit(limit).all()
    # Eagerly load owner to avoid N+1 queries if owner details are needed
    return messages

# Updated WebSocket endpoint to save messages
@app.websocket("/ws/{client_id}") # Renamed for clarity, original was fine too
async def websocket_chat(websocket: WebSocket, client_id: str, db: Session = Depends(get_db)): # Add db dependency
    # NOTE: For production, authenticate WebSocket connection with JWT too.
    # This example skips that for simplicity in message persistence, but it's crucial.
    # You would parse a token from a query parameter or header and validate it here.

    await manager.connect(websocket)
    welcome_message = f"Client #{client_id} joined the chat."
    await manager.broadcast(welcome_message)

    try:
        while True:
            data = await websocket.receive_text()
            full_message = f"Client #{client_id} says: {data}"

            # Save message to database
            new_message = models.Message(content=data, owner_id=1) # For now, hardcode owner_id=1 (e.g., testuser)
                                                                 # In next chapter, we'll use actual authenticated user
            db.add(new_message)
            db.commit()
            db.refresh(new_message)

            await manager.broadcast(full_message)
    except WebSocketDisconnect:
        manager.disconnect(websocket)
        await manager.broadcast(f"Client #{client_id} left the chat.")