All Downloads are FREE. Search and download functionalities are using the official Maven repository.

org.elasticsearch.script.mustache.CustomMustacheFactory Maven / Gradle / Ivy

There is a newer version: 8.14.1
Show newest version
/*
 * 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.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.IterableCode;
import com.github.mustachejava.codes.WriteCode;
import com.google.common.base.Function;
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.util.Collections;
import java.util.Iterator;
import java.util.Map;
import java.util.Objects;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class CustomMustacheFactory extends DefaultMustacheFactory {

    private final Encoder encoder;

    public CustomMustacheFactory(boolean escaping) {
        super();
        setObjectHandler(new CustomReflectionObjectHandler());
        if (escaping) {
            this.encoder = new JsonEscapeEncoder();
        } else {
            this.encoder = new NoEscapeEncoder();
        }
    }

    @Override
    public void encode(String value, Writer writer) {
        encoder.accept(value, writer);
    }

    @Override
    public MustacheVisitor createMustacheVisitor() {
        return new CustomMustacheVisitor(this);
    }

    class CustomMustacheVisitor extends DefaultMustacheVisitor {

        public 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 {
                list.add(new IterableCode(templateContext, df, mustache, variable));
            }
        }
    }

    /**
     * Base class for custom Mustache functions
     */
    static abstract class CustomCode extends IterableCode {

        private final String code;

        public 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, Object[] scopes) {
            Object resolved = get(scopes);
            writer = handle(writer, createFunction(resolved), scopes);
            appendText(writer);
            return writer;
        }

        @Override
        protected void tag(Writer writer, String tag) throws IOException {
            writer.write(tc.startChars());
            writer.write(tag);
            writer.write(code);
            writer.write(tc.endChars());
        }

        protected abstract Function createFunction(Object resolved);

        /**
         * At compile time, this function extracts the name of the variable:
         * {{#toJson}}variable_name{{/toJson}}
         */
        protected static String extractVariableName(String fn, Mustache mustache, TemplateContext tc) {
            Code[] codes = mustache.getCodes();
            if (codes == null || codes.length != 1) {
                throw new MustacheException("Mustache function [" + fn + "] must contain one and only one identifier");
            }

            try (StringWriter capture = new StringWriter()) {
                // Variable name is in plain text and has type WriteCode
                if (codes[0] instanceof WriteCode) {
                    codes[0].execute(capture, Collections.emptyList());
                    return capture.toString();
                } else {
                    codes[0].identity(capture);
                    return capture.toString();
                }
            } catch (IOException e) {
                throw new MustacheException("Exception while parsing mustache function [" + fn + "] at line " + tc.line(), e);
            }
        }
    }

    /**
     * This function renders {@link Iterable} and {@link Map} as their JSON representation
     */
    static class ToJsonCode extends CustomCode {

        private static final String CODE = "toJson";

        public ToJsonCode(TemplateContext tc, DefaultMustacheFactory df, Mustache mustache, String variable) {
            super(tc, df, mustache, CODE);
            if (CODE.equalsIgnoreCase(variable) == false) {
                throw new MustacheException("Mismatch function code [" + CODE + "] cannot be applied to [" + variable + "]");
            }
        }

        @Override
        @SuppressWarnings("unchecked")
        protected Function  createFunction(final Object resolved) {
            return new Function() {
                @Override
                public String apply(Object s) {
                    if (resolved == null) {
                        return null;
                    }
                    try (XContentBuilder builder = XContentBuilder.builder(XContentType.JSON.xContent())) {
                        if (resolved instanceof Iterable) {
                            builder.startArray();
                            for (Object o : (Iterable) resolved) {
                                builder.value(o);
                            }
                            builder.endArray();
                        } else if (resolved instanceof Map) {
                            builder.map((Map) resolved);
                        } else {
                            // Do not handle as JSON
                            return oh.stringify(resolved);
                        }
                        return builder.string();
                    } catch (IOException e) {
                        throw new MustacheException("Failed to convert object to JSON", e);
                    }
                }
            };
        }
        static boolean match(String variable) {
            return CODE.equalsIgnoreCase(variable);
        }
    }

    /**
     * This function concatenates the values of an {@link Iterable} using a given delimiter
     */
    static class JoinerCode extends CustomCode {

        protected static final String CODE = "join";
        private static final String DEFAULT_DELIMITER = ",";

        private final String delimiter;

        public JoinerCode(TemplateContext tc, DefaultMustacheFactory df, Mustache mustache, String delimiter) {
            super(tc, df, mustache, CODE);
            this.delimiter = delimiter;
        }

        public JoinerCode(TemplateContext tc, DefaultMustacheFactory df, Mustache mustache) {
            this(tc, df, mustache, DEFAULT_DELIMITER);
        }

        @Override
        @SuppressWarnings("unchecked")
        protected Function createFunction(final Object resolved) {
            return new Function() {
                @Override
                public Object apply(Object s) {
                    if (s == null) {
                        return null;
                    } else if (resolved instanceof Iterable) {
                        StringBuilder joiner = new StringBuilder();
                        Iterator it = ((Iterable) resolved).iterator();
                        while (it.hasNext()) {
                            joiner.append(oh.stringify(it.next()));
                            if (it.hasNext()) {
                                joiner.append(delimiter);
                            }
                        }
                        return joiner.toString();
                    }
                    return s;
                }
            };
        }

        static boolean match(String variable) {
            return CODE.equalsIgnoreCase(variable);
        }
    }

    static class CustomJoinerCode extends JoinerCode {

        private static final Pattern PATTERN = Pattern.compile("^(?:" + CODE + " delimiter='(.*)')$");

        public CustomJoinerCode(TemplateContext tc, DefaultMustacheFactory df, Mustache mustache, String variable) {
            super(tc, df, mustache, extractDelimiter(variable));
        }

        private static String extractDelimiter(String variable) {
            Matcher matcher = PATTERN.matcher(variable);
            if (matcher.find()) {
                return matcher.group(1);
            }
            throw new MustacheException("Failed to extract delimiter for join function");
        }

        static boolean match(String variable) {
            return PATTERN.matcher(variable).matches();
        }
    }

    class NoEscapeEncoder implements Encoder {

        @Override
        public void accept(String s, Writer writer) {
            try {
                writer.write(s);
            } catch (IOException e) {
                throw new MustacheException("Failed to encode value: " + s);
            }
        }
    }

    class JsonEscapeEncoder implements Encoder {

        @Override
        public void accept(String s, Writer writer) {
            try {
                writer.write(JsonStringEncoder.getInstance().quoteAsString(s));
            } catch (IOException e) {
                throw new MustacheException("Failed to escape and encode value: " + s);
            }
        }
    }

    interface Encoder {
        void accept(String s, Writer writer);
    }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy