machine_learn/util/mongodb.py

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

import pymongo
import datetime
import time
import random
import logging
from util.config import config


class MongoConfig(object):
    HOST = config.get('mongodb', 'host')
    USER = config.get('mongodb', 'user')
    PORT = int(config.get('mongodb', 'port'))
    PASSWORD = config.get('mongodb', 'password')
    DB = config.get('mongodb', 'db')


def _get_default_mongodb_instance():
    mongo_client = pymongo.MongoClient(MongoConfig.HOST, MongoConfig.PORT)
    db = mongo_client[MongoConfig.DB]
    if MongoConfig.PASSWORD is not None and MongoConfig.PASSWORD != '':
        db.authenticate(MongoConfig.USER, MongoConfig.PASSWORD)
    return db


def get_mongodb_tablenames():
    db = _get_default_mongodb_instance()
    return db.collection_names(include_system_collections=False)


def get_mongo_table_instance(tablename):
    db = _get_default_mongodb_instance()
    return db[tablename]


def del_mongodb_table(tablename):
    try:
        db = _get_default_mongodb_instance()
        db.drop_collection(tablename)
        return True
    except:
        return False


def rename_mongodb_table(from_tablename, to_tablename):
    try:
        get_mongo_table_instance(from_tablename).rename(to_tablename)
        return True
    except:
        return False


def get_mongodb_table_indexes(tablename):
    try:
        ret = []

        ftable = get_mongo_table_instance(tablename)
        findex = ftable.list_indexes()
        listindex = list(findex)
        if listindex is None or len(listindex) == 0:
            return ret

        for i in listindex:
            sonobj = i['key']
            indexes = []

            for k, v in sonobj.iteritems():
                v1 = pymongo.ASCENDING if v > 0 else pymongo.DESCENDING
                temp = (k, v1)
                indexes.append(temp)

            ret.append(indexes)

        return ret
    except Exception as e:
        return None


def set_mongodb_table_indexes(tablename, indexes):
    try:
        ftable = get_mongo_table_instance(tablename)

        if indexes is None or len(indexes) == 0:
            return True
        for i in indexes:
            ftable.ensure_index(i)

        return True
    except Exception as e:
        return False


def copy_mongodb_table_indexes(from_tablename, to_tablename):
    indexes = get_mongodb_table_indexes(from_tablename)
    set_mongodb_table_indexes(to_tablename, indexes)