add support for new json object list output format (#58)

fulups · web-flow · commit e82b101a66c0 · 2020-09-26T14:51:27.000+09:00
diff --git a/README.md b/README.md
@@ -4,7 +4,7 @@ Elasticsearch Data Format Plugin
 ## Overview
 
 Elasticsearch Data Format Plugin provides a feature to allow you to download a response of a search result as several formats other than JSON.
-The supported formats are CSV, Excel and JSON(Bulk).
+The supported formats are CSV, Excel, JSON(Bulk) and JSON(Object List).
 
 ## Version
 
@@ -67,3 +67,11 @@ If not, it's as scan query(all data are stored.).
 | bulk.index        | string  | Index name in Bulk file |
 | bulk.type         | string  | Type name in Bulk file |
 
+### JSON (Object List format)
+
+    $ curl -o /tmp/data.json -XGET "localhost:9200/{index}/{type}/_data?format=jsonlist&source=..."
+
+| Request Parameter |  Type  | Description                                                  |
+| :---------------- | :----: | :----------------------------------------------------------- |
+| source            | string | [Query DSL](http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/query-dsl.html) |
+
diff --git a/src/main/java/org/codelibs/elasticsearch/df/content/ContentType.java b/src/main/java/org/codelibs/elasticsearch/df/content/ContentType.java
@@ -2,6 +2,7 @@
 
 import org.codelibs.elasticsearch.df.content.csv.CsvContent;
 import org.codelibs.elasticsearch.df.content.json.JsonContent;
+import org.codelibs.elasticsearch.df.content.json.JsonListContent;
 import org.codelibs.elasticsearch.df.content.xls.XlsContent;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.rest.RestRequest;
@@ -95,6 +96,27 @@ public String fileName(final RestRequest request) {
             }
             return index + ".xlsx";
         }
+    },
+    JSONLIST(50) {
+        @Override
+        public String contentType() {
+            return "application/json";
+        }
+
+        @Override
+        public DataContent dataContent(final Client client,
+                final RestRequest request) {
+            return new JsonListContent(client, request, this);
+        }
+
+        @Override
+        public String fileName(final RestRequest request) {
+            final String index = request.param("index");
+            if (index == null) {
+                return "_all.json";
+            }
+            return index + ".json";
+        }
     };
 
     private int index;
diff --git a/src/main/java/org/codelibs/elasticsearch/df/content/json/JsonListContent.java b/src/main/java/org/codelibs/elasticsearch/df/content/json/JsonListContent.java
@@ -0,0 +1,134 @@
+package org.codelibs.elasticsearch.df.content.json;
+
+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 org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.codelibs.elasticsearch.df.content.ContentType;
+import org.codelibs.elasticsearch.df.content.DataContent;
+import org.codelibs.elasticsearch.df.util.RequestUtil;
+import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.client.Client;
+import org.elasticsearch.common.xcontent.XContentHelper;
+import org.elasticsearch.common.xcontent.XContentType;
+import org.elasticsearch.rest.RestChannel;
+import org.elasticsearch.rest.RestRequest;
+import org.elasticsearch.search.SearchHit;
+import org.elasticsearch.search.SearchHits;
+
+public class JsonListContent extends DataContent {
+    private static final Logger logger = LogManager.getLogger(JsonListContent.class);
+
+    public JsonListContent(final Client client, final RestRequest request, final ContentType contentType) {
+        super(client, request, contentType);
+    }
+
+    @Override
+    public void write(final File outputFile, final SearchResponse response, final RestChannel channel,
+            final ActionListener<Void> listener) {
+        try {
+            final OnLoadListener onLoadListener = new OnLoadListener(
+                    outputFile, listener);
+            onLoadListener.onResponse(response);
+        } catch (final Exception e) {
+            listener.onFailure(new ElasticsearchException("Failed to write data.",
+                    e));
+        }
+    }
+
+    protected class OnLoadListener implements ActionListener<SearchResponse> {
+        protected ActionListener<Void> listener;
+
+        protected Writer writer;
+
+        protected File outputFile;
+
+        private long currentCount = 0;
+
+        private boolean firstLine = true;
+
+        protected OnLoadListener(final File outputFile, final ActionListener<Void> listener) {
+            this.outputFile = outputFile;
+            this.listener = listener;
+            try {
+                writer = new BufferedWriter(new OutputStreamWriter(
+                        new FileOutputStream(outputFile), "UTF-8"));
+            } catch (final Exception e) {
+                throw new ElasticsearchException("Could not open "
+                        + outputFile.getAbsolutePath(), e);
+            }
+            try {
+                writer.append('[');
+            }catch (final Exception e) {
+                onFailure(e);
+            }
+        }
+
+        @Override
+        public void onResponse(final SearchResponse response) {
+            final String scrollId = response.getScrollId();
+            final SearchHits hits = response.getHits();
+            final int size = hits.getHits().length;
+            currentCount += size;
+            if (logger.isDebugEnabled()) {
+                logger.debug("scrollId: {}, totalHits: {}, hits: {}, current: {}",
+                        scrollId, hits.getTotalHits(), size, currentCount);
+            }
+            try {
+                for (final SearchHit hit : hits) {
+                    final String source = XContentHelper.convertToJson(
+                            hit.getSourceRef(), true, false, XContentType.JSON);
+                    if (!firstLine){
+                        writer.append(',');
+                    }else{
+                        firstLine = false;
+                    }
+                    writer.append('\n').append(source);
+                }
+
+                if (size == 0 || scrollId == null) {
+                    // end
+                    writer.append('\n').append(']');
+                    writer.flush();
+                    close();
+                    listener.onResponse(null);
+                } else {
+                    client.prepareSearchScroll(scrollId)
+                            .setScroll(RequestUtil.getScroll(request))
+                            .execute(this);
+                }
+            } catch (final Exception e) {
+                onFailure(e);
+            }
+        }
+
+        @Override
+        public void onFailure(final Exception e) {
+            try {
+                close();
+            } catch (final Exception e1) {
+                // ignore
+            }
+            listener.onFailure(new ElasticsearchException("Failed to write data.",
+                    e));
+        }
+
+        private void close() {
+            if (writer != null) {
+                try {
+                    writer.close();
+                } catch (final IOException e) {
+                    throw new ElasticsearchException("Could not close "
+                            + outputFile.getAbsolutePath(), e);
+                }
+            }
+        }
+    }
+}
diff --git a/src/main/java/org/codelibs/elasticsearch/df/rest/RestDataAction.java b/src/main/java/org/codelibs/elasticsearch/df/rest/RestDataAction.java
@@ -129,6 +129,9 @@ private ContentType getContentType(final RestRequest request) {
                 || "application/json".equals(contentType)
                 || "json".equalsIgnoreCase(contentType)) {
             return ContentType.JSON;
+        } else if ("application/list+json".equals(contentType)
+                || "jsonlist".equals(contentType)) {
+            return ContentType.JSONLIST;
         }
 
         return null;
