com.hishidama.embulk.tester.EmbulkPluginTester Maven / Gradle / Ivy
package com.hishidama.embulk.tester;
import java.io.Closeable;
import java.io.File;
import java.net.URISyntaxException;
import java.net.URL;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.text.MessageFormat;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import org.embulk.EmbulkEmbed;
import org.embulk.EmbulkEmbed.Bootstrap;
import org.embulk.config.ConfigLoader;
import org.embulk.config.ConfigSource;
import org.embulk.formatter.csv.CsvFormatterPlugin;
import org.embulk.input.file.LocalFileInputPlugin;
import org.embulk.parser.csv.CsvParserPlugin;
import org.embulk.spi.DecoderPlugin;
import org.embulk.spi.EncoderPlugin;
import org.embulk.spi.ExecutorPlugin;
import org.embulk.spi.FileInputPlugin;
import org.embulk.spi.FileOutputPlugin;
import org.embulk.spi.FilterPlugin;
import org.embulk.spi.FormatterPlugin;
import org.embulk.spi.GuessPlugin;
import org.embulk.spi.InputPlugin;
import org.embulk.spi.OutputPlugin;
import org.embulk.spi.ParserPlugin;
import com.hishidama.embulk.tester.EmbulkTestOutputPlugin.OutputRecord;
// https://github.com/embulk/embulk-output-jdbc/blob/master/embulk-output-jdbc/src/test/java/org/embulk/output/jdbc/tester/EmbulkPluginTester.java
public class EmbulkPluginTester implements Closeable {
private static class PluginDefinition {
@SuppressWarnings("unused")
public final Class iface;
public final String name;
public final Class impl;
public PluginDefinition(Class iface, String name, Class impl) {
this.iface = iface;
this.name = name;
this.impl = impl;
}
@SuppressWarnings("unchecked")
public Class impl() {
return (Class) impl;
}
}
private final Map, Map> plugins = new HashMap<>();
private EmbulkEmbed embulk;
private ConfigLoader configLoader;
private int inputTaskSize = 1;
private int outputMinTaskSize = 4;
public EmbulkPluginTester() {
}
public EmbulkPluginTester(Class iface, String name, Class impl) {
addPlugin(iface, name, impl);
}
public void addPlugin(Class iface, String name, Class impl) {
if (!iface.isAssignableFrom(impl)) {
throw new IllegalArgumentException(MessageFormat.format("name={1}. {2} is not {0}", iface.getSimpleName(), name, impl.getName()));
}
plugins.computeIfAbsent(iface, k -> new LinkedHashMap<>()).put(name, new PluginDefinition(iface, name, impl));
}
public Collection getPlugins(Class iface) {
return plugins.getOrDefault(iface, Collections.emptyMap()).values();
}
public void addInputPlugin(String name, Class impl) {
addPlugin(InputPlugin.class, name, impl);
}
public void addFileInputPlugin(String name, Class impl) {
addPlugin(FileInputPlugin.class, name, impl);
}
public void addParserPlugin(String name, Class impl) {
addPlugin(ParserPlugin.class, name, impl);
}
public void addOutputPlugin(String name, Class impl) {
addPlugin(OutputPlugin.class, name, impl);
}
public void addFileOutputPlugin(String name, Class impl) {
addPlugin(FileOutputPlugin.class, name, impl);
}
public void addFormatterPlugin(String name, Class impl) {
addPlugin(FormatterPlugin.class, name, impl);
}
public void setInputTaskSize(int taskSize) {
this.inputTaskSize = taskSize;
}
public void setOutputMinTaskSize(int taskSize) {
this.outputMinTaskSize = taskSize;
}
protected synchronized EmbulkEmbed getEmbulkEmbed() {
if (this.embulk == null) {
addDefaultPlugin();
Bootstrap bootstrap = new EmbulkEmbed.Bootstrap();
for (PluginDefinition plugin : getPlugins(DecoderPlugin.class)) {
bootstrap.builtinDecoderPlugin(plugin.name, plugin.impl());
}
for (PluginDefinition plugin : getPlugins(EncoderPlugin.class)) {
bootstrap.builtinEncoderPlugin(plugin.name, plugin.impl());
}
for (PluginDefinition plugin : getPlugins(ExecutorPlugin.class)) {
bootstrap.builtinExecutorPlugin(plugin.name, plugin.impl());
}
for (PluginDefinition plugin : getPlugins(FileInputPlugin.class)) {
bootstrap.builtinFileInputPlugin(plugin.name, plugin.impl());
}
for (PluginDefinition plugin : getPlugins(FileOutputPlugin.class)) {
bootstrap.builtinFileOutputPlugin(plugin.name, plugin.impl());
}
for (PluginDefinition plugin : getPlugins(FilterPlugin.class)) {
bootstrap.builtinFilterPlugin(plugin.name, plugin.impl());
}
for (PluginDefinition plugin : getPlugins(FormatterPlugin.class)) {
bootstrap.builtinFormatterPlugin(plugin.name, plugin.impl());
}
for (PluginDefinition plugin : getPlugins(GuessPlugin.class)) {
bootstrap.builtinGuessPlugin(plugin.name, plugin.impl());
}
for (PluginDefinition plugin : getPlugins(InputPlugin.class)) {
bootstrap.builtinInputPlugin(plugin.name, plugin.impl());
}
for (PluginDefinition plugin : getPlugins(OutputPlugin.class)) {
bootstrap.builtinOutputPlugin(plugin.name, plugin.impl());
}
for (PluginDefinition plugin : getPlugins(ParserPlugin.class)) {
bootstrap.builtinParserPlugin(plugin.name, plugin.impl());
}
this.embulk = bootstrap.initialize();
}
return this.embulk;
}
protected void addDefaultPlugin() {
addFileInputPlugin("file", LocalFileInputPlugin.class);
addFileInputPlugin(EmbulkTestFileInputPlugin.TYPE, EmbulkTestFileInputPlugin.class);
addParserPlugin("csv", CsvParserPlugin.class);
addOutputPlugin(EmbulkTestOutputPlugin.TYPE, EmbulkTestOutputPlugin.class);
addOutputPlugin(EmbulkTestOutputBinaryPlugin.TYPE, EmbulkTestOutputBinaryPlugin.class);
addFormatterPlugin("csv", CsvFormatterPlugin.class);
}
public synchronized ConfigLoader getConfigLoader() {
if (this.configLoader == null) {
this.configLoader = getEmbulkEmbed().newConfigLoader();
}
return this.configLoader;
}
public ConfigSource newConfigSource(String type) {
ConfigSource config = newConfigSource();
config.set("type", type);
return config;
}
public ConfigSource newConfigSource() {
return getConfigLoader().newConfigSource();
}
public EmbulkTestParserConfig newParserConfig(String type) {
EmbulkTestParserConfig parser = new EmbulkTestParserConfig();
parser.setType(type);
return parser;
}
public List runParser(URL inFile, EmbulkTestParserConfig parser) {
Path file;
try {
file = Paths.get(inFile.toURI());
} catch (URISyntaxException e) {
throw new RuntimeException(e);
}
return runParser(file, parser);
}
public List runParser(File inFile, EmbulkTestParserConfig parser) {
return runParser(inFile.toPath(), parser);
}
public List runParser(Path inFile, EmbulkTestParserConfig parser) {
ConfigSource in = newConfigSource("file");
in.set("path_prefix", inFile.toAbsolutePath().toString());
in.set("parser", parser);
return runInput(in);
}
public List runParser(List inList, EmbulkTestParserConfig parser) {
ConfigSource in = newConfigSourceTestFileInput(inList, parser);
return runInput(in);
}
protected ConfigSource newConfigSourceTestFileInput(List inList, EmbulkTestParserConfig parser) {
ConfigSource in = newConfigSource(EmbulkTestFileInputPlugin.TYPE);
in.set("textList", inList);
in.set("taskSize", this.inputTaskSize);
in.set("parser", parser);
return in;
}
public List runInput(ConfigSource in) {
ConfigSource out = newConfigSource(EmbulkTestOutputPlugin.TYPE);
EmbulkTestOutputPlugin.clearResult();
run(in, out);
return EmbulkTestOutputPlugin.getResult();
}
public void runOutput(List inList, EmbulkTestParserConfig parser, ConfigSource out) {
ConfigSource in = newConfigSourceTestFileInput(inList, parser);
run(in, out);
}
public List runFormatterToBinary(List inList, EmbulkTestParserConfig parser, ConfigSource formatter) {
ConfigSource in = newConfigSourceTestFileInput(inList, parser);
return runFormatterToBinary(in, formatter);
}
public List runFormatterToBinary(ConfigSource in, ConfigSource formatter) {
ConfigSource out = newConfigSource(EmbulkTestOutputBinaryPlugin.TYPE);
out.set("formatter", formatter);
EmbulkTestOutputBinaryPlugin.clearResult();
run(in, out);
return EmbulkTestOutputBinaryPlugin.getResult();
}
public void run(ConfigSource in, ConfigSource out) {
ConfigSource config = newConfigSource();
config.set("in", in);
config.set("out", out);
run(config);
}
public void run(ConfigSource config) {
// https://github.com/embulk/embulk/blob/master/embulk-core/src/main/java/org/embulk/exec/LocalExecutorPlugin.java
ConfigSource exec = newConfigSource();
exec.set("min_output_tasks", this.outputMinTaskSize);
config.set("exec", exec);
getEmbulkEmbed().run(config);
}
@Override
public void close() {
if (this.embulk != null) {
// embulk.destroy(); // unsupported
this.embulk = null;
}
}
}
© 2015 - 2025 Weber Informatics LLC | Privacy Policy