cdplib/cdplib/db_handlers/MongodbHandler.py

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

"""
Created on Mon Sep 16 13:27:44 2019

@author: oskar
@description: Implementation of a database handler for abstraction of the mongodb.
"""



import simplejson
import sys
import os

import pymongo
from pymongo import MongoClient
import pandas as pd
import numpy as np

sys.path.append(os.getcwd())
from cdplib.log import Log
from cdplib.db_migration.ParseJsonSchema import ParseJsonSchema


#class MongodbHandlerPool(metaclass=SingletonThreadsafe):
class MongodbHandlerPool():
    '''
    '''

    def __init__(self, size: int = 1):
        self._size = size
        self._mongodb_handlers = [MongodbHandler() for _ in range(size)]

    def aquire(self):
        while not self._mongodb_handlers:
            self._mongodb_handlers = [MongodbHandler() for _ in range(self._size)]
            log.warning("Ran out of Mongodb handlers, 10 more have been added. Are you sure you've returned yours?")
        return self._mongodb_handlers.pop()

    def release(self, mongodb_handler):
        if len(self._mongodb_handlers) < self._size:
            self._mongodb_handlers.append(mongodb_handler)


class MongodbHandler:

    '''

    '''
    pass
    def __init__(self, database_url: str = None,
                 database_name: str = None):
        '''
        :param str database_url: Url for the mongodb database
        :param str database_name: Name of the database the database handler should handle
        '''
        if database_url is None:

            from libraries.configuration import default as cfg

            database_url = "mongodb://{0}:{1}@{2}:{3}"\
                           .format(cfg["MONGO"]["MONGO_USER"],
                                   cfg["MONGO"]["MONGO_PASSWORD"],
                                   cfg["MONGO"]["MONGO_HOST"],
                                   cfg["MONGO"]["MONGO_PORT"])

            if database_name is None:

                database_name = cfg["MONGO"]["MONGO_DATABASE_NAME"]

        assert(isinstance(database_url, str)),\
            "Parameter 'database_url' must be a string type"
        assert(isinstance(database_name, str)),\
            "Parameter 'database_name' must be a string type"

        self._log = Log("Mongodb Handler")

        # Connect to the MongoDB
        self._client = MongoClient(database_url)
        # Connect to the oebb_db database, or create it if it doesnt exist.
        self._database = self._client[database_name]

        self._database_name = database_name

    def set_database(self, database_name):
        self._database = self._client[database_name]

    def drop_database(self):
        '''
        '''
        self._client.drop_database(self._database_name)

    def drop_collection(self, collection_name: str):
        '''
        '''
        self._database[collection_name].drop()

    def create_index(self, collection_name: str, key: str, direction: (str, int)='text'):
        '''
        :param str collection_name: name on the collection for which the schema will be set.
        :param str key: Which value should be used as the index.
        :param str direction: see https://api.mongodb.com/python/current/api/pymongo/collection.html for reference.
        '''

        allowed_directions = [1, -1, '2d', 'geoHaystack', '2dsphere', 'hashed', 'text']

        assert(isinstance(collection_name, str)),\
            "Parameter 'collection_name' must be a string type"
        assert(isinstance(key, (str, int)) & key in allowed_directions),\
            "Parameter 'key' must be a string or integer type and must be one of these values: 1, -1, '2d', 'geoHaystack', '2dsphere', 'hashed', 'text' "
        assert(isinstance(direction, str)),\
            "Parameter 'direction' must be a string type"
            
        self._database[collection_name].create_index([key, direction])

    def set_collection_schema(self, collection_name: str, schema_path: str,
                              validation_level: str = 'moderate',validation_action: str = 'error'):
        '''
        :param str collection_name: name on the collection for which the schema will be set.
        :param str schema_path: path to the schema file.
        :param str validation_level: level of validation done by the mongodb.
        :param str validation_action: what will happen upon validation error, warning or error message.
        '''
        assert(isinstance(collection_name, str)),\
            "Parameter 'collection_name' must be a string type"
        assert(isinstance(schema_path, str)),\
            "Parameter 'schema_path' must be a string type"
        assert(isinstance(validation_level, str)),\
            "Parameter 'validation_level' must be a string type"
        assert(isinstance(validation_action, str)),\
            "Parameter 'validation_action' must be a string type"
            
        parse_obj = ParseJsonSchema(schema_paths=schema_path)

        schema = parse_obj.read_schema_and_parse_for_mongodb(schema_path)

        command = {
                    'collMod': collection_name,
                    'validator': {
                        '$jsonSchema': schema
                    },
                    'validationLevel': validation_level,
                    'validationAction': validation_action
                    }

        try:
            self._database.command(command)

        except Exception as error:
            self._log.log_and_raise_error(('An error occured when trying to set a schema for the collection: {}. \nError: {}').format(collection_name, error))

    def create_collection(self, collection_name):
        '''
        :param str collection_name: name of the collection to be created.
        '''

        assert(isinstance(collection_name, str)),\
            "Parameter 'collection_name' must be a string type"

        if collection_name not in self._database.list_collection_names():
            try:
                self._log.info(("Collection '{}' has been created").format(collection_name))
                return self._database.create_collection(collection_name)

            except Exception as error:
                self._log.log_and_raise_error(('An error occured while creating the new collection: {}. \nError: {}').format(collection_name, error))
        else:
            self._log.info(("Collection '{}' already exists").format(collection_name))
            return self._database[collection_name]

    def insert_data_into_collection(self, data: (dict, list, np.ndarray, pd.DataFrame, pd.Series),
                                    collection_name: str,
                                    ordered: bool = False):
        '''
        :param dict data: dictionary containing the data to be inserted in the collection
        :param pymongo.database.Collection collection: The collection the data will be added to.
        '''

        allowed_types = (dict, list, np.ndarray, pd.DataFrame, pd.Series)

        assert(isinstance(data, allowed_types)),\
            "Parameter 'data' is of invalid type"

        if isinstance(data, np.ndarray):
            data = pd.DataFrame(data)

        if isinstance(data, pd.DataFrame):

            data = simplejson.loads(data.to_json(orient="records",
                                                 date_format="iso"))

        elif isinstance(data, pd.Series):

            data = simplejson.loads(data.to_json(date_format="iso"))

        try:
            if (len(data) == 1) or (isinstance(data, dict)):

                if isinstance(data, pd.DataFrame) and (len(data) == 1):
                    data = data.iloc[0]
                elif type(data) is list:
                    data = data[0]

                self._database[collection_name].insert_one(data)
            else:
                self._database[collection_name].insert_many(data, ordered=ordered)

        except Exception as error:
            self._log.log_and_raise_error(('An error occured when trying to insert data into {}, {}. \nError: {}').format(self._database_name, collection_name, error))

        self._log.info(('Data has been inserted into the {} collection').format(collection_name))

    def create_collection_and_set_schema(self, collection_name: str, schema_path: str):
        '''
        :param str collection_name: name of the collection to be created.
        :param str schema_path: path to the schema file.
        '''
        assert(isinstance(collection_name, str)),\
            "Parameter 'collection_name' must be a string type"
        assert(isinstance(schema_path, str)),\
            "Parameter 'schema_path' must be a string type"

        self.create_collection(collection_name)
        self.set_collection_schema(collection_name=collection_name, schema_path=schema_path)

    def query_data_and_generate_dataframe(self, collection_name: str, attribute: str = None,
                                          attribute_value: str = None, comparison_operator: str = '$eq'):
        '''

        '''
        try:
            if attribute == None or attribute_value == None:
                data = self._database[collection_name].find()
            else:
                data = self._database[collection_name].find({attribute: {comparison_operator: attribute_value}})

        except Exception as error:
            self._log.log_and_raise_error(('An error occured trying to query data from {}, with query {}: {}:{}. \nError:{}').format(collection_name, attribute_value, comparison_operator, attribute_value, error))

        return self.convert_mongo_data_into_dataframe(data)

    def aggregate_data_and_generate_dataframe(self, collection_name: str, aggregation_pipeline: list):

        try:
            data = self._database[collection_name].aggregate(pipeline=aggregation_pipeline, allowDiskUse=True)

        except Exception as error:
            self._log.log_and_raise_error(('A problem occured when aggregating the collection {} with the pipeline {}. \nError: {}').format(collection_name, aggregation_pipeline, error))

        return self.convert_mongo_data_into_dataframe(data)

    def convert_mongo_data_into_dataframe(self, data) -> pd.DataFrame():

        data = list(data)
        try:
            if len(data)> 0:
                df = pd.DataFrame(data)
                df.set_index('radsatznummer', inplace=True)
                return df
            else:
                self._log.warning(('No data for the query was found').format())
        except Exception as error:
            self._log.log_and_raise_error(('An error occured trying to convert mongo data into pd.Dataframe. \nError: {} ').format(error))

    def update_data_in_collection(self, query_label: str, query_value: str, update_label:str, update_value: str, collection_name:str):
        self._database[collection_name].update_one({query_label:query_value}, {"$set": {update_label: update_value}})


if __name__ == "__main__":

    log = Log("Test MongodbHandler:")

    log.info('Script started')

    database_url = "mongodb://{0}:{1}@{2}:{3}"\
                           .format('root',
                                   'oebb',
                                   'localhost',
                                   27017)
    database_name = 'test_database'

    db_handler = MongodbHandler(database_name=database_name, database_url=database_url)

    # Create a colleciton for the wheelsets and give it its schema.
    for schema_path in [
            os.path.join(".", "mongo_schema", "schema_wheelsets.json"),
            os.path.join(".", "mongo_schema", "schema_process_instances.json"),
            os.path.join(".", "mongo_schema", "schema_componets.json")]:

        if os.path.isfile(schema_path):

            collection_name = os.path.basename(schema_path).lstrip("_schema").split(".")[0]

            db_handler.create_collection_and_set_schema(collection_name, schema_path)
        else:
            log.log_and_raise_warning(('No file exists at path: {}').format(schema_path))

    log.info(("Existing databases: {}, Collection in OEBB database {}")\
             .format(db_handler._client.list_database_names(), db_handler._database.list_collection_names()))