diff --git a/src/test/java/org/codelibs/elasticsearch/df/DataFormatPluginTest.java b/src/test/java/org/codelibs/elasticsearch/df/DataFormatPluginTest.java
@@ -52,18 +52,21 @@ public class DataFormatPluginTest {
     private static final File csvTempFile;
     private static final File xlsTempFile;
     private static final File jsonTempFile;
+    private static final File jsonListTempFile;
     private static final String path;
 
     private final Map<String, String> paramsCsv = new HashMap<>();
     private final Map<String, String> paramsXls = new HashMap<>();
     private final Map<String, String> paramsJson = new HashMap<>();
+    private final Map<String, String> paramsJsonList = new HashMap<>();
 
     static {
         docNumber = 20;
 
         csvTempFile = createTempFile("csvtest", ".csv");
         xlsTempFile = createTempFile("xlstest", ".xls");
         jsonTempFile = createTempFile("jsontest", ".json");
+        jsonListTempFile = createTempFile("jsonlisttest", ".json");
         path = "/dataset0/_data";
     }
 
@@ -106,13 +109,15 @@ public void prepareParams() {
         paramsCsv.put("format", "csv");
         paramsXls.put("format", "xls");
         paramsJson.put("format", "json");
+        paramsJsonList.put("format", "jsonlist");
     }
 
     @After
     public void clearParams() {
         paramsCsv.clear();
         paramsXls.clear();
         paramsJson.clear();
+        paramsJsonList.clear();
     }
 
     @Test
@@ -421,6 +426,106 @@ public void dumpJsonInFile() throws IOException {
         }
     }
 
