MongoDB Atlas Search Client

MongoDB Atlas Search Client

Problem Statement

Full-text search is one of the most common requirements in modern applications, but implementing it correctly is surprisingly hard. MongoDB's standard $text operator only supports basic stemming and language-specific stop words — it can't do fuzzy matching, autocomplete, phrase proximity, or relevance scoring with custom weights. Meanwhile, standing up a separate Elasticsearch cluster just to search a collection you already have in MongoDB adds significant operational overhead.

MongoDB Atlas Search solves this by embedding a Lucene-based search engine directly into Atlas. You define an index on your collection, and the $search aggregation stage gives you fuzzy matching, autocomplete, highlighting, and custom scoring — all without leaving the MongoDB query language. This script wraps that capability in a reusable CLI tool so you can prototype queries, test index definitions, and validate relevance tuning from the terminal before wiring it into an application.

How It Works

The script connects to an Atlas cluster using pymongo, then runs an aggregation pipeline where the first stage is always $search. The $search stage accepts a query operator (text, phrase, compound, autocomplete, etc.) along with highlight and count options. Subsequent stages ($project, $limit, $addFields for score) shape the output.

Key design decisions:

• Operator selection via CLI flags — --mode text|phrase|autocomplete|compound maps to the right Lucene operator, so you don't need to remember the JSON shape.
• Score projection — Atlas attaches a searchScore to each document; we surface it and sort by it so relevance is visible.
• Highlight extraction — Instead of dumping raw HTML fragments, we extract the matched terms so terminal output stays readable.
• Graceful degradation — If the index doesn't exist or the URI is malformed, we raise clear exceptions instead of tracebacks.

The script uses pymongo.MongoClient with serverSelectionTimeoutMS set low so connection failures surface quickly, and it never writes to the database — it's read-only by design.

Full Script

