import logging
import datetime
from django.db.models.query import QuerySet
from django.db.models.base import ModelBase
from .settings import ADD_LOG_FAILURES
from .settings import ELASTICSEARCH_CONN
from .settings import ELASTICSEARCH_TTL

from elasticsearch import Elasticsearch
from elasticsearch.exceptions import TransportError, NotFoundError, ConflictError
from elastic_search.settings import ELASTIC_MAPPINGS, FIELDS_MAP, ES_SETTINGS
from elasticsearch.helpers import bulk

LOGGER = logging.getLogger(__name__)

default_op = ' AND '
default_facet_op = ' OR '


class ElasticSearch(object):
    """
    Syncs any data stream sent into ES using single save of bulk.
    """

    def __init__(self, conn_str=ELASTICSEARCH_CONN):
        """
        create ES connection for configured host.
        """
        self.es = Elasticsearch(conn_str)

    def build_mapping(self, *args, **kwargs):
        """
        Method to create initial mapping of properties for the index we
        want to create. Is called for the first time while index is being
        created. Takes key 'mapping' in kwargs for mapping type.
        """
        index = kwargs.get("index")
        doc_type = kwargs.get("_type")
        body = kwargs.get("mapping", {})
        if kwargs.get("settings"):
            body.update({'settings': kwargs.get("settings")})
        try:
            self.es.indices.create(index=index, body=body)
            self.es.indices.put_mapping(index=index, body=body,
                                        doc_type=doc_type)
            return True, "Index Mapping successful."
        except TransportError, e:
            # fail silently - just log and die ...
            message = 'Error in indexing, host: {}, unable to create mapping'.format(
                ELASTICSEARCH_CONN)
            if ADD_LOG_FAILURES:
                LOGGER.exception(e)
            return False, message


    def save(self, *args, **kwargs):
        """
        takes as kwargs-
        {
            "data": json data,
            "ts": timestamp,
            "index": index name,
            "type": index type,
        }
        """
        # create the index string
        message = str()
        data = kwargs.get("data")
        ts = kwargs.get("ts")
        # prepare the doc for indexing
        doc = dict()
        doc['_type'] = kwargs.get("type")
        doc['info'] = data
        doc['@timestamp'] = datetime.datetime(*ts[:6])
        try:
            exists = self.es.indices.exists(kwargs.get("index"))
            if not exists:
                map_st, map_msg = self.build_mapping(**kwargs)
                if not map_st:
                    return map_st, map_msg
            res = self.es.index(
                    index=kwargs.get("index"),
                    doc_type=doc['_type'],
                    body=doc,  # message
                    timestamp=datetime.datetime.utcnow(),  # set to current time
                    consistency='one',  # do not wait for quorum / all shards
                    replication='async',  # async
                    ttl=ELASTICSEARCH_TTL)  # as defined in settings
            return True, res
        except TransportError, e:
            # fail silently - just log and die ...
            message = 'Error in indexing, host: {}, unable to index'.format(
                ELASTICSEARCH_CONN)
            if ADD_LOG_FAILURES:
                LOGGER.exception(e)
            return False, message

    def bulk(self, *args, **kwargs):
        '''
        Takes as kwargs-
        {
            "data_iter": iterator object for data
            "index": index name to be indexed to
            "_type": index type
            "mapping": incase new index, must have index mappings
        }
        '''
        data_iter = kwargs.get("data_iter")
        index = kwargs.get("index")
        try:
            exists = self.es.indices.exists(index)
            if not exists:
                map_st, map_msg = self.build_mapping(**kwargs)
                if not map_st:
                    return map_st, map_msg
            res = bulk(self.es, data_iter, index=index)
            return True, res
        except TransportError, e:
            # fail silently - just log and die ...
            message = 'Error in indexing, host: {}, unable to index'.format(
                ELASTICSEARCH_CONN)
            if ADD_LOG_FAILURES:
                LOGGER.exception(e)
            return False, message


