NoSQL & Other Data Stores

import os
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Annotated, Any, TypedDict

from bson import ObjectId
from bson.errors import InvalidId
from fastapi.responses import JSONResponse
from fastapi import Depends, FastAPI, HTTPException, Request
from pydantic import BaseModel, BeforeValidator, Field
from pymongo import DESCENDING, AsyncMongoClient, ReturnDocument
from pymongo.asynchronous.database import AsyncDatabase
from pymongo.errors import PyMongoError


class State(TypedDict):
    mongo: AsyncMongoClient[dict[str, Any]]


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[State]:
    # The URI comes from the environment (a secret in production), never from the code.
    client: AsyncMongoClient[dict[str, Any]] = AsyncMongoClient(
        os.environ.get("MONGODB_URI", "mongodb://localhost:27017"),
        serverSelectionTimeoutMS=2000,
    )
    db = client[os.environ.get("MONGODB_DB", "bookshelf")]
    await db["reviews"].create_index("book_isbn")
    yield {"mongo": client}  # becomes request.state.mongo
    await client.close()


app = FastAPI(title="Bookshelf API", lifespan=lifespan)


@app.exception_handler(PyMongoError)
async def mongo_unavailable(request: Request, exc: PyMongoError) -> JSONResponse:
    return JSONResponse(status_code=503, content={"detail": "Review store unavailable"})


def get_db(request: Request) -> AsyncDatabase[dict[str, Any]]:
    return request.state.mongo[os.environ.get("MONGODB_DB", "bookshelf")]


DbDep = Annotated[AsyncDatabase[dict[str, Any]], Depends(get_db)]

# MongoDB's _id is an ObjectId: expose it as a string called "id".
PyObjectId = Annotated[str, BeforeValidator(str)]


class ReviewIn(BaseModel):
    reader: str = Field(min_length=1, max_length=60)
    rating: int = Field(ge=1, le=5)
    text: str = Field(default="", max_length=2000)
    tags: list[str] = []


class Review(ReviewIn):
    id: PyObjectId = Field(validation_alias="_id")
    book_isbn: str


def parse_object_id(review_id: str) -> ObjectId:
    try:
        return ObjectId(review_id)
    except InvalidId:
        raise HTTPException(status_code=404, detail="Review not found")


@app.post("/books/{isbn}/reviews", status_code=201)
async def add_review(isbn: str, review: ReviewIn, db: DbDep) -> Review:
    doc = {"book_isbn": isbn, **review.model_dump()}
    result = await db["reviews"].insert_one(doc)
    return Review.model_validate({**doc, "_id": result.inserted_id})


@app.get("/books/{isbn}/reviews")
async def list_reviews(isbn: str, db: DbDep, limit: int = 20) -> list[Review]:
    cursor = db["reviews"].find({"book_isbn": isbn}).sort("_id", DESCENDING).limit(limit)
    return [Review.model_validate(doc) for doc in await cursor.to_list(length=limit)]


@app.get("/reviews/{review_id}")
async def get_review(review_id: str, db: DbDep) -> Review:
    doc = await db["reviews"].find_one({"_id": parse_object_id(review_id)})
    if doc is None:
        raise HTTPException(status_code=404, detail="Review not found")
    return Review.model_validate(doc)


@app.patch("/reviews/{review_id}")
async def edit_review(review_id: str, patch: ReviewIn, db: DbDep) -> Review:
    doc = await db["reviews"].find_one_and_update(
        {"_id": parse_object_id(review_id)},
        {"$set": patch.model_dump()},
        return_document=ReturnDocument.AFTER,
    )
    if doc is None:
        raise HTTPException(status_code=404, detail="Review not found")
    return Review.model_validate(doc)


@app.delete("/reviews/{review_id}", status_code=204)
async def delete_review(review_id: str, db: DbDep) -> None:
    result = await db["reviews"].delete_one({"_id": parse_object_id(review_id)})
    if result.deleted_count == 0:
        raise HTTPException(status_code=404, detail="Review not found")


class RatingSummary(BaseModel):
    isbn: str = Field(validation_alias="_id")
    average: float
    count: int


@app.get("/books/{isbn}/rating")
async def rating(isbn: str, db: DbDep) -> RatingSummary:
    pipeline = [
        {"$match": {"book_isbn": isbn}},
        {"$group": {"_id": "$book_isbn", "average": {"$avg": "$rating"}, "count": {"$sum": 1}}},
    ]
    cursor = await db["reviews"].aggregate(pipeline)  # aggregate() is a coroutine in the async API
    docs = await cursor.to_list(length=1)
    if not docs:
        raise HTTPException(status_code=404, detail="No reviews")
    return RatingSummary.model_validate(docs[0])

export MONGODB_URI="mongodb://localhost:27017"
fastapi dev mongo_app.py
curl -X POST http://127.0.0.1:8000/books/978-0441172719/reviews \
  -H 'Content-Type: application/json' \
  -d '{"reader":"ana","rating":5,"text":"Great","tags":["sf"]}'