#!/usr/bin/env python3
"""
MongoDB Atlas Search Client

A command-line tool for querying MongoDB Atlas Search indexes with support
for text, phrase, autocomplete, and compound queries. Prints ranked results
with relevance scores and matched highlights.

Usage:
    python3 atlas_search.py --uri 'mongodb+srv://...' --index products \
        --query 'wireless keyboard' --limit 20
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from typing import Any

from pymongo import MongoClient
from pymongo.errors import (
    ConfigurationError,
    OperationFailure,
    PyMongoError,
    ServerSelectionTimeoutError,
)

# ---------------------------------------------------------------------------
# Configuration (overridable via environment variables)
# ---------------------------------------------------------------------------
DEFAULT_URI = os.environ.get("MONGODB_URI", "mongodb://localhost:27017")
DEFAULT_INDEX = os.environ.get("ATLAS_SEARCH_INDEX", "default")
DEFAULT_DB = os.environ.get("MONGODB_DB", "test")
DEFAULT_COLLECTION = os.environ.get("MONGODB_COLLECTION", "items")
DEFAULT_FIELD = os.environ.get("ATLAS_SEARCH_FIELD", "description")
SERVER_SELECTION_TIMEOUT_MS = int(os.environ.get("MONGO_TIMEOUT_MS", "5000"))

# ---------------------------------------------------------------------------
# Query builders — one per supported search mode
# ---------------------------------------------------------------------------

def build_text_query(query: str, field: str, fuzzy: bool) -> dict[str, Any]:
    """Standard full-text search with optional fuzzy matching."""
    op: dict[str, Any] = {
        "query": query,
        "path": field,
    }
    if fuzzy:
        # maxEdits=1 tolerates a single-char typo; prefixLength=1 keeps
        # the first character exact which dramatically reduces false hits.
        op["fuzzy"] = {"maxEdits": 1, "prefixLength": 1}
    return {"text": op}


def build_phrase_query(query: str, field: str, slop: int) -> dict[str, Any]:
    """Exact phrase match with configurable slop (word distance)."""
    return {
        "phrase": {
            "query": query,
            "path": field,
            "slop": slop,
        }
    }


def build_autocomplete_query(query: str, field: str) -> dict[str, Any]:
    """Autocomplete operator — requires the field to be indexed as 'autocomplete'."""
    return {
        "autocomplete": {
            "query": query,
            "path": field,
            "tokenOrder": "sequential",
        }
    }


def build_compound_query(query: str, field: str) -> dict[str, Any]:
    """Compound query: must-match the field, boost exact phrase hits."""
    return {
        "compound": {
            "must": [{"text": {"query": query, "path": field}}],
            "should": [
                {
                    "phrase": {
                        "query": query,
                        "path": field,
                        "score": {"boost": {"value": 2}},
                    }
                }
            ],
            "minimumShouldMatch": 0,
        }
    }


QUERY_BUILDERS = {
    "text": lambda q, f, a: build_text_query(q, f, fuzzy=a.fuzzy),
    "phrase": lambda q, f, a: build_phrase_query(q, f, slop=a.slop),
    "autocomplete": lambda q, f, a: build_autocomplete_query(q, f),
    "compound": lambda q, f, a: build_compound_query(q, f),
}

# ---------------------------------------------------------------------------
# Core execution
# ---------------------------------------------------------------------------

def run_search(
    uri: str,
    db_name: str,
    coll_name: str,
    index_name: str,
    query: str,
    field: str,
    mode: str,
    limit: int,
    args: argparse.Namespace,
) -> list[dict[str, Any]]:
    """Execute the Atlas Search aggregation and return ranked documents."""
    client = MongoClient(uri, serverSelectionTimeoutMS=SERVER_SELECTION_TIMEOUT_MS)

    try:
        # Force a round-trip so connection errors surface here, not later.
        client.admin.command("ping")
    except ServerSelectionTimeoutError as exc:
        raise ConnectionError(
            f"Could not reach MongoDB at {uri!r} within "
            f"{SERVER_SELECTION_TIMEOUT_MS}ms. Check the URI and your IP allowlist."
        ) from exc

    collection = client[db_name][coll_name]

    search_stage = {
        "$search": {
            "index": index_name,
            "highlight": {"path": field},
        }
    }
    # Merge the operator-specific payload into the $search stage.
    search_stage["$search"].update(QUERY_BUILDERS[mode](query, field, args))

    pipeline = [
        search_stage,
        # Expose the Lucene relevance score so we can sort and display it.
        {"$addFields": {"score": {"$meta": "searchScore"}}},
        {"$limit": limit},
        {"$project": {"_id": 0, "score": 1, "highlights": {"$meta": "searchHighlights"},
                      "doc": "$$ROOT"}},
    ]

    try:
        return list(collection.aggregate(pipeline))
    except OperationFailure as exc:
        # Most common cause: the named Atlas Search index doesn't exist yet.
        if "index not found" in str(exc).lower():
            raise LookupError(
                f"Atlas Search index {index_name!r} not found on "
                f"{db_name}.{coll_name}. Create it in the Atlas UI first."
            ) from exc
        raise
    finally:
        client.close()


# ---------------------------------------------------------------------------
# Output formatting
# ---------------------------------------------------------------------------

def extract_highlights(highlights: list[dict[str, Any]] | None) -> list[str]:
    """Flatten Atlas highlight fragments into a list of plain matched phrases."""
    if not highlights:
        return []
    phrases: list[str] = []
    for hl in highlights:
        for text in hl.get("texts", []):
            # Strip the <em>...</em> markers Atlas injects.
            cleaned = text["value"].replace("<em>", "").replace("</em>", "")
            phrases.append(cleaned.strip())
    return phrases


def print_results(results: list[dict[str, Any]], field: str) -> None:
    """Pretty-print results to stdout."""
    if not results:
        print("No results found.")
        return

    for i, hit in enumerate(results, start=1):
        doc = hit.get("doc", {})
        score = hit.get("score", 0.0)
        print(f"\n[{i}] score={score:.4f}")
        # Show the searched field first, then any other scalar fields.
        if field in doc:
            print(f"    {field}: {doc[field]}")
        for key, value in doc.items():
            if key == field or key.startswith("$"):
                continue
            if isinstance(value, (str, int, float, bool)):
                print(f"    {key}: {value}")
        for phrase in extract_highlights(hit.get("highlights")):
            print(f"    ↳ match: {phrase}")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Query a MongoDB Atlas Search index from the command line.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--uri", default=DEFAULT_URI, help="MongoDB connection string")
    parser.add_argument("--db", default=DEFAULT_DB, help="Database name")
    parser.add_argument("--collection", default=DEFAULT_COLLECTION, help="Collection name")
    parser.add_argument("--index", default=DEFAULT_INDEX, help="Atlas Search index name")
    parser.add_argument("--field", default=DEFAULT_FIELD, help="Field to search")
    parser.add_argument("--query", required=True, help="Search query string")
    parser.add_argument(
        "--mode",
        choices=list(QUERY_BUILDERS.keys()),
        default="text",
        help="Search operator to use",
    )
    parser.add_argument("--limit", type=int, default=10, help="Max results to return")
    parser.add_argument("--fuzzy", action="store_true", help="Enable fuzzy matching (text mode)")
    parser.add_argument("--slop", type=int, default=0, help="Word distance for phrase mode")
    parser.add_argument("--json", action="store_true", help="Emit raw JSON instead of text")
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)

    if args.limit < 1:
        print("error: --limit must be >= 1", file=sys.stderr)
        return 2

    try:
        results = run_search(
            uri=args.uri,
            db_name=args.db,
            coll_name=args.collection,
            index_name=args.index,
            query=args.query,
            field=args.field,
            mode=args.mode,
            limit=args.limit,
            args=args,
        )
    except (ConnectionError, LookupError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1
    except ConfigurationError as exc:
        print(f"error: invalid configuration — {exc}", file=sys.stderr)
        return 1
    except PyMongoError as exc:
        print(f"error: MongoDB operation failed — {exc}", file=sys.stderr)
        return 1

    if args.json:
        print(json.dumps(results, default=str, indent=2))
    else:
        print_results(results, args.field)
    return 0


if __name__ == "__main__":
    sys.exit(main())

Configuration

Usage Examples

Basic full-text with fuzzy matching:

python3 atlas_search.py \
  --uri 'mongodb+srv://app:pw@cluster0.abc.mongodb.net' \
  --db shop --collection products \
  --index products_idx --field description \
  --query 'wireless keyboard' --mode text --fuzzy --limit 5

[1] score=2.8413
    description: Ergonomic wireless keyboard with backlit keys
    price: 79.99
    ↳ match: wireless keyboard
[2] score=1.9204
    description: Compact wireless mouse and keyboard combo
    price: 45.00

Phrase search with slop (matching "red shoes" within 2 words):

python3 atlas_search.py --query 'red shoes' --mode phrase --slop 2 --limit 3

Autocomplete mode with JSON output for piping into jq:

python3 atlas_search.py --query 'wire' --mode autocomplete --json | jq '.[].doc.name'

Error Handling

• ServerSelectionTimeoutError → wrapped as ConnectionError with a hint about IP allowlists (the #1 cause of Atlas connection failures).
• OperationFailure: index not found → wrapped as LookupError telling the user to create the index in the Atlas UI.
• ConfigurationError (bad URI scheme, missing SRV record) → caught and reported without a traceback.
• PyMongoError (catch-all for network drops, auth failures) → reported cleanly.
• --limit < 1 → validated up front, exits with code 2.
• The finally: client.close() block ensures sockets are released even when the aggregation raises.

Exit codes: 0 success, 1 runtime error, 2 bad arguments.

Extending Ideas

• Index management: add a --create-index subcommand that calls the Atlas Admin API to POST a JSON index definition, so the tool can bootstrap itself against a fresh cluster.
• Faceted search: extend the pipeline with $searchMeta + $facet to return category/price-bucket counts alongside hits — useful for building search UIs.
• Vector search: Atlas supports $vectorSearch on the same cluster; add a --mode vector that embeds the query via an embedding API and searches a knnVector field.
• Interactive REPL: wrap run_search in a cmd.Cmd subclass so users can iterate on queries without re-invoking the process (avoids reconnect cost).