<commit_before>package whelk.rest.api

import groovy.transform.CompileStatic
import groovy.transform.TypeCheckingMode
import groovy.util.logging.Log4j2 as Log
import org.codehaus.jackson.map.ObjectMapper

import whelk.JsonLd
import whelk.Whelk


@CompileStatic
@Log
class ESQuery {
    private Whelk whelk
    private JsonLd jsonld
    private Set keywordFields
    private static final ObjectMapper mapper = new ObjectMapper()
    private static final int DEFAULT_PAGE_SIZE = 50
    private static final List RESERVED_PARAMS = [
        'q', '_limit', '_offset', '_sort', '_statsrepr', '_site_base_uri'
    ]

    ESQuery() {
        // NOTE: For unit tests only!
    }

    ESQuery(Whelk whelk) {
        this.whelk = whelk
        this.jsonld = whelk.jsonld
        this.keywordFields = getKeywords(this.whelk)
    }

    Set getKeywords(Whelk whelk) {
        if (whelk.elastic) {
            Map mappings = whelk.elastic.getMappings()
            return getKeywordFields(mappings)
        } else {
            return [] as Set
        }
    }

    void setKeywords(Set keywordFields) {
        // NOTE: For unit tests only!
        this.keywordFields = keywordFields
    }

    @CompileStatic(TypeCheckingMode.SKIP)
    Map doQuery(Map<String, String[]> queryParameters, String dataset) {
        Map esQuery = getESQuery(queryParameters)
        Map esResponse = hideKeywordFields(whelk.elastic.query(esQuery, dataset))
        return esResponse
    }

    @CompileStatic(TypeCheckingMode.SKIP)
    Map getESQuery(Map<String, String[]> queryParameters0) {
        // Legit params and their uses:
        //   q - simple_query_string
        String q
        //   _limit, _offset - pagination
        int limit
        int offset
        //   _sort - keyword sort
        List sortBy
        //   _statsrepr - aggregates
        Map aggQuery
        //   _site_base_uri - Filter id.kb.se resources
        List siteFilter
        //   any k=v param - FILTER query (same key => OR, different key => AND)
        List filters

        def (String[] originalTypeParam, queryParameters) =
            expandTypeParam(queryParameters0, whelk.jsonld)

        q = getQueryString(queryParameters)
        (limit, offset) = getPaginationParams(queryParameters)
        sortBy = getSortClauses(queryParameters)
        siteFilter = getSiteFilter(queryParameters)
        aggQuery = getAggQuery(queryParameters)
        filters = getFilters(queryParameters)

        queryParameters = maybeResetTypeParam(queryParameters, originalTypeParam)

        Map queryString = [
            'simple_query_string': [
                'query': q,
                'default_operator':  'AND'
            ]
        ]

        Map query = [
            'query': [
                'bool': [
                    'must': [
                        queryString
                    ]
                ]
            ]
        ]

        if (limit >= 0) {
            query['size'] = limit
        }

        if (offset) {
            query['from'] = offset
        }

        if (filters && siteFilter) {
            query['query']['bool']['filter'] = filters + siteFilter
        } else if (filters) {
            query['query']['bool']['filter'] = filters
        } else if (siteFilter) {
            query['query']['bool']['filter'] = siteFilter
        }

        if (sortBy) {
            query['sort'] = sortBy
        }

        if (aggQuery) {
            query['aggs'] = aggQuery
        }

        return query
    }

    /**
     * Expand `@type` query parameter with subclasses
     *
     * This also removes superclasses, since we only care about the most
     * specific class.
     *
     * Public for test only - don't call outside this class!
     *
     */
    public Tuple2<String[], Map<String, String[]>> expandTypeParam(Map<String, String[]> queryParameters,
                                                                   JsonLd jsonld) {
        // Filter out all @types that have (more specific) subclasses that are also in the list
        // So for example [Instance, Electronic] should be reduced to just [Electronic].
        // Afterwards, include all subclasses of the remaining @types
        String[] types = queryParameters.get('@type')
        if (types != null) {
            // Select types to prune
            Set<String> toBeRemoved = []
            for (String c1 : types) {
                ArrayList<String> c1SuperClasses = []
                jsonld.getSuperClasses(c1, c1SuperClasses)
                toBeRemoved.addAll(c1SuperClasses)
            }
            // Make a new pruned list without the undesired superclasses
            List<String> prunedTypes = []
            for (String type : types) {
                if (!toBeRemoved.contains(type))
                    prunedTypes.add(type)
            }
            // Add all subclasses of the remaining types
            ArrayList<String> subClasses = []
            for (String type : prunedTypes) {
                List<String> result = []
                jsonld.getSuperClasses('bar', result)
                jsonld.getSubClasses(type, subClasses)
                subClasses.add(type)
            }

            queryParameters.put('@type', (String[]) subClasses.toArray())
        }

        return new Tuple2(types, queryParameters)
    }

    private Map<String, String[]> maybeResetTypeParam(Map<String, String[]> queryParameters,
                                                      String[] originalTypeParam) {
        if (originalTypeParam != null) {
            queryParameters.put('@type', originalTypeParam)
        }

        return queryParameters
    }

