com.redislabs.riot.cli.file.FileExportOptions Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of riot Show documentation
Redis Input/Output Tool
There is a newer version: 1.11.2
package com.redislabs.riot.cli.file;

import java.io.IOException;
import java.io.Writer;
import java.util.Arrays;
import java.util.Locale;
import java.util.Map;

import org.springframework.batch.item.ItemWriter;
import org.springframework.batch.item.file.FlatFileHeaderCallback;
import org.springframework.batch.item.file.FlatFileItemWriter;
import org.springframework.batch.item.json.JacksonJsonObjectMarshaller;
import org.springframework.batch.item.support.AbstractItemStreamItemWriter;
import org.springframework.core.io.Resource;
import org.springframework.core.io.WritableResource;
import org.springframework.util.Assert;

import com.redislabs.riot.file.FlatResourceItemWriterBuilder;
import com.redislabs.riot.file.JsonResourceItemWriterBuilder;

import picocli.CommandLine.Option;

public class FileExportOptions extends FileOptions {

	@Option(names = "--append", description = "Append to file if it exists")
	private boolean append;
	@Option(names = "--force-sync", description = "Force-sync changes to disk on flush")
	private boolean forceSync;
	@Option(names = "--line-sep", description = "String to separate lines (default: system default)", paramLabel = "")
	private String lineSeparator = FlatFileItemWriter.DEFAULT_LINE_SEPARATOR;
	@Option(names = "--format", description = "Format string used to aggregate items", paramLabel = "")
	private String format;
	@Option(names = "--locale", description = "Locale", paramLabel = "")
	private Locale locale = Locale.ENGLISH;
	@Option(names = "--max-length", description = "Max length of the formatted string", paramLabel = "")
	private Integer maxLength;
	@Option(names = "--min-length", description = "Min length of the formatted string", paramLabel = "")
	private Integer minLength;

	public ItemWriter> writer() throws IOException {
		WritableResource resource = getResource();
		switch (getType()) {
		case JSON:
			return jsonWriter(resource);
		case FIXED:
			return formattedWriter(resource);
		default:
			return delimitedWriter(resource);
		}
	}

	protected WritableResource getResource() throws IOException {
		Resource resource = getOutputResource();
		Assert.isInstanceOf(WritableResource.class, resource);
		WritableResource writable = (WritableResource) resource;
		if (isGzip()) {
			return new GZIPOutputStreamResource(writable.getOutputStream(), writable.getDescription());
		}
		return writable;
	}

	private FlatResourceItemWriterBuilder> flatWriterBuilder(Resource resource, String headerLine) {
		FlatResourceItemWriterBuilder> builder = new FlatResourceItemWriterBuilder<>();
		builder.append(append);
		builder.encoding(getEncoding());
		builder.lineSeparator(lineSeparator);
		builder.resource(resource);
		builder.saveState(false);
		if (headerLine != null) {
			builder.headerCallback(new FlatFileHeaderCallback() {

				@Override
				public void writeHeader(Writer writer) throws IOException {
					writer.write(headerLine);
				}
			});
		}
		return builder;
	}

	private AbstractItemStreamItemWriter> jsonWriter(Resource resource) {
		JsonResourceItemWriterBuilder> builder = new JsonResourceItemWriterBuilder<>();
		builder.name("json-s3-writer");
		builder.append(append);
		builder.encoding(getEncoding());
		builder.jsonObjectMarshaller(new JacksonJsonObjectMarshaller<>());
		builder.lineSeparator(lineSeparator);
		builder.resource(resource);
		builder.saveState(false);
		return builder.build();
	}

	private AbstractItemStreamItemWriter> delimitedWriter(Resource resource) {
		String headerLine = null;
		if (isHeader()) {
			headerLine = String.join(getDelimiter(), getNames());
		}
		FlatResourceItemWriterBuilder> builder = flatWriterBuilder(resource, headerLine);
		builder.name("delimited-s3-writer");
		com.redislabs.riot.file.FlatResourceItemWriterBuilder.DelimitedBuilder> delimited = builder
				.delimited();
		delimited.delimiter(getDelimiter());
		delimited.fieldExtractor(new MapFieldExtractor(getNames()));
		if (getNames().length > 0) {
			delimited.names(getNames());
		}
		return builder.build();
	}

	private AbstractItemStreamItemWriter> formattedWriter(Resource resource) {
		String headerLine = null;
		if (isHeader()) {
			headerLine = String.format(locale, format, Arrays.asList(getNames()).toArray());
		}
		FlatResourceItemWriterBuilder> builder = flatWriterBuilder(resource, headerLine);
		FlatResourceItemWriterBuilder.FormattedBuilder> formatted = builder.formatted();
		builder.name("formatted-s3-writer");
		formatted.fieldExtractor(new MapFieldExtractor(getNames()));
		if (getNames().length > 0) {
			formatted.names(getNames());
		}
		formatted.format(format);
		formatted.locale(locale);
		if (minLength != null) {
			formatted.minimumLength(minLength);
		}
		if (maxLength != null) {
			formatted.maximumLength(maxLength);
		}
		return builder.build();
	}

}