org.apache.solr.update.processor.StatelessScriptUpdateProcessorFactory.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.solr.update.processor.StatelessScriptUpdateProcessorFactory.java

Source

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF 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.apache.solr.update.processor;

import org.apache.solr.common.SolrException;
import org.apache.solr.common.util.NamedList;
import org.apache.solr.common.params.ModifiableSolrParams;
import org.apache.solr.core.SolrCore;
import org.apache.solr.core.SolrResourceLoader;
import org.apache.solr.request.SolrQueryRequest;
import org.apache.solr.request.LocalSolrQueryRequest;
import org.apache.solr.response.SolrQueryResponse;
import org.apache.solr.update.*;
import org.apache.solr.util.plugin.SolrCoreAware;
import org.apache.commons.lang.StringUtils;
import org.apache.commons.io.IOUtils;
import org.apache.commons.io.FilenameUtils;

import javax.script.Invocable;
import javax.script.ScriptEngine;
import javax.script.ScriptEngineManager;
import javax.script.ScriptEngineFactory;
import javax.script.ScriptException;

import java.io.IOException;
import java.io.InputStream;
import java.io.Reader;
import java.lang.invoke.MethodHandles;
import java.nio.charset.StandardCharsets;
import java.util.Set;
import java.util.LinkedHashSet;
import java.util.ArrayList;
import java.util.List;
import java.util.Collection;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * <p>
 * An update request processor factory that enables the use of update 
 * processors implemented as scripts which can be loaded by the 
 * {@link SolrResourceLoader} (usually via the <code>conf</code> dir for 
 * the SolrCore).
 * </p>
 * <p>
 * This factory requires at least one configuration parameter named
 * <code>script</code> which may be the name of a script file as a string, 
 * or an array of multiple script files.  If multiple script files are 
 * specified, they are executed sequentially in the order specified in the 
 * configuration -- as if multiple factories were configured sequentially
 * </p>
 * <p>
 * Each script file is expected to declare functions with the same name 
 * as each method in {@link UpdateRequestProcessor}, using the same 
 * arguments.  One slight deviation is in the optional return value from 
 * these functions: If a script function has a <code>boolean</code> return 
 * value, and that value is <code>false</code> then the processor will 
 * cleanly terminate processing of the command and return, without forwarding 
 * the command on to the next script or processor in the chain.
 * Due to limitations in the {@link ScriptEngine} API used by 
 * this factory, it can not enforce that all functions exist on initialization,
 * so errors from missing functions will only be generated at runtime when
 * the chain attempts to use them.
 * </p>
 * <p>
 * The factory may also be configured with an optional "params" argument, 
 * which can be an {@link NamedList} (or array, or any other simple Java 
 * object) which will be put into the global scope for each script.
 * </p>
 * <p>
 * The following variables are define as global variables for each script:
 * <ul>
 *  <li>req - The {@link SolrQueryRequest}</li>
 *  <li>rsp - The {@link SolrQueryResponse}</li>
 *  <li>logger - A {@link Logger} that can be used for logging purposes in the script</li>
 *  <li>params - The "params" init argument in the factory configuration (if any)</li>
 * </ul>
 * <p>
 * Internally this update processor uses JDK 6 scripting engine support, 
 * and any {@link Invocable} implementations of <code>ScriptEngine</code> 
 * that can be loaded using the Solr Plugin ClassLoader may be used.  
 * By default, the engine used for each script is determined by the filed 
 * extension (ie: a *.js file will be treated as a JavaScript script) but 
 * this can be overridden by specifying an explicit "engine" name init 
 * param for the factory, which identifies a registered name of a 
 * {@link ScriptEngineFactory}. 
 * (This may be particularly useful if multiple engines are available for 
 * the same scripting language, and you wish to force the usage of a 
 * particular engine because of known quirks)
 * </p>
 * <p>
 * A new {@link ScriptEngineManager} is created for each 
 * <code>SolrQueryRequest</code> defining a "global" scope for the script(s) 
 * which is request specific.  Separate <code>ScriptEngine</code> instances 
 * are then used to evaluate the script files, resulting in an "engine" scope 
 * that is specific to each script.
 * </p>
 * <p>
 * A simple example...
 * </p>
 * <pre class="prettyprint">
 * &lt;processor class="solr.StatelessScriptUpdateProcessorFactory"&gt;
 *   &lt;str name="script"&gt;updateProcessor.js&lt;/str&gt;
 * &lt;/processor&gt;
 * </pre>
 * <p>
 * A more complex example involving multiple scripts in different languages, 
 * and a "params" <code>NamedList</code> that will be put into the global 
 * scope of each script...
 * </p>
 * <pre class="prettyprint">
 * &lt;processor class="solr.StatelessScriptUpdateProcessorFactory"&gt;
 *   &lt;arr name="script"&gt;
 *     &lt;str name="script"&gt;first-processor.js&lt;/str&gt;
 *     &lt;str name="script"&gt;second-processor.py&lt;/str&gt;
 *   &lt;/arr&gt;
 *   &lt;lst name="params"&gt;
 *     &lt;bool name="a_bool_value"&gt;true&lt;/bool&gt;
 *     &lt;int name="and_int_value"&gt;3&lt;/int&gt;
 *   &lt;/lst&gt;
 * &lt;/processor&gt;
 * </pre>
 * <p>
 * An example where the script file extensions are ignored, and an 
 * explicit script engine is used....
 * </p>
 * <pre class="prettyprint">
 * &lt;processor class="solr.StatelessScriptUpdateProcessorFactory"&gt;
 *   &lt;arr name="script"&gt;
 *     &lt;str name="script"&gt;first-processor.txt&lt;/str&gt;
 *     &lt;str name="script"&gt;second-processor.txt&lt;/str&gt;
 *   &lt;/arr&gt;
 *   &lt;str name="engine"&gt;rhino&lt;/str&gt;
 * &lt;/processor&gt;
 * </pre>
 * 
 */
