org.codelibs.elasticsearch.velocity.service.VelocityScriptEngineService Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of elasticsearch-lang-velocity Show documentation
Show all versions of elasticsearch-lang-velocity Show documentation
This plugin adds Velocity language to Elasticsearch.
package org.codelibs.elasticsearch.velocity.service;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.Writer;
import java.lang.ref.SoftReference;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.AccessController;
import java.security.PrivilegedAction;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Queue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import org.apache.logging.log4j.Logger;
import org.apache.velocity.Template;
import org.apache.velocity.VelocityContext;
import org.apache.velocity.app.VelocityEngine;
import org.apache.velocity.context.Context;
import org.apache.velocity.exception.VelocityException;
import org.codelibs.elasticsearch.velocity.util.ContextProperties;
import org.elasticsearch.ElasticsearchException;
import org.elasticsearch.common.bytes.BytesReference;
import org.elasticsearch.common.component.AbstractComponent;
import org.elasticsearch.common.io.UTF8StreamWriter;
import org.elasticsearch.common.io.stream.BytesStreamOutput;
import org.elasticsearch.common.settings.Setting;
import org.elasticsearch.common.settings.Setting.Property;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.util.concurrent.ThreadContext;
import org.elasticsearch.env.Environment;
import org.elasticsearch.script.CompiledScript;
import org.elasticsearch.script.ExecutableScript;
import org.elasticsearch.script.ScriptEngineService;
import org.elasticsearch.script.SearchScript;
import org.elasticsearch.search.lookup.SearchLookup;
public class VelocityScriptEngineService extends AbstractComponent implements ScriptEngineService {
public static final Setting SETTING_SCRIPT_VELOCITY_PROPS =
Setting.groupSetting("script.velocity.props.", Property.NodeScope);
public static final Setting SETTING_SCRIPT_VELOCITY_CONTEXT_PROPS =
Setting.groupSetting("script.velocity.context.props.", Property.NodeScope);
public static final String NAME = "velocity";
/**
* Thread local UTF8StreamWriter to store template execution results in,
* thread local to save object creation.
*/
private static ThreadLocal> utf8StreamWriter = new ThreadLocal<>();
private VelocityEngine velocityEngine;
private File workDir;
private final Queue templateFileQueue = new ConcurrentLinkedQueue<>();
private final Map contextPropMap = new ConcurrentHashMap<>();
private ThreadContext threadContext;
/**
* If exists, reset and return, otherwise create, reset and return a writer.
*/
private static UTF8StreamWriter utf8StreamWriter() {
final SoftReference ref = utf8StreamWriter.get();
UTF8StreamWriter writer = ref == null ? null : ref.get();
if (writer == null) {
writer = new UTF8StreamWriter(1024 * 4);
utf8StreamWriter.set(new SoftReference<>(writer));
}
writer.reset();
return writer;
}
public VelocityScriptEngineService(final Settings settings) {
super(settings);
workDir = findWorkDir(settings);
final Path configPath = Paths.get(Environment.PATH_CONF_SETTING.get(settings));
final Map contextPropSettings = SETTING_SCRIPT_VELOCITY_CONTEXT_PROPS.get(settings).getAsMap();
for (final Map.Entry entry : contextPropSettings.entrySet()) {
final String key = entry.getKey();
if (key.indexOf('.') == -1) {
final Path path = configPath.resolve(entry.getValue());
if (exists(path)) {
final ContextProperties properties = new ContextProperties(path.toFile());
contextPropMap.put(key, properties);
final String interval = contextPropSettings.get(key + ".interval");
if (interval != null) {
try {
properties.checkInterval = Long.parseLong(interval);
} catch (final NumberFormatException e) {
logger.warn("{} is not long type.", e, interval);
}
}
} else {
logger.warn("{} is not found.", path);
}
}
}
final Properties props = new Properties();
for (final Map.Entry entry : SETTING_SCRIPT_VELOCITY_PROPS.get(settings).getAsMap().entrySet()) {
props.put(entry.getKey(), entry.getValue());
}
final String resourceLoader = (String) props.get("resource.loader");
if (resourceLoader != null) {
props.put("resource.loader", "WORK_TMPL,ES_TMPL," + resourceLoader);
} else {
props.put("resource.loader", "WORK_TMPL,ES_TMPL");
}
initPropertyValue(props, "WORK_TMPL.resource.loader.class", "org.apache.velocity.runtime.resource.loader.FileResourceLoader");
initPropertyValue(props, "WORK_TMPL.resource.loader.path", workDir.getAbsolutePath());
initPropertyValue(props, "WORK_TMPL.resource.loader.cache", "true");
initPropertyValue(props, "WORK_TMPL.resource.loader.modificationCheckInterval", "0");
initPropertyValue(props, "ES_TMPL.resource.loader.class", "org.apache.velocity.runtime.resource.loader.FileResourceLoader");
initPropertyValue(props, "ES_TMPL.resource.loader.path", configPath.resolve("scripts").toFile().getAbsolutePath());
initPropertyValue(props, "ES_TMPL.resource.loader.cache", "true");
initPropertyValue(props, "ES_TMPL.resource.loader.modificationCheckInterval", "60");
final Path logsFile = Paths.get(Environment.PATH_LOGS_SETTING.get(settings));
initPropertyValue(props, "velocimacro.library.autoreload", "false");
initPropertyValue(props, "input.encoding", "UTF-8");
initPropertyValue(props, "output.encoding", "UTF-8");
initPropertyValue(props, "runtime.log", logsFile.resolve("velocity.log").toFile().getAbsolutePath());
velocityEngine = AccessController.doPrivileged((PrivilegedAction) () -> {
final VelocityEngine engine = new VelocityEngine(props);
engine.init();
return engine;
});
}
private boolean exists(final Path path) {
return AccessController.doPrivileged((PrivilegedAction) () -> {
return Files.exists(path);
});
}
private File findWorkDir(final Settings settings) {
final List lookupPathList = new ArrayList<>();
List pathList = Environment.PATH_DATA_SETTING.get(settings);
if (pathList.isEmpty()) {
pathList = Arrays.asList(new File(Environment.PATH_HOME_SETTING.get(settings), "data").getAbsolutePath());
}
for (final String path : pathList) {
final File vmCacheDir = Paths.get(path, "vm_cache").toFile();
if (vmCacheDir.isDirectory()) {
return vmCacheDir;
} else if (vmCacheDir.exists()) {
continue;
} else if (vmCacheDir.mkdirs()) {
return vmCacheDir;
}
lookupPathList.add(vmCacheDir.getAbsolutePath());
}
throw new ElasticsearchException(
"Could not create a working directory: " + String.join(", ", lookupPathList.toArray(new String[lookupPathList.size()])));
}
private boolean initPropertyValue(final Properties props, final String key, final String value) {
if (!props.containsKey(key)) {
props.put(key, value);
return true;
}
return false;
}
@Override
public String getType() {
return NAME;
}
@Override
public String getExtension() {
return "vm";
}
@Override
public boolean isInlineScriptEnabled() {
return true;
}
@Override
public Object compile(final String templateName, final String templateSource, final Map params) {
final VelocityScriptTemplate scriptTemplate = new VelocityScriptTemplate(velocityEngine, workDir, templateSource);
final File templateFile = scriptTemplate.getTemplateFile();
if (templateFile != null) {
if (logger.isDebugEnabled()) {
logger.debug("Adding {}", templateFile.getAbsolutePath());
}
templateFileQueue.add(templateFile);
}
return scriptTemplate;
}
@Override
public ExecutableScript executable(final CompiledScript compiledScript, final Map vars) {
final Map scriptVars;
if (!contextPropMap.isEmpty()) {
scriptVars = new HashMap<>(contextPropMap);
} else {
scriptVars = new HashMap<>();
}
scriptVars.putAll(vars);
scriptVars.put("threadContext", threadContext);
return new VelocityExecutableScript((VelocityScriptTemplate) compiledScript.compiled(), scriptVars, logger);
}
@Override
public SearchScript search(final CompiledScript compiledScript, final SearchLookup lookup, final Map vars) {
throw new UnsupportedOperationException();
}
@Override
public void close() {
for (final File templateFile : templateFileQueue) {
if (logger.isDebugEnabled()) {
logger.debug("Deleting {}", templateFile.getAbsolutePath());
}
if (!templateFile.delete()) {
logger.warn("Failed to delete {}.", templateFile.getAbsolutePath());
}
}
}
public static class VelocityScriptTemplate {
private Template template;
private File templateFile;
private String script;
private final VelocityEngine velocityEngine;
public VelocityScriptTemplate(final VelocityEngine velocityEngine, final File workDir, final String script) {
this.velocityEngine = velocityEngine;
if (script.startsWith("##cache")) {
String encoding = (String) velocityEngine.getProperty("input.encoding");
if (encoding == null) {
encoding = "UTF-8";
}
if (!workDir.exists() && !workDir.mkdirs()) {
throw new VelocityException("Could not create a working directory: " + workDir.getAbsolutePath());
}
templateFile = null;
BufferedWriter bw = null;
try {
templateFile = File.createTempFile("templ", ".vm", workDir);
bw = new BufferedWriter(new OutputStreamWriter(new FileOutputStream(templateFile), encoding));
bw.write(script);
bw.flush();
} catch (final IOException e) {
throw new VelocityException("Failed to create a template file.", e);
} finally {
if (bw != null) {
try {
bw.close();
} catch (final IOException e) {
// ignore
}
}
}
template = AccessController
.doPrivileged((PrivilegedAction) () -> {
return velocityEngine
.getTemplate(templateFile.getName());
});
} else {
this.script = script;
}
}
public File getTemplateFile() {
return templateFile;
}
public void merge(final Context context, final Writer writer) {
if (script != null) {
final String logTag = Integer.toString(script.hashCode());
AccessController
.doPrivileged((PrivilegedAction) () -> {
return velocityEngine.evaluate(context, writer,
logTag, script);
});
} else {
AccessController.doPrivileged((PrivilegedAction) () -> {
template.merge(context, writer);
return null;
});
}
}
}
private static class VelocityExecutableScript implements ExecutableScript {
/** Compiled template object. */
private final VelocityScriptTemplate context;
/** Parameters to fill above object with. */
private Map vars;
private final Logger logger;
/**
* @param template the compiled template object
* @param vars the parameters to fill above object with
* @param logger
**/
public VelocityExecutableScript(final VelocityScriptTemplate context, final Map vars, final Logger logger) {
this.context = context;
this.logger = logger;
if (vars == null) {
this.vars = Collections.emptyMap();
} else {
this.vars = vars;
}
}
@Override
public void setNextVar(final String name, final Object value) {
vars.put(name, value);
}
@Override
public Object run() {
final BytesStreamOutput result = new BytesStreamOutput();
final UTF8StreamWriter writer = utf8StreamWriter().setOutput(result);
try {
context.merge(new VelocityContext(vars), writer);
writer.flush();
} catch (final Exception e) {
throw new ElasticsearchException("Could not execute query template: ", e);
} finally {
try {
writer.close();
} catch (final IOException e) {
logger.error("Could not execute query template (failed to close writer): ", e);
}
}
final BytesReference bytes = result.bytes();
if (logger.isDebugEnabled()) {
logger.debug("output: {}", bytes.utf8ToString());
}
return bytes;
}
@Override
public Object unwrap(final Object value) {
return value;
}
}
public void setThreadContext(final ThreadContext threadContext) {
this.threadContext = threadContext;
}
}