Generic indexes in mongo test

generic_mongo_indexes.py

from functools import partial
from pprint import pprint

from pymongo import Connection, ASCENDING


DATABASE_NAME = 'index_test'


def setup():
    get_pool().testcol.create_index([('entity', ASCENDING), ('indexes', ASCENDING)])


def teardown():
    get_pool().connection.drop_database(DATABASE_NAME)


def get_pool():
    return Connection()[DATABASE_NAME]


def add_indexes(document, indexes):
    if 'indexes' not in document:
        document['indexes'] = []

    for index in indexes:
        document['indexes'].append([index, document[index]])

    return document


def main():
    try:
        setup()

        docs = [{'entity': 'person', 'firstname': 'bob', 'lastname': 'smith', 'age': 100},
                {'entity': 'person', 'firstname': 'fred', 'lastname': 'smith'},
                {'entity': 'pet',    'firstname': 'fluffy', 'lastname': 'smith'},
                {'entity': 'person', 'firstname': 'adam', 'lastname': 'wagner'},
                {'entity': 'person', 'firstname': 'james', 'lastname': 'bond'}]
        docs = map(partial(add_indexes, indexes=['firstname', 'lastname']), docs)
        
        pprint(docs)

        print '-' * 100

        db = get_pool()
        db.testcol.insert(docs)
        query = db.testcol.find({
            'entity': 'pet',
            'indexes': {
                '$gte': ['lastname', 'smith']
            }
        })
        pprint(list(query))

        print '-' * 100

        pprint(query.explain())


    finally:
        teardown()


if __name__ == '__main__':
    main()