ufmongo merger

ufmongo.py

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

import sys, os, argparse
import logging

import csv
import json
from functools import partial

import pymongo
from pymongo import MongoClient

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Created Dict from Mongo results
def collection_to_dict(mongo_db, coll_name, key_name):
    collection=mongo_db[coll_name]
    cursor= collection.find({})
    return dict((k, v)
         for (k, v) in
                [(doc.get(key_name), doc)
                  for doc in cursor])

def process(fk_set1, fk_set2, doc):
    try:
        fk_obj1 = fk_set1.get(doc["fk_name1"])
        fk_obj2 = fk_set2.get(doc["fk_name2"])

        # Merge objs create Vector of CSV ROW
        csv_row = [
            doc.get(""),
            doc.get(""),
            fk_obj1.get(""),
            fk_obj2.get("")
        ]
        return csv_row 
    except Exception, e:
        logger.error(str(e), exc_info=True)
        

if __name__ == "__main__":
    
    desc = '''
examples:
    ufmongo 
    '''
    parser = argparse.ArgumentParser(
        description=" ufmongo ", 
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=desc)
    # parser.add_argument("input_file", help="input file")
    parser.add_argument("output_file", help="output file")
    # parser.add_argument("-d", "--dump_exceptions", action="store_true", help="dump exception rows to stdout, Default=False")

    args = parser.parse_args()
    
    # Mongo Setup 
    client = MongoClient("localhost", 27017, maxPoolSize=50)
    DB= client.localhost

    MAIN_COLLECTION="main_coll_name"
    FK_COLL_NAME1="coll_name1"
    FK_KEY_NAME1="fk_name1"
    
    FK_COLL_NAME2="coll_name2"
    FK_COLL_NAME2="fk_name2"    

    # Populate Set 1
    fk_set1 = collection_to_dict(DB, FK_COLL_NAME1, FK_KEY_NAME1)
    
    # Populate Set 2
    fk_set2 = collection_to_dict(DB, FK_COLL_NAME2, FK_KEY_NAME2)
    
    row_processor = partial(process, fk_set1, fk_set2)

    with open(args.output_file, 'wb') as out_file:
        csvwriter = csv.writer(out_file)
        collection=mongo_db[MAIN_COLLECTION]
        cursor= collection.find({})
        
        seen= set()
        for document in curosr:
            try:
                result = row_process(document)
                
                k= result[0] + result[1]
                if k in seen:
                    continue
                else:
                    csvwriter.writerow(result)
                    seen.add(k)
                    
            except Exception, e:
                logger.error(str(e), exc_info=True)