@Override
public FileOutput open(TaskSource taskSource, final FileOutput fileOutput) {
final PluginTask task = taskSource.loadTask(PluginTask.class);
final FileOutputOutputStream output =
new FileOutputOutputStream(
fileOutput, task.getBufferAllocator(), FileOutputOutputStream.CloseMode.FLUSH);
return new OutputStreamFileOutput(
new OutputStreamFileOutput.Provider() {
public OutputStream openNext() throws IOException {
output.nextFile();
return new GZIPOutputStream(output) {
{
this.def.setLevel(task.getLevel());
}
};
}
public void finish() throws IOException {
fileOutput.finish();
}
public void close() throws IOException {
fileOutput.close();
}
});
}
private List<ColumnGetter> newColumnGetters(
PluginTask task, JdbcSchema querySchema, PageBuilder pageBuilder) throws SQLException {
ColumnGetterFactory factory = newColumnGetterFactory(pageBuilder, task.getDefaultTimeZone());
ImmutableList.Builder<ColumnGetter> getters = ImmutableList.builder();
for (JdbcColumn c : querySchema.getColumns()) {
JdbcColumnOption columnOption = columnOptionOf(task.getColumnOptions(), c);
getters.add(factory.newColumnGetter(c, columnOption));
}
return getters.build();
}
@Override
public void transaction(ConfigSource config, Schema inputSchema, FilterPlugin.Control control) {
PluginTask task = config.loadConfig(PluginTask.class);
try {
task.setJsonTable(buildJsonTable(task, loadJsonFile(task.getJsonFilePath())));
} catch (IOException e) {
logger.error(e.getMessage());
throw new RuntimeException(e);
}
Schema outputSchema = buildOutputSchema(task, inputSchema, task.getJsonColumns());
control.run(task.dump(), outputSchema);
}
@Override
public ConfigDiff transaction(ConfigSource config, InputPlugin.Control control) {
PluginTask task = config.loadConfig(getTaskClass());
// if (task.getLastValue().isPresent() && !task.getOrderBy().isPresent()) {
// throw new ConfigException("order_by parameter must be set if last_value parameter is
// set");
// }
Schema schema;
try (JdbcInputConnection con = newConnection(task)) {
schema = setupTask(con, task);
} catch (SQLException ex) {
throw Throwables.propagate(ex);
}
return buildNextConfigDiff(task, control.run(task.dump(), schema, 1));
}
private Schema setupTask(JdbcInputConnection con, PluginTask task) throws SQLException {
// build SELECT query and gets schema of its result
String query = getQuery(task, con);
logger.info("SQL: " + query);
JdbcSchema querySchema = con.getSchemaOfQuery(query);
task.setQuerySchema(querySchema);
// validate column_options
newColumnGetters(task, querySchema, null);
ColumnGetterFactory factory = newColumnGetterFactory(null, task.getDefaultTimeZone());
ImmutableList.Builder<Column> columns = ImmutableList.builder();
for (int i = 0; i < querySchema.getCount(); i++) {
JdbcColumn column = querySchema.getColumn(i);
JdbcColumnOption columnOption = columnOptionOf(task.getColumnOptions(), column);
columns.add(
new Column(
i, column.getName(), factory.newColumnGetter(column, columnOption).getToType()));
}
return new Schema(columns.build());
}
private HashMap<String, HashMap<String, String>> buildJsonTable(
final PluginTask task, List<HashMap<String, String>> jsonData) {
HashMap<String, HashMap<String, String>> jsonTable = Maps.newHashMap();
for (HashMap<String, String> json : jsonData) {
HashMap<String, String> record = Maps.newHashMap();
for (ColumnConfig columnConfig : task.getJsonColumns()) {
String columnKey = task.getJoinedColumnsPrefix() + columnConfig.getName();
String value = json.get(columnConfig.getName());
record.put(columnKey, value);
}
String rowKey = json.get(task.getCounterColumn().getName());
jsonTable.put(rowKey, record);
}
return jsonTable;
}
@Override
public TaskReport run(TaskSource taskSource, Schema schema, int taskIndex, PageOutput output) {
PluginTask task = taskSource.loadTask(getTaskClass());
JdbcSchema querySchema = task.getQuerySchema();
BufferAllocator allocator = task.getBufferAllocator();
PageBuilder pageBuilder = new PageBuilder(allocator, schema, output);
try {
List<ColumnGetter> getters = newColumnGetters(task, querySchema, pageBuilder);
try (JdbcInputConnection con = newConnection(task)) {
try (BatchSelect cursor =
con.newSelectCursor(
getQuery(task, con), task.getFetchRows(), task.getSocketTimeout())) {
while (true) {
// TODO run fetch() in another thread asynchronously
// TODO retry fetch() if it failed (maybe order_by is required and unique_column(s)
// option is also required)
boolean cont = fetch(cursor, getters, pageBuilder);
if (!cont) {
break;
}
}
}
}
} catch (SQLException ex) {
throw Throwables.propagate(ex);
}
pageBuilder.finish();
TaskReport report = Exec.newTaskReport();
// TODO
// if (orderByColumn != null) {
// report.set("last_value", lastValue);
// }
return report;
}
private Schema buildOutputSchema(
final PluginTask task, Schema inputSchema, List<ColumnConfig> jsonColumns) {
ImmutableList.Builder<Column> builder = ImmutableList.builder();
int i = 0;
for (Column inputColumn : inputSchema.getColumns()) {
Column outputColumn = new Column(i++, inputColumn.getName(), inputColumn.getType());
builder.add(outputColumn);
}
for (ColumnConfig columnConfig : jsonColumns) {
String columnName = task.getJoinedColumnsPrefix() + columnConfig.getName();
builder.add(new Column(i++, columnName, columnConfig.getType()));
}
return new Schema(builder.build());
}
@Override
public PageOutput open(
TaskSource taskSource,
final Schema inputSchema,
final Schema outputSchema,
final PageOutput output) {
final PluginTask task = taskSource.loadTask(PluginTask.class);
// create jsonColumns/baseColumn
final List<Column> outputColumns = outputSchema.getColumns();
final List<Column> inputColumns = inputSchema.getColumns();
Map<String, Column> inputColumnMap = Maps.newHashMap();
final List<Column> jsonColumns = new ArrayList<>();
for (Column column : outputColumns) {
if (!inputColumns.contains(column)) {
jsonColumns.add(column);
} else {
inputColumnMap.put(column.getName(), column);
}
}
final Column baseColumn = inputColumnMap.get(task.getBaseColumn().getName());
// create timestampParserMap
final HashMap<String, TimestampParser> timestampParserMap = Maps.newHashMap();
for (ColumnConfig jsonColumnConfig : task.getJsonColumns()) {
if (Types.TIMESTAMP.equals(jsonColumnConfig.getType())) {
String format = jsonColumnConfig.getOption().get(String.class, "format");
DateTimeZone timezone = DateTimeZone.forID(task.getTimeZone());
TimestampParser parser = new TimestampParser(task.getJRuby(), format, timezone);
timestampParserMap.put(task.getJoinedColumnsPrefix() + jsonColumnConfig.getName(), parser);
}
}
// get jsonTable
final HashMap<String, HashMap<String, String>> jsonTable = task.getJsonTable();
return new LeftOuterJoinJsonTableFilterFilteredPageOutput(
inputSchema, outputSchema, baseColumn, jsonTable, jsonColumns, timestampParserMap, output);
}
private String getQuery(PluginTask task, JdbcInputConnection con) {
if (task.getQuery().isPresent()) {
if (task.getTable().isPresent()
|| task.getSelect().isPresent()
|| task.getWhere().isPresent()
|| task.getOrderBy().isPresent()) {
throw new ConfigException(
"'table', 'select', 'where' and 'order_by' parameters are unnecessary if 'query' parameter is set.");
}
return task.getQuery().get();
} else if (task.getTable().isPresent()) {
return con.buildSelectQuery(
task.getTable().get(), task.getSelect(), task.getWhere(), task.getOrderBy());
} else {
throw new ConfigException("'table' parameter is required (if 'query' parameter is not set)");
}
}
@Override
public void transaction(ConfigSource config, Control control) {
PluginTask task = config.loadConfig(PluginTask.class);
control.run(task.dump(), task.getSchemaConfig().toSchema());
}
public void transaction(ConfigSource config, EncoderPlugin.Control control) {
PluginTask task = config.loadConfig(PluginTask.class);
control.run(task.dump());
}