~jcsackett/charmworld/bac-tag-constraints

# Copyright 2013 Canonical Ltd.  This software is licensed under the
# GNU Affero General Public License version 3 (see the file LICENSE).

__metaclass__ = type

from contextlib import contextmanager
from logging import getLogger
import random
import re

import pyelasticsearch
from pyelasticsearch.exceptions import (
    ElasticHttpError,
    ElasticHttpNotFoundError,
)
from requests.exceptions import ConnectionError

from charmworld.utils import (
    get_ini,
    Timer,
)

# Where wildcards are used, the final var name must be used, to avoid
# selecting any fields that may be integers, such as config.options.*.default,
# provides.*.interface.limit, or changes.revno.  Selecting an integer field
# causes a SearchParseException/NumberFormatException.
charm_exact_fields = [
    'owner',
    'series',
]
charm_free_text_fields = {
    'name': 10,
    'summary': 5,
    'description': 3,
    'config.options.description': None,
    'relations': None,
    'services': None,
}

bundle_exact_fields = [
    'owner',
    'data.series',
]
bundle_free_text_fields = {
    'name': 10,
    'basket_name': 5,
    'description': 3,
    'title': None,
    'data.services.*.charm': None,
}

CHARM = 'charm'
BUNDLE = 'bundle'


class InvalidCharmType(Exception):
    """Raised when an invalid charm type is specified."""


class NegativeLimit(Exception):
    """Raised when a negative limit is specified."""


class IncompatibleMapping(Exception):
    """Raised when desired mapping is not compatible with active mapping."""


class IndexMissing(Exception):
    """Raised when specified index does not exist."""


class IndexNotReady(Exception):
    """Raised when specified index is not ready yet."""


class SearchServiceNotAvailable(Exception):
    """Raised when the specified search server is not ready yet."""
    def __init__(self, message, orig_error):
        super(SearchServiceNotAvailable, self).__init__(message)
        self.orig_error = orig_error


@contextmanager
def translate_error():
    try:
        yield
    except ElasticHttpNotFoundError as e:
        if 'IndexMissingException' not in e.error:
            raise
        raise IndexMissing('Index does not exist', e.error)
    except ElasticHttpError as e:
        if e.status_code != 400:
            raise
        if 'MergeMappingException' not in e.error:
            raise
        raise IncompatibleMapping('Mapping not compatible', e.error)
    except ConnectionError as e:
        raise SearchServiceNotAvailable('Could not connect to search server.',
                                        e)


_BAD_CHARS = u'()/[]^{}'
_BAD_CHAR_TABLE = [
    unichr(c) if unichr(c) not in _BAD_CHARS else u' ' for c in range(128)]
_BAD_CHAR_TABLE = u''.join(_BAD_CHAR_TABLE)
_BAD_TRAILING_CHARS = r'\!:'