class ESMigration(ElasticSearch):
    """
    child class of ElasticSearch to handle the migration of postgres model to
    elastic search index.
    """
    def __init__(self, index_name, doc_type='string'):
        """
        create ES connection for configured host.
        """
        super(ESMigration, self).__init__()
        self.index_name = index_name
        self.doc_type = doc_type

    def create(self, model_or_queryset, **filter_args):
        """
        :arg model_or_queryset accept the model name or queryset. if passed
        model
        :arg filter_args args on which the model_or_queryset needs to further
         filtered
        :arg index index_name in elastic search
        :return success(boolean), res(message)
        """
        # TODO: Should pass the field list and configuration should be switchable
        # validating the arguments and transforming it in a parse-able object
        if isinstance(model_or_queryset, ModelBase):
            data_iter = model_or_queryset.objects.filter(**filter_args
            ).select_related().values(*FIELDS_MAP[self.index_name].keys())
        elif isinstance(model_or_queryset, QuerySet):
            data_iter = model_or_queryset.filter(**filter_args)\
                .select_related().values(*FIELDS_MAP[self.index_name].keys())
        else:
            # throw the validation issue
            raise ValueError('model_or_queryset should be django model or '
                             'queryset')
        # add the _id in data to set the postgres and elastic search same index
        for obj in data_iter:
            obj['_id'] = obj['id']
            obj['_type'] = self.doc_type
        mapping = ELASTIC_MAPPINGS.get(self.index_name, {})
        settings = ES_SETTINGS.get(self.index_name, {})

        data = {"data_iter": data_iter, "index": self.index_name,
                "_type": self.doc_type, "mapping": mapping, "settings":settings}

        return self.bulk(**data)

    def update(self, instances):
        """
        idea is to have a functionality which will update the existing documents
         in es
        TODO: if document doesn't exist there should be configurable flag to
        throw the exception or create the new document
        :arg instances queryset
        :arg index_name
        :return
        """
        if isinstance(instances, QuerySet):
            docs = instances.select_related().values(
                *FIELDS_MAP[self.index_name].keys())
        else:
            raise ValueError('Instance should be model queryset only.')
        try:
            # ideally it should be one but
            for doc in docs:
                try:
                    self.es.update(index=self.index_name, doc_type=self.doc_type,
                        id=doc['id'], body={"doc": doc})
                except NotFoundError as e:
                    self.create(instances.filter(id=doc['id']))
                except ConflictError as e:
                    LOGGER.info(e)
                    pass

        except KeyError as e:
            if ADD_LOG_FAILURES:
                LOGGER.exception(e)
            # log the issue and raise the same and caller decide.
            raise e
        except TransportError, e:
            message = 'Error in indexing, host: {}, unable to index'.format(
                ELASTICSEARCH_CONN)
            if ADD_LOG_FAILURES:
                LOGGER.exception(e)
            # raise the custom exception and let the caller decide if it want to
            # ignore the failure
            raise Exception(message)
        else:
            return True, 'updated the instance'

    def delete(self, instances):
        """
        idea is to have a functionality which will delete the existing documents
         in es
        :arg instances queryset
        :return
        """
        if isinstance(instances, QuerySet):
            docs = instances.select_related().values(
                *FIELDS_MAP[self.index_name].keys())
        else:
            raise ValueError('Instance should be model queryset only.')
        try:
            # ideally it should be one but ignore if doesn't exist
            for doc in docs:
                self.es.delete(index=self.index_name, doc_type=self.doc_type,
                               id=doc['id'], ignore=[400, 404])
        except KeyError as e:
            if ADD_LOG_FAILURES:
                LOGGER.exception(e)
            # log the issue and raise the same and caller decide.
            raise e
        except TransportError, e:
            message = 'Error in indexing, host: {}, unable to index'.format(
                ELASTICSEARCH_CONN)
            if ADD_LOG_FAILURES:
                LOGGER.exception(e)
            # raise the custom exception and let the caller decide if it want to
            # ignore the failure
            raise Exception(message)
        else:
            return True, 'deleted the instances.'

    @staticmethod
    def get_es_key(idx, ps_key):
        """
        return the corresponding es_key
        :param ps_key:
        :return:
        """
        for es_k, ps_k in FIELDS_MAP[idx].iteritems():
            if ps_key == ps_k:
                return es_k

    def get_list(self, *fields, **query_data):
        """
        interface to search the client warehouse
        :arg query_data query strings as dict. {model_field_name:value}
        :return list of dict of keys for which search was executed.
        """
        from utils import ElasticSearchBase
        kwargs = {}
        if query_data.get('size'):
            kwargs['size'] = query_data.pop('size')
        # mapping the data. creating the data which can parsed by query builder.
        data = dict()
        # creating fields list which we expect elastic search should return
        fields = [self.get_es_key(self.index_name, field) for field in fields] \
            if fields else FIELDS_MAP[self.index_name].keys()
        for k, v in query_data.iteritems():
            # mapping the keys with the one in elastic search
            data[self.get_es_key(self.index_name, k)] = v

        esb = ElasticSearchBase(query=data, fields=fields, **kwargs)
        search_result = esb.fetch_es_response(index_name=self.index_name,
                                              lucene_query=True)
        # translating the result readable to client (keys expected by api)
        if search_result.get('error'):
            return []
        results = []
        try:
            for res in search_result['hits']['hits']:
                # now preparing the list of dict with the keys which user passes
                # initially
                results.append({FIELDS_MAP[self.index_name][field]:
                                res['fields'][field][0] for field in fields})
        except KeyError as e:
            LOGGER.exception(e)
            # the required keys are not there.Issue could be configuration error
            raise EnvironmentError('Issue in fetch or parsing the result. '
                                   'Please report this to admin!')
        return results

