embulk-parser-grok 0.1.6

data/src/main/java/org/embulk/parser/grok/GrokParserPlugin.java

@@ -0,0 +1,94 @@
+package org.embulk.parser.grok;
+
+
+import com.google.common.base.Optional;
+import org.embulk.config.*;
+import org.embulk.spi.*;
+import org.embulk.spi.time.TimestampParser;
+import org.embulk.spi.util.LineDecoder;
+import org.slf4j.Logger;
+
+import java.util.List;
+import java.util.Map;
+
+public class GrokParserPlugin implements ParserPlugin {
+
+    public final Logger logger = Exec.getLogger(GrokParserPlugin.class.getName());
+
+    public interface PluginTask
+            extends Task, LineDecoder.DecoderTask, TimestampParser.Task {
+
+        @Config("grok_pattern")
+        String getGrokPattern();
+
+        @Config("first_line_pattern")
+        @ConfigDefault("null")
+        Optional<String> getFirstLinePattern();
+
+        @Config("grok_pattern_files")
+        List<String> getGrokPatternFiles();
+
+        @Config("timestamp_parser")
+        @ConfigDefault("\"ruby\"")
+        String getTimestampParser();
+
+        @Config("columns")
+        SchemaConfig getColumns();
+
+        @Config("stop_on_invalid_record")
+        @ConfigDefault("false")
+        boolean getStopOnInvalidRecord();
+    }
+
+    @Override
+    public void transaction(ConfigSource config, ParserPlugin.Control control) {
+        PluginTask task = config.loadConfig(PluginTask.class);
+        Schema schema = task.getColumns().toSchema();
+
+        control.run(task.dump(), schema);
+    }
+
+    @Override
+    public void run(TaskSource taskSource, Schema schema, FileInput input, PageOutput output) {
+        GrokParserPlugin.PluginTask task = taskSource.loadTask(GrokParserPlugin.PluginTask.class);
+
+        LineDecoder decoder;
+        if (task.getFirstLinePattern().isPresent()) {
+            decoder = new MultipleLineDecoder(input, task);
+        } else {
+            decoder = new LineDecoder(input, task);
+        }
+
+        try (GrokRecordIterator iterator = new GrokRecordIterator(decoder, task)) {
+            final List<DateParser> timestampParsers = TimestampParserFactory.create(task);
+            PageBuilder pageBuilder = new PageBuilder(Exec.getBufferAllocator(), schema, output);
+            while (true) {
+                if (!iterator.nextFile()) {
+                    break;
+                }
+                while (true) {
+                    if (!iterator.nextLine()) {
+                        break;
+                    }
+
+                    try {
+                        Map<String, Object> record = iterator.getCurrentRecord();
+                        if (record.keySet().size() != 0) {
+                            schema.visitColumns(new GrokColumnVisitor(record, pageBuilder, timestampParsers));
+                            pageBuilder.addRecord();
+                        }
+                    } catch (GrokRecordValidateException e) {
+                        String skippedLine = iterator.getCurrentLine();
+                        long lineNumber = iterator.getCurrentLineNumber();
+                        if (task.getStopOnInvalidRecord()) {
+                            throw new DataException(String.format("Invalid record at line %d: %s", lineNumber, skippedLine), e);
+                        }
+                        logger.warn(String.format("Skipped line %d (%s): %s", lineNumber, e.getMessage(), skippedLine));
+                    }
+                }
+            }
+            pageBuilder.finish();
+        }
+    }
+
+}

data/src/main/java/org/embulk/parser/grok/GrokRecordIterator.java

