sdg/mongo_example.py

"""
Quick smoke-test: write a document to MongoDB and read it back.

Usage:
    python -m sdg.test_mongo                          # localhost:27017
    python -m sdg.test_mongo --mongo-uri mongodb://HOST:PORT
"""

import argparse
import sys
from pymongo import MongoClient


def main() -> None:
    parser = argparse.ArgumentParser(description="Test MongoDB read/write")
    parser.add_argument("--mongo-uri", default="mongodb://john5.stanford.edu:27017",
                        help="MongoDB connection URI")
    args = parser.parse_args()

    print(f"Connecting to {args.mongo_uri} ...")
    client = MongoClient(args.mongo_uri, serverSelectionTimeoutMS=5000)

    # Verify connectivity
    client.admin.command("ping")
    print("Ping OK")

    db = client["sdg_test"]
    col = db["smoke_test"]

    # Write
    doc = {"key": "hello", "value": "world"}
    col.replace_one({"key": "hello"}, doc, upsert=True)
    print(f"Wrote: {doc}")

    # Read
    result = col.find_one({"key": "hello"})
    assert result is not None, "Read back None — write may have failed"
    print(f"Read:  key={result['key']}  value={result['value']}")

    # Cleanup
    col.delete_one({"key": "hello"})
    print("Cleaned up test document")

    client.close()
    print("All good!")


if __name__ == "__main__":
    main()