# result = e.update(index=index, doc_type='string', id=index, body={"doc":document})

c4778c66c18b434b854758f0a99a4433402c3843
import logging
import datetime
from django.db.models.query import QuerySet
from django.db.models.base import ModelBase
from .settings import ADD_LOG_FAILURES
from .settings import ELASTICSEARCH_CONN
from .settings import ELASTICSEARCH_TTL

from elasticsearch import Elasticsearch
from elasticsearch.exceptions import TransportError, NotFoundError, ConflictError
from elastic_search.settings import ELASTIC_MAPPINGS, FIELDS_MAP, ES_SETTINGS
from elasticsearch.helpers import bulk

LOGGER = logging.getLogger(__name__)

default_op = ' AND '
default_facet_op = ' OR '


class ElasticSearch(object):
    """
    Syncs any data stream sent into ES using single save of bulk.
    """

    def __init__(self, conn_str=ELASTICSEARCH_CONN):
        """
        create ES connection for configured host.
        """
        self.es = Elasticsearch(conn_str)

    def build_mapping(self, *args, **kwargs):
        """
        Method to create initial mapping of properties for the index we
        want to create. Is called for the first time while index is being
        created. Takes key 'mapping' in kwargs for mapping type.
        """
        index = kwargs.get("index")
        doc_type = kwargs.get("_type")
        body = kwargs.get("mapping", {})
        if kwargs.get("settings"):
            body.update({'settings': kwargs.get("settings")})
        try:
            self.es.indices.create(index=index, body=body)
            self.es.indices.put_mapping(index=index, body=body,
                                        doc_type=doc_type)
            return True, "Index Mapping successful."
        except TransportError, e:
            # fail silently - just log and die ...
            message = 'Error in indexing, host: {}, unable to create mapping'.format(
                ELASTICSEARCH_CONN)
            if ADD_LOG_FAILURES:
                LOGGER.exception(e)
            return False, message


    def save(self, *args, **kwargs):
        """
        takes as kwargs-
        {
            "data": json data,
            "ts": timestamp,
            "index": index name,
            "type": index type,
        }
        """
        # create the index string
        message = str()
        data = kwargs.get("data")
        ts = kwargs.get("ts")
        # prepare the doc for indexing
        doc = dict()
        doc['_type'] = kwargs.get("type")
        doc['info'] = data
        doc['@timestamp'] = datetime.datetime(*ts[:6])
        try:
            exists = self.es.indices.exists(kwargs.get("index"))
            if not exists:
                map_st, map_msg = self.build_mapping(**kwargs)
                if not map_st:
                    return map_st, map_msg
            res = self.es.index(
                    index=kwargs.get("index"),
                    doc_type=doc['_type'],
                    body=doc,  # message
                    timestamp=datetime.datetime.utcnow(),  # set to current time
                    consistency='one',  # do not wait for quorum / all shards
                    replication='async',  # async
                    ttl=ELASTICSEARCH_TTL)  # as defined in settings
            return True, res
        except TransportError, e:
            # fail silently - just log and die ...
            message = 'Error in indexing, host: {}, unable to index'.format(
                ELASTICSEARCH_CONN)
            if ADD_LOG_FAILURES:
                LOGGER.exception(e)
            return False, message

    def bulk(self, *args, **kwargs):
        '''
        Takes as kwargs-
        {
            "data_iter": iterator object for data
            "index": index name to be indexed to
            "_type": index type
            "mapping": incase new index, must have index mappings
        }
        '''
        data_iter = kwargs.get("data_iter")
        index = kwargs.get("index")
        try:
            exists = self.es.indices.exists(index)
            if not exists:
                map_st, map_msg = self.build_mapping(**kwargs)
                if not map_st:
                    return map_st, map_msg
            res = bulk(self.es, data_iter, index=index)
            return True, res
        except TransportError, e:
            # fail silently - just log and die ...
            message = 'Error in indexing, host: {}, unable to index'.format(
                ELASTICSEARCH_CONN)
            if ADD_LOG_FAILURES:
                LOGGER.exception(e)
            return False, message


