def default_index(f):
    def decorator(*args, **kwargs):
def author_index(f):
    def decorator(*args, **kwargs):
def search(query,
           index=None,
           filters=list(),
           size=10,
           include="*",
           exclude="authors",
           offset=0,
           sort_field=None,
           sort_order='',
           post_filter=None):
def search_authors(name, size=20, author_index=None):
def reindex_all(index=None, author_index=None, recreate=False, batch=5, start=-1, end=-1, synchronous=False):
def reindex_batch(hepsubmission_record_ids, index):
def get_record(record_id, index=None):
def get_records_matching_field(field, id, index=None, doc_type=None, source=None):
def delete_item_from_index(id, index, doc_type, parent=None):
def push_data_keywords(pub_ids=None, index=None):
def index_record_ids(record_ids, index=None):
def index_record_dict(record_dict, doc_type, recid, index=None, parent=None):
def recreate_index(index=None):
def fetch_record(record_id, doc_type, index=None):
def get_n_latest_records(n_latest, field="last_updated", index=None):
def get_count_for_collection(doc_type, index=None):
def get_all_ids(index=None, id_field='recid', last_updated=None, latest_first=False): def add_data_keywords(doc):
    # Aggregate and filter out invalid keywords
    valid_keywords = list(es_mapping['data_keywords']['properties'].keys())
    agg_keywords = defaultdict(list)
    for kw in doc["keywords"]:
        if kw['name'] in valid_keywords:
            agg_keywords[kw['name']].append(kw['value'])

    # Remove duplicates
    for k, v in agg_keywords.items():
        agg_keywords[k] = list(set(v))

    agg_keywords = process_cmenergies(agg_keywords)
    doc['data_keywords'] = dict(agg_keywords) def process_cmenergies(keywords):
    cmenergies = []
    if keywords['cmenergies']:
        for cmenergy in keywords['cmenergies']:
            cmenergy = cmenergy.strip(" gevGEV")
            try:
                cmenergy_val = float(cmenergy)
                cmenergies.append({"gte": cmenergy_val, "lte": cmenergy_val})
            except ValueError:
                m = re.match(r'^(-?\d+(?:\.\d+)?)[ \-+andAND]+(-?\d+(?:\.\d+)?)$', cmenergy)
                if m:
                    cmenergy_range = [float(m.group(1)), float(m.group(2))]
                    cmenergy_range.sort()
                    cmenergies.append({"gte": cmenergy_range[0], "lte": cmenergy_range[1]})
                else:
                    log.warning("Invalid value for cmenergies: %s" % cmenergy)

        keywords['cmenergies'] = cmenergies

    return keywords class QueryBuilder:

    @staticmethod
    def add_filters(search, filters):
        from .config.es_config import get_filter_field

        for name, value in filters:
            if name == "author":
                search = search.query('nested', path='authors', query=Q('match', authors__full_name = value))
            else:
                (filter_type, field, value) = get_filter_field(name, value)
                search = search.filter(filter_type, **{field: value})

        return search class HEPDataQueryParser(object):
    @staticmethod
    def parse_query(query_string):
        # query should be something like 'observable:ASYM' which
        # would translate to data_keywords.observables:ASYM
        mapping = {
            "keys": {
                "observables": "data_keywords.observables",
                "cmenergies": "data_keywords.cmenergies",
                "phrases": "data_keywords.phrases",
                "reactions": "data_keywords.reactions"
            }
        }

        new_query_string = query_string

        for query_part in re.split("AND|OR", query_string):
            query_part = query_part.strip()
            if ':' in query_part:
                _key, _value = query_part.split(':', maxsplit=1)
                _key = mapping['keys'].get(_key, _key)
                _value = HEPDataQueryParser._quote_phrase(_value)
                new_query_string = new_query_string.replace(query_part, f"{_key}:{_value}")
            else:
                new_query_string = new_query_string.replace(
                    query_part, HEPDataQueryParser()._quote_phrase(query_part)
                )

        return new_query_string

    @staticmethod
    def _quote_phrase(phrase):
        # Match phrases containing a reaction (including "-->") or a doi (word
        # chars with / in the middle) and quote them
        pattern = re.compile("(.*-->.*|[\w\.]+\/[\w\.]+)")

        if '"' not in phrase and pattern.fullmatch(phrase):
            return f'"{phrase}"'
        return phrase def flip_sort_order(order):
    """ Flip the sort order from descending to ascending or vice-versa. """
    if order == 'asc':
        return 'desc'
    else:
        return 'asc' def parse_and_format_date(datestring):
    """ Parse a date string into a nice readable format. """
    if not datestring:
        return None
    import dateutil.parser
    timedate = dateutil.parser.parse(datestring)
    return timedate.strftime('%d %b %Y') def prepare_author_for_indexing(document):
    """ Extract the author list from the document and index it in a separate
     index. """

    index = current_app.config['AUTHOR_INDEX']

    author_data = []
    authors = document.get('authors', None)

    if authors is not None:
        for author in authors:
            data_dict = author

            op_dict = {
                "index": {
                    "_index": index,
                    "_id": author['full_name']
                }
            }
            author_data.append(op_dict)
            author_data.append(data_dict)

    return author_data def calculate_sort_order(is_reversed, sorting_field):
    """ Take the default sort order for a given field and an information
     whether to flip it and compute the final sorting order. """
    from .config.es_config import default_sort_order_for_field
    default_sort_order = default_sort_order_for_field(sorting_field)
    if is_reversed == 'rev':
        return flip_sort_order(default_sort_order)
    else:
        return default_sort_order def push_keywords(docs):
    """
        Add keywords from datatables to the corresponding publication record
    """
    from hepdata.utils.miscellaneous import splitter
    datatables, publications = splitter(docs,
                                        lambda d: 'related_publication' in d)
    if len(publications) == 0 and len(datatables) == 0:
        raise ValueError("Documents provided are not appropriate " +
                         "for pushing keywords")

    # check the related publication field

    for pub in publications:
        data = filter(lambda table:
                      table['related_publication'] == pub['recid'],
                      datatables)

        keywords = reduce(lambda acc, d: acc + d['keywords'], data, [])

        agg_keywords = defaultdict(list)
        for kw in keywords:
            agg_keywords[kw['name']].append(kw['value'])

        # Remove duplicates
        for k, v in agg_keywords.items():
            agg_keywords[k] = list(set(v))

        pub['data_keywords'] = agg_keywords

    return publications + datatables def import_records(inspire_ids, synchronous=False, update_existing=False,
                   base_url='https://hepdata.net', send_email=False):
    """
    Import records from hepdata.net

    :param inspire_ids: array of inspire ids to load (in the format insXXX).
    :param synchronous: if should be run immediately rather than via celery
    :param update_existing: whether to update records that already exist
    :param base_url: override default base URL
    :param send_email: whether to send emails on finalising submissions
    :return: None
    """
    for index, inspire_id in enumerate(inspire_ids):
        _cleaned_id = str(inspire_id).replace("ins", "")
        if synchronous:
            _import_record(_cleaned_id, update_existing=update_existing,
                           base_url=base_url, send_email=send_email)
        else:
            log.info("Sending import_record task to celery for id %s"
                     % _cleaned_id)
            _import_record.delay(_cleaned_id, update_existing=update_existing,
                                 base_url=base_url, send_email=send_email) def check_max_results(args):
    """
    Get the size query parameter from the URL and if it doesn't exist
    assign a default value.
    """
    max_results = args.get('size', HEPDATA_CFG_DEFAULT_RESULTS_PER_PAGE)
    try:
        max_results = int(max_results)
    except ValueError:
        max_results = HEPDATA_CFG_DEFAULT_RESULTS_PER_PAGE

    if max_results < 1:
        max_results = HEPDATA_CFG_DEFAULT_RESULTS_PER_PAGE
    elif max_results > LIMIT_MAX_RESULTS_PER_PAGE:
        max_results = LIMIT_MAX_RESULTS_PER_PAGE

    args['size'] = max_results
