Examples with TaskSource org.embulk.config.TaskSource used on opensource projects

Example 11 with TaskSource

use of org.embulk.config.TaskSource in project embulk by embulk.

the class TestRenameFilterPlugin method renameAndCheckSchema.

private void renameAndCheckSchema(ConfigSource config, final String[] original, final String[] expected) {
    final Schema originalSchema = makeSchema(original);
    filter.transaction(config, originalSchema, new FilterPlugin.Control() {

        @Override
        public void run(TaskSource task, Schema renamedSchema) {
            assertEquals(originalSchema.getColumnCount(), renamedSchema.getColumnCount());
            assertEquals(expected.length, renamedSchema.getColumnCount());
            for (int i = 0; i < renamedSchema.getColumnCount(); ++i) {
                assertEquals(originalSchema.getColumnType(i), renamedSchema.getColumnType(i));
                assertEquals(expected[i], renamedSchema.getColumnName(i));
            }
        }
    });
}

Example 12 with TaskSource

use of org.embulk.config.TaskSource in project embulk by embulk.

the class TestRenameFilterPlugin method checkRenaming.

@Test
public void checkRenaming() {
    ConfigSource pluginConfig = Exec.newConfigSource().set("columns", ImmutableMap.of("_c0", "_c0_new"));
    filter.transaction(pluginConfig, SCHEMA, new FilterPlugin.Control() {

        @Override
        public void run(TaskSource task, Schema newSchema) {
            // _c0 -> _c0_new
            Column old0 = SCHEMA.getColumn(0);
            Column new0 = newSchema.getColumn(0);
            assertEquals("_c0_new", new0.getName());
            assertEquals(old0.getType(), new0.getType());
            // _c1 is not changed
            Column old1 = SCHEMA.getColumn(1);
            Column new1 = newSchema.getColumn(1);
            assertEquals("_c1", new1.getName());
            assertEquals(old1.getType(), new1.getType());
        }
    });
}

Example 13 with TaskSource

use of org.embulk.config.TaskSource in project embulk by embulk.

the class TestFileOutputRunner method testMockFormatterIteration.

@Test
public void testMockFormatterIteration() {
    MockFileOutputPlugin fileOutputPlugin = new MockFileOutputPlugin();
    final FileOutputRunner runner = new FileOutputRunner(fileOutputPlugin);
    ImmutableList<ImmutableMap<String, Object>> columns = ImmutableList.of(ImmutableMap.<String, Object>of("name", "col1", "type", "boolean", "option", ImmutableMap.of()), ImmutableMap.<String, Object>of("name", "col2", "type", "long", "option", ImmutableMap.of()), ImmutableMap.<String, Object>of("name", "col3", "type", "double", "option", ImmutableMap.of()), ImmutableMap.<String, Object>of("name", "col4", "type", "string", "option", ImmutableMap.of()), ImmutableMap.<String, Object>of("name", "col5", "type", "timestamp", "option", ImmutableMap.of()), ImmutableMap.<String, Object>of("name", "col6", "type", "json", "option", ImmutableMap.of()));
    ConfigSource config = Exec.newConfigSource().set("type", "unused?").set("formatter", ImmutableMap.of("type", "mock", "columns", columns));
    final Schema schema = config.getNested("formatter").loadConfig(MockParserPlugin.PluginTask.class).getSchemaConfig().toSchema();
    runner.transaction(config, schema, 1, new OutputPlugin.Control() {

        public List<TaskReport> run(final TaskSource outputTask) {
            TransactionalPageOutput tran = runner.open(outputTask, schema, 1);
            boolean committed = false;
            try {
                ImmutableMapValue jsonValue = newMap(newString("_c1"), newBoolean(true), newString("_c2"), newInteger(10), newString("_c3"), newString("embulk"), newString("_c4"), newMap(newString("k"), newString("v")));
                for (Page page : PageTestUtils.buildPage(runtime.getBufferAllocator(), schema, true, 2L, 3.0D, "45", Timestamp.ofEpochMilli(678L), jsonValue, true, 2L, 3.0D, "45", Timestamp.ofEpochMilli(678L), jsonValue)) {
                    tran.add(page);
                }
                tran.commit();
                committed = true;
            } finally {
                if (!committed) {
                    tran.abort();
                }
                tran.close();
            }
            return new ArrayList<TaskReport>();
        }
    });
    assertEquals(true, fileOutputPlugin.transactionCompleted);
    assertEquals(2, MockFormatterPlugin.records.size());
    for (List<Object> record : MockFormatterPlugin.records) {
        assertEquals(Boolean.TRUE, record.get(0));
        assertEquals(2L, record.get(1));
        assertEquals(3.0D, (Double) record.get(2), 0.1D);
        assertEquals("45", record.get(3));
        assertEquals(678L, ((Timestamp) record.get(4)).toEpochMilli());
        assertEquals("{\"_c1\":true,\"_c2\":10,\"_c3\":\"embulk\",\"_c4\":{\"k\":\"v\"}}", record.get(5).toString());
    }
}

Example 14 with TaskSource

use of org.embulk.config.TaskSource in project embulk by embulk.

the class TestFileOutputRunner method testTransactionAborted.

