Skip to content
Permalink
Browse files
better header handling for CSVFormatFactory
  • Loading branch information
Romain Manni-Bucau committed Nov 24, 2015
1 parent ef9bfa9 commit e915a6c8cd00c2253b8338f9aebd2cb41a765be0
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 3 deletions.
@@ -17,8 +17,15 @@
package org.apache.batchee.csv;

import org.apache.commons.csv.CSVFormat;
import org.apache.commons.csv.CSVRecord;
import org.apache.commons.csv.QuoteMode;

import java.io.IOException;
import java.io.StringReader;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

class CSVFormatFactory {
private CSVFormatFactory() {
// no-op
@@ -75,7 +82,17 @@ static CSVFormat newFormat(final String format,
out = out.withHeader();
}
if (header != null && !header.trim().isEmpty()) {
out = out.withHeader(header.split(" *, *"));
try { // headers can have CSV header names so parse it there
final Iterator<CSVRecord> iterator = out.withHeader(new String[0]).parse(new StringReader(header + '\n' + header)).iterator();
final CSVRecord record = iterator.next();
final List<String> list = new ArrayList<String>(record.size());
for (final String h : record) {
list.add(h);
}
out = out.withHeader(list.toArray(new String[record.size()]));
} catch (final IOException e) { // can't occur actually
out = out.withHeader(header.split(" *, *"));
}
}
if (skipHeaderRecord != null) {
out = out.withSkipHeaderRecord(Boolean.parseBoolean(skipHeaderRecord));
@@ -122,7 +122,7 @@ public void open(final Serializable checkpoint) throws Exception {
(defaultMapper != null ? new BeanLocator.LocatorInstance<CsvWriterMapper>(defaultMapper, null) : null) :
BeanLocator.Finder.get(locator).newInstance(CsvWriterMapper.class, mapper);

if ((header == null || header.isEmpty()) && Boolean.parseBoolean(writeHeaders) && DefaultMapper.class.isInstance(mapperInstance.getValue())) {
if ((header == null || header.isEmpty()) && Boolean.parseBoolean(writeHeaders) && DefaultMapper.class.isInstance(mapperInstance.getValue()) && checkpoint == null) {
header = toListString(DefaultMapper.class.cast(mapperInstance.getValue()).getHeaders());
}
final CSVFormat format = newFormat();
@@ -138,7 +138,15 @@ public void open(final Serializable checkpoint) throws Exception {
private String toListString(final Iterable<String> headers) {
final StringBuilder b = new StringBuilder();
for (final String s : headers) {
b.append(s).append(",");
final boolean quote = s.contains(",");
if (quote) {
b.append('"');
}
b.append(s);
if (quote) {
b.append('"');
}
b.append(",");
}
if (b.length() > 0) {
b.setLength(b.length() - 1);

0 comments on commit e915a6c

Please sign in to comment.