from elasticsearch.exceptions import NotFoundError
from elasticsearch_dsl import Search
from invenio_db import db
from unittest.mock import call
from hepdata.ext.elasticsearch.config.es_config import \
    add_default_aggregations, sort_fields_mapping
from hepdata.ext.elasticsearch import api as es_api
from hepdata.ext.elasticsearch.config.es_config import get_filter_field
from hepdata.ext.elasticsearch.document_enhancers import add_data_keywords, process_cmenergies
from hepdata.ext.elasticsearch.process_results import merge_results, match_tables_to_papers, \
    get_basic_record_information, is_datatable
from hepdata.ext.elasticsearch.query_builder import QueryBuilder, HEPDataQueryParser
from hepdata.ext.elasticsearch.utils import flip_sort_order, parse_and_format_date, prepare_author_for_indexing, \
    calculate_sort_order, push_keywords
from hepdata.modules.records.importer.api import import_records
from hepdata.modules.submission.models import HEPSubmission
from invenio_search import current_search_client as es
from hepdata.modules.search.config import LIMIT_MAX_RESULTS_PER_PAGE, \
    HEPDATA_CFG_DEFAULT_RESULTS_PER_PAGE
from hepdata.modules.search.views import check_max_results
import datetime
import pytest
                                    "AND unknown_field:hello")

    _test_query4 = 'reactions:P P --> LQ LQ X AND doi:10.1007/s100520000432'
    parsed_query_string4 = HEPDataQueryParser.parse_query(_test_query4)

    assert (parsed_query_string4 == 'data_keywords.reactions:"P P --> LQ LQ X"'
                                    ' AND doi:"10.1007/s100520000432"')

    _test_query5 = 'P P --> LQ LQ X'
    parsed_query_string5 = HEPDataQueryParser.parse_query(_test_query5)

    assert (parsed_query_string5 == '"P P --> LQ LQ X"')