class ESMigration(ElasticSearch):
    """
    child class of ElasticSearch to handle the migration of postgres model to
    elastic search index.
    """
    def __init__(self, index_name, doc_type='string'):
        """
        create ES connection for configured host.
        """
        super(ESMigration, self).__init__()
        self.index_name = index_name
        self.doc_type = doc_type

    def create(self, model_or_queryset, **filter_args):
        """
        :arg model_or_queryset accept the model name or queryset. if passed
        model
        :arg filter_args args on which the model_or_queryset needs to further
         filtered
        :arg index index_name in elastic search
        :return success(boolean), res(message)
        """
        # TODO: Should pass the field list and configuration should be switchable
        # validating the arguments and transforming it in a parse-able object
        if isinstance(model_or_queryset, ModelBase):
            data_iter = model_or_queryset.objects.filter(**filter_args
            ).select_related().values(*FIELDS_MAP[self.index_name].keys())
        elif isinstance(model_or_queryset, QuerySet):
            data_iter = model_or_queryset.filter(**filter_args)\
                .select_related().values(*FIELDS_MAP[self.index_name].keys())
        else:
            # throw the validation issue
            raise ValueError('model_or_queryset should be django model or '
                             'queryset')
        # add the _id in data to set the postgres and elastic search same index
        for obj in data_iter:
            obj['_id'] = obj['id']
            obj['_type'] = self.doc_type
        mapping = ELASTIC_MAPPINGS.get(self.index_name, {})
        settings = ES_SETTINGS.get(self.index_name, {})

        data = {"data_iter": data_iter, "index": self.index_name,
                "_type": self.doc_type, "mapping": mapping, "settings":settings}

        return self.bulk(**data)

    def update(self, instances):
        """
        idea is to have a functionality which will update the existing documents
         in es
        TODO: if document doesn't exist there should be configurable flag to
        throw the exception or create the new document
        :arg instances queryset
        :arg index_name
        :return
        """
        if isinstance(instances, QuerySet):
            docs = instances.select_related().values(
                *FIELDS_MAP[self.index_name].keys())
        else:
            raise ValueError('Instance should be model queryset only.')
        try:
            # ideally it should be one but
            for doc in docs:
                try:
                    self.es.update(index=self.index_name, doc_type=self.doc_type,
                        id=doc['id'], body={"doc": doc})
                except NotFoundError as e:
                    self.create(instances.filter(id=doc['id']))
                except ConflictError as e:
                    LOGGER.info(e)
                    pass

        except KeyError as e:
            if ADD_LOG_FAILURES:
                LOGGER.exception(e)
            # log the issue and raise the same and caller decide.
            raise e
        except TransportError, e:
            message = 'Error in indexing, host: {}, unable to index'.format(
                ELASTICSEARCH_CONN)
            if ADD_LOG_FAILURES:
                LOGGER.exception(e)
            # raise the custom exception and let the caller decide if it want to
            # ignore the failure
            raise Exception(message)
        else:
            return True, 'updated the instance'

    def delete(self, instances):
        """
        idea is to have a functionality which will delete the existing documents
         in es
        :arg instances queryset
        :return
        """
        if isinstance(instances, QuerySet):
            docs = instances.select_related().values(
                *FIELDS_MAP[self.index_name].keys())
        else:
            raise ValueError('Instance should be model queryset only.')
        try:
            # ideally it should be one but ignore if doesn't exist
            for doc in docs:
                self.es.delete(index=self.index_name, doc_type=self.doc_type,
                               id=doc['id'], ignore=[400, 404])
        except KeyError as e:
            if ADD_LOG_FAILURES:
                LOGGER.exception(e)
            # log the issue and raise the same and caller decide.
            raise e
        except TransportError, e:
            message = 'Error in indexing, host: {}, unable to index'.format(
                ELASTICSEARCH_CONN)
            if ADD_LOG_FAILURES:
                LOGGER.exception(e)
            # raise the custom exception and let the caller decide if it want to
            # ignore the failure
            raise Exception(message)
        else:
            return True, 'deleted the instances.'

    @staticmethod
    def get_es_key(idx, ps_key):
        """
        return the corresponding es_key
        :param ps_key:
        :return:
        """
        for es_k, ps_k in FIELDS_MAP[idx].iteritems():
            if ps_key == ps_k:
                return es_k

    def get_list(self, *fields, **query_data):
        """
        interface to search the client warehouse
        :arg query_data query strings as dict. {model_field_name:value}
        :return list of dict of keys for which search was executed.
        """
        from utils import ElasticSearchBase
        kwargs = {}
        if query_data.get('size'):
            kwargs['size'] = query_data.pop('size')
        # mapping the data. creating the data which can parsed by query builder.
        data = dict()
        # creating fields list which we expect elastic search should return
        fields = [self.get_es_key(self.index_name, field) for field in fields] \
            if fields else FIELDS_MAP[self.index_name].keys()
        for k, v in query_data.iteritems():
            # mapping the keys with the one in elastic search
            data[self.get_es_key(self.index_name, k)] = v

        esb = ElasticSearchBase(query=data, fields=fields, **kwargs)
        search_result = esb.fetch_es_response(index_name=self.index_name,
                                              lucene_query=True)
        # translating the result readable to client (keys expected by api)
        if search_result.get('error'):
            return []
        results = []
        try:
            for res in search_result['hits']['hits']:
                # now preparing the list of dict with the keys which user passes
                # initially
                results.append({FIELDS_MAP[self.index_name][field]:
                                res['fields'][field][0] for field in fields})
        except KeyError as e:
            LOGGER.exception(e)
            # the required keys are not there.Issue could be configuration error
            raise EnvironmentError('Issue in fetch or parsing the result. '
                                   'Please report this to admin!')
        return results

# result = e.update(index=index, doc_type='string', id=index, body={"doc":document})
