persistencehandler.py

'''
Copyright (C) e2its - All Rights Reserved
 * Unauthorized copying of this file, via any medium is strictly prohibited
 * Proprietary and confidential
 *
 * This file is part of gDayF project.
 *
 * Written by Jose L. Sanchez <e2its.es@gmail.com>, 2016-2019
'''

from json import dump, dumps, load, loads
from collections import OrderedDict
import codecs
from os import path as ospath, chmod, remove
from os import path
from shutil import rmtree, copy2
from pathlib import Path
from gdayf.common.utils import hash_key
from gdayf.conf.loadconfig import LoadConfig
from gdayf.common.utils import get_model_fw
from gdayf.common.storagemetadata import StorageMetadata
import gzip
import mimetypes
from hdfs import InsecureClient as Client, HdfsError
from pymongo import MongoClient
from pymongo.errors import *
#from pymongo import *
from copy import deepcopy
import bson
from bson.codec_options import CodecOptions



class PersistenceHandler(object):
    
    def __init__(self, e_c):
        self._ec = e_c
        self._config = self._ec.config.get_config()['storage']
        self._persistence = self._ec.config.get_config()['persistence']

    
    def store_file(self, storage_json, filename):
        global_op = 0

        '''try:
            file = open(filename, 'rb')
            mmap_ = mmap.mmap(file.fileno(), 0, access=mmap.ACCESS_READ)
        except IOError:
            return 1'''

        for each_storage_type in storage_json:
            if each_storage_type['type'] == 'localfs':
                result, each_storage_type['hash_value'] = self._store_file_to_localfs(each_storage_type, filename)
                global_op += result
            elif each_storage_type['type'] == 'hdfs':
                result, each_storage_type['hash_value'] = self._store_file_to_hdfs(each_storage_type, filename)
                global_op += result
        '''mmap_.close()'''
        return global_op

    
    def _store_file_to_hdfs(self, storage_json, filename):
        try:
            client = Client(url=self._config['hdfs']['url'])
        except HdfsError as hexecution_error:
            print(repr(hexecution_error))
            return 1, None
        except IOError as iexecution_error:
            print(repr(iexecution_error))
            return 1, None
        except OSError as oexecution_error:
            print(repr(oexecution_error))
            return 1, None
        try:
            self._mkdir_hdfs(path=path.dirname(storage_json['value']),
                             grants=self._config['grants'],
                             client=client)
            #client.write(storage_json['value'], data=mmap_, encoding='utf-8', overwrite=True)
            client.upload(storage_json['value'], filename, overwrite=True)

            '''with client.write(storage_json['value'],overwrite=True) as wfile:
                mmap_.seek(0)
                iterator = 0
                while iterator < mmap_.size():
                    wfile.write(mmap_.read())
                    wfile.flush()
                    iterator += 1
                wfile.close()'''

        except HdfsError as hexecution_error:
            print(repr(hexecution_error))
            return 1, None
        except IOError as iexecution_error:
            print(repr(iexecution_error))
            return 1, None
        except OSError as oexecution_error:
            print(repr(oexecution_error))
            return 1, None
        finally:
            del client

        return 0, None

    
    def _store_file_to_localfs(self, storage_json, filename):
        if not ospath.exists(path=storage_json['value']):
            try:
                self._mkdir_localfs(path=path.dirname(storage_json['value']), grants=int(self._config['grants'], 8))
                copy2(filename, storage_json['value'])
                '''with open(storage_json['value'], 'wb') as wfile:
                    mmap_.seek(0)
                    iterator = 0
                    while iterator < mmap_.size():
                        wfile.write(mmap_.read())
                        wfile.flush()
                        iterator += 1
                    wfile.close()'''
                chmod(storage_json['value'], int(self._config['grants'], 8))
            except IOError:
                return 1, None

        return 0, hash_key(hash_type=storage_json['hash_type'], filename=storage_json['value'])

    
    def remove_file(self, load_path):
        global_op = 0

        storage_metadata = StorageMetadata(self._ec)
        for each_storage_type in load_path:
            if each_storage_type['type'] == 'localfs':
                result, storage = self._remove_file_to_localfs(each_storage_type)
                if storage is not None:
                    storage_metadata.append(storage)
                global_op += result
            elif each_storage_type['type'] == 'hdfs':
                result, storage = self._remove_file_to_hdfs(each_storage_type)
                if storage is not None:
                    storage_metadata.append(storage)
                global_op += result
        return global_op, storage_metadata.copy()

    
    def _remove_file_to_hdfs(self,storage_json):
        path = storage_json['value']
        url_beginning = path.find('//') + 2
        url_ending = path.find('/', url_beginning)
        path = path[url_ending:]

        try:
            client = Client(url=self._config['hdfs']['url'])
        except HdfsError as hexecution_error:
            print(repr(hexecution_error))
            return 1, None
        except IOError as iexecution_error:
            print(repr(iexecution_error))
            return 1, None
        except OSError as oexecution_error:
            print(repr(oexecution_error))
            return 1, None
        try:
            if client.delete(hdfs_path=path, recursive=True):
                return 0, None
            else:
                return 1, storage_json

        except HdfsError as hexecution_error:
            print(repr(hexecution_error))
            return 1, storage_json
        except IOError as iexecution_error:
            print(repr(iexecution_error))
            return 1, storage_json
        except OSError as oexecution_error:
            print(repr(oexecution_error))
            return 1, storage_json
        finally:
            del client


    
    @staticmethod
    def _remove_file_to_localfs(storage_json):
        if not ospath.exists(path=storage_json['value']):
            return 0, None
        else:
            try:
                if ospath.isdir(storage_json['value']):
                    rmtree(storage_json['value'])
                else:
                    remove(storage_json['value'])
                return 0, None
            except OSError:
                return 1, storage_json

    
    def store_json(self, storage_json, ar_json, other=None):
        '''assert isinstance(storage_json, StorageMetadata)'''
        global_op = 0
        for each_storage_type in storage_json:
            if each_storage_type['type'] == 'localfs':
                if other is None:
                    global_op += self._store_json_to_localfs(each_storage_type, ar_json)
                else:
                    global_op += self._store_json_to_localfs(each_storage_type, other)
            elif each_storage_type['type'] == 'hdfs':
                if other is None:
                    global_op += self._store_json_to_hdfs(each_storage_type, ar_json)
                else:
                    global_op += self._store_json_to_hdfs(each_storage_type, other)
            elif each_storage_type['type'] == 'mongoDB':
                if other is None:
                    global_op += self._store_json_to_mongoDB(each_storage_type, ar_json)
                else:
                    global_op += self._store_other_to_mongoDB(each_storage_type, other)
        return global_op

    
    def _store_json_to_localfs(self, storage_json, ar_json):
        compress = self._persistence['compress_json']
        #if not ospath.exists(storage_json['value']):
        try:
            self._mkdir_localfs(path=path.dirname(storage_json['value']), grants=int(self._config['grants'], 8))
            if compress:
                file = gzip.GzipFile(storage_json['value'], 'w')
                json_str = dumps(ar_json, indent=4)
                json_bytes = json_str.encode('utf-8')
                file.write(json_bytes)
            else:
                # file = open(storage_json['value'], 'w')
                file = codecs.open(storage_json['value'], 'w', encoding='utf-8')
                dump(ar_json, file, indent=4, ensure_ascii=False)
            file.close()
            chmod(storage_json['value'], int(self._config['grants'], 8))
            return 0
        except IOError as iexecution_error:
            print(repr(iexecution_error))
            return 1

    
    def _store_json_to_hdfs(self, storage_json, ar_json, client=None):
        remove_client = False
        if client is None:
            client = Client(url=self._config['hdfs']['url'])
            remove_client = True
        compress = self._persistence['compress_json']
        #if not ospath.exists(storage_json['value']):
        try:
            self._mkdir_hdfs(path=path.dirname(storage_json['value']),
                             grants=self._config['grants'],
                             client=client)
            if compress:
                json_str = dumps(ar_json, indent=4)
                json_bytes = json_str.encode('utf-8')
                client.write(storage_json['value'],
                             data=gzip.compress(json_bytes),
                             overwrite=True)
            else:
                client.write(storage_json['value'], data=dumps(ar_json, indent=4, ensure_ascii=False),
                             encoding='utf-8', overwrite=True)
            return 0
        except HdfsError as hexecution_error:
            print(repr(hexecution_error))
            return 1
        except IOError as iexecution_error:
            print(repr(iexecution_error))
            return 1
        except OSError as oexecution_error:
            print(repr(oexecution_error))
            return 1
        finally:
            if remove_client:
                del client

    
    def _store_json_to_mongoDB(self, storage_json, ar_json, client=None):
        remove_client = False
        if client is None or not isinstance(client(MongoClient)):
            try:
                client = MongoClient(host=self._config['mongoDB']['url'],
                                     port=int(self._config['mongoDB']['port']),
                                     document_class=OrderedDict)
                remove_client = True
            except ConnectionFailure as cexecution_error:
                print(repr(cexecution_error))
                return 1
        try:
            description = Path(storage_json['value']).parts
            db = client[self._config['mongoDB']['value']]
            collection = db[description[1]]
            model_id = ar_json['model_parameters'][get_model_fw(ar_json)]['parameters']['model_id']['value']
            filter_cond = "model_parameters." + get_model_fw(ar_json) + ".parameters.model_id.value"
            cond = [{filter_cond: model_id}, {"type": ar_json['type']},
                    {"model_id": ar_json['model_id']},  {"timestamp": ar_json['timestamp']}]
            query = {"$and": cond}

            count = collection.find(query).count()
            new_ar_json = OrderedDict(ar_json)
            if count == 1:
                collection.delete_one(query)
                collection.insert(new_ar_json, check_keys=False)
                return 0
            elif count == 0:
                collection.insert(new_ar_json, check_keys=False)
                return 0
            else:
                print("Trace: Duplicate Model %s" % model_id)
                return 1
        finally:
            if remove_client:
                client.close()

    
    def _store_other_to_mongoDB(self, storage_json, other, client=None):

        remove_client = False
        if client is None or not isinstance(client(MongoClient)):
            try:
                client = MongoClient(host=self._config['mongoDB']['url'],
                                     port=int(self._config['mongoDB']['port']),
                                     document_class=OrderedDict)
                remove_client = True
            except ConnectionFailure as cexecution_error:
                print(repr(cexecution_error))
                return 1
        try:
            db = client[self._config['mongoDB']['value']]
            collection = db[self._ec.get_id_user() + '_' + storage_json['value']]
            new_ar_json = OrderedDict(other)
            try:
                collection.insert(new_ar_json, check_keys=False)
                return 0
            except Exception as execution_error:
                print(repr(execution_error))
                return 1
        finally:
            if remove_client:
                client.close()


    
    def recover_experiment_mongoDB(self, client=None):
        execution_list = list()
        remove_client = False
        if client is None or not isinstance(client(MongoClient)):
            try:
                client = MongoClient(host=self._config['mongoDB']['url'],
                                     port=int(self._config['mongoDB']['port']),
                                     document_class=OrderedDict)
                remove_client = True
            except ConnectionFailure as cexecution_error:
                print(repr(cexecution_error))
                return execution_list
        try:
            db = client[self._config['mongoDB']['value']]
            collection = db[self._ec.get_id_user()]
            query = {"$and": [{"model_id": self._ec.get_id_analysis()}, {"type": "train"}]}
            for element in collection.find(query):
                execution_list.append(element)
            for element in execution_list:
                element.pop('_id')
            #print(execution_list)
        except PyMongoError as pexecution_error:
            print(repr(pexecution_error))
        finally:
            if remove_client:
                client.close()
            return deepcopy(execution_list)

    
    def mkdir(self, type, path, grants):
        if type == 'localfs':
            return self._mkdir_localfs(path=path, grants=int(grants, 8))
        elif type == 'hdfs':
            return self._mkdir_hdfs(path=path, grants=grants)
        elif type == 'mongoDB':
            return self._mkdir_mongoDB(path=path, grants=grants)

    
    @staticmethod
    def _mkdir_localfs(path, grants):
        try:
            Path(path).mkdir(mode=grants, parents=True, exist_ok=True)
            return 0
        except IOError:
            return 1

    
    def _mkdir_hdfs(self, path, grants, client=None):
        remove_client = False
        if client is None:
            client = Client(url=self._config['hdfs']['url'])
            remove_client = True
        try:
            if client.status(hdfs_path=path, strict=False) is None:
                client.makedirs(hdfs_path=path, permission=grants)
            return 0
        except HdfsError as hexecution_error:
            print(repr(hexecution_error))
            return 1
        except IOError as iexecution_error:
            print(repr(iexecution_error))
            return 1
        except OSError as oexecution_error:
            print(repr(oexecution_error))
            return 1
        finally:
            if remove_client:
                del client

    
    @staticmethod
    def _mkdir_mongoDB(path, grants):
        try:
            return 0
        except IOError:
            return 1

    
    def get_ar_from_engine(self, path):
        found = False
        for storage in ['localfs', 'hdfs', 'mongoDB']:
            if storage == 'localfs' and not found:
                if ospath.exists(path):
                    _, type = mimetypes.guess_type(path)
                    if type == 'gzip':
                        file = gzip.GzipFile(filename=path, mode='r')
                        json_bytes = file.read()
                        json_str = json_bytes.decode('utf-8')
                        ar_metadata = loads(json_str, object_hook=OrderedDict)
                    else:
                        #Bug pending
                        file = open(path, 'r')
                        ar_metadata = load(file, object_hook=OrderedDict)
                    file.close()
                    return 0, ar_metadata
            elif storage == 'hdfs' and not found:
                url = self._config[storage]['url']
                client = Client(url=url)
                remove_client = True
                try:
                    if client.status(hdfs_path=path, strict=False) is not None:
                        _, type = mimetypes.guess_type(path)
                        if type == 'gzip':
                            with client.read(path) as file_hdfs:
                                file = gzip.GzipFile(fileobj=file_hdfs)
                                json_bytes = file.read()
                                json_str = json_bytes.decode('utf-8')
                                ar_metadata = loads(json_str, object_hook=OrderedDict)
                                file.close()
                        else:
                            with client.read(path) as file_hdfs:
                                json_bytes = file_hdfs.read()
                                json_str = json_bytes.decode('utf-8')
                                ar_metadata = loads(json_str, object_hook=OrderedDict)
                        return 0, ar_metadata
                except HdfsError as hexecution_error:
                    print(repr(hexecution_error))
                    return 1, repr(hexecution_error)
                except IOError as iexecution_error:
                    print(repr(iexecution_error))
                    return 1, repr(iexecution_error)
                except OSError as oexecution_error:
                    print(repr(oexecution_error))
                    return 1, repr(oexecution_error)
                finally:
                    if remove_client:
                        del client
            elif storage == 'mongoDB' and not found:
                try:
                    client = MongoClient(host=self._config['mongoDB']['url'],
                                         port=int(self._config['mongoDB']['port']),
                                         document_class=OrderedDict)
                    remove_client = True
                except ConnectionFailure as cexecution_error:
                    print(repr(cexecution_error))
                    return 1, repr(cexecution_error)
                try:
                    db = client[self._config['mongoDB']['value']]
                    description = Path(path).parts
                    if description[1] is not None \
                            and description[2] is not None \
                            and description[3] is not None \
                            and description[4] is not None \
                            and description[1] in db.collection_names():

                        collection = db[description[1]]
                        query1 = {"$and": [{"model_id": description[3]},
                                           {"workflow_id": description[2]},
                                           {'type': 'train'}]
                                  }
                        for element in collection.find(query1):
                            # Return the first element found
                            if element['model_parameters'][get_model_fw(element)]['parameters']['model_id']['value'] \
                                    == description[4]:
                                element.pop('_id')
                                print(element)
                                return 0, element
                        return 1, None

                    else:
                        return 1, None
                except PyMongoError as pexecution_error:
                    print(repr(pexecution_error))
                    return 1, repr(pexecution_error)
                finally:
                    if remove_client:
                        client.close()

        return 1, None