@@ -0,0 +1,71 @@
+package org.embulk.parser.grok;
+
+import java.io.Closeable;
+import java.util.Map;
+
+import oi.thekraken.grok.api.Grok;
+import oi.thekraken.grok.api.Match;
+import oi.thekraken.grok.api.exception.GrokException;
+
+import org.embulk.spi.Exec;
+import org.embulk.spi.util.LineDecoder;
+import org.slf4j.Logger;
+
+public class GrokRecordIterator implements Closeable {
+
+    private LineDecoder decoder;
+    private Grok grok = null;
+
+    private String currentLine = null;
+    private long currentLineNumber = 0;
+
+    private final Logger logger = Exec.getLogger(GrokRecordIterator.class.getName());
+
+    public GrokRecordIterator(LineDecoder decoder, GrokParserPlugin.PluginTask task) {
+        this.decoder = decoder;
+        try {
+            this.grok = new Grok();
+            for (String file : task.getGrokPatternFiles()) {
+                this.grok.addPatternFromFile(file);
+            }
+            this.grok.compile(task.getGrokPattern());
+        } catch (GrokException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    public boolean nextFile() {
+        currentLine = null;
+        currentLineNumber = 0;
+        return decoder.nextFile();
+    }
+
+    public boolean nextLine() {
+        currentLine = decoder.poll();
+        currentLineNumber++;
+        return currentLine != null;
+    }
+
+    @Override
+    public void close() {
+        decoder.close();
+    }
+
+    public Map<String, Object> getCurrentRecord() {
+        Match gm = grok.match(currentLine);
+        gm.captures();
+        if (gm.isNull()) {
+            throw new GrokRecordValidateException("Couldn't parse line");
+        }
+        return gm.toMap();
+    }
+
+    public String getCurrentLine() {
+        return currentLine;
+    }
+
+    public long getCurrentLineNumber() {
+        return currentLineNumber;
+    }
+}
+

data/src/main/java/org/embulk/parser/grok/GrokRecordValidateException.java

@@ -0,0 +1,11 @@
+package org.embulk.parser.grok;
+
+public class GrokRecordValidateException extends RuntimeException {
+    GrokRecordValidateException(String message) {
+        super(message);
+    }
+
+    GrokRecordValidateException(Throwable cause) {
+        super(cause);
+    }
+}

data/src/main/java/org/embulk/parser/grok/MultipleLineDecoder.java

@@ -0,0 +1,66 @@
+package org.embulk.parser.grok;
+
+import com.google.common.base.Strings;
+import oi.thekraken.grok.api.Grok;
+import oi.thekraken.grok.api.Match;
+import oi.thekraken.grok.api.exception.GrokException;
+import org.embulk.spi.FileInput;
+import org.embulk.spi.util.LineDecoder;
+
+public class MultipleLineDecoder extends LineDecoder {
+
+    private StringBuilder buildLines;
+    private Grok grok = null;
+
+    public MultipleLineDecoder(FileInput in, GrokParserPlugin.PluginTask task) {
+        super(in, task);
+
+        try {
+            this.grok = new Grok();
+            for (String file : task.getGrokPatternFiles()) {
+                this.grok.addPatternFromFile(file);
+            }
+            this.grok.compile(task.getFirstLinePattern().get());
+        } catch (GrokException e) {
+            throw new RuntimeException(e);
+        }
+        this.buildLines = new StringBuilder();
+    }
+
+    @Override
+    public boolean nextFile() {
+        return super.nextFile();
+    }
+
+    @Override
+    public void close() {
+        super.close();
+    }
+
+    @Override
+    public String poll() {
+        String currentLine;
+        while ((currentLine = super.poll()) != null) {
+            Match gm = grok.match(currentLine);
+            gm.captures();
+
+            if (!gm.isNull()) {
+                String fullLog = this.buildLines.toString();
+                if (!Strings.isNullOrEmpty(fullLog)) {
+                    this.buildLines.setLength(0);
+                    this.buildLines.append(currentLine).append(System.lineSeparator());
+                    return fullLog;
+                }
+            }
+            this.buildLines.append(currentLine).append(System.lineSeparator());
+        }
+
+        String fullLog = this.buildLines.toString();
+        if (!Strings.isNullOrEmpty(fullLog)) {
+            this.buildLines.setLength(0);
+            return fullLog;
+        }
+
+        return null;
+    }
+}

data/src/main/java/org/embulk/parser/grok/TimestampParserFactory.java

@@ -0,0 +1,85 @@
+package org.embulk.parser.grok;
+
+import com.google.common.collect.ImmutableMap;
+import org.embulk.config.Task;
+import org.embulk.spi.ColumnConfig;
+import org.embulk.spi.time.Timestamp;
+import org.embulk.spi.time.TimestampParser;
+import org.embulk.spi.type.TimestampType;
+import org.embulk.spi.util.Timestamps;
+import org.joda.time.DateTimeZone;
+
+import java.text.ParseException;
+import java.text.SimpleDateFormat;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public class TimestampParserFactory {
+    private static Map<String, String> formatMap = ImmutableMap.<String, String>builder()
+            .put("%Y", "yyyy")
+            .put("%m", "MM")
+            .put("%d", "dd")
+            .put("%H", "HH")
+            .put("%M", "mm")
+            .put("%S", "ss")
+            .put("%z", "z")
+            .put("%T", "HH:mm:ss")
+            .put("%b", "MMM")
+            .put("%N", "SSS")
+            .put("%6N", "SSS")
+            .put("T", "'T'")
+            .build();
+
+    private interface TimestampColumnOption extends Task, TimestampParser.TimestampColumnOption {
+    }
+
+    private static String convertToJavaDateFormat(String rubyFormat) {
+        String current = rubyFormat;
+
+        for (Map.Entry<String, String> entry : formatMap.entrySet()) {
+            current = current.replace(entry.getKey(), entry.getValue());
+        }
+
+        return current;
+    }
+
+    public static List<DateParser> create(GrokParserPlugin.PluginTask task) {
+        switch (task.getTimestampParser().toLowerCase()) {
+            case "ruby":
+                TimestampParser[] ps = Timestamps.newTimestampColumnParsers(task, task.getColumns());
+                return Arrays.stream(ps)
+                        .map(parser -> (DateParser) (text) -> parser.parse(text))
+                        .collect(Collectors.toList());
+            case "epoch":
+                return task.getColumns().getColumns().stream()
+                        .map(x -> (DateParser) (text) -> Timestamp.ofEpochMilli(Long.parseLong(text)))
+                        .collect(Collectors.toList());
+            case "sdf":
+            case "simpledateformat":
+            default:
+                SimpleDateFormat[] parsers = new SimpleDateFormat[task.getColumns().getColumnCount()];
+                int i = 0;
+                for (ColumnConfig column : task.getColumns().getColumns()) {
+                    if (column.getType() instanceof TimestampType) {
+                        TimestampColumnOption option = column.getOption().loadConfig(TimestampColumnOption.class);
+                        String format = convertToJavaDateFormat(option.getFormat().or("yyyy-MM-dd HH:MM:ss.SSS z"));
+                        SimpleDateFormat sdf = new SimpleDateFormat(format, Locale.ENGLISH);
+                        sdf.setTimeZone(option.getTimeZone().or(DateTimeZone.UTC).toTimeZone());
+                        parsers[i] = sdf;
+                    }
+                    i++;
+                }
+                return Arrays.stream(parsers).map(parser ->
+                        (DateParser) (String date) -> {
+                            try {
+                                return Timestamp.ofEpochMilli(parser.parse(date).getTime());
+                            } catch (ParseException e) {
+                                throw new GrokRecordValidateException(e);
+                            }
+                        }).collect(Collectors.toList());
+        }
+    }
+}

data/src/test/java/org/embulk/parser/TestGrokGuessPlugin.java

@@ -0,0 +1,50 @@
+package org.embulk.parser;
+
+import com.google.common.collect.ImmutableMap;
+import org.junit.Test;
+
+import java.util.Map;
+
+import static org.hamcrest.core.Is.is;
+import static org.junit.Assert.assertThat;
+
+public class TestGrokGuessPlugin extends TestGrokPluginBase {
+
+    public void testGuess(String template, String input, String expected) throws Exception {
+
+        String inputPath = TestGrokParserPlugin.class.getClassLoader().getResource(input).getPath();
+
+        Map<String, String> params = ImmutableMap.of(
+                "__INPUT_PATH__", inputPath,
+                "__PROJECT_DIR__", System.getProperty("user.dir")
+        );
+        String yamlPath = generateConfigYaml(template, params);
+        String diff = tester.guess(yamlPath);
+
+        assertThat(diff, is(expected));
+    }
+
+    @Test
+    public void testApacheGuess() throws Exception {
+        testGuess("guess.yml", "apache.log", "in:\n" +
+                "  parser:\n" +
+                "    grok_pattern: '%{COMBINEDAPACHELOG}'\n" +
+                "    columns:\n" +
+                "    - {name: request, type: string}\n" +
+                "    - {name: agent, type: string}\n" +
+                "    - {name: COMMONAPACHELOG, type: string}\n" +
+                "    - {name: auth, type: string}\n" +
+                "    - {name: ident, type: string}\n" +
+                "    - {name: verb, type: string}\n" +
+                "    - {name: referrer, type: string}\n" +
+                "    - {name: bytes, type: long}\n" +
+                "    - {name: response, type: long}\n" +
+                "    - {name: clientip, type: string}\n" +
+                "    - {name: COMBINEDAPACHELOG, type: string}\n" +
+                "    - {name: httpversion, type: string}\n" +
+                "    - {name: rawrequest, type: string}\n" +
+                "    - {name: timestamp, format: '%d/%b/%Y:%T %z', type: timestamp}\n");
+
+    }
+
+}

data/src/test/java/org/embulk/parser/TestGrokParserPlugin.java

@@ -0,0 +1,55 @@
+package org.embulk.parser;
+
+import com.google.common.collect.ImmutableMap;
+import org.embulk.util.StreamUtil;
+import org.junit.Test;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+import static junit.framework.Assert.assertEquals;
+
+public class TestGrokParserPlugin extends TestGrokPluginBase{
+
+
+    public void testRun(String template, String input, String expected) throws Exception {
+
+        String inputPath = TestGrokParserPlugin.class.getClassLoader().getResource(input).getPath();
+        Path outputPath = Paths.get(outputDirectoryPath.toString(), input);
+        String expectedPath = TestGrokParserPlugin.class.getClassLoader().getResource(expected).getPath();
+
+        Map<String, String> params = ImmutableMap.of(
+                "__INPUT_PATH__", inputPath,
+                "__OUTPUT_PATH__", outputPath.toString(),
+                "__PROJECT_DIR__", System.getProperty("user.dir")
+        );
+        String yamlPath = generateConfigYaml(template, params);
+        tester.run(yamlPath);
+
+        Path outputFilePath = Paths.get(outputDirectoryPath.toString(), input + "000.00.output.csv");
+        List<Integer> counter = new ArrayList<>();
+        counter.add(0);
+        StreamUtil.zip(Files.lines(outputFilePath), Files.lines(Paths.get(expectedPath)))
+                .forEach(tuple -> {
+                    Integer line = counter.get(0);
+                    line++;
+                    counter.set(0, line);
+                    assertEquals(tuple.b, tuple.a);
+                });
+    }
+
+    @Test
+    public void testApacheLog() throws Exception {
+        testRun("apache.yml", "apache.log", "expected_apache.csv");
+    }
+
+    @Test
+    public void testMultiLineLog() throws Exception {
+        testRun("multiline.yml", "multiline.log", "expected_multiline.csv");
+    }
+
+}

data/src/test/java/org/embulk/parser/TestGrokPluginBase.java

@@ -0,0 +1,69 @@
+package org.embulk.parser;
+
+import org.embulk.parser.grok.GrokGuessPlugin;
+import org.embulk.parser.grok.GrokParserPlugin;
+import org.embulk.spi.GuessPlugin;
+import org.embulk.spi.ParserPlugin;
+import org.embulk.util.EmbulkPluginTester;
+import org.junit.After;
+import org.junit.Before;
+
+import java.io.*;
+import java.nio.file.FileVisitResult;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.SimpleFileVisitor;
+import java.nio.file.attribute.BasicFileAttributes;
+import java.util.Map;
+
+public class TestGrokPluginBase {
+
+    protected String generateConfigYaml(String template, Map<String, String> params) throws IOException {
+        File generatedFile = File.createTempFile("generated", "yaml");
+
+        String templatePath = TestGrokParserPlugin.class.getClassLoader().getResource(template).getPath();
+
+        try (BufferedReader br = new BufferedReader(new FileReader(new File(templatePath)));
+             BufferedWriter writer = new BufferedWriter(new FileWriter(generatedFile))) {
+            String line;
+            while ((line = br.readLine()) != null) {
+
+                for (Map.Entry<String, String> entry : params.entrySet()) {
+                    line = line.replaceAll(entry.getKey(), entry.getValue());
+                }
+                writer.write(line + "\n");
+            }
+        }
+
+        return generatedFile.getAbsolutePath();
+    }
+
+    protected Path outputDirectoryPath;
+    protected EmbulkPluginTester tester;
+
+    @Before
+    public void before() throws IOException {
+        outputDirectoryPath = Files.createTempDirectory(null);
+        tester = new EmbulkPluginTester();
+        tester.addPlugin(ParserPlugin.class, "grok", GrokParserPlugin.class);
+        tester.addPlugin(GuessPlugin.class, "grok", GrokGuessPlugin.class);
+    }
+
+    @After
+    public void after() throws IOException {
+        Files.walkFileTree(outputDirectoryPath, new SimpleFileVisitor<Path>() {
+            @Override
+            public FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOException {
+                Files.deleteIfExists(dir);
+                return super.postVisitDirectory(dir, exc);
+            }
+
+            @Override
+            public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
+                Files.deleteIfExists(file);
+                return super.visitFile(file, attrs);
+            }
+        });
+    }
+
+}