ExportFields.java

package crate.elasticsearch.export;

import org.elasticsearch.common.bytes.BytesReference;
import org.elasticsearch.common.xcontent.*;
import org.elasticsearch.index.mapper.internal.TTLFieldMapper;
import org.elasticsearch.index.mapper.internal.TimestampFieldMapper;
import org.elasticsearch.search.SearchHitField;
import org.elasticsearch.search.internal.InternalSearchHit;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class ExportFields implements ToXContent {

    private final List<String> fields;
    private InternalSearchHit hit;
    private final List<FieldExtractor> fieldExtractors;

    static final class Fields {
        static final XContentBuilderString _SOURCE = new XContentBuilderString("_source");
        static final XContentBuilderString _TYPE = new XContentBuilderString("_type");
        static final XContentBuilderString _INDEX = new XContentBuilderString("_index");
        static final XContentBuilderString _ID = new XContentBuilderString("_id");
        static final XContentBuilderString _VERSION = new XContentBuilderString("_version");
        static final XContentBuilderString _TIMESTAMP = new XContentBuilderString(TimestampFieldMapper.NAME);
        static final XContentBuilderString _TTL = new XContentBuilderString(TTLFieldMapper.NAME);
    }

    abstract class FieldExtractor implements ToXContent {


    }


    class SourceFieldExtractor extends FieldExtractor {

        @Override
        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
            BytesReference source = hit.sourceRef();
            XContentType contentType = XContentFactory.xContentType(source);
            XContentParser parser = XContentFactory.xContent(contentType).createParser(source);
            try {
                parser.nextToken();
                builder.field(Fields._SOURCE);
                builder.copyCurrentStructure(parser);
            } finally {
                parser.close();
            }
            return builder;
        }
    }

    class HitFieldExtractor extends FieldExtractor {

        private final String fieldName;

        public HitFieldExtractor(String fieldName) {
            this.fieldName = fieldName;
        }

        @Override
        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
            SearchHitField field = hit.getFields().get(fieldName);
            if (field != null && !field.values().isEmpty()) {
                if (field.values().size() == 1) {
                    builder.field(field.name(), field.values().get(0));
                } else {
                    builder.field(field.name());
                    builder.startArray();
                    for (Object value : field.values()) {
                        builder.value(value);
                    }
                    builder.endArray();
                }
            }
            return builder;
        }
    }

    public void hit(InternalSearchHit hit) {
        this.hit = hit;
    }

    public ExportFields(List<String> fields) {
        this.fields = fields;
        this.fieldExtractors = getFieldExtractors();
    }

    private List<FieldExtractor> getFieldExtractors() {
        List<FieldExtractor> extractors = new ArrayList<FieldExtractor>(fields.size());
        for (String fn : fields) {
            FieldExtractor fc = null;
            if (fn.startsWith("_")) {
                if (fn.equals("_source")) {
                    fc = new SourceFieldExtractor();
                } else if (fn.equals("_id")) {
                    fc = new FieldExtractor() {
                        @Override
                        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
                            return builder.field(Fields._ID, hit.getId());
                        }
                    };
                } else if (fn.equals("_version")) {
                    fc = new FieldExtractor() {
                        @Override
                        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
                            return builder.field(Fields._VERSION, hit.getVersion());
                        }
                    };
                } else if (fn.equals("_index")) {
                    fc = new FieldExtractor() {
                        @Override
                        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
                            return builder.field(Fields._INDEX, hit.getIndex());
                        }
                    };
                } else if (fn.equals("_type")) {
                    fc = new FieldExtractor() {
                        @Override
                        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
                            return builder.field(Fields._TYPE, hit.getType());
                        }
                    };
                } else {
                    fc = new HitFieldExtractor(fn);
                }
            } else {
                fc = new HitFieldExtractor(fn);
            }
            extractors.add(fc);
        }
        return extractors;
    }

    @Override
    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
        builder.startObject();
        for (FieldExtractor fc : fieldExtractors) {
            fc.toXContent(builder, params);
        }
        builder.endObject();
        return builder;
    }
}