+    @Test
+    public void dumpJsonList() throws IOException {
+
+        // Download All as JSON
+        try (CurlResponse curlResponse = EcrCurl.get(node, "/dataset0/_data")
+                .header("Content-Type", "application/json")
+                .param("format", "jsonlist").execute()) {
+            final String content = curlResponse.getContentAsString();
+            final String[] lines = content.split("\n");
+            assertEquals(docNumber + 2, lines.length);
+            assertTrue(lines[0].equals("["));
+            assertTrue(lines[1].startsWith("{" + "\"aaa\":\"test"));
+            assertTrue(lines[docNumber + 1].equals("]"));
+        }
+
+        final String query = "{\"query\":{\"bool\":{\"must\":[{\"range\":{\"bbb\":{\"from\":\"1\",\"to\":\"10\"}}}],\"must_not\":[],\"should\":[]}},\"sort\":[\"bbb\"]}";
+
+        // Download 10 docs as JSON with Query
+        try (CurlResponse curlResponse = EcrCurl.get(node, "/dataset0/_data")
+                .header("Content-Type", "application/json")
+                .param("format", "jsonlist")
+                .param("search_type", "query_then_fetch").body(query)
+                .execute()) {
+            final String content = curlResponse.getContentAsString();
+            final String[] lines = content.split("\n");
+            assertEquals(10 + 2, lines.length);
+            assertTrue(lines[0].startsWith("["));
+            assertTrue(lines[1].startsWith("{" + "\"aaa\":\"test"));
+        }
+
+        // Download 10 docs as JSON
+        try (CurlResponse curlResponse = EcrCurl.get(node, "/dataset0/_data")
+                .header("Content-Type", "application/json").param("q", "*:*")
+                .param("format", "jsonlist").param("from", "5").execute()) {
+            final String content = curlResponse.getContentAsString();
+            final String[] lines = content.split("\n");
+            assertEquals(15 + 2, lines.length);
+        }
+
+        // Download all the docs from the 5th as JSON
+        try (CurlResponse curlResponse = EcrCurl.get(node, "/dataset0/_data")
+                .header("Content-Type", "application/json").param("q", "*:*")
+                .param("format", "jsonlist").param("from", "5")
+                .param("size", String.valueOf(docNumber)).execute()) {
+            final String content = curlResponse.getContentAsString();
+            final String[] lines = content.split("\n");
+            assertEquals((docNumber - 5) + 2, lines.length);
+        }
+
+        final String queryWithFrom = "{\"query\":{\"match_all\":{}},\"from\":5,\"size\":" + String.valueOf(docNumber) + ",\"sort\":[\"bbb\"]}";
+
+        // Download All as JSON with Query and from
+        try (CurlResponse curlResponse = EcrCurl.get(node, "/dataset0/_data")
+                .header("Content-Type", "application/json")
+                .param("format", "jsonlist").body(queryWithFrom).execute()) {
+            final String content = curlResponse.getContentAsString();
+            final String[] lines = content.split("\n");
+            assertEquals((docNumber - 5) + 2, lines.length);
+        }
+
+        // Download All as JSON with Query and from
+        try (CurlResponse curlResponse = EcrCurl.get(node, "/dataset0/_data")
+                .header("Content-Type", "application/json")
+                .param("format", "jsonlist").param("source", queryWithFrom)
+                .param("source_content_type", "application/json")
+                .execute()) {
+            final String content = curlResponse.getContentAsString();
+            final String[] lines = content.split("\n");
+            assertEquals((docNumber - 5) + 2, lines.length);
+        }
+
+        // Download All as JSON with search_type
+        try (CurlResponse curlResponse = EcrCurl.get(node, "/dataset0/_data")
+                .header("Content-Type", "application/json")
+                .param("search_type", "query_then_fetch")
+                .param("format", "jsonlist").execute()) {
+            final String content = curlResponse.getContentAsString();
+            final String[] lines = content.split("\n");
+            assertEquals(docNumber + 2, lines.length);
+            assertTrue(lines[0].equals("["));
+            assertTrue(lines[1].startsWith("{" + "\"aaa\":\"test"));
+            assertTrue(lines[docNumber + 1].equals("]"));            
+        }
+    }
+
+    @Test
+    public void dumpJsonListInFile() throws IOException {
+        paramsJsonList.put("file", jsonListTempFile.getAbsolutePath());
+
+        try (CurlResponse curlResponse = createRequest(node, path, paramsJsonList).execute()) {
+            assertAcknowledged(curlResponse, jsonListTempFile);
+            final List<String> lines = Files.readAllLines(jsonListTempFile.toPath(), Charsets.UTF_8);
+            assertEquals(docNumber + 2, lines.size());
+            assertTrue(lines.get(0).equals("["));
+            assertTrue(lines.get(1).startsWith("{" + "\"aaa\":\"test"));
+            assertTrue(lines.get(docNumber).startsWith("{" + "\"aaa\":\"test"));
+            assertTrue(lines.get(docNumber + 1).equals("]"));   
+        }
+    }
+
     @Test
     public void dumpSizeLimit() throws IOException {
 

Original file line number	Diff line number	Diff line change
`@@ -129,6 +129,9 @@ private ContentType getContentType(final RestRequest request) {`
`129`	`129`	`\|\| "application/json".equals(contentType)`
`130`	`130`	`\|\| "json".equalsIgnoreCase(contentType)) {`
`131`	`131`	`return ContentType.JSON;`
	`132`	`+ } else if ("application/list+json".equals(contentType)`
	`133`	`+ \|\| "jsonlist".equals(contentType)) {`
	`134`	`+ return ContentType.JSONLIST;`
`132`	`135`	`}`
`133`	`136`
`134`	`137`	`return null;`