pymongo_aug.py

from pymongo import MongoClient
from pymongo.collection import Collection
from pymongo.cursor import Cursor


def update_d(d, add=None, rm=None):
    if add is not None:
        d.update(add)
    if rm is not None:
        for k in rm:
            del d[k]
    return d


def get_collection(db_name, collection, **conn_kwargs):
    """
    Gets collection

    :keyword db_name: database name
    :type db_name: ``str``

    :keyword collection: collection name
    :type collection: ``str``

    :return: Collection
    :rtype: ``Collection``
    """
    client = MongoClient(**conn_kwargs)
    db = client[db_name]
    return db[collection]


def get_top_ancestors(modulestore):
    """
    Gets top ancestors

    :keyword modulestore: modulestore containing the blocks
    :type modulestore: ``Collection``

    :return: Cursor
    :rtype: ``Cursor``
    """
    return modulestore.find({'exam': {'$type': 10}})


def update_descendants(modulestore, blocks, ancestor_fields):
    """
    Gets top ancestors

    :keyword modulestore: modulestore containing the blocks
    :type modulestore: ``Collection``

    :keyword blocks: iterator over the blocks (collections within modulestore)
    :type blocks: ``Cursor`` | `tuple`

    :keyword ancestor_fields: fields of the top most ancestor
    :type ancestor_fields: ``dict``
    """
    for block in blocks:
        modulestore.replace_one({'block_id': block['block_id'],
                                 'block_type': block['block_type']},
                                update_d(block, add={'ancestor_fields': ancestor_fields},
                                         rm=('_id',)))
        update_descendants.counter += 1
        print 'Iteration:', update_descendants.counter

        if 'children' in block and block['children']:
            for block_type, block_id in block['children']:
                update_descendants(modulestore,
                                   modulestore.find({'block_id': block_id,
                                                     'block_type': block_type,
                                                     # 'ancestor_fields': {'$exists': False}
                                                     }),
                                   ancestor_fields)


update_descendants.counter = 0

if __name__ == '__main__':
    ms = get_collection('edxapp', 'modulestore.mapped0')
    for _block in get_top_ancestors(ms):
        update_descendants(ms, (_block,), _block['fields'])