curl http://127.0.0.1:8000/books/978-0441172719/rating

{"reader":"ana","rating":5,"text":"Great","tags":["sf"],"id":"000000000000000000000001","book_isbn":"978-0441172719"}
{"isbn":"978-0441172719","average":3.5,"count":2}

import os
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Annotated, TypedDict

from fastapi import Depends, FastAPI, HTTPException, Request, Response
from pydantic import BaseModel
from redis import asyncio as redis


class State(TypedDict):
    redis: redis.Redis


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[State]:
    # One client = one connection pool, shared by every request.
    client = redis.Redis.from_url(
        os.environ.get("REDIS_URL", "redis://localhost:6379/0"),
        decode_responses=True,  # str in, str out
        max_connections=20,
    )
    await client.ping()  # fail fast at startup if Redis is unreachable
    yield {"redis": client}
    await client.aclose()  # redis-py >= 5.0.1; close() is deprecated


app = FastAPI(title="Bookshelf API", lifespan=lifespan)


def get_redis(request: Request) -> redis.Redis:
    return request.state.redis


RedisDep = Annotated[redis.Redis, Depends(get_redis)]


class Book(BaseModel):
    isbn: str
    title: str
    rating_average: float


DB: dict[str, Book] = {"978-0441172719": Book(isbn="978-0441172719", title="Dune", rating_average=4.5)}
DB_READS = 0  # counts how often the "database" is really hit


@app.get("/books/{isbn}")
async def get_book(isbn: str, r: RedisDep, response: Response) -> Book:
    global DB_READS
    key = f"book:{isbn}"
    cached = await r.get(key)
    if cached is not None:
        response.headers["X-Cache"] = "hit"
        return Book.model_validate_json(cached)
    DB_READS += 1
    book = DB.get(isbn)
    if book is None:
        raise HTTPException(status_code=404, detail="Book not found")
    await r.set(key, book.model_dump_json(), ex=60)  # expire after 60 seconds
    response.headers["X-Cache"] = "miss"
    return book


@app.put("/books/{isbn}")
async def update_book(isbn: str, book: Book, r: RedisDep) -> Book:
    DB[isbn] = book
    await r.delete(f"book:{isbn}")  # invalidate: the next read repopulates the cache
    return book


class RateLimit:
    """Fixed window: at most `limit` requests per `window` seconds per client and route."""

    def __init__(self, limit: int, window: int) -> None:
        self.limit = limit
        self.window = window

    async def __call__(self, request: Request, response: Response, r: RedisDep) -> None:
        client = request.client.host if request.client else "unknown"
        key = f"ratelimit:{request.url.path}:{client}"
        async with r.pipeline(transaction=True) as pipe:  # MULTI/EXEC
            pipe.incr(key)
            pipe.expire(key, self.window, nx=True)  # start the window only once
            pipe.ttl(key)
            count, _, ttl = await pipe.execute()
        response.headers["X-RateLimit-Remaining"] = str(max(self.limit - count, 0))
        if count > self.limit:
            raise HTTPException(
                status_code=429,
                detail="Too many requests",
                headers={"Retry-After": str(max(ttl, 1))},
            )


@app.get("/search", dependencies=[Depends(RateLimit(limit=3, window=60))])
async def search(q: str) -> dict[str, str]:
    return {"q": q}

curl -i http://127.0.0.1:8000/books/978-0441172719   # X-Cache: miss
curl -i http://127.0.0.1:8000/books/978-0441172719   # X-Cache: hit
curl -X PUT http://127.0.0.1:8000/books/978-0441172719 \
  -H 'Content-Type: application/json' \
  -d '{"isbn":"978-0441172719","title":"Dune","rating_average":4.6}'
curl -i http://127.0.0.1:8000/books/978-0441172719   # X-Cache: miss, rating_average 4.6

{"q":"dune"}
{"q":"dune"}
{"q":"dune"}
{"detail":"Too many requests"}

import hashlib
import math
import os
import re
import uuid
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Annotated, Protocol, TypedDict

from fastapi import Depends, FastAPI, Query, Request
from pydantic import BaseModel, Field
from qdrant_client import AsyncQdrantClient, models

COLLECTION = "reviews"
DIM = 64


class Embedder(Protocol):
    dimension: int

    async def embed(self, texts: list[str]) -> list[list[float]]: ...


class HashingEmbedder:
    """Toy embedder so the example runs offline: it matches shared words, not meaning.
    Replace it with a real embedding model (an API call or a local model) in production."""

    dimension = DIM

    async def embed(self, texts: list[str]) -> list[list[float]]:
        vectors = []
        for text in texts:
            vec = [0.0] * DIM
            for word in re.findall(r"[a-z']+", text.lower()):
                bucket = int(hashlib.sha256(word.encode()).hexdigest(), 16) % DIM
                vec[bucket] += 1.0
            norm = math.sqrt(sum(v * v for v in vec)) or 1.0
            vectors.append([v / norm for v in vec])
        return vectors


