NoSQL, Redis, MongoDB & Search

from flask import Blueprint, jsonify

stats_bp = Blueprint("stats", __name__)


@stats_bp.post("/books/<int:book_id>/views")
def record_view(book_id):
    pipe = get_redis().pipeline()
    pipe.incr(f"book:{book_id}:views")
    pipe.zincrby("books:popular", 1, book_id)
    views, _ = pipe.execute()
    return jsonify(book_id=book_id, views=views)


@stats_bp.get("/popular")
def popular():
    top = get_redis().zrange("books:popular", 0, 2, desc=True, withscores=True)
    return jsonify([{"book_id": int(member), "views": int(score)} for member, score in top])


app = create_app(
    {"TESTING": True},
    redis_client=fakeredis.FakeRedis(decode_responses=True),
    mongo_client=mongomock.MongoClient(),
    search_client=object(),
    blueprints=[stats_bp],
)
client = app.test_client()
for book_id in (1, 2, 2, 3, 2):
    client.post(f"/books/{book_id}/views")
assert client.post("/books/1/views").get_json() == {"book_id": 1, "views": 2}
assert client.get("/popular").get_json() == [
    {"book_id": 2, "views": 3},
    {"book_id": 1, "views": 2},
    {"book_id": 3, "views": 1},
]
print(client.get("/popular").get_json())

import secrets

from flask import abort, url_for

tokens_bp = Blueprint("tokens", __name__)
DOWNLOAD_TTL = 300


@tokens_bp.post("/books/<int:book_id>/download-links")
def create_download_link(book_id):
    token = secrets.token_urlsafe(16)
    get_redis().set(f"dl:{token}", book_id, ex=DOWNLOAD_TTL)
    return jsonify(url=url_for("tokens.download", token=token), expires_in=DOWNLOAD_TTL), 201


@tokens_bp.get("/download/<token>")
def download(token):
    book_id = get_redis().getdel(f"dl:{token}")
    if book_id is None:
        abort(404)
    return jsonify(book_id=int(book_id), note="the file would be streamed here")


app = create_app(
    {"TESTING": True},
    redis_client=fakeredis.FakeRedis(decode_responses=True),
    mongo_client=mongomock.MongoClient(),
    search_client=object(),
    blueprints=[tokens_bp],
)
client = app.test_client()
link = client.post("/books/7/download-links").get_json()
with app.app_context():
    ttl = get_redis().ttl("dl:" + link["url"].rsplit("/", 1)[1])
assert 0 < ttl <= DOWNLOAD_TTL
assert client.get(link["url"]).get_json()["book_id"] == 7
assert client.get(link["url"]).status_code == 404  # already redeemed
print("token redeemed once, ttl was", ttl)

import json

BOOK_TTL = 60


def get_book_cached(book_id, loader):
    key = f"cache:book:{book_id}"
    redis_conn = get_redis()
    if (hit := redis_conn.get(key)) is not None:
        return json.loads(hit)
    book = loader(book_id)
    if book is not None:
        redis_conn.set(key, json.dumps(book), ex=BOOK_TTL)
    return book


def invalidate_book(book_id):
    get_redis().delete(f"cache:book:{book_id}")


calls = []


def load_from_database(book_id):
    calls.append(book_id)
    return {"id": book_id, "title": "Flask Web Development"}


with app.app_context():
    assert get_book_cached(1, load_from_database)["title"] == "Flask Web Development"
    assert get_book_cached(1, load_from_database)["id"] == 1
    assert calls == [1]  # the second call was a cache hit
    invalidate_book(1)
    get_book_cached(1, load_from_database)
    assert calls == [1, 1]
print("loader calls:", calls)

import functools
import logging

from flask import make_response, request
from redis.exceptions import RedisError

log = logging.getLogger(__name__)


def rate_limit(limit: int, window: int):
    def decorator(view):
        @functools.wraps(view)
        def wrapper(*args, **kwargs):
            key = f"rl:{request.endpoint}:{request.remote_addr}"
            try:
                pipe = get_redis().pipeline()
                pipe.set(key, 0, ex=window, nx=True)
                pipe.incr(key)
                pipe.ttl(key)
                _, count, ttl = pipe.execute()
            except RedisError:
                log.exception("rate limiter unavailable, letting the request through")
                return view(*args, **kwargs)
            if count > limit:
                response = make_response(jsonify(error="rate limit exceeded"), 429)
                response.headers["Retry-After"] = str(max(ttl, 1))
                return response
            return view(*args, **kwargs)
        return wrapper
    return decorator


limited_bp = Blueprint("limited", __name__)