    /**
     * Get query string from query params, or default if not found.
     *
     * Public for test only - don't call outside this class!
     *<commit_msg>Adapt rest component to change in JsonLd.getSubClasses signature<commit_after>package whelk.rest.api

import groovy.transform.CompileStatic
import groovy.transform.TypeCheckingMode
import groovy.util.logging.Log4j2 as Log
import org.codehaus.jackson.map.ObjectMapper

import whelk.JsonLd
import whelk.Whelk


@CompileStatic
@Log
class ESQuery {
    private Whelk whelk
    private JsonLd jsonld
    private Set keywordFields
    private static final ObjectMapper mapper = new ObjectMapper()
    private static final int DEFAULT_PAGE_SIZE = 50
    private static final List RESERVED_PARAMS = [
        'q', '_limit', '_offset', '_sort', '_statsrepr', '_site_base_uri'
    ]

    ESQuery() {
        // NOTE: For unit tests only!
    }

    ESQuery(Whelk whelk) {
        this.whelk = whelk
        this.jsonld = whelk.jsonld
        this.keywordFields = getKeywords(this.whelk)
    }

    Set getKeywords(Whelk whelk) {
        if (whelk.elastic) {
            Map mappings = whelk.elastic.getMappings()
            return getKeywordFields(mappings)
        } else {
            return [] as Set
        }
    }

    void setKeywords(Set keywordFields) {
        // NOTE: For unit tests only!
        this.keywordFields = keywordFields
    }

    @CompileStatic(TypeCheckingMode.SKIP)
    Map doQuery(Map<String, String[]> queryParameters, String dataset) {
        Map esQuery = getESQuery(queryParameters)
        Map esResponse = hideKeywordFields(whelk.elastic.query(esQuery, dataset))
        return esResponse
    }

    @CompileStatic(TypeCheckingMode.SKIP)
    Map getESQuery(Map<String, String[]> queryParameters0) {
        // Legit params and their uses:
        //   q - simple_query_string
        String q
        //   _limit, _offset - pagination
        int limit
        int offset
        //   _sort - keyword sort
        List sortBy
        //   _statsrepr - aggregates
        Map aggQuery
        //   _site_base_uri - Filter id.kb.se resources
        List siteFilter
        //   any k=v param - FILTER query (same key => OR, different key => AND)
        List filters

        def (String[] originalTypeParam, queryParameters) =
            expandTypeParam(queryParameters0, whelk.jsonld)

        q = getQueryString(queryParameters)
        (limit, offset) = getPaginationParams(queryParameters)
        sortBy = getSortClauses(queryParameters)
        siteFilter = getSiteFilter(queryParameters)
        aggQuery = getAggQuery(queryParameters)
        filters = getFilters(queryParameters)

        queryParameters = maybeResetTypeParam(queryParameters, originalTypeParam)

        Map queryString = [
            'simple_query_string': [
                'query': q,
                'default_operator':  'AND'
            ]
        ]

        Map query = [
            'query': [
                'bool': [
                    'must': [
                        queryString
                    ]
                ]
            ]
        ]

        if (limit >= 0) {
            query['size'] = limit
        }

        if (offset) {
            query['from'] = offset
        }

        if (filters && siteFilter) {
            query['query']['bool']['filter'] = filters + siteFilter
        } else if (filters) {
            query['query']['bool']['filter'] = filters
        } else if (siteFilter) {
            query['query']['bool']['filter'] = siteFilter
        }

        if (sortBy) {
            query['sort'] = sortBy
        }

        if (aggQuery) {
            query['aggs'] = aggQuery
        }

        return query
    }

    /**
     * Expand `@type` query parameter with subclasses
     *
     * This also removes superclasses, since we only care about the most
     * specific class.
     *
     * Public for test only - don't call outside this class!
     *
     */
    public Tuple2<String[], Map<String, String[]>> expandTypeParam(Map<String, String[]> queryParameters,
                                                                   JsonLd jsonld) {
        // Filter out all @types that have (more specific) subclasses that are also in the list
        // So for example [Instance, Electronic] should be reduced to just [Electronic].
        // Afterwards, include all subclasses of the remaining @types
        String[] types = queryParameters.get('@type')
        if (types != null) {
            // Select types to prune
            Set<String> toBeRemoved = []
            for (String c1 : types) {
                ArrayList<String> c1SuperClasses = []
                jsonld.getSuperClasses(c1, c1SuperClasses)
                toBeRemoved.addAll(c1SuperClasses)
            }
            // Make a new pruned list without the undesired superclasses
            List<String> prunedTypes = []
            for (String type : types) {
                if (!toBeRemoved.contains(type))
                    prunedTypes.add(type)
            }
            // Add all subclasses of the remaining types
            Set<String> subClasses = []
            for (String type : prunedTypes) {
                subClasses += jsonld.getSubClasses(type)
                subClasses.add(type)
            }

            queryParameters.put('@type', (String[]) subClasses.toArray())
        }

        return new Tuple2(types, queryParameters)
    }

    private Map<String, String[]> maybeResetTypeParam(Map<String, String[]> queryParameters,
                                                      String[] originalTypeParam) {
        if (originalTypeParam != null) {
            queryParameters.put('@type', originalTypeParam)
        }

        return queryParameters
    }

    /**
     * Get query string from query params, or default if not found.
     *
     * Public for test only - don't call outside this class!
     *
{'fraction_non_alphanumeric': 0.09167098669431484, 'fraction_numerical': 0.002419215482979091, 'mean_word_length': 1.8060606060606061, 'pattern_counts': {'":': 0, '<': 26, '<?xml version=': 0, '>': 32, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6927', 'n_tokens_mistral': 3157, 'n_tokens_neox': 3002, 'n_words': 799}