@Test
public void testTransactionAborted() {
    MockFileOutputPlugin fileOutputPlugin = new MockFileOutputPlugin();
    final FileOutputRunner runner = new FileOutputRunner(fileOutputPlugin);
    ImmutableList<ImmutableMap<String, Object>> columns = ImmutableList.of(ImmutableMap.<String, Object>of("name", "col1", "type", "boolean", "option", ImmutableMap.of()), ImmutableMap.<String, Object>of("name", "col2", "type", "long", "option", ImmutableMap.of()), ImmutableMap.<String, Object>of("name", "col3", "type", "double", "option", ImmutableMap.of()), ImmutableMap.<String, Object>of("name", "col4", "type", "string", "option", ImmutableMap.of()), ImmutableMap.<String, Object>of("name", "col5", "type", "timestamp", "option", ImmutableMap.of()), ImmutableMap.<String, Object>of("name", "col6", "type", "json", "option", ImmutableMap.of()));
    ConfigSource config = Exec.newConfigSource().set("type", "unused?").set("formatter", ImmutableMap.of("type", "mock", "columns", columns));
    final Schema schema = config.getNested("formatter").loadConfig(MockParserPlugin.PluginTask.class).getSchemaConfig().toSchema();
    try {
        runner.transaction(config, schema, 1, new OutputPlugin.Control() {

            public List<TaskReport> run(final TaskSource outputTask) {
                TransactionalPageOutput tran = runner.open(outputTask, schema, 1);
                boolean committed = false;
                try {
                    tran.add(null);
                    tran.commit();
                    committed = true;
                } finally {
                    if (!committed) {
                        tran.abort();
                    }
                    tran.close();
                }
                return new ArrayList<TaskReport>();
            }
        });
    } catch (NullPointerException npe) {
    // Just passing through.
    }
    assertEquals(false, fileOutputPlugin.transactionCompleted);
}

Example 15 with TaskSource

use of org.embulk.config.TaskSource in project embulk by embulk.

the class BulkLoader method doResume.

private ExecutionResult doResume(ConfigSource config, final ResumeState resume) {
    final BulkLoaderTask task = config.loadConfig(BulkLoaderTask.class);
    final ExecutorPlugin exec = newExecutorPlugin(task);
    final ProcessPluginSet plugins = new ProcessPluginSet(task);
    final LoaderState state = newLoaderState(Exec.getLogger(BulkLoader.class), plugins);
    state.setTransactionStage(TransactionStage.INPUT_BEGIN);
    try {
        @SuppressWarnings("checkstyle:LineLength") ConfigDiff inputConfigDiff = plugins.getInputPlugin().resume(resume.getInputTaskSource(), resume.getInputSchema(), resume.getInputTaskReports().size(), new InputPlugin.Control() {

            public List<TaskReport> run(final TaskSource inputTask, final Schema inputSchema, final int inputTaskCount) {
                // TODO validate inputTask?
                // TODO validate inputSchema
                state.setInputTaskSource(inputTask);
                state.setTransactionStage(TransactionStage.FILTER_BEGIN);
                Filters.transaction(plugins.getFilterPlugins(), task.getFilterConfigs(), inputSchema, new Filters.Control() {

                    public void run(final List<TaskSource> filterTasks, final List<Schema> schemas) {
                        state.setSchemas(schemas);
                        state.setFilterTaskSources(filterTasks);
                        state.setTransactionStage(TransactionStage.EXECUTOR_BEGIN);
                        exec.transaction(task.getExecConfig(), last(schemas), inputTaskCount, new ExecutorPlugin.Control() {

                            public void transaction(final Schema executorSchema, final int outputTaskCount, final ExecutorPlugin.Executor executor) {
                                // TODO validate executorSchema
                                state.setExecutorSchema(executorSchema);
                                state.setTransactionStage(TransactionStage.OUTPUT_BEGIN);
                                @SuppressWarnings("checkstyle:LineLength") ConfigDiff outputConfigDiff = plugins.getOutputPlugin().resume(resume.getOutputTaskSource(), executorSchema, outputTaskCount, new OutputPlugin.Control() {

                                    public List<TaskReport> run(final TaskSource outputTask) {
                                        // TODO validate outputTask?
                                        state.setOutputTaskSource(outputTask);
                                        restoreResumedTaskReports(resume, state);
                                        state.setTransactionStage(TransactionStage.RUN);
                                        if (!state.isAllTasksCommitted()) {
                                            execute(task, executor, state);
                                        }
                                        if (!state.isAllTasksCommitted()) {
                                            throw new RuntimeException(String.format("%d input tasks and %d output tasks failed", state.countUncommittedInputTasks(), state.countUncommittedOutputTasks()));
                                        }
                                        state.setTransactionStage(TransactionStage.OUTPUT_COMMIT);
                                        return state.getAllOutputTaskReports();
                                    }
                                });
                                state.setOutputConfigDiff(outputConfigDiff);
                                state.setTransactionStage(TransactionStage.EXECUTOR_COMMIT);
                            }
                        });
                        state.setTransactionStage(TransactionStage.FILTER_COMMIT);
                    }
                });
                state.setTransactionStage(TransactionStage.INPUT_COMMIT);
                return state.getAllInputTaskReports();
            }
        });
        state.setInputConfigDiff(inputConfigDiff);
        state.setTransactionStage(TransactionStage.CLEANUP);
        cleanupCommittedTransaction(config, state);
        return state.buildExecuteResult();
    } catch (Throwable ex) {
        if (isSkippedTransaction(ex)) {
            ConfigDiff configDiff = ((SkipTransactionException) ex).getConfigDiff();
            return state.buildExecuteResultOfSkippedExecution(configDiff);
        } else if (state.isAllTasksCommitted() && state.isAllTransactionsCommitted()) {
            // ignore the exception
            return state.buildExecuteResultWithWarningException(ex);
        }
        throw state.buildPartialExecuteException(ex, Exec.session());
    }
}