mongo_connect.py

import pandas as pd
from pymongo import MongoClient
import settings

def _connect_mongo(host=settings.MONGO_HOST, port=settings.MONGO_PORT,
        username=settings.MONGO_USER, password=settings.MONGO_PASSWORD, db=settings.MONGO_DB):
    """ A util for making a connection to mongo """
    if username and password:
        mongo_uri = 'mongodb://%s:%s@%s:%s/%s' % (username, password, host, port, db)
        conn = MongoClient(mongo_uri)
    else:
        conn = MongoClient(host, port)
    return conn[db]


def get_collection(collection, host=settings.MONGO_HOST, port=settings.MONGO_PORT, 
                   username=settings.MONGO_USER, password=settings.MONGO_PASSWORD, db=settings.MONGO_DB):
    db = _connect_mongo(host=host, port=port, username=username, password=password, db=db)
    return db[collection]

def read_mongo(collection, filter_=None, projection=None, host=settings.MONGO_HOST, port=settings.MONGO_PORT, 
                   username=settings.MONGO_USER, password=settings.MONGO_PASSWORD, db=settings.MONGO_DB, no_id=True):
    """ Read from Mongo and Store into DataFrame """

    my_collection = get_collection(collection, host=host, port=port, username=username, password=password, 
                                   db=settings.MONGO_DB)

    # Make a query to the specific DB and Collection
    cursor = my_collection.find(filter_, projection)

    # Expand the cursor and construct the DataFrame
    df =  pd.DataFrame(list(cursor))

    # Delete the _id
    if no_id:
        del df['_id']
    
    if '__v' in df:
        del df['__v']

    return df

def write_mongo(dataframe, collection):
    col_ = get_collection(collection)
    col_.delete_many( {} )
    json = dataframe.to_dict(orient='index')
    values = list(json.values())
    results = col_.insert_many(values)
    print('{} records inserted in {}.'.format(len(results.inserted_ids), collection))