org.elasticsearch.script.MockScriptEngine Maven / Gradle / Ivy

Go to download
/*
 * Licensed to Elasticsearch under one or more contributor
 * license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright
 * ownership. Elasticsearch licenses this file to you under
 * the Apache License, Version 2.0 (the "License"); you may
 * not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package org.elasticsearch.script;

import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.Scorer;
import org.elasticsearch.common.Nullable;
import org.elasticsearch.common.bytes.BytesArray;
import org.elasticsearch.search.lookup.LeafSearchLookup;
import org.elasticsearch.search.lookup.SearchLookup;

import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.function.Function;

import static java.util.Collections.emptyMap;

/**
 * A mocked script engine that can be used for testing purpose.
 *
 * This script engine allows to define a set of predefined scripts that basically a combination of a key and a
 * function:
 *
 * The key can be anything as long as it is a {@link String} and is used to resolve the scripts
 * at compilation time. For inline scripts, the key can be a description of the script. For stored and file scripts,
 * the source must match a key in the predefined set of scripts.
 *
 * The function is used to provide the result of the script execution and can return anything.
 */
public class MockScriptEngine implements ScriptEngineService {

    public static final String NAME = "mockscript";

    private final String type;
    private final Map, Object>> scripts;

    public MockScriptEngine(String type, Map, Object>> scripts) {
        this.type = type;
        this.scripts = Collections.unmodifiableMap(scripts);
    }

    public MockScriptEngine() {
        this(NAME, Collections.emptyMap());
    }

    @Override
    public String getType() {
        return type;
    }

    @Override
    public String getExtension() {
        return getType();
    }

    @Override
    public Object compile(String name, String source, Map params) {
        // Scripts are always resolved using the script's source. For inline scripts, it's easy because they don't have names and the
        // source is always provided. For stored and file scripts, the source of the script must match the key of a predefined script.
        Function, Object> script = scripts.get(source);
        if (script == null) {
            throw new IllegalArgumentException("No pre defined script matching [" + source + "] for script with name [" + name + "], " +
                    "did you declare the mocked script?");
        }
        return new MockCompiledScript(name, params, source, script);
    }

    @Override
    public ExecutableScript executable(CompiledScript compiledScript, @Nullable Map vars) {
        MockCompiledScript compiled = (MockCompiledScript) compiledScript.compiled();
        return compiled.createExecutableScript(vars);
    }

    @Override
    public SearchScript search(CompiledScript compiledScript, SearchLookup lookup, @Nullable Map vars) {
        MockCompiledScript compiled = (MockCompiledScript) compiledScript.compiled();
        return compiled.createSearchScript(vars, lookup);
    }

    @Override
    public void close() throws IOException {
    }

    @Override
    public boolean isInlineScriptEnabled() {
        return true;
    }


    public class MockCompiledScript {

        private final String name;
        private final String source;
        private final Map params;
        private final Function, Object> script;

        public MockCompiledScript(String name, Map params, String source, Function, Object> script) {
            this.name = name;
            this.source = source;
            this.params = params;
            this.script = script;
        }

        public String getName() {
            return name;
        }

        public ExecutableScript createExecutableScript(Map vars) {
            Map context = new HashMap<>();
            if (params != null) {
                context.putAll(params);
            }
            if (vars != null) {
                context.putAll(vars);
            }
            return new MockExecutableScript(context, script != null ? script : ctx -> new BytesArray(source));
        }

        public SearchScript createSearchScript(Map vars, SearchLookup lookup) {
            Map context = new HashMap<>();
            if (params != null) {
                context.putAll(params);
            }
            if (vars != null) {
                context.putAll(vars);
            }
            return new MockSearchScript(lookup, context, script != null ? script : ctx -> source);
        }
    }

    public class MockExecutableScript implements ExecutableScript {

        private final Function, Object> script;
        private final Map vars;

        public MockExecutableScript(Map vars, Function, Object> script) {
            this.vars = vars;
            this.script = script;
        }

        @Override
        public void setNextVar(String name, Object value) {
            vars.put(name, value);
        }

        @Override
        public Object run() {
            return script.apply(vars);
        }
    }

    public class MockSearchScript implements SearchScript {

        private final Function, Object> script;
        private final Map vars;
        private final SearchLookup lookup;

        public MockSearchScript(SearchLookup lookup, Map vars, Function, Object> script) {
            this.lookup = lookup;
            this.vars = vars;
            this.script = script;
        }

        @Override
        public LeafSearchScript getLeafSearchScript(LeafReaderContext context) throws IOException {
            LeafSearchLookup leafLookup = lookup.getLeafSearchLookup(context);

            Map ctx = new HashMap<>(leafLookup.asMap());
            if (vars != null) {
                ctx.putAll(vars);
            }

            AbstractSearchScript leafSearchScript = new AbstractSearchScript() {

                @Override
                public Object run() {
                    return script.apply(ctx);
                }

                @Override
                public void setNextVar(String name, Object value) {
                    ctx.put(name, value);
                }

                @Override
                public void setScorer(Scorer scorer) {
                    super.setScorer(scorer);
                    ctx.put("_score", new ScoreAccessor(scorer));
                }
            };
            leafSearchScript.setLookup(leafLookup);
            return leafSearchScript;
        }

        @Override
        public boolean needsScores() {
            return true;
        }
    }

    public static Script mockInlineScript(final String script) {
        return new Script(ScriptType.INLINE, "mock", script, emptyMap());
    }

}