Cache MongoDB Document IDs in Redis SET with Python

cache_mongo_ids.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from pymongo import MongoClient
from bson import ObjectId
import redis


def cache_ids(db, cache, cache_key, reset=True, limit=1000):
    if reset:
        # O(M) where M = number of items in key
        cache.delete(cache_key)

    # adjust your query, projection and other parameters as needed
    cursor = db.find({}, {'_id': 1}).limit(limit)

    # O(N) where N is the number of IDs to be cached
    for doc in cursor:
        # if you're on Python 3, do str(doc['_id']) instead
        cache.sadd(cache_key, unicode(doc['_id']))

if __name__ == '__main__':
    db = MongoClient('mongodb://localhost:27017')['some_db_name']
    cache = redis.from_url('redis://localhost:6379', db=0)
    cache_key = 'id_set'

    cache_ids(db, cache, cache_key)

    # Random ID determination? Redis SRANDMEMBER!
    _id = cache.srandmember(cache_key)
    
    # Random doc
    doc = db[COLLECTION].find_one({'_id': ObjectId(_id)})
    
    # Need N random IDs? Redis SRANDMEMBER still!
    _ids = [ObjectId(_id) for _id in cache.srandmember(cache_key, N)]
    
    # Random docs
    docs = db[COLLECTION].find({'_id': {'$in': _ids}})