@limited_bp.get("/search-suggestions")
@rate_limit(limit=3, window=60)
def suggestions():
    return jsonify(suggestions=["flask", "flask web development"])


app = create_app(
    {"TESTING": True},
    redis_client=fakeredis.FakeRedis(decode_responses=True),
    mongo_client=mongomock.MongoClient(),
    search_client=object(),
    blueprints=[limited_bp],
)
client = app.test_client()
statuses = [client.get("/search-suggestions").status_code for _ in range(5)]
assert statuses == [200, 200, 200, 429, 429]
blocked = client.get("/search-suggestions")
assert 1 <= int(blocked.headers["Retry-After"]) <= 60
print(statuses, "Retry-After:", blocked.headers["Retry-After"])

from flask_limiter import Limiter
from flask_limiter.util import get_remote_address

app = Flask(__name__)
app.config["RATELIMIT_STORAGE_URI"] = os.environ.get("RATELIMIT_STORAGE_URI", "memory://")
app.config["RATELIMIT_HEADERS_ENABLED"] = True
limiter = Limiter(get_remote_address, app=app, default_limits=["200 per hour"])


@app.get("/ping")
@limiter.limit("2 per minute")
def ping():
    return {"pong": True}


client = app.test_client()
statuses = [client.get("/ping").status_code for _ in range(3)]
assert statuses == [200, 200, 429]
assert "Retry-After" in client.get("/ping").headers
print(statuses)

from datetime import datetime, timezone

from bson import ObjectId
from bson.errors import InvalidId
from flask.json.provider import DefaultJSONProvider
from werkzeug.routing import BaseConverter, ValidationError


class ObjectIdConverter(BaseConverter):
    regex = r"[0-9a-fA-F]{24}"

    def to_python(self, value):
        try:
            return ObjectId(value)
        except InvalidId:
            raise ValidationError()

    def to_url(self, value):
        return str(value)


class MongoJSONProvider(DefaultJSONProvider):
    @staticmethod
    def default(o):
        if isinstance(o, ObjectId):
            return str(o)
        return DefaultJSONProvider.default(o)


reviews_bp = Blueprint("reviews", __name__)


def serialise(review):
    document = dict(review)
    document["id"] = str(document.pop("_id"))
    return document


@reviews_bp.post("/books/<isbn>/reviews")
def add_review(isbn):
    data = request.get_json(silent=True) or {}
    rating, text, tags = data.get("rating"), data.get("text"), data.get("tags", [])
    if not (isinstance(rating, int) and not isinstance(rating, bool) and 1 <= rating <= 5):
        return jsonify(error="rating must be an integer from 1 to 5"), 422
    if not isinstance(text, str) or not 1 <= len(text) <= 2000:
        return jsonify(error="text must be a string of 1 to 2000 characters"), 422
    if not (isinstance(tags, list) and all(isinstance(tag, str) for tag in tags)):
        return jsonify(error="tags must be a list of strings"), 422
    document = {
        "isbn": isbn, "rating": rating, "text": text, "tags": tags,
        "created_at": datetime.now(timezone.utc),
    }
    result = get_mongo_db().reviews.insert_one(document)
    return jsonify(id=result.inserted_id), 201  # an ObjectId, serialised by the provider


@reviews_bp.get("/books/<isbn>/reviews")
def list_reviews(isbn):
    cursor = get_mongo_db().reviews.find({"isbn": isbn}).sort("created_at", -1).limit(20)
    return jsonify([serialise(review) for review in cursor])


@reviews_bp.get("/reviews/<objectid:review_id>")
def get_review(review_id):
    review = get_mongo_db().reviews.find_one({"_id": review_id})
    if review is None:
        abort(404)
    return jsonify(serialise(review))


def create_reviews_app():
    app = create_app(
        {"TESTING": True},
        redis_client=fakeredis.FakeRedis(decode_responses=True),
        mongo_client=mongomock.MongoClient(tz_aware=True),
        search_client=object(),
    )
    app.url_map.converters["objectid"] = ObjectIdConverter
    app.json = MongoJSONProvider(app)
    app.register_blueprint(reviews_bp)
    return app


app = create_reviews_app()
client = app.test_client()
created = client.post("/books/9781098135195/reviews", json={"rating": 5, "text": "Clear and current", "tags": ["flask"]})
assert created.status_code == 201
review_id = created.get_json()["id"]
assert client.get(f"/reviews/{review_id}").get_json()["rating"] == 5
assert client.get("/reviews/not-an-object-id").status_code == 404
assert client.post("/books/x/reviews", json={"rating": 9, "text": "?"}).status_code == 422
print(client.get("/books/9781098135195/reviews").get_json())