class ElasticSearchClient:

    reviewed_clause = {'not': {'prefix': {'store_url': 'cs:~'}}}

    valid_charms_clauses = [{
        'or': [{'and': [{'missing': {'field': 'data.store_data.errors'}}]},
               {'type': {'value': 'bundle'}},
               ]
    }]

    def __init__(self, client, index_name):
        self._client = client
        self.index_name = index_name
        self.stable_status_timeout = '30s'

    def exists(self):
        try:
            self.get_status()
        except IndexMissing:
            return False
        else:
            return True

    def create_index(self):
        settings = get_ini()
        self._client.create_index(
            self.index_name, settings={
                "index": {
                    "number_of_shards": settings['es_shards'],
                    "number_of_replicas": settings['es_replicas'],
                }
            })
        # The ES server may need some time to actually create the index
        # after the request has been issued. Waiting until the server is
        # in "green" status again prevents IndexMissing exceptions.
        self.wait_for_green_status()

    def update_aliased(self, to_add, to_remove):
        """Update the set of indices that are aliased to this one.

        :param to_add: (a single index name) Add this to the list of indices.
        :param to_remove: (a list of index names) Remove these from the list
            of indices.
        """
        actions = [
            {'remove': {'index': remove, 'alias': self.index_name}}
            for remove in to_remove]
        if to_add is not None:
            actions.append({'add': {
                'index': to_add, 'alias': self.index_name}})
        if len(actions) > 0:
            self._client.update_aliases({'actions': actions})

    def get_aliased(self):
        """Return a list of indices that are aliased to this one."""
        return [index for index, info in self._client.aliases().items()
                if self.index_name in info['aliases']]

    def wait_for_green_status(self):
        """Wait until the ElasticSearch server is in a stable state.

        Some operations, like index creation or a change of an index mapping,
        are not finished when the client method returns. Waiting for the
        server status "green" before any other operations are performed
        help to prevent spurious failures.
        """
        self._client.health(self.index_name, wait_for_status='green',
                            timeout=self.stable_status_timeout)

    def delete_index(self):
        with translate_error():
            self._client.delete_index(self.index_name)

    @classmethod
    def from_settings(cls, settings, index_name='charms'):
        urls = settings['es_urls'].split(' ')
        return cls(pyelasticsearch.ElasticSearch(urls), index_name)

    @staticmethod
    def escape_query(query):
        return re.sub(r'[+\-&|!(){}\[\]\^"~*?:\\/]',
                      lambda m: r'\%s' % m.group(0), query)

    def put_mapping(self):
        """Merge the current mapping with the existing mapping.

        :raises: IncompatibleMapping if the existing mapping is not compatible
            with the current mapping.
        """
        charm_exact_index = [
            'categories',
            'name',
            'owner',
            'i_provides',
            'i_requires',
            'series',
            'store_url',
            '_id',
        ]
        inner_charm_properties = dict(
            (name, {'type': 'string', 'index': 'not_analyzed'})
            for name in charm_exact_index)

        CHARM_TEXT_INDEX = (
            'summary',
            'description',
        )
        inner_charm_properties.update(dict(
            (name, {'type': 'string'}) for name in CHARM_TEXT_INDEX))

        CHARM_INTEGER_INDEX = (
            'downloads',
            'downloads_in_past_30_days',
            'downloads_in_past_7_days',
            'downloads_in_past_half_year',
        )
        inner_charm_properties.update(dict(
            (name, {'type': 'long'}) for name in CHARM_INTEGER_INDEX))

        inner_charm_properties['date_created'] = {
            'type': 'date',
        }
        inner_charm_properties['is_featured'] = {
            'type': 'boolean',
        }
        inner_charm_properties['store_data'] = {
            'properties': {
                'errors': {'type': 'string'},
            }
        }
        inner_charm_properties['config'] = {
            'properties': {
                'options': {
                    'properties': {
                        'default': {'type': 'string'},
                        'description': {'type': 'string'},
                    }
                }
            }
        }

        searchable_interfaces = {
            'type': 'object',
            'dynamic': False,
            'properties': {
                'relation_name': {'type': 'string', 'index': 'not_analyzed'},
                'interface': {'type': 'string', 'index': 'not_analyzed'},
                'scope': {'type': 'string', 'index': 'not_analyzed'},
            },
        }

        charm_properties = {
            'data': {
                'properties': inner_charm_properties,
            },
            'requires': searchable_interfaces,
            'provides': searchable_interfaces,
        }

        bundle_exact_index = [
            'name',
            'owner',
            'series',
            'title',
        ]

        bundle_properties = {
            'data': {
                'properties':  # XXX The linter won't let me indent this well.
                dict((name, {'type': 'string', 'index': 'not_analyzed'})
                for name in bundle_exact_index)
            }}

        with translate_error():
            try:
                for (name, properties, dynamic) in (
                        (CHARM, charm_properties, False),
                        (BUNDLE, bundle_properties, True)):
                    self._client.put_mapping(self.index_name, name, {
                        name: {
                            'type': 'object',
                            'dynamic': dynamic,
                            'properties': properties,
                        }
                    })
            finally:
                    # Successful as well as failing calls of put_mapping()
                    # may put the index for some time into an instable
                    # state. Attempts to use it in either case during this
                    # phase may lead to IndexMissing exceptions. Waiting
                    # for status "green" prevents them.
                    self.wait_for_green_status()

    def get_mapping(self, doctype=CHARM):
        """Return the current mapping for 'charm'."""
        return self._client.get_mapping(self.index_name, doctype)

    def _index_item(self, item, doctype):
        id_ = item['_id']
        with translate_error():
            self._client.index(
                self.index_name, doctype, item, id_, refresh=True)

    @staticmethod
    def _wrap_charm(charm_data):
        # Avoid circular import.
        from charmworld.models import construct_charm_id
        result = {
            '_id': construct_charm_id(charm_data, use_revision=False),
            'data': charm_data,
        }
        for relation_type in ('requires', 'provides'):
            if relation_type in charm_data:
                searchable_infos = []
                for name, info in charm_data[relation_type].items():
                    info = info.copy()
                    info['relation_name'] = name
                    searchable_infos.append(info)
                result[relation_type] = searchable_infos
        return result

    def index_charm(self, charm):
        self._index_item(self._wrap_charm(charm), CHARM)

    @staticmethod
    def _wrap_bundle(bundle_data):
        # Avoid circular import.
        from charmworld.models import Bundle
        return {
            '_id': Bundle.construct_id(
                bundle_data['owner'], bundle_data['basket_name'],
                bundle_data['name']),
            'data': bundle_data,
        }

    def index_bundle(self, bundle):
        self._index_item(self._wrap_bundle(bundle), BUNDLE)

    def _delete_item(self, id_, doctype):
        with translate_error():
            self._client.delete(self.index_name, doctype, id_, refresh=True)

    def delete_charm(self, charm):
        from charmworld.models import construct_charm_id
        charm_id = construct_charm_id(charm, use_revision=False)
        self._delete_item(charm_id, CHARM)

    def delete_bundle(self, bundle_data):
        # Avoid circular import.
        from charmworld.models import Bundle
        search_id = Bundle(bundle_data).search_id
        self._delete_item(search_id, BUNDLE)

    def _index_items(self, items, doctype):
        if len(items) == 0:
            return
        self._client.bulk_index(
            self.index_name, doctype, items, '_id', refresh=True)

    def index_charms(self, charms):
        self._index_items([self._wrap_charm(charm) for charm in charms], CHARM)

    def index_bundles(self, bundles):
        self._index_items(
            [self._wrap_bundle(bundle) for bundle in bundles], BUNDLE)

    @staticmethod
    def _get_text_query(text, autocomplete=False):
        if isinstance(text, basestring):
            if text.count('"') % 2 != 0:
                # Quotation marks must be "balanced".
                text = text + '"'
            text = text.translate(_BAD_CHAR_TABLE)
            while text and text[-1] in _BAD_TRAILING_CHARS:
                text = text[:-1]
        if autocomplete:
            return {'prefix': {'data.name': text}}

        charm_fields = [field + ('' if boost is None else '^%d' % boost)
                        for field, boost in charm_free_text_fields.items()]
        charm_fields.extend(charm_exact_fields)

        bundle_fields = [field + ('' if boost is None else '^%d' % boost)
                         for field, boost in bundle_free_text_fields.items()]
        bundle_fields.extend(bundle_exact_fields)

        if text != '':
            def make_query(data_fields, other_fields=[]):
                fields = ['data.' + field for field in data_fields]
                fields.extend(other_fields)
                return {
                    'query_string': {
                    'query': text,
                    'fields': fields,
                    }}

            charm_dsl = {'filtered': {
                'query': make_query(
                    charm_fields, ['requires.*', 'provides.*']),
                'filter': {'type': {'value': CHARM}}
            }}
            bundle_dsl = {'filtered': {
                'query': make_query(bundle_fields),
                'filter': {'type': {'value': BUNDLE}}
            }}
            # Union the bundle_dsl and charm_dsl results
            dsl = {'bool': {'should': [bundle_dsl, charm_dsl]}}
        else:
            dsl = {'match_all': {}}
        return dsl

    @classmethod
    def _get_official_boost(cls, dsl):
        dsl = {
            'custom_filters_score': {
                'query': dsl,
                'filters': [
                    {
                        'filter': cls.reviewed_clause,
                        'boost': 10,
                    },
                ],
            }
        }
        return dsl

    @classmethod
    def _get_filtered(cls, dsl, filters, type_, valid_only, doctype=CHARM):
        and_clause = {
            'and': [{'terms': {'data.' + key: value}}
                    for key, value in filters.items()],
        }

        if valid_only:
            and_clause['and'].extend(cls.valid_charms_clauses)

        if type_ is not None:
            typeclauses = []
            for value in type_:
                if value == 'approved':
                    typeclauses.append(cls.reviewed_clause)
                elif value == 'community':
                    typeclauses.append({
                        'not': cls.reviewed_clause,
                    })
                elif value == 'environment':
                    typeclauses.append({'term': {'_id': ''}})
                else:
                    raise InvalidCharmType(value)

            if len(typeclauses) > 0:
                and_clause['and'].append({'or': typeclauses})

        if doctype is not None:
            and_clause['and'].append({'type': {'value': doctype}})

        if len(and_clause['and']) == 0:
            and_clause = {'match_all': {}}
        dsl = {
            'query': {
                'filtered': {
                    'filter': and_clause,
                    'query': dsl,
                }
            }
        }
        return dsl

    def _get_item(self, id_, doctype):
        result = self._client.get(self.index_name, doctype, id_)
        return result['_source']

    def get(self, charm_id, doctype=CHARM):
        if doctype != CHARM:
            try:
                return self._get_item(charm_id, doctype)
            except ElasticHttpNotFoundError:
                return None
        with translate_error():
            hits = self._client.search({
                'query': {'filtered': {'filter': {
                    'term': {'data._id': charm_id}
                }}}
            }, index=self.index_name, doc_type=CHARM)
            results = hits['hits']['hits']
            if len(results) == 1:
                return results[0]['_source']['data']

    def get_items(self, typed_ids):
        """Return documents and their doctype that match the id and types.

        :param typed_ids: a list of dicts that specify the _type and _id of
            the wanted docs. eg. [{_id: charm_id, _type=charm}]
        :return: A list of dicts that contain the doctype and source data.
            eg. [{doctype: charm, data: {}}]
        """
        if typed_ids == []:
            return []
        result = self._client.multi_get(
            typed_ids, self.index_name, doc_type=None)
        return [
            dict(doctype=doc['_type'], data=doc['_source']['data'])
            for doc in result['docs']]

    def _unlimited_search(self, dsl, limit):
        """Thin interface to search that permits unlimited results.

        If limit is None, no limit is specified to search.  If the default
        limit is exceeded, the search is repeated with the expected number of
        hits.
        """
        kwargs = {'index': self.index_name}
        if limit is not None:
            kwargs['size'] = limit
        with translate_error():
            hits = self._client.search(dsl, **kwargs)
        if limit is None and len(hits['hits']['hits']) < hits['hits']['total']:
            kwargs['size'] = hits['hits']['total']
            hits = self._client.search(dsl, **kwargs)
        return hits

    def api_search(self, text='', filters=None, type_=None,
                   limit=None, valid_only=True, sort=None,
                   autocomplete=False, doctype=CHARM):
        """Search charms and bundles (as used by the API).

        :param text: A string to search for-- will be analyzed (i.e. split
            into stems).
        :param filters: A dict of variable to list of acceptable values.  If
            the field contents match any of these values, the result is
            permitted.  The meaning of "match" depends on the how the field
            was indexed.  For not-analyzed fields, this must be an exact
            match, for analyzed fields, it must match a stem.
        :param type_: Either 'approved' or 'community'.  If 'approved' is
            supplied, only "reviewed" charms will be returned.  If 'community'
            is supplied, charms need not be "reviewed".
        :param limit: The maximum number of charms to return.  If None, all
            possible charms will be returned, but this is not efficient.
        :param valid_only: If True, only charms and bundles that were
            successfully parsed are returned.
        :param sort: If None, charms are sorted by score.  If 'downloaded',
            charms are in order of most to least downloaded.  If 'new', charms
            are in order of most-recently to least-recently created.
        :param autocomplete: If true use a prefix search for autocompletion.
        :param doctype: Only return entities of the given type.  May be
            'charm', 'bundle', or None.  If None, do not filter.
        return: A list of {doctype=<charm|bundle>, data=entity}.
        """
        if filters is None:
            filters = {}
        dsl = self._get_text_query(text, autocomplete)
        dsl = self._get_official_boost(dsl)
        dsl = self._get_filtered(dsl, filters, type_, valid_only, doctype)
        if sort == 'downloaded':
            dsl['sort'] = [{
                'downloads': {'order': 'desc'},
            }]
        elif sort == 'new':
            dsl['sort'] = [{
                'date_created': {'order': 'desc'},
            }]
        elif sort is not None:
            raise ValueError('Invalid sort: %s.' % sort)
        if limit is not None:
            if limit < 0:
                raise NegativeLimit()
        hits = self._unlimited_search(dsl, limit)
        return [
            dict(doctype=hit['_type'], data=hit['_source']['data'])
            for hit in hits['hits']['hits']]

    def get_status(self):
        with translate_error():
            return self._client.status(index=self.index_name)['indices']

    def search(self, terms, valid_only=True):
        # Avoid circular imports.
        from charmworld.models import (
            Bundle,
            Charm,
        )
        size = None
        if terms == '':
            size = 0

        dsl = self._get_text_query(terms)
        dsl = self._get_official_boost(dsl)
        dsl = self._get_filtered(dsl, {}, None, valid_only, doctype=None)

        with Timer() as timer:
            status = self.get_status()
            try:
                real_index = self.get_aliased()[0]
            except IndexError:
                real_index = self.index_name
            docs = status[real_index].get('docs')
            if docs is None:
                log = getLogger('search')
                message = (
                    'Cannot find "docs" in status of index "%s".' %
                    real_index)
                log.info(message + '  Raising IndexNotReady.')
                raise IndexNotReady(message)
            result_total = docs['num_docs']
            if size is None:
                size = result_total
            hits = self._client.search(dsl, index=self.index_name, size=size)
        hitlist = hits['hits']['hits']
        charms = [{'data': Charm(hit['_source']['data']),
                   'weight': hit['_score']}
                  for hit in hitlist if hit['_type'] == CHARM]
        bundles = [{'data': Bundle(hit['_source']['data']),
                    'weight': hit['_score']}
                   for hit in hitlist if hit['_type'] == BUNDLE]

        matches = len(charms) + len(bundles)
        return {
            'results': {CHARM: charms, BUNDLE: bundles},
            'search_time': u"%.5f" % timer.duration(),
            'result_total': result_total,
            'matches': matches,
            'matches_human_readable_estimate': '%d' % matches,
        }

    def related_charms(self, requires, provides, series=None,
                       exclude_name=None, doctype=CHARM):
        """Determine the charms related to the specified interfaces.

        Results are of the form:
        {'data': {/*actual schema*/}, 'weight': 5.6}

        :param requires: A set of interfaces that charms may require.
        :param provides: A set of interfaces that charms may provide.
        :param series: If provided, limit results to the specified series.
        :param exclude_name: If not None, exclude results with the specified
            name.
        :param doctype: Only return entities of the given type.  May be
            'charm', 'bundle', or None.  If None, do not filter.
        :return: requires, provides, each of which is a dict mapping
            interfaces to lists of results.
        """
        and_clauses = [{'or': [
            {'terms': {'i_provides': provides}},
            {'terms': {'i_requires': requires}},
        ]}]
        if series is not None:
            and_clauses.append({'term': {'data.series': series}})
        if exclude_name is not None:
            and_clauses.append({'not': {'term': {'data.name': exclude_name}}})
        and_clauses.extend(self.valid_charms_clauses)
        if doctype is not None:
            and_clauses.append({'type': {'value': doctype}})
        dsl = {
            'filtered': {
                'filter': {
                    'and': and_clauses,
                }
            }
        }
        dsl = {'query': self._get_official_boost(dsl)}
        hits = self._unlimited_search(dsl, None)
        result_p = {}
        result_r = {}
        for hit in hits['hits']['hits']:
            charm = hit['_source']['data']
            payload = {'data': charm, 'weight': hit['_score']}
            for interface in provides.intersection(charm['i_provides']):
                result_p.setdefault(interface, []).append(payload)
            for interface in requires.intersection(charm['i_requires']):
                result_r.setdefault(interface, []).append(payload)
        return result_r, result_p,

    def get_related_client(self, name):
        return self.__class__(self._client, name)

    def get_aliasable_client(self, name=None):
        """Return a client that can be aliased to this one.

        By default, the new client uses a name based on this index's name,
        with a random suffix.
        """
        if name is None:
            name = '%s-%d' % (self.index_name, random.randint(1, 99999))
        return self.get_related_client(name)

    @contextmanager
    def replacement(self, name=None):
        copy = self.get_aliasable_client(name)
        copy.create_index()
        try:
            copy.put_mapping()
            yield copy
        except:
            copy.delete_index()
            raise

    def create_replacement(self, name=None, charms=None, bundles=None):
        """Create a replacement for this index.

        The mapping will not be this index's mapping, it will be the current
        mapping (e.g. put_mapping).

        :param name: If not supplied, a name based on this index's name will
            be used.
        :param charms: A list of charms.  If None, the charms in the current
            index will be used.
        :return: The ElasticSearchClient for the supplied mapping.
        """
        with self.replacement(name) as copy:
            if charms is None:
                try:
                    result = self.api_search(valid_only=False, doctype=CHARM)
                    charms = [r['data'] for r in result]
                except IndexMissing:
                    charms = []
            copy.index_charms(charms)

            if bundles is None:
                try:
                    result = self.api_search(valid_only=False, doctype=BUNDLE)
                    bundles = [r['data'] for r in result]
                except IndexMissing:
                    bundles = []
            copy.index_bundles(bundles)
            return copy

    def health(self):
        """The health status of the ES server."""
        return self._client.health(self.index_name)