class State(TypedDict):
    qdrant: AsyncQdrantClient


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[State]:
    url = os.environ.get("QDRANT_URL")  # e.g. http://localhost:6333
    client = AsyncQdrantClient(url=url, api_key=os.environ.get("QDRANT_API_KEY")) if url else AsyncQdrantClient(":memory:")
    if not await client.collection_exists(COLLECTION):
        await client.create_collection(
            COLLECTION,
            vectors_config=models.VectorParams(size=DIM, distance=models.Distance.COSINE),
        )
    yield {"qdrant": client}
    await client.close()


app = FastAPI(title="Bookshelf API", lifespan=lifespan)


def get_qdrant(request: Request) -> AsyncQdrantClient:
    return request.state.qdrant


def get_embedder() -> Embedder:
    return HashingEmbedder()


QdrantDep = Annotated[AsyncQdrantClient, Depends(get_qdrant)]
EmbedderDep = Annotated[Embedder, Depends(get_embedder)]


class ReviewIn(BaseModel):
    book_isbn: str
    rating: int = Field(ge=1, le=5)
    text: str = Field(min_length=3)


class ReviewHit(BaseModel):
    id: str
    score: float
    book_isbn: str
    rating: int
    text: str


@app.post("/semantic/reviews", status_code=201)
async def index_review(review: ReviewIn, qdrant: QdrantDep, embedder: EmbedderDep) -> dict[str, str]:
    [vector] = await embedder.embed([review.text])
    point_id = str(uuid.uuid4())
    await qdrant.upsert(
        COLLECTION,
        points=[models.PointStruct(id=point_id, vector=vector, payload=review.model_dump())],
    )
    return {"id": point_id}


@app.get("/semantic/reviews")
async def search_reviews(
    q: Annotated[str, Query(min_length=2)],
    qdrant: QdrantDep,
    embedder: EmbedderDep,
    isbn: str | None = None,
    min_rating: Annotated[int, Query(ge=1, le=5)] = 1,
    limit: Annotated[int, Query(ge=1, le=20)] = 5,
) -> list[ReviewHit]:
    [vector] = await embedder.embed([q])
    must = [models.FieldCondition(key="rating", range=models.Range(gte=min_rating))]
    if isbn:
        must.append(models.FieldCondition(key="book_isbn", match=models.MatchValue(value=isbn)))
    response = await qdrant.query_points(
        COLLECTION, query=vector, query_filter=models.Filter(must=must), limit=limit, with_payload=True
    )
    return [ReviewHit(id=str(p.id), score=p.score, **(p.payload or {})) for p in response.points]

[
  {"id": "<uuid>", "score": 0.655, "book_isbn": "978-0441172719", "rating": 5, "text": "Epic desert planet story with giant sandworms"}
]

from typing import Any

from pgvector.sqlalchemy import Vector
from sqlalchemy import Column
from sqlalchemy.dialects import postgresql
from sqlmodel import Field, SQLModel, select


class ReviewEmbedding(SQLModel, table=True):
    __tablename__ = "review_embedding"
    id: int | None = Field(default=None, primary_key=True)
    book_isbn: str = Field(index=True)
    text: str
    embedding: list[float] = Field(sa_column=Column(Vector(64)))  # needs CREATE EXTENSION vector


def nearest(query_vector: list[float], limit: int = 5):
    distance = ReviewEmbedding.embedding.cosine_distance(query_vector)
    return select(ReviewEmbedding, distance.label("distance")).order_by(distance).limit(limit)


if __name__ == "__main__":
    stmt = nearest([0.1] * 64)
    print(stmt.compile(dialect=postgresql.dialect()))
    print(ReviewEmbedding.__table__.c.embedding.type)
    from sqlalchemy.schema import CreateTable
    print(CreateTable(ReviewEmbedding.__table__).compile(dialect=postgresql.dialect()))

    from sqlalchemy import Index
    from sqlalchemy.schema import CreateIndex
    hnsw = Index(
        "review_embedding_hnsw",
        ReviewEmbedding.embedding,
        postgresql_using="hnsw",
        postgresql_with={"m": 16, "ef_construction": 64},
        postgresql_ops={"embedding": "vector_cosine_ops"},
    )
    print(CreateIndex(hnsw).compile(dialect=postgresql.dialect()))

SELECT review_embedding.id, review_embedding.book_isbn, review_embedding.text, review_embedding.embedding,
       review_embedding.embedding <=> %(embedding_1)s AS distance
FROM review_embedding ORDER BY review_embedding.embedding <=> %(embedding_1)s LIMIT %(param_1)s

CREATE TABLE review_embedding (
    id SERIAL NOT NULL,
    book_isbn VARCHAR NOT NULL,
    text VARCHAR NOT NULL,
    embedding VECTOR(64),
    PRIMARY KEY (id)
)

CREATE INDEX review_embedding_hnsw ON review_embedding USING hnsw (embedding vector_cosine_ops)
  WITH (m = 16, ef_construction = 64)