import click
from pymongo import ASCENDING, DESCENDING


reviews_tools_bp = Blueprint("reviews_tools", __name__, cli_group=None)


@reviews_tools_bp.cli.command("ensure-indexes")
def ensure_indexes():
    """Create the indexes the reviews queries rely on."""
    reviews = get_mongo_db().reviews
    reviews.create_index([("isbn", ASCENDING), ("created_at", DESCENDING)])
    click.echo("indexes ready")


@reviews_tools_bp.get("/books/<isbn>/rating")
def rating(isbn):
    pipeline = [
        {"$match": {"isbn": isbn}},
        {"$group": {"_id": "$isbn", "average": {"$avg": "$rating"}, "count": {"$sum": 1}}},
    ]
    result = next(get_mongo_db().reviews.aggregate(pipeline), None)
    if result is None:
        return jsonify(average=None, count=0)
    return jsonify(average=round(result["average"], 2), count=result["count"])


app = create_reviews_app()
app.register_blueprint(reviews_tools_bp)
runner = app.test_cli_runner()
result = runner.invoke(args=["ensure-indexes"])
assert result.exit_code == 0 and "indexes ready" in result.output, result.output
client = app.test_client()
for stars in (5, 4, 4):
    client.post("/books/9781098135195/reviews", json={"rating": stars, "text": "ok"})
assert client.get("/books/9781098135195/rating").get_json() == {"average": 4.33, "count": 3}
print(result.output.strip(), client.get("/books/9781098135195/rating").get_json())

users = mongomock.MongoClient().bookshelf.users
users.insert_one({"name": "admin", "password_hash": secrets.token_hex(16)})

# Vulnerable shape: the client controls the whole value of "password_hash".
attack = {"name": "admin", "password_hash": {"$ne": None}}
assert users.find_one(attack) is not None  # matches without knowing the hash


def safe_find_user(name):
    if not isinstance(name, str):
        return None
    return users.find_one({"name": name})


assert safe_find_user({"$ne": None}) is None
assert safe_find_user("admin")["name"] == "admin"
print("operator injection blocked by type checks")

from flask_pymongo import PyMongo

mongo = PyMongo()


def create_flask_pymongo_app():
    app = Flask(__name__)
    app.config["MONGO_URI"] = os.environ.get("MONGO_URI", "mongodb://localhost:27017/bookshelf")
    mongo.init_app(app)

    @app.get("/reviews/<ObjectId:review_id>")
    def review(review_id):
        return jsonify(mongo.db.reviews.find_one_or_404({"_id": review_id}))

    return app


app = create_flask_pymongo_app()
assert mongo.db.name == "bookshelf"
assert "ObjectId" in app.url_map.converters
assert app.test_client().get("/reviews/zzz").status_code == 404  # rejected by the converter, no database access
print("Flask-PyMongo configured; database:", mongo.db.name)

from mongoengine import Document, IntField, ListField, StringField, connect, disconnect


class Review(Document):
    isbn = StringField(required=True)
    rating = IntField(min_value=1, max_value=5, required=True)
    text = StringField(max_length=2000)
    tags = ListField(StringField(max_length=30))
    meta = {"indexes": ["isbn"]}


def init_mongoengine(app, **options):
    disconnect()
    return connect(host=app.config["MONGO_URI"], db=app.config["MONGO_DBNAME"], **options)


app = create_app({"TESTING": True}, redis_client=object(), mongo_client=object(), search_client=object())
init_mongoengine(app, mongo_client_class=mongomock.MongoClient)

mongoengine_bp = Blueprint("mongoengine_reviews", __name__)


@mongoengine_bp.post("/books/<isbn>/reviews")
def add_review_odm(isbn):
    data = request.get_json(silent=True) or {}
    review = Review(isbn=isbn, rating=data.get("rating"), text=data.get("text"), tags=data.get("tags", []))
    try:
        review.validate()
    except Exception as error:  # mongoengine.ValidationError
        return jsonify(error=str(error)), 422
    review.save()
    return jsonify(id=str(review.id)), 201


app.register_blueprint(mongoengine_bp)
client = app.test_client()
assert client.post("/books/x/reviews", json={"rating": 5, "text": "good"}).status_code == 201
assert client.post("/books/x/reviews", json={"rating": 6}).status_code == 422
assert Review.objects(isbn="x").count() == 1
print("documents:", Review.objects.count())

import re


class FakeIndices:
    def __init__(self, store):
        self.store = store

    def exists(self, *, index):
        return index in self.store.docs

    def create(self, *, index, mappings=None, **_):
        self.store.docs.setdefault(index, {})

    def refresh(self, *, index):
        pass