def test_search(app, load_default_data, identifiers):
    """
    Test the search functions work correctly, also with the new query syntax.
    Testing both standard and authors search here to save loading the data multiple times
    :return:
    """
    index = app.config.get('ELASTICSEARCH_INDEX')

    # Test searching with an empty query
    results = es_api.search('', index=index)
    assert(results['total'] == len(identifiers))
    assert(len(results['facets']) == 8)
    assert(len(results['results']) == len(identifiers))

    for i in range(len(results['results'])):
        assert(results['results'][i]['title'] == identifiers[i]['title'])
        assert(results['results'][i]['inspire_id'] == identifiers[i]['inspire_id'])
        assert(len(results['results'][i]['data']) == identifiers[i]['data_tables'])

    # Test pagination (1 item per page as we only have 2; get 2nd page)
    results = es_api.search('', index=index, size=1, offset=1)
    assert(results['total'] == len(identifiers))
    assert(len(results['results']) == 1)
    assert(results['results'][0]['title'] == identifiers[1]['title'])

    # Test a simple search query from the second test submission
    # The search matches the publication but not the data tables
    results = es_api.search('charmonium', index=index)
    assert(results['total'] == 1)
    assert(results['results'][0]['inspire_id'] == identifiers[1]['inspire_id'])
    assert(len(results['results'][0]['data']) == 0)

    # Test the authors search (fuzzy)
    results = es_api.search_authors('Bal')
    expected = [
        {'affiliation': 'Texas U., Arlington', 'full_name': 'Pal, Arnab'},
        {'affiliation': 'Panjab U.', 'full_name': 'Bala, A.'}
    ]

    assert(len(results) == len(expected))
    for author in expected:
        assert(author in results)

    # Test search queries that ES can't parse
    results = es_api.search('/', index=index)
    assert results == {'error': 'Failed to parse query [/]'}

    results = es_api.search('cmenergies:[1.3%20TO%201.4]', index=index)
    assert results == {'error': 'Failed to parse query [data_keywords.cmenergies:[1.3%20TO%201.4]]'}

    results = es_api.search('(SELECT (CHR(113)||CHR(122)||CHR(122)||CHR(122)||CHR(113))||(SELECT (CASE WHEN (6242=6242) THEN 1 ELSE 0 END))::text||(CHR(113)||CHR(120)||CHR(107)||CHR(98)||CHR(113)))', index=index)
    assert results == {'error': 'Failed to parse query [(SELECT (CHR(113)||CHR(122)||CHR(122)||CHR(122)||CHR(113))||(SELECT (CASE WHEN (6242=6242) THEN 1 ELSE 0 END))::text||(CHR(113)||CHR(120)||CHR(107)||CHR(98)||CHR(113)))]'}

    # Test a search query to an invalid index
    results = es_api.search('hello', index='thisisnotanindex')
    assert results == {'error': 'An unexpected error occurred: index_not_found_exception'}


def test_merge_results():
    pub_result = {
        "hits": {
            "hits": [{"key": "testa"}],
            "total": { "value": 2, "relation": "eq" }
        }
    }
    data_result = {
        "hits": {
            "hits": [{"key": "testb"}],
            "total": { "value": 1, "relation": "eq" }
        }
    }

    merged = merge_results(pub_result, data_result)
    assert ("hits" in merged)
    assert (len(merged["hits"]) == 2)
    assert (merged["total"] == 2)


def test_flip_sort_order():
    order = "desc"
    order = flip_sort_order(order=order)
    assert (order is "asc")
    order = flip_sort_order(order=order)
    assert (order is "desc")


def test_parse_date():
    date = parse_and_format_date("09-01-2016")
    assert (date == "01 Sep 2016")

    date = parse_and_format_date(None)
    assert (date is None)


def test_calculate_sort_order():
    sort_title = calculate_sort_order(None, "title")

    assert (sort_title is "asc")

    sort_title_rev = calculate_sort_order("rev", "title")
    assert (sort_title_rev is "desc")

    sort_collaborations = calculate_sort_order(None, "collaborations")
    assert (sort_collaborations is "asc")


    sort_date = calculate_sort_order(None, "date")