org.embulk.spi.Exec Java Examples

@Test
public void checkExtraConfigurations()
        throws NoSuchMethodException, InvocationTargetException, IllegalAccessException
{
    ConfigSource map = Exec.newConfigSource()
            .set("foo", "bar");

    ConfigSource config = Exec.newConfigSource()
            .set("path_prefix", "test")
            .setNested("extra_configurations", map);

    ParquetOutputPlugin.PluginTask task = config.loadConfig(ParquetOutputPlugin.PluginTask.class);

    Map<String, String> extra = task.getExtraConfigurations();
    assertTrue(extra.containsKey("foo"));
    assertEquals("bar", extra.get("foo"));

    ParquetOutputPlugin plugin = new ParquetOutputPlugin();
    Method method = ParquetOutputPlugin.class.getDeclaredMethod("createConfiguration", Map.class, List.class);
    method.setAccessible(true);
    Configuration conf = (Configuration) method.invoke(plugin, extra, new ArrayList<String>());
    assertEquals("bar", conf.get("foo"));
}
 

@Override
public PageOutput open(final TaskSource taskSource, final Schema inputSchema,
        final Schema outputSchema, final PageOutput output)
{
    final PluginTask task = taskSource.loadTask(PluginTask.class);

    return new PageOutput() {
        private PageReader pageReader = new PageReader(inputSchema);
        private PageBuilder pageBuilder = new PageBuilder(Exec.getBufferAllocator(), outputSchema, output);
        private ColumnVisitorImpl visitor = new ColumnVisitorImpl(task, inputSchema, outputSchema, pageReader, pageBuilder);

        @Override
        public void finish()
        {
            pageBuilder.finish();
        }

        @Override
        public void close()
        {
            pageBuilder.close();
        }

        @Override
        public void add(Page page)
        {
            pageReader.setPage(page);

            while (pageReader.nextRecord()) {
                outputSchema.visitColumns(visitor);
                pageBuilder.addRecord();
            }
        }
    };
}
 

private PluginTask taskFromYamlString(String... lines)
{
    StringBuilder builder = new StringBuilder();
    for (String line : lines) {
        builder.append(line).append("\n");
    }
    String yamlString = builder.toString();

    ConfigLoader loader = new ConfigLoader(Exec.getModelManager());
    ConfigSource config = loader.fromYamlString(yamlString);
    return config.loadConfig(PluginTask.class);
}
 

private ConfigSource configFromYamlString(String... lines)
{
    StringBuilder builder = new StringBuilder();
    for (String line : lines) {
        builder.append(line).append("\n");
    }
    String yamlString = builder.toString();

    ConfigLoader loader = new ConfigLoader(Exec.getModelManager());
    return loader.fromYamlString(yamlString);
}
 

private PluginTask taskFromYamlString(String... lines)
{
    StringBuilder builder = new StringBuilder();
    for (String line : lines) {
        builder.append(line).append("\n");
    }
    String yamlString = builder.toString();

    ConfigLoader loader = new ConfigLoader(Exec.getModelManager());
    ConfigSource config = loader.fromYamlString(yamlString);
    return config.loadConfig(PluginTask.class);
}
 

private static JsonlColumnOption columnOptionOf(Map<String, JsonlColumnOption> columnOptions, String columnName)
{
    return Optional.fromNullable(columnOptions.get(columnName)).or(
            // default column option
            new Supplier<JsonlColumnOption>()
            {
                public JsonlColumnOption get()
                {
                    return Exec.newConfigSource().loadConfig(JsonlColumnOption.class);
                }
            });
}
 

public ConfigDiff transaction(ConfigSource config,
                              Schema schema, int processorCount,
                              OutputPlugin.Control control)
{
    PluginTask task = config.loadConfig(PluginTask.class);

    //TODO


    try (@SuppressWarnings("unchecked") ClassLoaderSwap clswp = new ClassLoaderSwap(this.getClass())) {
        control.run(task.dump());
    }
    return Exec.newConfigDiff();
}
 

@Test
public void checkDefaultValues()
{
    ConfigSource config = Exec.newConfigSource()
            .set("path_prefix", "test");

    ParquetOutputPlugin.PluginTask task = config.loadConfig(ParquetOutputPlugin.PluginTask.class);
    assertEquals(".parquet", task.getFileNameExtension());
    assertEquals(".%03d", task.getSequenceFormat());
    assertEquals(134217728, task.getBlockSize());
    assertEquals(1048576, task.getPageSize());
    assertEquals("UNCOMPRESSED", task.getCompressionCodec());
    assertFalse(task.getOverwrite());
    assertEquals("false", task.getSignature());
}
 

@Test(expected = ConfigException.class)
public void checkColumnsRequired()
{
    ConfigSource config = Exec.newConfigSource();

    config.loadConfig(ParquetOutputPlugin.PluginTask.class);
}
 

@Override
public ConfigDiff resume(TaskSource taskSource, int taskCount,
        Control control)
{
    control.run(taskSource);
    return Exec.newConfigDiff();
}
 

public JsonlParserPlugin()
{
    this.log = Exec.getLogger(JsonlParserPlugin.class);
}
 

@Override
public void run(TaskSource taskSource, Schema schema, FileInput input, PageOutput output)
{
    PluginTask task = taskSource.loadTask(PluginTask.class);

    setColumnNameValues(schema);

    final SchemaConfig schemaConfig = getSchemaConfig(task);
    final TimestampParser[] timestampParsers = Timestamps.newTimestampColumnParsers(task, schemaConfig);
    final LineDecoder decoder = newLineDecoder(input, task);
    final JsonParser jsonParser = newJsonParser();
    final boolean stopOnInvalidRecord = task.getStopOnInvalidRecord();

    try (final PageBuilder pageBuilder = new PageBuilder(Exec.getBufferAllocator(), schema, output)) {
        ColumnVisitorImpl visitor = new ColumnVisitorImpl(task, schema, pageBuilder, timestampParsers);

        while (decoder.nextFile()) { // TODO this implementation should be improved with new JsonParser API on Embulk v0.8.3
            lineNumber = 0;

            while ((line = decoder.poll()) != null) {
                lineNumber++;

                try {
                    Value value = jsonParser.parse(line);

                    if (!value.isMapValue()) {
                        throw new JsonRecordValidateException("Json string is not representing map value.");
                    }

                    final Map<Value, Value> record = value.asMapValue().map();
                    for (Column column : schema.getColumns()) {
                        Value v = record.get(getColumnNameValue(column));
                        visitor.setValue(v);
                        column.visit(visitor);
                    }

                    pageBuilder.addRecord();
                }
                catch (JsonRecordValidateException | JsonParseException e) {
                    if (stopOnInvalidRecord) {
                        throw new DataException(String.format("Invalid record at line %d: %s", lineNumber, line), e);
                    }
                    log.warn(String.format("Skipped line %d (%s): %s", lineNumber, e.getMessage(), line));
                }
            }
        }

        pageBuilder.finish();
    }
}
 

private ConfigSource getConfigFromYamlFile(File yamlFile)
        throws IOException
{
    ConfigLoader loader = new ConfigLoader(Exec.getModelManager());
    return loader.fromYamlFile(yamlFile);
}
 

@Override
public TaskReport commit()
{
    return Exec.newTaskReport();
    //TODO
}
 

@Override
public TaskReport commit()
{
    TaskReport report = Exec.newTaskReport();
    return report;
}