class FakeSearch:
    """Just enough of the Elasticsearch client for the tests on this page: not a search engine."""

    def __init__(self):
        self.docs = {}
        self.indices = FakeIndices(self)

    def index(self, *, index, document, id=None, **_):
        self.docs.setdefault(index, {})[id] = dict(document)
        return {"result": "created", "_id": id}

    def delete(self, *, index, id, **_):
        self.docs.get(index, {}).pop(id, None)
        return {"result": "deleted"}

    def bulk(self, *, operations, **_):
        for action, document in zip(operations[::2], operations[1::2]):
            meta = action["index"]
            self.index(index=meta["_index"], id=meta["_id"], document=document)
        return {"errors": False}

    def search(self, *, index, query, from_=0, size=10, **_):
        spec = query["multi_match"]
        terms = re.findall(r"\w+", spec["query"].lower())
        scored = []
        for doc_id, doc in self.docs.get(index, {}).items():
            score = 0.0
            for field in spec["fields"]:
                name, _, boost = field.partition("^")
                words = re.findall(r"\w+", str(doc.get(name, "")).lower())
                score += sum(term in words for term in terms) * float(boost or 1)
            if score:
                scored.append((score, doc_id, doc))
        scored.sort(key=lambda row: (-row[0], str(row[1])))
        page = scored[from_:from_ + size]
        return {"hits": {
            "total": {"value": len(scored), "relation": "eq"},
            "hits": [{"_id": i, "_score": s, "_source": d} for s, i, d in page],
        }}


fake = FakeSearch()
fake.index(index="books", id="1", document={"title": "Flask Web Development", "author": "Miguel Grinberg"})
assert fake.search(index="books", query={"multi_match": {"query": "flask", "fields": ["title^2", "author"]}})["hits"]["total"]["value"] == 1
print("fake search client ready")

import hmac
import sqlite3
import tempfile
from pathlib import Path

from elastic_transport import TransportError
from elasticsearch import ApiError
from flask import g

BOOKS_INDEX = "books"
BOOKS_MAPPINGS = {"properties": {
    "title": {"type": "text"},
    "author": {"type": "text"},
    "isbn": {"type": "keyword"},
    "year": {"type": "integer"},
}}


def ensure_books_index(es):
    if not es.indices.exists(index=BOOKS_INDEX):
        es.indices.create(index=BOOKS_INDEX, mappings=BOOKS_MAPPINGS)


def index_book(es, book):
    es.index(index=BOOKS_INDEX, id=book["isbn"], document={
        "title": book["title"], "author": book["author"],
        "isbn": book["isbn"], "year": book["year"],
    })


def search_books(es, text, page=1, per_page=10):
    response = es.search(
        index=BOOKS_INDEX,
        query={"multi_match": {"query": text, "fields": ["title^2", "author"]}},
        from_=(page - 1) * per_page,
        size=per_page,
    )
    hits = response["hits"]
    return {
        "total": hits["total"]["value"],
        "results": [{"score": h["_score"], **h["_source"]} for h in hits["hits"]],
    }


def get_db():
    if "db" not in g:
        g.db = sqlite3.connect(current_app.config["DATABASE"])
        g.db.row_factory = sqlite3.Row
    return g.db


def admin_required(view):
    @functools.wraps(view)
    def wrapper(*args, **kwargs):
        expected = current_app.config.get("ADMIN_TOKEN")
        supplied = request.headers.get("Authorization", "").removeprefix("Bearer ")
        if not expected or not hmac.compare_digest(supplied, expected):
            abort(401)
        return view(*args, **kwargs)
    return wrapper


catalogue_bp = Blueprint("catalogue", __name__)


@catalogue_bp.teardown_app_request
def close_db(_exc):
    if db := g.pop("db", None):
        db.close()


@catalogue_bp.get("/search")
def search():
    text = request.args.get("q", "").strip()
    if not text:
        return jsonify(error="q is required"), 400
    page = request.args.get("page", 1, type=int)
    try:
        return jsonify(search_books(get_search(), text, page=max(page, 1)))
    except (ApiError, TransportError):
        current_app.logger.exception("search backend unavailable")
        return jsonify(error="search is temporarily unavailable"), 503