public class StatelessScriptUpdateProcessorFactory extends UpdateRequestProcessorFactory implements SolrCoreAware {

    private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());

    private final static String SCRIPT_ARG = "script";
    private final static String PARAMS_ARG = "params";
    private final static String ENGINE_NAME_ARG = "engine";

    private List<ScriptFile> scriptFiles;

    /** if non null, this is an override for the engine for all scripts */
    private String engineName = null;

    private Object params = null;

    private SolrResourceLoader resourceLoader;

    private ScriptEngineCustomizer scriptEngineCustomizer;

    @Override
    public void init(NamedList args) {
        Collection<String> scripts = args.removeConfigArgs(SCRIPT_ARG);
        if (scripts.isEmpty()) {
            throw new SolrException(SolrException.ErrorCode.SERVER_ERROR,
                    "StatelessScriptUpdateProcessorFactory must be " + "initialized with at least one "
                            + SCRIPT_ARG);
        }
        scriptFiles = new ArrayList<>();
        for (String script : scripts) {
            scriptFiles.add(new ScriptFile(script));
        }

        params = args.remove(PARAMS_ARG);

        Object engine = args.remove(ENGINE_NAME_ARG);
        if (engine != null) {
            if (engine instanceof String) {
                engineName = (String) engine;
            } else {
                throw new SolrException(SolrException.ErrorCode.SERVER_ERROR,
                        "'" + ENGINE_NAME_ARG + "' init param must be a String (found: " + engine.getClass() + ")");
            }
        }

        super.init(args);

    }

    @Override
    public UpdateRequestProcessor getInstance(SolrQueryRequest req, SolrQueryResponse rsp,
            UpdateRequestProcessor next) {
        List<EngineInfo> scriptEngines = null;

        scriptEngines = initEngines(req, rsp);

        return new ScriptUpdateProcessor(req, rsp, scriptEngines, next);
    }

    // TODO: Make this useful outside of tests, such that a ScriptEngineCustomizer could be looked up through the resource loader
    void setScriptEngineCustomizer(ScriptEngineCustomizer scriptEngineCustomizer) {
        this.scriptEngineCustomizer = scriptEngineCustomizer;
    }

    @Override
    public void inform(SolrCore core) {
        resourceLoader = core.getResourceLoader();

        // test that our engines & scripts are valid

        SolrQueryResponse rsp = new SolrQueryResponse();
        SolrQueryRequest req = new LocalSolrQueryRequest(core, new ModifiableSolrParams());
        try {
            initEngines(req, rsp);
        } catch (Exception e) {
            String msg = "Unable to initialize scripts: " + e.getMessage();
            log.error(msg, e);
            throw new SolrException(SolrException.ErrorCode.SERVER_ERROR, msg, e);
        } finally {
            req.close();
        }

    }

    //================================================ Helper Methods ==================================================

    /**
     * Initializes a list of script engines - an engine per script file.
     *
     * @param req The solr request.
     * @param rsp The solr response
     * @return The list of initialized script engines.
     */
    private List<EngineInfo> initEngines(SolrQueryRequest req, SolrQueryResponse rsp) throws SolrException {

        List<EngineInfo> scriptEngines = new ArrayList<>();

        ScriptEngineManager scriptEngineManager = new ScriptEngineManager(resourceLoader.getClassLoader());

        scriptEngineManager.put("logger", log);
        scriptEngineManager.put("req", req);
        scriptEngineManager.put("rsp", rsp);
        if (params != null) {
            scriptEngineManager.put("params", params);
        }

        for (ScriptFile scriptFile : scriptFiles) {
            ScriptEngine engine = null;
            if (null != engineName) {
                engine = scriptEngineManager.getEngineByName(engineName);
                if (engine == null) {
                    String details = getSupportedEngines(scriptEngineManager, false);
                    throw new SolrException(SolrException.ErrorCode.SERVER_ERROR, "No ScriptEngine found by name: "
                            + engineName + (null != details ? " -- supported names: " + details : ""));
                }
            } else {
                engine = scriptEngineManager.getEngineByExtension(scriptFile.getExtension());
                if (engine == null) {
                    String details = getSupportedEngines(scriptEngineManager, true);
                    throw new SolrException(SolrException.ErrorCode.SERVER_ERROR,
                            "No ScriptEngine found by file extension: " + scriptFile.getFileName()
                                    + (null != details ? " -- supported extensions: " + details : ""));

                }
            }

            if (!(engine instanceof Invocable)) {
                String msg = "Engine "
                        + ((null != engineName) ? engineName : ("for script " + scriptFile.getFileName()))
                        + " does not support function invocation (via Invocable): " + engine.getClass().toString()
                        + " (" + engine.getFactory().getEngineName() + ")";
                log.error(msg);
                throw new SolrException(SolrException.ErrorCode.SERVER_ERROR, msg);
            }

            if (scriptEngineCustomizer != null) {
                scriptEngineCustomizer.customize(engine);
            }

            scriptEngines.add(new EngineInfo((Invocable) engine, scriptFile));
            try {
                Reader scriptSrc = scriptFile.openReader(resourceLoader);

                try {
                    engine.eval(scriptSrc);
                } catch (ScriptException e) {
                    throw new SolrException(SolrException.ErrorCode.SERVER_ERROR,
                            "Unable to evaluate script: " + scriptFile.getFileName(), e);
                } finally {
                    IOUtils.closeQuietly(scriptSrc);
                }
            } catch (IOException ioe) {
                throw new SolrException(SolrException.ErrorCode.SERVER_ERROR,
                        "Unable to evaluate script: " + scriptFile.getFileName(), ioe);
            }
        }
        return scriptEngines;
    }

    /**
     * For error messages - returns null if there are any exceptions of any 
     * kind building the string (or of the list is empty for some unknown reason).
     * @param ext - if true, list of extensions, otherwise a list of engine names
     */
    private static String getSupportedEngines(ScriptEngineManager mgr, boolean ext) {
        String result = null;
        try {
            List<ScriptEngineFactory> factories = mgr.getEngineFactories();
            if (null == factories)
                return result;

            Set<String> engines = new LinkedHashSet<>(factories.size());
            for (ScriptEngineFactory f : factories) {
                if (ext) {
                    engines.addAll(f.getExtensions());
                } else {
                    engines.addAll(f.getNames());
                }
            }
            result = StringUtils.join(engines, ", ");
        } catch (RuntimeException e) {
            /* :NOOP: */
        }
        return result;
    }

    //================================================= Inner Classes ==================================================

    /**
     * The actual update processor. All methods delegate to scripts.
     */
    private static class ScriptUpdateProcessor extends UpdateRequestProcessor {

        private List<EngineInfo> engines;

        private ScriptUpdateProcessor(SolrQueryRequest req, SolrQueryResponse res, List<EngineInfo> engines,
                UpdateRequestProcessor next) {
            super(next);
            this.engines = engines;
        }

        @Override
        public void processAdd(AddUpdateCommand cmd) throws IOException {
            if (invokeFunction("processAdd", cmd)) {
                super.processAdd(cmd);
            }
        }

        @Override
        public void processDelete(DeleteUpdateCommand cmd) throws IOException {
            if (invokeFunction("processDelete", cmd)) {
                super.processDelete(cmd);
            }

        }

        @Override
        public void processMergeIndexes(MergeIndexesCommand cmd) throws IOException {
            if (invokeFunction("processMergeIndexes", cmd)) {
                super.processMergeIndexes(cmd);
            }
        }

        @Override
        public void processCommit(CommitUpdateCommand cmd) throws IOException {
            if (invokeFunction("processCommit", cmd)) {
                super.processCommit(cmd);
            }
        }

        @Override
        public void processRollback(RollbackUpdateCommand cmd) throws IOException {
            if (invokeFunction("processRollback", cmd)) {
                super.processRollback(cmd);
            }
        }

        @Override
        public void finish() throws IOException {
            if (invokeFunction("finish")) {
                super.finish();
            }
        }

        /**
         * returns true if processing should continue, or false if the 
         * request should be ended now.  Result value is computed from the return 
         * value of the script function if: it exists, is non-null, and can be 
         * cast to a java Boolean.
         */
        private boolean invokeFunction(String name, Object... cmd) {

            for (EngineInfo engine : engines) {
                try {
                    Object result = engine.getEngine().invokeFunction(name, cmd);
                    if (null != result && result instanceof Boolean) {
                        if (!((Boolean) result).booleanValue()) {
                            return false;
                        }
                    }

                } catch (ScriptException | NoSuchMethodException e) {
                    throw new SolrException(
                            SolrException.ErrorCode.SERVER_ERROR, "Unable to invoke function " + name
                                    + " in script: " + engine.getScriptFile().getFileName() + ": " + e.getMessage(),
                            e);
                }
            }

            return true;
        }
    }

    /**
     * Holds the script engine and its associated script file.
     */
    private static class EngineInfo {

        private final Invocable engine;
        private final ScriptFile scriptFile;

        private EngineInfo(Invocable engine, ScriptFile scriptFile) {
            this.engine = engine;
            this.scriptFile = scriptFile;
        }

        public Invocable getEngine() {
            return engine;
        }

        public ScriptFile getScriptFile() {
            return scriptFile;
        }
    }

    /**
     * Represents a script file.
     */
    private static class ScriptFile {

        private final String fileName;
        private final String extension;

        private ScriptFile(String fileName) {
            this.fileName = fileName;
            this.extension = FilenameUtils.getExtension(fileName);
        }

        public String getFileName() {
            return fileName;
        }

        public String getExtension() {
            return extension;
        }

        public Reader openReader(SolrResourceLoader resourceLoader) throws IOException {
            InputStream input = resourceLoader.openResource(fileName);
            return org.apache.lucene.util.IOUtils.getDecodingReader(input, StandardCharsets.UTF_8);
        }
    }
}