Created FormatInputStreat to unify reading from file or stdin

3 files changed, 220 insertions, 0 deletions

diff --git a/vespa-http-client/pom.xml b/vespa-http-client/pom.xml
index ad352cfe4cd..9e2325e8016 100644
--- a/vespa-http-client/pom.xml
+++ b/vespa-http-client/pom.xml
@@ -74,6 +74,11 @@
       <version>${jackson2.version}</version>
     </dependency>
     <dependency>
+      <groupId>com.fasterxml.jackson.dataformat</groupId>
+      <artifactId>jackson-dataformat-xml</artifactId>
+      <version>${jackson2.version}</version>
+    </dependency>
+    <dependency>
       <groupId>io.airlift</groupId>
       <artifactId>airline</artifactId>
       <version>0.6</version>
diff --git a/vespa-http-client/src/main/java/com/yahoo/vespa/http/client/runner/FormatInputStream.java b/vespa-http-client/src/main/java/com/yahoo/vespa/http/client/runner/FormatInputStream.java
new file mode 100644
index 00000000000..b0c3ba3c1ac
--- /dev/null
+++ b/vespa-http-client/src/main/java/com/yahoo/vespa/http/client/runner/FormatInputStream.java
@@ -0,0 +1,102 @@
+package com.yahoo.vespa.http.client.runner;
+
+import com.fasterxml.jackson.core.JsonFactory;
+import com.fasterxml.jackson.core.format.DataFormatDetector;
+import com.fasterxml.jackson.core.format.DataFormatMatcher;
+import com.fasterxml.jackson.core.format.MatchStrength;
+import com.fasterxml.jackson.dataformat.xml.XmlFactory;
+
+import java.io.BufferedInputStream;
+import java.io.ByteArrayInputStream;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.SequenceInputStream;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Optional;
+
+/**
+ * @author valerijf
+ */
+public class FormatInputStream {
+    private InputStream inputStream;
+    private Format format;
+
+    /**
+     * Creates a single data input stream from either file or InputStream depending on which one is present. Preference
+     * for file if both present. Additionally also detects input data format of the result stream, throws
+     * IllegalArgumentException if unable to determine data format.
+     *
+     * @param stream              InputStream of the data if present
+     * @param inputFile           Path to file to use as input
+     * @param addRootElementToXml To add vespafeed root element around the input data stream
+     * @throws IOException
+     */
+    public FormatInputStream(InputStream stream, Optional<String> inputFile, boolean addRootElementToXml)
+            throws IOException {
+        final DataFormatDetector dataFormatDetector = new DataFormatDetector(new JsonFactory(), new XmlFactory());
+        final DataFormatMatcher formatMatcher;
+
+        if (inputFile.isPresent()) {
+            try (FileInputStream fileInputStream = new FileInputStream(inputFile.get())) {
+                formatMatcher = dataFormatDetector.findFormat(fileInputStream);
+            }
+            inputStream = new FileInputStream(inputFile.get());
+
+        } else {
+            if (stream.available() == 0) {
+                System.out.println("No data in stream yet and no file specified, waiting for data.");
+            }
+
+            inputStream = stream.markSupported() ? stream : new BufferedInputStream(stream);
+            inputStream.mark(DataFormatDetector.DEFAULT_MAX_INPUT_LOOKAHEAD);
+            formatMatcher = dataFormatDetector.findFormat(inputStream);
+            inputStream.reset();
+        }
+
+        if (addRootElementToXml) {
+            inputStream = addVespafeedTag(inputStream);
+            format = Format.XML;
+            return;
+        }
+
+        if (formatMatcher.getMatchStrength() == MatchStrength.INCONCLUSIVE ||
+                formatMatcher.getMatchStrength() == MatchStrength.NO_MATCH) {
+            throw new IllegalArgumentException("Could not detect input format");
+        }
+
+        switch (formatMatcher.getMatchedFormatName().toLowerCase()) {
+            case "json":
+                format = Format.JSON;
+                break;
+
+            case "xml":
+                format = Format.XML;
+                break;
+
+            default:
+                throw new IllegalArgumentException("Unknown data format");
+        }
+    }
+
+    private static InputStream addVespafeedTag(InputStream inputStream) {
+        return new SequenceInputStream(Collections.enumeration(Arrays.asList(
+                new ByteArrayInputStream("<vespafeed>".getBytes()),
+                inputStream,
+                new ByteArrayInputStream("</vespafeed>".getBytes())))
+        );
+    }
+
+    public InputStream getInputStream() {
+        return inputStream;
+    }
+
+    public Format getFormat() {
+        return format;
+    }
+
+    public enum Format {
+        JSON, XML
+    }
+}
diff --git a/vespa-http-client/src/test/java/com/yahoo/vespa/http/client/runner/FormatInputStreamTest.java b/vespa-http-client/src/test/java/com/yahoo/vespa/http/client/runner/FormatInputStreamTest.java
new file mode 100644
index 00000000000..f16ac239b0d
--- /dev/null
+++ b/vespa-http-client/src/test/java/com/yahoo/vespa/http/client/runner/FormatInputStreamTest.java
@@ -0,0 +1,113 @@
+package com.yahoo.vespa.http.client.runner;
+
+import org.junit.Test;
+
+import java.io.ByteArrayInputStream;
+import java.io.File;
+import java.io.FileWriter;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Optional;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+/**
+ * @author valerijf
+ */
+public class FormatInputStreamTest {
+    @Test(expected=IllegalArgumentException.class)
+    public void testWithGarbageText() throws IOException {
+        String streamString = "This is neither XML nor JSON!";
+        InputStream jsonStream = getInputStreamOf(streamString);
+        FormatInputStream formatInputStream = new FormatInputStream(jsonStream, Optional.empty(), false);
+    }
+
+    @Test
+    public void testWithFileInput() throws IOException {
+        String fileString = "{\"format\": \"json\"}";
+        File file = File.createTempFile("feeddata", "json");
+        file.deleteOnExit();
+        try (FileWriter writer = new FileWriter(file)) {
+            writer.write(fileString);
+        }
+
+        FormatInputStream formatInputStream = new FormatInputStream(null, Optional.of(file.getAbsolutePath()), false);
+        assertThat(fileString, is(convertStreamToString(formatInputStream.getInputStream())));
+        assertThat(formatInputStream.getFormat(), is(FormatInputStream.Format.JSON));
+    }
+
+    @Test
+    public void testPreferenceFileOverStream() throws IOException {
+        String streamString = "something entirely different";
+        String fileString = "{\"format\": \"json\"}";
+
+        InputStream jsonStream = getInputStreamOf(streamString);
+        File file = File.createTempFile("feeddata", "json");
+        file.deleteOnExit();
+        try (FileWriter writer = new FileWriter(file)) {
+            writer.write(fileString);
+        }
+
+        FormatInputStream formatInputStream = new FormatInputStream(jsonStream, Optional.of(file.getAbsolutePath()), false);
+        assertThat(fileString, is(convertStreamToString(formatInputStream.getInputStream())));
+        assertThat(formatInputStream.getFormat(), is(FormatInputStream.Format.JSON));
+    }
+
+    @Test
+    public void testSimpleJsonInputStream() throws IOException {
+        String streamString = "{\"format\": \"json\"}";
+        InputStream jsonStream = getInputStreamOf(streamString);
+        FormatInputStream formatInputStream = new FormatInputStream(jsonStream, Optional.empty(), false);
+
+        assertThat(streamString, is(convertStreamToString(formatInputStream.getInputStream())));
+        assertThat(formatInputStream.getFormat(), is(FormatInputStream.Format.JSON));
+    }
+
+    @Test
+    public void testSimpleXmlInputStream() throws IOException {
+        String streamString = "<scope><tag>format</tag><value>xml</value></scope>";
+        InputStream jsonStream = getInputStreamOf(streamString);
+        FormatInputStream formatInputStream = new FormatInputStream(jsonStream, Optional.empty(), false);
+
+        assertThat(streamString, is(convertStreamToString(formatInputStream.getInputStream())));
+        assertThat(formatInputStream.getFormat(), is(FormatInputStream.Format.XML));
+    }
+
+    @Test
+    public void testSparselyFormattedXml() throws IOException {
+        String streamString = "               \t\t\n<scope>\n\n\n<tag>format</tag><value>xml</value></scope>";
+        InputStream jsonStream = getInputStreamOf(streamString);
+        FormatInputStream formatInputStream = new FormatInputStream(jsonStream, Optional.empty(), false);
+
+        assertThat(streamString, is(convertStreamToString(formatInputStream.getInputStream())));
+        assertThat(formatInputStream.getFormat(), is(FormatInputStream.Format.XML));
+    }
+
+    @Test
+    public void testAddRootToXml() throws IOException {
+        String streamString = "some random text";
+        InputStream textStream = getInputStreamOf(streamString);
+        FormatInputStream formatInputStream = new FormatInputStream(textStream, Optional.empty(), true);
+
+        assertThat("<vespafeed>" + streamString + "</vespafeed>",
+                is(convertStreamToString(formatInputStream.getInputStream())));
+        assertThat(formatInputStream.getFormat(), is(FormatInputStream.Format.XML));
+    }
+
+    private static String convertStreamToString(InputStream inputStream) throws IOException {
+        StringBuilder builder = new StringBuilder();
+        while (true) {
+            int character = inputStream.read();
+            if (character == -1) {
+                inputStream.close();
+                return builder.toString();
+            }
+            builder.append((char)character);
+        }
+    }
+
+    private static InputStream getInputStreamOf(String text) {
+        return new ByteArrayInputStream(text.getBytes());
+    }
+}