@catalogue_bp.post("/books")
@admin_required
def add_book():
    data = request.get_json(silent=True) or {}
    try:
        book = {"isbn": str(data["isbn"]), "title": str(data["title"]),
                "author": str(data["author"]), "year": int(data["year"])}
    except (KeyError, TypeError, ValueError):
        return jsonify(error="isbn, title, author and year are required"), 422
    db = get_db()
    db.execute("INSERT OR REPLACE INTO books (isbn, title, author, year) VALUES (:isbn, :title, :author, :year)", book)
    db.commit()                      # the database is the source of truth: commit first ...
    try:
        index_book(get_search(), book)  # ... then update the derived index
    except (ApiError, TransportError):
        current_app.logger.exception("indexing failed; run the reindex command to repair")
    return jsonify(book), 201


@catalogue_bp.delete("/books/<isbn>")
@admin_required
def delete_book(isbn):
    db = get_db()
    db.execute("DELETE FROM books WHERE isbn = ?", (isbn,))
    db.commit()
    try:
        get_search().delete(index=BOOKS_INDEX, id=isbn)
    except (ApiError, TransportError):
        current_app.logger.exception("removing from the index failed; run the reindex command")
    return "", 204


workdir = Path(tempfile.mkdtemp())
admin_token = secrets.token_hex(16)
fake = FakeSearch()
app = create_app(
    {"TESTING": True, "DATABASE": str(workdir / "bookshelf.db"), "ADMIN_TOKEN": admin_token},
    redis_client=fakeredis.FakeRedis(decode_responses=True),
    mongo_client=mongomock.MongoClient(),
    search_client=fake,
    blueprints=[catalogue_bp],
)
with sqlite3.connect(app.config["DATABASE"]) as connection:
    connection.execute("CREATE TABLE books (isbn TEXT PRIMARY KEY, title TEXT, author TEXT, year INTEGER)")
ensure_books_index(fake)

client = app.test_client()
auth = {"Authorization": f"Bearer {admin_token}"}
assert client.post("/books", json={"isbn": "1", "title": "Flask Web Development", "author": "Miguel Grinberg", "year": 2018}).status_code == 401
for isbn, title, author, year in [
    ("1", "Flask Web Development", "Miguel Grinberg", 2018),
    ("2", "Flask Framework Cookbook", "Shalabh Aggarwal", 2023),
    ("3", "Designing Data-Intensive Applications", "Martin Kleppmann", 2017),
]:
    response = client.post("/books", json={"isbn": isbn, "title": title, "author": author, "year": year}, headers=auth)
    assert response.status_code == 201
found = client.get("/search", query_string={"q": "flask cookbook"}).get_json()
assert found["total"] == 2 and found["results"][0]["isbn"] == "2"
assert client.get("/search").status_code == 400
assert client.delete("/books/2", headers=auth).status_code == 204
assert client.get("/search", query_string={"q": "cookbook"}).get_json()["total"] == 0
print(found["results"][0]["title"], "ranked first")

from whoosh.fields import ID, TEXT, Schema
from whoosh.index import create_in
from whoosh.qparser import MultifieldParser

schema = Schema(isbn=ID(stored=True, unique=True), title=TEXT(stored=True), author=TEXT(stored=True))
index_dir = Path(tempfile.mkdtemp())
ix = create_in(index_dir, schema)

writer = ix.writer()
for isbn, title, author in [
    ("1", "Flask Web Development", "Miguel Grinberg"),
    ("2", "Flask Framework Cookbook", "Shalabh Aggarwal"),
    ("3", "Designing Data-Intensive Applications", "Martin Kleppmann"),
]:
    writer.update_document(isbn=isbn, title=title, author=author)  # upsert by the unique field
writer.commit()

with ix.searcher() as searcher:
    query = MultifieldParser(["title", "author"], ix.schema).parse("cookbook")
    titles = [hit["title"] for hit in searcher.search(query, limit=10)]
assert titles == ["Flask Framework Cookbook"]
print(titles)

connection = sqlite3.connect(":memory:")
connection.execute("CREATE VIRTUAL TABLE book_fts USING fts5(title, author, tokenize='porter')")
connection.executemany(
    "INSERT INTO book_fts (title, author) VALUES (?, ?)",
    [
        ("Flask Web Development", "Miguel Grinberg"),
        ("Flask Framework Cookbook", "Shalabh Aggarwal"),
        ("Designing Data-Intensive Applications", "Martin Kleppmann"),
    ],
)


def fts_search(text):
    phrase = '"' + text.replace('"', '""') + '"'   # a quoted phrase is always valid FTS5 syntax
    return connection.execute(
        "SELECT title FROM book_fts WHERE book_fts MATCH ? ORDER BY bm25(book_fts)", (phrase,)
    ).fetchall()


assert [row[0] for row in fts_search("developments")] == ["Flask Web Development"]  # the porter stemmer matches
assert fts_search('"unbalanced') == []
print(fts_search("flask"))