com.o19s.es.template.mustache.CustomMustacheFactory Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of elasticsearch-learning-to-rank Show documentation
Show all versions of elasticsearch-learning-to-rank Show documentation
Learing to Rank Query w/ RankLib Models
/*
* Copyright [2017] Wikimedia Foundation
*
* Licensed 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 com.o19s.es.template.mustache;
import com.fasterxml.jackson.core.io.JsonStringEncoder;
import com.github.mustachejava.Code;
import com.github.mustachejava.DefaultMustacheFactory;
import com.github.mustachejava.DefaultMustacheVisitor;
import com.github.mustachejava.Mustache;
import com.github.mustachejava.MustacheException;
import com.github.mustachejava.MustacheVisitor;
import com.github.mustachejava.TemplateContext;
import com.github.mustachejava.codes.DefaultMustache;
import com.github.mustachejava.codes.IterableCode;
import com.github.mustachejava.codes.WriteCode;
import org.elasticsearch.common.Strings;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentType;
import java.io.IOException;
import java.io.StringWriter;
import java.io.Writer;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.StringJoiner;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
/**
* XXX: shamelessly copied from the mustache module
*/
public class CustomMustacheFactory extends DefaultMustacheFactory {
static final String JSON_MIME_TYPE_WITH_CHARSET = "application/json; charset=UTF-8";
static final String JSON_MIME_TYPE = "application/json";
static final String PLAIN_TEXT_MIME_TYPE = "text/plain";
static final String X_WWW_FORM_URLENCODED_MIME_TYPE = "application/x-www-form-urlencoded";
private static final String DEFAULT_MIME_TYPE = JSON_MIME_TYPE;
private static final Map> ENCODERS;
static {
Map> encoders = new HashMap<>();
encoders.put(JSON_MIME_TYPE_WITH_CHARSET, JsonEscapeEncoder::new);
encoders.put(JSON_MIME_TYPE, JsonEscapeEncoder::new);
encoders.put(PLAIN_TEXT_MIME_TYPE, DefaultEncoder::new);
encoders.put(X_WWW_FORM_URLENCODED_MIME_TYPE, UrlEncoder::new);
ENCODERS = Collections.unmodifiableMap(encoders);
}
private final Encoder encoder;
public CustomMustacheFactory(String mimeType) {
super();
setObjectHandler(new CustomReflectionObjectHandler());
this.encoder = createEncoder(mimeType);
}
public CustomMustacheFactory() {
this(DEFAULT_MIME_TYPE);
}
@Override
public void encode(String value, Writer writer) {
try {
encoder.encode(value, writer);
} catch (IOException e) {
throw new MustacheException("Unable to encode value", e);
}
}
static Encoder createEncoder(String mimeType) {
Supplier supplier = ENCODERS.get(mimeType);
if (supplier == null) {
throw new IllegalArgumentException("No encoder found for MIME type [" + mimeType + "]");
}
return supplier.get();
}
@Override
public MustacheVisitor createMustacheVisitor() {
return new CustomMustacheVisitor(this);
}
class CustomMustacheVisitor extends DefaultMustacheVisitor {
CustomMustacheVisitor(DefaultMustacheFactory df) {
super(df);
}
@Override
public void iterable(TemplateContext templateContext, String variable, Mustache mustache) {
if (ToJsonCode.match(variable)) {
list.add(new ToJsonCode(templateContext, df, mustache, variable));
} else if (JoinerCode.match(variable)) {
list.add(new JoinerCode(templateContext, df, mustache));
} else if (CustomJoinerCode.match(variable)) {
list.add(new CustomJoinerCode(templateContext, df, mustache, variable));
} else if (UrlEncoderCode.match(variable)) {
list.add(new UrlEncoderCode(templateContext, df, mustache, variable));
} else {
list.add(new IterableCode(templateContext, df, mustache, variable));
}
}
}
/**
* Base class for custom Mustache functions
*/
abstract static class CustomCode extends IterableCode {
private final String code;
CustomCode(TemplateContext tc, DefaultMustacheFactory df, Mustache mustache, String code) {
super(tc, df, mustache, extractVariableName(code, mustache, tc));
this.code = Objects.requireNonNull(code);
}
@Override
public Writer execute(Writer writer, final List