diff --git a/README.md b/README.md index 9e927a3..40d4f41 100644 --- a/README.md +++ b/README.md @@ -17,9 +17,10 @@ playground's README states what it needs. | --------------------------------- | ---------------------- | ------------------------------------------------------------------------- | | [mongoose](playgrounds/mongoose/) | Node.js — Mongoose ODM | Express REST API + a CRUD/compatibility test suite using the Mongoose ODM. | | [beanie](playgrounds/beanie/) | Python — Beanie ODM | FastAPI REST API + a CRUD/compatibility test suite using the Beanie ODM. | +| [pymongo](playgrounds/pymongo/) | Python — PyMongo driver | Flask REST API + a CRUD/compatibility test suite using the raw PyMongo driver. | -More playgrounds are planned (for example **PyMongo** and other MongoDB -drivers). Contributions are welcome. +More playgrounds are planned (for example the **MongoDB Node.js native driver** +and other MongoDB drivers). Contributions are welcome. ## Getting Started @@ -45,6 +46,14 @@ cd playgrounds/beanie ./scripts/run-app.sh # or run the demo REST API ``` +To try the PyMongo playground: + +```bash +cd playgrounds/pymongo +./scripts/run-test.sh # start DocumentDB locally and run the compatibility suite +./scripts/run-app.sh # or run the demo REST API +``` + ## Repository Layout ``` @@ -53,7 +62,8 @@ documentdb-playground/ ├── LICENSE └── playgrounds/ ├── mongoose/ # Node.js + Mongoose ODM - └── beanie/ # Python + Beanie ODM + ├── beanie/ # Python + Beanie ODM + └── pymongo/ # Python + PyMongo driver ``` ## License diff --git a/playgrounds/beanie/README.md b/playgrounds/beanie/README.md index 023eba3..718db61 100644 --- a/playgrounds/beanie/README.md +++ b/playgrounds/beanie/README.md @@ -80,7 +80,7 @@ Stop the database when you are done: ./scripts/stop-documentdb.sh ``` -The suite should end with `Passed: 13 Failed: 0`. +The suite should end with `Passed: 16 Failed: 0`. ## Trying the API @@ -175,7 +175,7 @@ Verified against `documentdb-local:latest` (release `0.114`): | Index creation via `Settings.indexes` | ✅ Supported | Built asynchronously by the engine; `createIndexes` returns in ~2s. Avoid `collation`. | | Unique indexes | ✅ Supported | Duplicate keys raise `DuplicateKeyError` (code `11000`). | | Aggregation pipelines | ✅ Common stages | `$match`, `$group`, `$unwind`, `$sort`, etc. Atlas-only stages differ. | -| `$vectorSearch` | ❌ Not supported | Atlas-only operator. | +| `$vectorSearch` / vector search | ✅ Supported | DocumentDB supports a `cosmosSearch` vector index (e.g. `vector-ivf`) queried via the `$vectorSearch` stage. The test suite creates one and runs a nearest-neighbor query. | | Index `collation` | ❌ Not supported | `createIndex.collation is not implemented yet`; omit it. | | Change streams / transactions | ⚠️ Check version | Verify against your DocumentDB version before relying on them. | @@ -212,9 +212,12 @@ Beanie DocumentDB compatibility test ✅ aggregation ($unwind/$group) ✅ unique index enforcement (duplicate sku rejected) ✅ delete_one + ✅ vector index + insert (cosmosSearch vector-ivf) + ✅ $vectorSearch returns nearest neighbor + ✅ vector cleanup (drop collection) ✅ cleanup (drop collection) ==================================== -Passed: 13 Failed: 0 +Passed: 16 Failed: 0 ``` ## What the Scripts Do @@ -229,7 +232,7 @@ Passed: 13 Failed: 0 ## Verification -- `./scripts/run-test.sh` ends with `Passed: 13 Failed: 0`. +- `./scripts/run-test.sh` ends with `Passed: 16 Failed: 0`. - With `./scripts/run-app.sh` running, `curl http://localhost:3000/health` returns `{"status":"healthy","db":"connected"}`, `POST /books` returns `201` with the created document, and `GET /stats/genres` returns per-genre counts. diff --git a/playgrounds/beanie/app/beanie_crud_test.py b/playgrounds/beanie/app/beanie_crud_test.py index 4fb1683..24e5887 100644 --- a/playgrounds/beanie/app/beanie_crud_test.py +++ b/playgrounds/beanie/app/beanie_crud_test.py @@ -217,6 +217,70 @@ async def _delete_one() -> None: await step("delete_one", _delete_one) + # Vector search: DocumentDB supports a `cosmosSearch` vector index and the + # `$vectorSearch` aggregation stage. The index is created with the raw + # `createIndexes` command since Beanie does not model vector indexes. + db_ref = client[DB_NAME] + vec_name = f"vectors_{int(time.time() * 1000)}" + vectors = db_ref[vec_name] + + async def _vector_index_insert() -> None: + await vectors.insert_many( + [ + {"name": "a", "v": [1, 0, 0]}, + {"name": "b", "v": [0.9, 0.1, 0]}, + {"name": "c", "v": [0, 0, 1]}, + ] + ) + res = await db_ref.command( + { + "createIndexes": vec_name, + "indexes": [ + { + "name": "v_ivf", + "key": {"v": "cosmosSearch"}, + "cosmosSearchOptions": { + "kind": "vector-ivf", + "numLists": 1, + "similarity": "COS", + "dimensions": 3, + }, + } + ], + } + ) + if res.get("ok") != 1: + raise RuntimeError("createIndexes did not return ok:1") + + await step("vector index + insert (cosmosSearch vector-ivf)", _vector_index_insert) + + async def _vector_search() -> None: + hits = await vectors.aggregate( + [ + { + "$vectorSearch": { + "index": "v_ivf", + "path": "v", + "queryVector": [1, 0, 0], + "numCandidates": 10, + "limit": 2, + } + }, + {"$project": {"name": 1, "_id": 0}}, + ] + ).to_list(length=10) + if not hits: + raise RuntimeError("vector search returned no results") + if hits[0].get("name") != "a": + raise RuntimeError(f"expected nearest 'a', got {hits[0].get('name')!r}") + + await step("$vectorSearch returns nearest neighbor", _vector_search) + + async def _vector_cleanup() -> None: + await vectors.drop() + + await step("vector cleanup (drop collection)", _vector_cleanup) + async def _drop() -> None: await coll.drop() diff --git a/playgrounds/pymongo/README.md b/playgrounds/pymongo/README.md new file mode 100644 index 0000000..f9c05ed --- /dev/null +++ b/playgrounds/pymongo/README.md @@ -0,0 +1,307 @@ +# PyMongo with DocumentDB (local) + +This playground shows how to use [PyMongo](https://pymongo.readthedocs.io/), the +official **synchronous Python driver** for MongoDB, against DocumentDB — running +**entirely on your machine**. Unlike the Mongoose and Beanie playgrounds (which +use ODMs), this one talks to DocumentDB at the **raw driver level**: no schema +classes, just collections and BSON documents. It includes: + +- a small **Flask + PyMongo REST API** (`app/`), and +- a standalone **PyMongo CRUD/compatibility test suite** + (`app/pymongo_crud_test.py`) that exercises connect, index creation, insert, + query, update, aggregation, unique-index enforcement, and delete. + +There is **no Kubernetes and no cloud**. DocumentDB runs as the +[`documentdb-local`](https://github.com/documentdb/documentdb) emulator in a +single Docker container, and the app/test run as local Python processes that +connect straight to it. + +> **What is PyMongo?** PyMongo is the **official MongoDB driver for Python** — a +> library your application imports to talk to a MongoDB-compatible database. It +> is *not* an ODM: you work directly with databases, collections, and dict-like +> BSON documents. Here it is used by the demo **app** +> ([`app/main.py`](app/main.py)) and the standalone **test script** +> ([`app/pymongo_crud_test.py`](app/pymongo_crud_test.py)). + +## Architecture + +Everything is local. The emulator container exposes the MongoDB wire protocol on +`localhost:10260`; the Python processes connect to it directly. + +``` + Your machine (WSL / Linux / macOS) +┌──────────────────────────────────────────────────────────────────┐ +│ ┌────────────────────┐ ┌──────────────────────────────┐ │ +│ │ pymongo app / │ TLS, │ documentdb-local (Docker) │ │ +│ │ test script │ wire │ ┌────────────┐ ┌─────────┐ │ │ +│ │ (Python + PyMongo)│────────▶│ │ Gateway │▶│Postgres │ │ │ +│ │ │ :10260 │ │ (10260) │ │ (engine)│ │ │ +│ └────────────────────┘ │ └────────────┘ └─────────┘ │ │ +│ └──────────────────────────────┘ │ +└──────────────────────────────────────────────────────────────────┘ +``` + +PyMongo talks to the emulator exactly as it would to a standalone `mongod`, with +a few required options (see [Connecting PyMongo to +DocumentDB](#connecting-pymongo-to-documentdb)). + +## Prerequisites + +- **Docker** (to run the `documentdb-local` emulator) +- **Python 3.10+** with `venv` (to run the app and test suite) + +The scripts create a Python virtualenv and install dependencies for you. On +Windows, run these from a **WSL** shell. + +## Quick Start + +From this directory (`playgrounds/pymongo/`). `run-test.sh` and `run-app.sh` +are **two independent operations** — each starts DocumentDB on its own if it +isn't already running. + +### Option A — run the test suite + +```bash +# Run the full CRUD/compatibility suite end-to-end. +# Starts DocumentDB in Docker (first run pulls the image), then runs the tests. +./scripts/run-test.sh +``` + +### Option B — run the demo REST API + +```bash +# Starts DocumentDB (if not already running) and serves the API on :3000. +# This stays in the foreground until you press Ctrl-C. +./scripts/run-app.sh +``` + +Stop the database when you are done: + +```bash +./scripts/stop-documentdb.sh +``` + +The suite should end with `Passed: 16 Failed: 0`. + +## Trying the API + +With `./scripts/run-app.sh` running, the API is on `http://localhost:3000`: + +```bash +# Health +curl -s http://localhost:3000/health +# {"status":"healthy","db":"connected"} + +# Create a book +curl -s -X POST http://localhost:3000/books \ + -H 'Content-Type: application/json' \ + -d '{"title":"Dune","author":"Herbert","genres":["sci-fi"],"pages":412,"rating":5}' + +# List books +curl -s http://localhost:3000/books | jq . + +# Count books per genre (aggregation) +curl -s http://localhost:3000/stats/genres | jq . +``` + +## Connecting PyMongo to DocumentDB + +The DocumentDB gateway speaks the MongoDB wire protocol but advertises itself as +a **standalone** server over **TLS** (with a self-signed cert). PyMongo +therefore needs these options (see [`app/db.py`](app/db.py)): + +```python +from pymongo import MongoClient + +client = MongoClient( + uri, + directConnection=True, # gateway is standalone, not a replica set + tls=True, # gateway only accepts TLS + tlsAllowInvalidCertificates=True, # emulator uses a self-signed cert +) +db = client["pymongo_demo"] +``` + +The connection string built by the scripts is: + +``` +mongodb://:@localhost:10260/?tls=true&tlsAllowInvalidCertificates=true&directConnection=true +``` + +If your connection string contains `replicaSet=rs0`, strip it — a direct +connection to the standalone gateway conflicts with it. Both +[`app/db.py`](app/db.py) and the test script strip it automatically. + +For production against a real (non-emulator) deployment, set `TLS_INSECURE=false` +and pass a CA bundle via `tlsCAFile` instead of `tlsAllowInvalidCertificates`. + +## Configuration Reference + +All settings are passed via environment variables; there is no config file. + +### Emulator + scripts (`scripts/`) + +Read by [`lib.sh`](scripts/lib.sh) and the `start`/`stop`/`run` scripts. + +| Variable | Default | Description | +| ---------------------- | -------------------------------------------------------- | -------------------------------------------------------- | +| `DOCUMENTDB_IMAGE` | `ghcr.io/documentdb/documentdb/documentdb-local:latest` | Emulator image to pull/run. | +| `DOCUMENTDB_CONTAINER` | `documentdb-local` | Docker container name. | +| `DOCUMENTDB_HOST` | `localhost` | Host the app/test connect to. | +| `DOCUMENTDB_PORT` | `10260` | Host port mapped to the gateway. | +| `DOCUMENTDB_USERNAME` | `docdbadmin` | Emulator admin username. **Do not use `documentdb`** (reserved — the gateway rejects it as "Username is invalid"). | +| `DOCUMENTDB_PASSWORD` | `Documentdb!Local1` | Emulator admin password. If you use special characters, URL-encode them in the connection string. | +| `PORT` | `3000` | Local port the Flask app listens on (`run-app.sh`). | + +### App + test script (`app/`) + +Read by [`app/db.py`](app/db.py), [`app/main.py`](app/main.py), and +[`app/pymongo_crud_test.py`](app/pymongo_crud_test.py). The scripts set +`MONGO_URI` for you from the variables above. + +| Variable | Default | Description | +| ----------------------------- | ---------------- | -------------------------------------------------------------------------------------------- | +| `MONGO_URI` | _(set by scripts)_ | DocumentDB connection string. `replicaSet=rs0` is stripped automatically. The test script also accepts it as the first CLI argument. | +| `MONGO_DB` | `pymongo_demo` (app), `pymongo_test` (test) | Database name PyMongo connects to. | +| `TLS_INSECURE` | `true` | When `true`, accepts the self-signed cert. Set `false` for CA-verified TLS. | +| `SERVER_SELECTION_TIMEOUT_MS` | `10000` | How long PyMongo waits to select a server before erroring. | +| `PORT` | `3000` | Port the Flask API listens on. | + +## DocumentDB Compatibility Notes + +| PyMongo feature | Status | Notes | +| ---------------------------------------- | ------------- | --------------------------------------------------------------------- | +| CRUD (`insert_*`/`find*`/`update_*`/`delete_*`) | ✅ Supported | Standard document operations work as expected. | +| `find_one` / `_id` point lookups | ✅ Supported | Works on the current `documentdb-local:latest` image. | +| `create_indexes` | ✅ Supported | Built asynchronously by the engine. Avoid `collation`. | +| Unique indexes | ✅ Supported | Duplicate keys raise `DuplicateKeyError` (code `11000`). | +| `find_one_and_update` (returns new) | ✅ Supported | `ReturnDocument.AFTER` returns the updated document. | +| Aggregation pipelines | ✅ Common stages | `$match`, `$group`, `$unwind`, `$sort`, etc. Atlas-only stages differ. | +| `$vectorSearch` / vector search | ✅ Supported | DocumentDB supports a `cosmosSearch` vector index (e.g. `vector-ivf`) queried via the `$vectorSearch` stage. The test suite creates one and runs a nearest-neighbor query. | +| Index `collation` | ❌ Not supported | `createIndex.collation is not implemented yet`; omit it. | +| Change streams / transactions | ⚠️ Check version | Verify against your DocumentDB version before relying on them. | + +The CRUD test suite ([`app/pymongo_crud_test.py`](app/pymongo_crud_test.py)) +covers the supported rows above and prints a pass/fail summary. + +## Running the Test Suite Manually + +`scripts/run-test.sh` sets `MONGO_URI` and runs the suite for you. To run it +directly against any reachable connection string: + +```bash +cd app +python3 -m venv .venv +.venv/bin/pip install -r requirements.txt +MONGO_URI="mongodb://docdbadmin:Documentdb!Local1@localhost:10260/?tls=true&tlsAllowInvalidCertificates=true&directConnection=true" \ + .venv/bin/python pymongo_crud_test.py +``` + +Expected output: + +``` +PyMongo DocumentDB compatibility test +===================================== + ✅ connect + ✅ create indexes + ✅ insert_one + ✅ insert_many + ✅ find_one by _id + ✅ find with filter + sort + limit + ✅ count_documents + ✅ update_one ($set) + ✅ find_one_and_update (returns new) + ✅ aggregation ($unwind/$group) + ✅ unique index enforcement (duplicate sku rejected) + ✅ delete_one + ✅ vector index + insert (cosmosSearch vector-ivf) + ✅ $vectorSearch returns nearest neighbor + ✅ vector cleanup (drop collection) + ✅ cleanup (drop collection) +===================================== +Passed: 16 Failed: 0 +``` + +## What the Scripts Do + +| Script | Purpose | +| ----------------------------- | ------------------------------------------------------------------------------------------ | +| `scripts/start-documentdb.sh` | Start the local emulator container and wait until the gateway is ready. | +| `scripts/run-test.sh` | Start DocumentDB (if needed), set up the venv, and run the PyMongo CRUD/compatibility suite. | +| `scripts/run-app.sh` | Start DocumentDB (if needed), set up the venv, and run the Flask + PyMongo demo app. | +| `scripts/stop-documentdb.sh` | Stop and remove the emulator container (full reset of its data). | +| `scripts/lib.sh` | Shared helpers: container lifecycle, readiness wait, connection-string builder, and venv setup. | + +## Verification + +- `./scripts/run-test.sh` ends with `Passed: 16 Failed: 0`. +- With `./scripts/run-app.sh` running, `curl http://localhost:3000/health` + returns `{"status":"healthy","db":"connected"}`, `POST /books` returns `201` + with the created document, and `GET /stats/genres` returns per-genre counts. + +## Cleanup + +- **App / test:** press `Ctrl-C` to stop the app; the test exits on its own. + Optionally remove the virtualenv: `rm -rf app/.venv`. +- **Emulator:** `./scripts/stop-documentdb.sh` removes the container and all its + data. + +## Troubleshooting + +### `AuthenticationFailed: Username is invalid.` + +The emulator rejects certain reserved usernames — notably `documentdb`. Use a +different admin username (the default here is `docdbadmin`). If you changed +`DOCUMENTDB_USERNAME`, recreate the container so the new credentials take effect: + +```bash +./scripts/stop-documentdb.sh && ./scripts/start-documentdb.sh +``` + +### `ServerSelectionTimeoutError` / TLS handshake failures + +The gateway requires TLS. Confirm the connection string includes `tls=true` and +`tlsAllowInvalidCertificates=true` (the scripts add these). Make sure the +emulator is running: `docker ps` should list `documentdb-local`, and +`docker logs documentdb-local` should show the gateway accepting connections. + +### Port `10260` already in use + +Another process (or a previous emulator) holds the port. Stop it, or run on a +different port: + +```bash +DOCUMENTDB_PORT=10261 ./scripts/start-documentdb.sh +DOCUMENTDB_PORT=10261 ./scripts/run-test.sh +``` + +### Credentials changed but auth still fails + +The username/password are baked into the container at creation time. Changing +`DOCUMENTDB_USERNAME`/`DOCUMENTDB_PASSWORD` only takes effect after you recreate +the container (`stop-documentdb.sh` then `start-documentdb.sh`). + +### `createIndex.collation is not implemented yet` + +An index uses `collation`. Remove it; DocumentDB does not support collation +indexes. The indexes in this playground intentionally avoid it. + +## Directory Layout + +``` +pymongo/ +├── README.md +├── app/ +│ ├── requirements.txt +│ ├── db.py # MongoClient connection (DocumentDB options) +│ ├── main.py # Flask REST API (/books, /health, /stats) +│ ├── models/ +│ │ └── book.py # Collection name, indexes, doc builder/serializer +│ └── pymongo_crud_test.py # Standalone CRUD/compatibility test suite +└── scripts/ + ├── lib.sh # Container lifecycle + connection-string builder + venv setup + ├── start-documentdb.sh # Start the local emulator (Docker) + ├── run-app.sh # Run the demo app locally + ├── run-test.sh # Run the test suite locally + └── stop-documentdb.sh # Stop + remove the emulator +``` diff --git a/playgrounds/pymongo/app/db.py b/playgrounds/pymongo/app/db.py new file mode 100644 index 0000000..db58ae0 --- /dev/null +++ b/playgrounds/pymongo/app/db.py @@ -0,0 +1,80 @@ +"""PyMongo connection helpers for DocumentDB. + +DocumentDB exposes a single gateway endpoint that speaks the MongoDB wire +protocol but advertises itself as a *standalone* server (not a replica set). +PyMongo therefore needs three tweaks when creating the ``MongoClient``: + + - directConnection=True -> don't attempt replica-set topology + discovery (the gateway is standalone). + - tls=True -> the gateway only accepts TLS. + - tlsAllowInvalidCertificates -> the default install uses a self-signed + cert. Set TLS_INSECURE=false and mount a + CA bundle (tlsCAFile) for production-grade + verification. + +The connection string itself (MONGO_URI) carries the credentials. If it still +contains ``replicaSet=rs0`` we strip it here so the driver does not try to match +a replica-set name the gateway never advertises. +""" + +from __future__ import annotations + +import os +import re + +from pymongo import MongoClient +from pymongo.database import Database + +_client: MongoClient | None = None + + +def sanitize_uri(uri: str | None) -> str: + """Remove ``replicaSet=...`` from a DocumentDB connection string. + + ``replicaSet`` is incompatible with a direct connection to the gateway; the + driver raises "client is configured to connect to a replica set named 'rs0' + but this node belongs to a set named 'None'" otherwise. + """ + if not uri: + raise ValueError("MONGO_URI is not set. Provide a DocumentDB connection string.") + return re.sub(r"[?&]replicaSet=[^&]*", "", uri) + + +def build_client_kwargs() -> dict: + """Return the MongoClient kwargs DocumentDB's gateway requires.""" + tls_insecure = os.environ.get("TLS_INSECURE", "true").lower() != "false" + return { + "directConnection": True, + "tls": True, + "tlsAllowInvalidCertificates": tls_insecure, + "serverSelectionTimeoutMS": int(os.environ.get("SERVER_SELECTION_TIMEOUT_MS", "10000")), + } + + +def connect(uri: str | None = None, db_name: str | None = None) -> Database: + """Create the shared MongoClient and return the target database.""" + global _client + clean_uri = sanitize_uri(uri or os.environ.get("MONGO_URI")) + database_name = db_name or os.environ.get("MONGO_DB", "pymongo_demo") + + _client = MongoClient(clean_uri, **build_client_kwargs()) + return _client[database_name] + + +def get_client() -> MongoClient | None: + return _client + + +def ping() -> bool: + """Return True if the gateway responds to an admin ``ping``.""" + if _client is None: + return False + result = _client.admin.command("ping") + return result.get("ok") == 1 + + +def close() -> None: + global _client + if _client is not None: + _client.close() + _client = None diff --git a/playgrounds/pymongo/app/main.py b/playgrounds/pymongo/app/main.py new file mode 100644 index 0000000..1a98b6d --- /dev/null +++ b/playgrounds/pymongo/app/main.py @@ -0,0 +1,131 @@ +"""Flask + PyMongo demo REST API for DocumentDB. + +Run locally with scripts/run-app.sh, or directly: + + MONGO_URI="mongodb://user:pass@localhost:10260/?tls=true&tlsAllowInvalidCertificates=true&directConnection=true" \ + python main.py + +Endpoints: + GET /health -> liveness/readiness (checks the gateway ping) + POST /books -> create a book + GET /books -> list books (optional ?author= filter) + GET /books/{id} -> fetch one book by id + PATCH /books/{id} -> update a book + DELETE /books/{id} -> delete a book + GET /stats/genres -> aggregation: count of books per genre +""" + +from __future__ import annotations + +import os +from datetime import datetime, timezone + +from bson import ObjectId +from bson.errors import InvalidId +from flask import Flask, jsonify, request + +import db +from models import book as book_model + +PORT = int(os.environ.get("PORT", "3000")) + +app = Flask(__name__) + +# Connect once at import/startup and ensure the demo indexes exist. +database = db.connect() +books = database[book_model.COLLECTION_NAME] +books.create_indexes(book_model.INDEXES) +print("Connected to DocumentDB via PyMongo") + + +def _parse_object_id(book_id: str) -> ObjectId | None: + try: + return ObjectId(book_id) + except (InvalidId, TypeError): + return None + + +@app.get("/health") +def health(): + """Return 200 only when the gateway responds to a ping.""" + try: + if db.ping(): + return jsonify(status="healthy", db="connected") + except Exception as exc: # noqa: BLE001 - report any connection error as unhealthy + return jsonify(status="unhealthy", db=str(exc)), 503 + return jsonify(status="unhealthy", db="disconnected"), 503 + + +@app.post("/books") +def create_book(): + payload = request.get_json(silent=True) or {} + if not payload.get("title") or not payload.get("author"): + return jsonify(error="title and author are required"), 400 + doc = book_model.build_document(payload) + result = books.insert_one(doc) + created = books.find_one({"_id": result.inserted_id}) + return jsonify(book_model.serialize(created)), 201 + + +@app.get("/books") +def list_books(): + author = request.args.get("author") + query = {"author": author} if author else {} + cursor = books.find(query).sort("created_at", -1).limit(100) + docs = [book_model.serialize(d) for d in cursor] + return jsonify(count=len(docs), books=docs) + + +@app.get("/books/") +def get_book(book_id: str): + oid = _parse_object_id(book_id) + if oid is None: + return jsonify(error="invalid id"), 400 + doc = books.find_one({"_id": oid}) + if doc is None: + return jsonify(error="not found"), 404 + return jsonify(book_model.serialize(doc)) + + +@app.patch("/books/") +def update_book(book_id: str): + oid = _parse_object_id(book_id) + if oid is None: + return jsonify(error="invalid id"), 400 + changes = request.get_json(silent=True) or {} + allowed = {"title", "author", "genres", "pages", "published", "in_stock", "rating"} + updates = {k: v for k, v in changes.items() if k in allowed} + if updates: + updates["updated_at"] = datetime.now(timezone.utc) + result = books.update_one({"_id": oid}, {"$set": updates}) + if result.matched_count == 0: + return jsonify(error="not found"), 404 + doc = books.find_one({"_id": oid}) + if doc is None: + return jsonify(error="not found"), 404 + return jsonify(book_model.serialize(doc)) + + +@app.delete("/books/") +def delete_book(book_id: str): + oid = _parse_object_id(book_id) + if oid is None: + return jsonify(error="invalid id"), 400 + result = books.delete_one({"_id": oid}) + if result.deleted_count == 0: + return jsonify(error="not found"), 404 + return "", 204 + + +@app.get("/stats/genres") +def genre_stats(): + pipeline = [ + {"$unwind": "$genres"}, + {"$group": {"_id": "$genres", "count": {"$sum": 1}}}, + {"$sort": {"count": -1}}, + ] + return jsonify(list(books.aggregate(pipeline))) + + +if __name__ == "__main__": + app.run(host="0.0.0.0", port=PORT) diff --git a/playgrounds/pymongo/app/models/__init__.py b/playgrounds/pymongo/app/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/playgrounds/pymongo/app/models/book.py b/playgrounds/pymongo/app/models/book.py new file mode 100644 index 0000000..2b3891a --- /dev/null +++ b/playgrounds/pymongo/app/models/book.py @@ -0,0 +1,63 @@ +"""Example "book" collection helpers for the PyMongo demo. + +PyMongo is a **raw driver**, not an ODM: there is no schema class. Instead this +module centralizes everything the app and tests need to treat a plain BSON +document as a "Book": + + - the collection name, + - the index definitions (no ``collation`` — DocumentDB does not implement it), + - a builder that applies defaults + timestamps to a create payload, and + - a JSON serializer that makes ``_id``/``datetime`` values response-friendly. + +It mirrors the Mongoose/Beanie ``Book`` models in the sibling playgrounds so the +three demos are directly comparable. +""" + +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Any, Mapping + +from pymongo import ASCENDING, IndexModel + +COLLECTION_NAME = "books" + +# Compound index: exercises DocumentDB index creation via PyMongo. +# No `collation` option here; DocumentDB does not implement it. +INDEXES = [ + IndexModel([("author", ASCENDING), ("title", ASCENDING)], name="author_title"), +] + + +def _utcnow() -> datetime: + return datetime.now(timezone.utc) + + +def build_document(payload: Mapping[str, Any]) -> dict: + """Apply defaults + timestamps to a create payload, returning a new dict.""" + now = _utcnow() + return { + "title": payload["title"], + "author": payload["author"], + "genres": list(payload.get("genres", [])), + "pages": payload.get("pages"), + "published": payload.get("published"), + "in_stock": payload.get("in_stock", True), + "rating": payload.get("rating"), + "created_at": now, + "updated_at": now, + } + + +def serialize(doc: Mapping[str, Any] | None) -> dict | None: + """Convert a BSON document into a JSON-serializable dict.""" + if doc is None: + return None + out = dict(doc) + if "_id" in out: + out["_id"] = str(out["_id"]) + for key in ("published", "created_at", "updated_at"): + value = out.get(key) + if isinstance(value, datetime): + out[key] = value.isoformat() + return out diff --git a/playgrounds/pymongo/app/pymongo_crud_test.py b/playgrounds/pymongo/app/pymongo_crud_test.py new file mode 100644 index 0000000..6e19821 --- /dev/null +++ b/playgrounds/pymongo/app/pymongo_crud_test.py @@ -0,0 +1,293 @@ +"""PyMongo CRUD/compatibility test against DocumentDB. + +Usage: + MONGO_URI="mongodb://user:pass@host:10260/?tls=true&tlsAllowInvalidCertificates=true&directConnection=true" \ + python pymongo_crud_test.py + +Or pass the URI as the first argument: + python pymongo_crud_test.py "mongodb://user:pass@host:10260/?..." + +Exercises connect, index creation, insert, find, update, aggregation, +unique-index enforcement, and delete using the raw PyMongo driver against the +DocumentDB gateway. Exits non-zero on the first real failure. +""" + +from __future__ import annotations + +import os +import re +import sys +import time +from datetime import datetime, timezone + +from pymongo import ASCENDING, DESCENDING, IndexModel, MongoClient, ReturnDocument +from pymongo.errors import DuplicateKeyError + +URI = sys.argv[1] if len(sys.argv) > 1 else os.environ.get("MONGO_URI") +DB_NAME = os.environ.get("MONGO_DB", "pymongo_test") +TLS_INSECURE = os.environ.get("TLS_INSECURE", "true").lower() != "false" + +passed = 0 +failed = 0 + + +def _ok(name: str) -> None: + global passed + passed += 1 + print(f" \u2705 {name}") + + +def _fail(name: str, err: object) -> None: + global failed + failed += 1 + print(f" \u274C {name}: {err}", file=sys.stderr) + + +def step(name: str, fn) -> None: + try: + fn() + _ok(name) + except Exception as err: # noqa: BLE001 - report any failure and continue + _fail(name, err) + + +def sanitize_uri(uri: str) -> str: + return re.sub(r"[?&]replicaSet=[^&]*", "", uri) + + +def run() -> int: + if not URI: + print("MONGO_URI not provided. Pass it as $1 or set MONGO_URI.", file=sys.stderr) + return 2 + + print("PyMongo DocumentDB compatibility test") + print("=====================================") + + # Fresh collection per run keeps the test idempotent. + coll_name = f"widgets_{int(time.time() * 1000)}" + + client: MongoClient | None = None + state: dict = {"coll": None, "created_id": None} + + def _connect() -> None: + nonlocal client + client = MongoClient( + sanitize_uri(URI), + directConnection=True, + tls=True, + tlsAllowInvalidCertificates=TLS_INSECURE, + serverSelectionTimeoutMS=15000, + ) + ping = client.admin.command("ping") + if ping.get("ok") != 1: + raise RuntimeError("ping did not return ok:1") + state["coll"] = client[DB_NAME][coll_name] + + step("connect", _connect) + + if client is None or state["coll"] is None: + print("\nConnection failed; aborting remaining steps.", file=sys.stderr) + return 1 + + coll = state["coll"] + + def _create_indexes() -> None: + coll.create_indexes( + [ + IndexModel([("sku", ASCENDING)], unique=True), + IndexModel([("name", ASCENDING), ("price", DESCENDING)]), + ] + ) + + step("create indexes", _create_indexes) + + def _insert_one() -> None: + res = coll.insert_one( + { + "sku": "SKU-001", + "name": "Gizmo", + "tags": ["alpha", "beta"], + "price": 9.99, + "active": True, + "created_at": datetime.now(timezone.utc), + } + ) + state["created_id"] = res.inserted_id + if state["created_id"] is None: + raise RuntimeError("no _id returned") + + step("insert_one", _insert_one) + + def _insert_many() -> None: + res = coll.insert_many( + [ + {"sku": "SKU-002", "name": "Gadget", "tags": ["beta"], "price": 19.5}, + {"sku": "SKU-003", "name": "Widget", "tags": ["alpha", "gamma"], "price": 4.25}, + ] + ) + if len(res.inserted_ids) != 2: + raise RuntimeError(f"expected 2 inserted, got {len(res.inserted_ids)}") + + step("insert_many", _insert_many) + + def _get_by_id() -> None: + doc = coll.find_one({"_id": state["created_id"]}) + if doc is None or doc.get("sku") != "SKU-001": + raise RuntimeError("document not found or mismatched") + + step("find_one by _id", _get_by_id) + + def _find_filter_sort_limit() -> None: + docs = list(coll.find({"price": {"$gte": 5}}).sort("price", DESCENDING).limit(10)) + if len(docs) != 2: + raise RuntimeError(f"expected 2 docs, got {len(docs)}") + if docs[0]["price"] < docs[1]["price"]: + raise RuntimeError("sort order incorrect") + + step("find with filter + sort + limit", _find_filter_sort_limit) + + def _count() -> None: + n = coll.count_documents({}) + if n != 3: + raise RuntimeError(f"expected 3 docs, got {n}") + + step("count_documents", _count) + + def _update_one() -> None: + res = coll.update_one({"sku": "SKU-002"}, {"$set": {"price": 21}}) + if res.modified_count != 1: + raise RuntimeError(f"expected 1 modified, got {res.modified_count}") + + step("update_one ($set)", _update_one) + + def _find_one_and_update() -> None: + doc = coll.find_one_and_update( + {"sku": "SKU-003"}, + {"$push": {"tags": "delta"}}, + return_document=ReturnDocument.AFTER, + ) + if not doc or "delta" not in doc.get("tags", []): + raise RuntimeError("update not applied") + + step("find_one_and_update (returns new)", _find_one_and_update) + + def _aggregate() -> None: + stats = list( + coll.aggregate( + [ + {"$unwind": "$tags"}, + {"$group": {"_id": "$tags", "count": {"$sum": 1}}}, + {"$sort": {"count": -1}}, + ] + ) + ) + if not stats: + raise RuntimeError("aggregation returned no results") + + step("aggregation ($unwind/$group)", _aggregate) + + def _unique_index() -> None: + try: + coll.insert_one({"sku": "SKU-001", "name": "Duplicate"}) + except DuplicateKeyError: + return + raise RuntimeError("duplicate insert was not rejected") + + step("unique index enforcement (duplicate sku rejected)", _unique_index) + + def _delete_one() -> None: + res = coll.delete_one({"sku": "SKU-002"}) + if res.deleted_count != 1: + raise RuntimeError(f"expected 1 deleted, got {res.deleted_count}") + + step("delete_one", _delete_one) + + # Vector search: DocumentDB supports a `cosmosSearch` vector index and the + # `$vectorSearch` aggregation stage. The index is created with the raw + # `createIndexes` command since it is not a standard MongoDB index type. + db_ref = client[DB_NAME] + vec_name = f"vectors_{int(time.time() * 1000)}" + vectors = db_ref[vec_name] + + def _vector_index_insert() -> None: + vectors.insert_many( + [ + {"name": "a", "v": [1, 0, 0]}, + {"name": "b", "v": [0.9, 0.1, 0]}, + {"name": "c", "v": [0, 0, 1]}, + ] + ) + res = db_ref.command( + { + "createIndexes": vec_name, + "indexes": [ + { + "name": "v_ivf", + "key": {"v": "cosmosSearch"}, + "cosmosSearchOptions": { + "kind": "vector-ivf", + "numLists": 1, + "similarity": "COS", + "dimensions": 3, + }, + } + ], + } + ) + if res.get("ok") != 1: + raise RuntimeError("createIndexes did not return ok:1") + + step("vector index + insert (cosmosSearch vector-ivf)", _vector_index_insert) + + def _vector_search() -> None: + hits = list( + vectors.aggregate( + [ + { + "$vectorSearch": { + "index": "v_ivf", + "path": "v", + "queryVector": [1, 0, 0], + "numCandidates": 10, + "limit": 2, + } + }, + {"$project": {"name": 1, "_id": 0}}, + ] + ) + ) + if not hits: + raise RuntimeError("vector search returned no results") + if hits[0].get("name") != "a": + raise RuntimeError(f"expected nearest 'a', got {hits[0].get('name')!r}") + + step("$vectorSearch returns nearest neighbor", _vector_search) + + def _vector_cleanup() -> None: + vectors.drop() + + step("vector cleanup (drop collection)", _vector_cleanup) + + def _drop() -> None: + coll.drop() + + step("cleanup (drop collection)", _drop) + + client.close() + + print("\n=====================================") + print(f"Passed: {passed} Failed: {failed}") + return 0 if failed == 0 else 1 + + +def main() -> None: + try: + code = run() + except Exception as err: # noqa: BLE001 + print(f"Unexpected error: {err}", file=sys.stderr) + sys.exit(1) + sys.exit(code) + + +if __name__ == "__main__": + main() diff --git a/playgrounds/pymongo/app/requirements.txt b/playgrounds/pymongo/app/requirements.txt new file mode 100644 index 0000000..9aa133e --- /dev/null +++ b/playgrounds/pymongo/app/requirements.txt @@ -0,0 +1,2 @@ +pymongo>=4.9 +flask>=3.0 diff --git a/playgrounds/pymongo/scripts/lib.sh b/playgrounds/pymongo/scripts/lib.sh new file mode 100755 index 0000000..c030dda --- /dev/null +++ b/playgrounds/pymongo/scripts/lib.sh @@ -0,0 +1,125 @@ +#!/usr/bin/env bash +# Shared helpers for the PyMongo local playground scripts. +# +# Everything runs on your machine: the DocumentDB local emulator runs in Docker +# and the app/test run as local Python processes that connect to it directly. +set -euo pipefail + +# Connection defaults. Override any of these via environment variables. +DOCUMENTDB_CONTAINER="${DOCUMENTDB_CONTAINER:-documentdb-local}" +DOCUMENTDB_IMAGE="${DOCUMENTDB_IMAGE:-ghcr.io/documentdb/documentdb/documentdb-local:latest}" +DOCUMENTDB_HOST="${DOCUMENTDB_HOST:-localhost}" +DOCUMENTDB_PORT="${DOCUMENTDB_PORT:-10260}" +# Note: the emulator rejects some reserved names (e.g. "documentdb"); use a +# distinct admin username. +DOCUMENTDB_USERNAME="${DOCUMENTDB_USERNAME:-docdbadmin}" +DOCUMENTDB_PASSWORD="${DOCUMENTDB_PASSWORD:-Documentdb!Local1}" + +# Build the MongoDB connection string for the local emulator. The gateway only +# speaks TLS and advertises itself as a standalone server, so we request TLS, +# accept its self-signed cert, and use a direct connection. +build_uri() { + echo "mongodb://${DOCUMENTDB_USERNAME}:${DOCUMENTDB_PASSWORD}@${DOCUMENTDB_HOST}:${DOCUMENTDB_PORT}/?tls=true&tlsAllowInvalidCertificates=true&directConnection=true" +} + +# Wait until a TCP port accepts connections. +# Args: [tries] +wait_for_port() { + local host="$1" port="$2" tries="${3:-60}" i + for i in $(seq 1 "$tries"); do + if (exec 3<>"/dev/tcp/${host}/${port}") 2>/dev/null; then + exec 3>&- 3<&- 2>/dev/null || true + return 0 + fi + sleep 1 + done + return 1 +} + +# Create (once) a Python virtualenv in the app directory and install +# requirements. Echoes the path to the venv's python interpreter. +# Args: +setup_venv() { + local app_dir="$1" + local venv_dir="$app_dir/.venv" + + if [ ! -d "$venv_dir" ]; then + echo "Creating Python virtualenv in $venv_dir ..." >&2 + python3 -m venv "$venv_dir" + fi + "$venv_dir/bin/pip" install --quiet --disable-pip-version-check \ + -r "$app_dir/requirements.txt" >&2 + echo "$venv_dir/bin/python" +} + +require_docker() { + command -v docker >/dev/null || { + echo "docker is required (start Docker Desktop / the Docker daemon first)" >&2 + exit 1 + } + docker info >/dev/null 2>&1 || { + echo "Cannot reach the Docker daemon. Is Docker running?" >&2 + exit 1 + } +} + +container_running() { + [ "$(docker inspect -f '{{.State.Running}}' "$DOCUMENTDB_CONTAINER" 2>/dev/null)" = "true" ] +} + +container_exists() { + docker inspect "$DOCUMENTDB_CONTAINER" >/dev/null 2>&1 +} + +ensure_documentdb() { + require_docker + + if container_running; then + echo "DocumentDB container '$DOCUMENTDB_CONTAINER' is already running." + else + if container_exists; then + docker rm -f "$DOCUMENTDB_CONTAINER" >/dev/null 2>&1 || true + fi + + echo "Starting DocumentDB container '$DOCUMENTDB_CONTAINER' on port ${DOCUMENTDB_PORT} ..." + docker run -dt \ + -p "127.0.0.1:${DOCUMENTDB_PORT}:10260" \ + --name "$DOCUMENTDB_CONTAINER" \ + "$DOCUMENTDB_IMAGE" \ + --username "$DOCUMENTDB_USERNAME" \ + --password "$DOCUMENTDB_PASSWORD" >/dev/null + fi + + wait_for_documentdb +} + +wait_for_documentdb() { + local uri + uri="$(build_uri)" + echo "Waiting for DocumentDB to accept connections ..." + + local i + for i in $(seq 1 60); do + if docker exec "$DOCUMENTDB_CONTAINER" mongosh "$uri" \ + --quiet --eval 'db.adminCommand({ ping: 1 })' >/dev/null 2>&1; then + echo "DocumentDB is ready." + return 0 + fi + sleep 2 + done + + echo "DocumentDB did not become ready in time." >&2 + echo "Check logs with: docker logs $DOCUMENTDB_CONTAINER" >&2 + return 1 +} + +stop_documentdb() { + require_docker + if container_exists; then + echo "Removing DocumentDB container '$DOCUMENTDB_CONTAINER' ..." + docker rm -f "$DOCUMENTDB_CONTAINER" >/dev/null + echo "Done." + else + echo "No DocumentDB container named '$DOCUMENTDB_CONTAINER' found." + fi +} diff --git a/playgrounds/pymongo/scripts/run-app.sh b/playgrounds/pymongo/scripts/run-app.sh new file mode 100755 index 0000000..495be2a --- /dev/null +++ b/playgrounds/pymongo/scripts/run-app.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +# Run the PyMongo demo API locally against a local DocumentDB container. +# +# Starts DocumentDB in Docker (if not already running), sets up the app's +# Python virtualenv/dependencies, then runs the Flask + PyMongo server. +# +# Prerequisites: docker, python3. +set -euo pipefail + +command -v python3 >/dev/null || { echo "python3 is required" >&2; exit 1; } + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +# shellcheck source=lib.sh +source "$SCRIPT_DIR/lib.sh" + +APP_DIR="$SCRIPT_DIR/../app" +PORT="${PORT:-3000}" + +ensure_documentdb + +echo "Setting up Python virtualenv and installing dependencies..." +VENV_PY=$(setup_venv "$APP_DIR") + +echo "" +echo "=== PyMongo demo API running locally ===" +echo "API: http://localhost:${PORT}" +echo "Health: curl http://localhost:${PORT}/health" +echo "Create: curl -X POST http://localhost:${PORT}/books -H 'Content-Type: application/json' \\" +echo " -d '{\"title\":\"Dune\",\"author\":\"Herbert\",\"genres\":[\"sci-fi\"],\"pages\":412}'" +echo "Press Ctrl-C to stop (DocumentDB keeps running; stop it with ./scripts/stop-documentdb.sh)." +echo "" + +cd "$APP_DIR" +MONGO_URI="$(build_uri)" PORT="$PORT" "$VENV_PY" main.py diff --git a/playgrounds/pymongo/scripts/run-test.sh b/playgrounds/pymongo/scripts/run-test.sh new file mode 100755 index 0000000..fe12e04 --- /dev/null +++ b/playgrounds/pymongo/scripts/run-test.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +# Run the PyMongo CRUD/compatibility test suite end-to-end against a local +# DocumentDB container. +# +# Starts DocumentDB in Docker (if not already running), then sets up the app +# virtualenv/dependencies and runs the standalone test suite locally. +# +# Prerequisites: docker, python3. +# +# Set KEEP_DB=0 to remove the DocumentDB container when the tests finish +# (default keeps it running for fast re-runs). +set -euo pipefail + +command -v python3 >/dev/null || { echo "python3 is required" >&2; exit 1; } + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +# shellcheck source=lib.sh +source "$SCRIPT_DIR/lib.sh" + +APP_DIR="$SCRIPT_DIR/../app" +KEEP_DB="${KEEP_DB:-1}" + +ensure_documentdb + +if [ "$KEEP_DB" != "1" ]; then + trap 'stop_documentdb' EXIT +fi + +echo "Setting up Python virtualenv and installing dependencies..." +VENV_PY=$(setup_venv "$APP_DIR") + +echo "" +cd "$APP_DIR" +MONGO_URI="$(build_uri)" "$VENV_PY" pymongo_crud_test.py diff --git a/playgrounds/pymongo/scripts/start-documentdb.sh b/playgrounds/pymongo/scripts/start-documentdb.sh new file mode 100755 index 0000000..7d62c16 --- /dev/null +++ b/playgrounds/pymongo/scripts/start-documentdb.sh @@ -0,0 +1,22 @@ +#!/usr/bin/env bash +# Start the DocumentDB local emulator in Docker on your machine. +# +# Idempotent: if the container is already running, this just verifies readiness. +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +# shellcheck source=lib.sh +source "$SCRIPT_DIR/lib.sh" + +ensure_documentdb + +echo "" +echo "=== DocumentDB local emulator is ready ===" +echo "Container: $DOCUMENTDB_CONTAINER (image: $DOCUMENTDB_IMAGE)" +echo "Connection string:" +echo " $(build_uri)" +echo "" +echo "Next:" +echo " ./scripts/run-test.sh # start DocumentDB if needed, then run tests" +echo " ./scripts/run-app.sh # start DocumentDB if needed, then run API" +echo " ./scripts/stop-documentdb.sh # stop + remove the emulator" diff --git a/playgrounds/pymongo/scripts/stop-documentdb.sh b/playgrounds/pymongo/scripts/stop-documentdb.sh new file mode 100755 index 0000000..d045e94 --- /dev/null +++ b/playgrounds/pymongo/scripts/stop-documentdb.sh @@ -0,0 +1,12 @@ +#!/usr/bin/env bash +# Stop and remove the DocumentDB local emulator container. +# +# This removes the container and all data it held (the emulator stores data +# inside the container, so this is a full reset). +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +# shellcheck source=lib.sh +source "$SCRIPT_DIR/lib.sh" + +stop_documentdb