Skip to content

Commit 4bfb4b0

Browse files
Backport to branch(3.16) : Fix --no-header option not working for CSV exports (#3060)
Co-authored-by: inv-jishnu <31100916+inv-jishnu@users.noreply.github.com>
1 parent 042b36d commit 4bfb4b0

File tree

2 files changed

+74
-3
lines changed

2 files changed

+74
-3
lines changed

data-loader/core/src/main/java/com/scalar/db/dataloader/core/dataexport/CsvExportManager.java

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -38,9 +38,11 @@ public CsvExportManager(
3838
@Override
3939
void processHeader(ExportOptions exportOptions, TableMetadata tableMetadata, Writer writer)
4040
throws IOException {
41-
String header = createCsvHeaderRow(exportOptions, tableMetadata);
42-
writer.append(header);
43-
writer.flush();
41+
if (!exportOptions.isExcludeHeaderRow()) {
42+
String header = createCsvHeaderRow(exportOptions, tableMetadata);
43+
writer.append(header);
44+
writer.flush();
45+
}
4446
}
4547

4648
/**

data-loader/core/src/test/java/com/scalar/db/dataloader/core/dataexport/CsvExportManagerTest.java

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import com.scalar.db.io.Column;
1515
import com.scalar.db.io.IntColumn;
1616
import com.scalar.db.io.Key;
17+
import java.io.BufferedReader;
1718
import java.io.BufferedWriter;
1819
import java.io.File;
1920
import java.io.IOException;
@@ -129,4 +130,72 @@ void startExport_givenPartitionKey_shouldGenerateOutputFile()
129130
Assertions.assertTrue(file.exists());
130131
Assertions.assertTrue(file.delete());
131132
}
133+
134+
@Test
135+
void startExport_givenNoHeaderRequired_shouldGenerateOutputFileWithoutHeader() throws Exception {
136+
String expectedFirstLine =
137+
"9007199254740992,2147483647,true,0.000000000000000000000000000000000000000000001401298464324817,0.0000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000049,test value,YmxvYiB0ZXN0IHZhbHVl,2000-01-01,01:01:01,2000-01-01T01:01,1970-01-21T03:20:41.740Z";
138+
139+
runExportAndAssertFirstLine(true, expectedFirstLine);
140+
}
141+
142+
@Test
143+
void startExport_givenHeaderRequired_shouldGenerateOutputFileWithHeader() throws Exception {
144+
String expectedFirstLine = "col1,col2,col3,col4,col5,col6,col7,col8,col9,col10,col11";
145+
146+
runExportAndAssertFirstLine(false, expectedFirstLine);
147+
}
148+
149+
private void runExportAndAssertFirstLine(boolean excludeHeader, String expectedFirstLine)
150+
throws Exception {
151+
// Arrange
152+
producerTaskFactory = new ProducerTaskFactory(",", false, false);
153+
exportManager = new CsvExportManager(storage, dao, producerTaskFactory);
154+
Scanner scanner = Mockito.mock(Scanner.class);
155+
String filePath = Paths.get("").toAbsolutePath() + "/output.csv";
156+
Map<String, Column<?>> values = UnitTestUtils.createTestValues();
157+
Result result = new ResultImpl(values, mockData);
158+
List<Result> results = Collections.singletonList(result);
159+
160+
ExportOptions exportOptions =
161+
ExportOptions.builder(
162+
"namespace",
163+
"table",
164+
Key.newBuilder().add(IntColumn.of("col1", 1)).build(),
165+
FileFormat.CSV)
166+
.sortOrders(Collections.emptyList())
167+
.scanRange(new ScanRange(null, null, false, false))
168+
.delimiter(",")
169+
.excludeHeaderRow(excludeHeader)
170+
.build();
171+
172+
Mockito.when(
173+
dao.createScanner(
174+
exportOptions.getNamespace(),
175+
exportOptions.getTableName(),
176+
exportOptions.getScanPartitionKey(),
177+
exportOptions.getScanRange(),
178+
exportOptions.getSortOrders(),
179+
exportOptions.getProjectionColumns(),
180+
exportOptions.getLimit(),
181+
storage))
182+
.thenReturn(scanner);
183+
Mockito.when(scanner.iterator()).thenReturn(results.iterator());
184+
try (BufferedWriter writer =
185+
Files.newBufferedWriter(
186+
Paths.get(filePath),
187+
Charset.defaultCharset(),
188+
StandardOpenOption.CREATE,
189+
StandardOpenOption.TRUNCATE_EXISTING)) {
190+
exportManager.startExport(exportOptions, mockData, writer);
191+
}
192+
File file = new File(filePath);
193+
Assertions.assertTrue(file.exists());
194+
try (BufferedReader br = Files.newBufferedReader(file.toPath(), Charset.defaultCharset())) {
195+
String firstLine = br.readLine();
196+
Assertions.assertEquals(expectedFirstLine, firstLine);
197+
} finally {
198+
Assertions.assertTrue(file.delete());
199+
}
200+
}
132201
}

0 commit comments

Comments
 (0)