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

Example 16 with TaskSource

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

the class PreviewExecutor method doPreview.

@SuppressWarnings("checkstyle:OverloadMethodsDeclarationOrder")
private PreviewResult doPreview(final PreviewTask task, final InputPlugin input, final List<FilterPlugin> filterPlugins) {
    try {
        input.transaction(task.getInputConfig(), new InputPlugin.Control() {

            public List<TaskReport> run(final TaskSource inputTask, Schema inputSchema, final int taskCount) {
                Filters.transaction(filterPlugins, task.getFilterConfigs(), inputSchema, new Filters.Control() {

                    public void run(final List<TaskSource> filterTasks, final List<Schema> filterSchemas) {
                        Schema inputSchema = filterSchemas.get(0);
                        Schema outputSchema = filterSchemas.get(filterSchemas.size() - 1);
                        PageOutput out = new SamplingPageOutput(task.getSampleRows(), outputSchema);
                        try {
                            for (int taskIndex = 0; taskIndex < taskCount; taskIndex++) {
                                try {
                                    out = Filters.open(filterPlugins, filterTasks, filterSchemas, out);
                                    input.run(inputTask, inputSchema, taskIndex, out);
                                } catch (NoSampleException ex) {
                                    if (taskIndex == taskCount - 1) {
                                        throw ex;
                                    }
                                }
                            }
                        } finally {
                            out.close();
                        }
                    }
                });
                // program never reaches here because SamplingPageOutput.finish throws an error.
                throw new NoSampleException("No input records to preview");
            }
        });
        throw new AssertionError("PreviewExecutor executor must throw PreviewedNoticeError");
    } catch (PreviewedNoticeError previewed) {
        return previewed.getPreviewResult();
    }
}

Example 17 with TaskSource

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

the class SamplingParserPlugin method runFileInputSampling.

public static Buffer runFileInputSampling(final FileInputRunner runner, ConfigSource inputConfig, ConfigSource sampleBufferConfig) {
    final SampleBufferTask sampleBufferTask = sampleBufferConfig.loadConfig(SampleBufferTask.class);
    // override in.parser.type so that FileInputRunner creates SamplingParserPlugin
    ConfigSource samplingInputConfig = inputConfig.deepCopy();
    samplingInputConfig.getNestedOrSetEmpty("parser").set("type", "system_sampling").set("sample_buffer_bytes", sampleBufferTask.getSampleBufferBytes());
    samplingInputConfig.set("decoders", null);
    try {
        runner.transaction(samplingInputConfig, new InputPlugin.Control() {

            public List<TaskReport> run(TaskSource taskSource, Schema schema, int taskCount) {
                if (taskCount == 0) {
                    throw new NoSampleException("No input files to read sample data");
                }
                int maxSize = -1;
                int maxSizeTaskIndex = -1;
                for (int taskIndex = 0; taskIndex < taskCount; taskIndex++) {
                    try {
                        runner.run(taskSource, schema, taskIndex, new PageOutput() {

                            @Override
                            public void add(Page page) {
                                // TODO exception class
                                throw new RuntimeException("Input plugin must be a FileInputPlugin to guess parser configuration");
                            }

                            public void finish() {
                            }

                            public void close() {
                            }
                        });
                    } catch (NotEnoughSampleError ex) {
                        if (maxSize < ex.getSize()) {
                            maxSize = ex.getSize();
                            maxSizeTaskIndex = taskIndex;
                        }
                        continue;
                    }
                }
                if (maxSize <= 0) {
                    throw new NoSampleException("All input files are empty");
                }
                taskSource.getNested("ParserTaskSource").set("force", true);
                try {
                    runner.run(taskSource, schema, maxSizeTaskIndex, new PageOutput() {

                        @Override
                        public void add(Page page) {
                            // TODO exception class
                            throw new RuntimeException("Input plugin must be a FileInputPlugin to guess parser configuration");
                        }

                        public void finish() {
                        }

                        public void close() {
                        }
                    });
                } catch (NotEnoughSampleError ex) {
                    throw new NoSampleException("All input files are smaller than minimum sampling size");
                }
                throw new NoSampleException("All input files are smaller than minimum sampling size");
            }
        });
        throw new AssertionError("SamplingParserPlugin must throw SampledNoticeError");
    } catch (SampledNoticeError error) {
        return error.getSample();
    }
}