-
Notifications
You must be signed in to change notification settings - Fork 24.2k
/
MustacheScriptEngineService.java
205 lines (184 loc) · 7.01 KB
/
MustacheScriptEngineService.java
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
/*
* 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.mustache;
import com.github.mustachejava.Mustache;
import org.elasticsearch.common.Nullable;
import org.elasticsearch.common.component.AbstractComponent;
import org.elasticsearch.common.inject.Inject;
import org.elasticsearch.common.io.FastStringReader;
import org.elasticsearch.common.io.UTF8StreamWriter;
import org.elasticsearch.common.io.stream.BytesStreamOutput;
import org.elasticsearch.common.settings.Settings;
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;
import java.io.IOException;
import java.lang.ref.SoftReference;
import java.util.Collections;
import java.util.Map;
/**
* Main entry point handling template registration, compilation and
* execution.
*
* Template handling is based on Mustache. Template handling is a two step
* process: First compile the string representing the template, the resulting
* {@link Mustache} object can then be re-used for subsequent executions.
*/
public class MustacheScriptEngineService extends AbstractComponent implements ScriptEngineService {
/** Thread local UTF8StreamWriter to store template execution results in, thread local to save object creation.*/
private static ThreadLocal<SoftReference<UTF8StreamWriter>> utf8StreamWriter = new ThreadLocal<>();
/** If exists, reset and return, otherwise create, reset and return a writer.*/
private static UTF8StreamWriter utf8StreamWriter() {
SoftReference<UTF8StreamWriter> 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;
}
/**
* @param settings automatically wired by Guice.
* */
@Inject
public MustacheScriptEngineService(Settings settings) {
super(settings);
}
/**
* Compile a template string to (in this case) a Mustache object than can
* later be re-used for execution to fill in missing parameter values.
*
* @param template
* a string representing the template to compile.
* @return a compiled template object for later execution.
* */
public Object compile(String template) {
/** Factory to generate Mustache objects from. */
return (new JsonEscapingMustacheFactory()).compile(new FastStringReader(template), "query-template");
}
/**
* Execute a compiled template object (as retrieved from the compile method)
* and fill potential place holders with the variables given.
*
* @param template
* compiled template object.
* @param vars
* map of variables to use during substitution.
*
* @return the processed string with all given variables substitued.
* */
public Object execute(Object template, Map<String, Object> vars) {
BytesStreamOutput result = new BytesStreamOutput();
UTF8StreamWriter writer = utf8StreamWriter().setOutput(result);
((Mustache) template).execute(writer, vars);
try {
writer.flush();
} catch (IOException e) {
logger.error("Could not execute query template (failed to flush writer): ", e);
} finally {
try {
writer.close();
} catch (IOException e) {
logger.error("Could not execute query template (failed to close writer): ", e);
}
}
return result.bytes();
}
@Override
public String[] types() {
return new String[] {"mustache"};
}
@Override
public String[] extensions() {
return new String[] {"mustache"};
}
@Override
public boolean sandboxed() {
return true;
}
@Override
public ExecutableScript executable(Object mustache,
@Nullable Map<String, Object> vars) {
return new MustacheExecutableScript((Mustache) mustache, vars);
}
@Override
public SearchScript search(Object compiledScript, SearchLookup lookup,
@Nullable Map<String, Object> vars) {
throw new UnsupportedOperationException();
}
@Override
public Object unwrap(Object value) {
return value;
}
@Override
public void close() {
// Nothing to do here
}
@Override
public void scriptRemoved(CompiledScript script) {
// Nothing to do here
}
/**
* Used at query execution time by script service in order to execute a query template.
* */
private class MustacheExecutableScript implements ExecutableScript {
/** Compiled template object. */
private Mustache mustache;
/** Parameters to fill above object with. */
private Map<String, Object> vars;
/**
* @param mustache the compiled template object
* @param vars the parameters to fill above object with
**/
public MustacheExecutableScript(Mustache mustache,
Map<String, Object> vars) {
this.mustache = mustache;
this.vars = vars == null ? Collections.EMPTY_MAP : vars;
}
@Override
public void setNextVar(String name, Object value) {
this.vars.put(name, value);
}
@Override
public Object run() {
BytesStreamOutput result = new BytesStreamOutput();
UTF8StreamWriter writer = utf8StreamWriter().setOutput(result);
((Mustache) mustache).execute(writer, vars);
try {
writer.flush();
} catch (IOException e) {
logger.error("Could not execute query template (failed to flush writer): ", e);
} finally {
try {
writer.close();
} catch (IOException e) {
logger.error("Could not execute query template (failed to close writer): ", e);
}
}
return result.bytes();
}
@Override
public Object unwrap(Object value) {
return value;
}
}
}