def reindex(index_client, charms=None):
    """Reindex documents with the current mapping.

    This works by creating a new index and then aliasing it to the old index.
    If the old index was already an alias, this is a simple swap.  If not, the
    old index is deleted.

    :param charms: A list of charms to use.  If None, the charms in the index
        will be used.
    """
    new_index = index_client.create_replacement(charms=charms)
    try:
        aliased = index_client.get_aliased()
        if aliased == []:
            try:
                index_client.delete_index()
            except IndexMissing:
                pass
        index_client.update_aliased(new_index.index_name, aliased)
    except:
        new_index.delete_index()
        raise
    if aliased != []:
        for name in aliased:
            ElasticSearchClient(index_client._client, name).delete_index()
    return new_index


def update(index_client, force_reindex=False):
    """Update the index to use the current mapping.

    If the index isn't present, a new index is created and aliased to it.
    If the index is present and compatible, it is simply upgraded.
    if the index is present and incompatible, its contents are reindexed.
    """
    actual_client = index_client
    try:
        if force_reindex:
            actual_client = reindex(index_client)
        else:
            index_client.put_mapping()
    except IncompatibleMapping:
        actual_client = reindex(index_client)
    except IndexMissing:
        actual_client = index_client.get_aliasable_client()
        actual_client.create_index()
        try:
            index_client.update_aliased(actual_client.index_name, [])
            actual_client.put_mapping()
        except:
            actual_client.delete_index()
            raise
    return actual_client


def update_main():
    """A main function for the es-update script."""
    update(ElasticSearchClient.from_settings(get_ini()))