Added escape string test

kuzudb · Aug 3, 2023 · d812ea4 · d812ea4
1 parent e962c36
commit d812ea4
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 39 deletions.
diff --git a/src/include/processor/operator/copy_to/csv_file_writer.h b/src/include/processor/operator/copy_to/csv_file_writer.h
@@ -20,8 +20,8 @@ class CSVFileWriter {
     std::unique_ptr<common::FileInfo> fileInfo;
 
     void escapeString(std::string& value);
+    void writeValue(common::ValueVector* vector);
     void flush();
-    void writeValue(common::ValueVector* vector, int64_t pos);
 
     template<typename T>
     void writeToBuffer(common::ValueVector* vector, int64_t pos, bool escapeStringValue = false);

diff --git a/src/processor/operator/copy_to/csv_file_writer.cpp b/src/processor/operator/copy_to/csv_file_writer.cpp
@@ -10,23 +10,6 @@ using namespace kuzu::common;
 namespace kuzu {
 namespace processor {
 
-template<typename T>
-void CSVFileWriter::writeListToBuffer(common::ValueVector* vector, int64_t pos) {
-    auto value = TypeUtils::toString(vector->getValue<T>(pos), vector);
-    escapeString(value);
-    writeToBuffer(value);
-}
-
-template<typename T>
-void CSVFileWriter::writeToBuffer(
-    common::ValueVector* vector, int64_t pos, bool escapeStringValue) {
-    auto value = TypeUtils::toString(vector->getValue<T>(pos));
-    if (escapeStringValue) {
-        escapeString(value);
-    }
-    writeToBuffer(value);
-}
-
 void CSVFileWriter::open(const std::string& filePath) {
     fileInfo = FileUtils::openFile(filePath, O_WRONLY | O_CREAT | O_TRUNC);
 }
@@ -44,20 +27,47 @@ void CSVFileWriter::writeHeader(const std::vector<std::string>& columnNames) {
     flush();
 }
 
-void CSVFileWriter::flush() {
-    const std::string str = buffer.str();
-    FileUtils::writeToFile(fileInfo.get(), (uint8_t*)str.data(), str.size(), fileOffset);
-    fileOffset += str.size();
-    buffer.str("");
+void CSVFileWriter::writeValues(std::vector<common::ValueVector*>& outputVectors) {
+    bool hasData = true;
+    if (outputVectors.size() == 0) {
+        return;
+    }
+    auto i = 0u;
+    for (; i < outputVectors.size() - 1; i++) {
+        assert(outputVectors[i]->state->isFlat);
+        writeValue(outputVectors[i]);
+        writeToBuffer(CopyConstants::DEFAULT_CSV_DELIMITER);
+    }
+    writeValue(outputVectors[i]);
+    writeToBuffer(CopyConstants::DEFAULT_CSV_LINE_BREAK);
+    flush();
+}
+
+template<typename T>
+void CSVFileWriter::writeToBuffer(
+    common::ValueVector* vector, int64_t pos, bool escapeStringValue) {
+    auto value = TypeUtils::toString(vector->getValue<T>(pos));
+    if (escapeStringValue) {
+        escapeString(value);
+    }
+    writeToBuffer(value);
+}
+
+template<typename T>
+void CSVFileWriter::writeListToBuffer(common::ValueVector* vector, int64_t pos) {
+    auto value = TypeUtils::toString(vector->getValue<T>(pos), vector);
+    escapeString(value);
+    writeToBuffer(value);
 }
 
 void CSVFileWriter::escapeString(std::string& value) {
     StringUtils::replaceAll(value, "\"", "\"\"");
     value = "\"" + value + "\"";
 }
 
-void CSVFileWriter::writeValue(common::ValueVector* vector, int64_t pos) {
-    auto selPos = vector->state->selVector->selectedPositions[pos];
+void CSVFileWriter::writeValue(common::ValueVector* vector) {
+    // vectors are always flat
+    auto selPos = vector->state->selVector->selectedPositions[0];
     switch (vector->dataType.getLogicalTypeID()) {
     case LogicalTypeID::BOOL:
         return writeToBuffer<int8_t>(vector, selPos);
@@ -92,20 +102,11 @@ void CSVFileWriter::writeValue(common::ValueVector* vector, int64_t pos) {
     }
 }
 
-void CSVFileWriter::writeValues(std::vector<common::ValueVector*>& outputVectors) {
-    bool hasData = true;
-    if (outputVectors.size() == 0) {
-        return;
-    }
-    auto i = 0u;
-    for (; i < outputVectors.size() - 1; i++) {
-        assert(outputVectors[i]->state->isFlat);
-        writeValue(outputVectors[i], 0);
-        writeToBuffer(CopyConstants::DEFAULT_CSV_DELIMITER);
-    }
-    writeValue(outputVectors[i], 0);
-    writeToBuffer(CopyConstants::DEFAULT_CSV_LINE_BREAK);
-    flush();
+void CSVFileWriter::flush() {
+    const std::string str = buffer.str();
+    FileUtils::writeToFile(fileInfo.get(), (uint8_t*)str.data(), str.size(), fileOffset);
+    fileOffset += str.size();
+    buffer.str("");
 }
 
 } // namespace processor

diff --git a/test/test_files/copy/copy_to.test b/test/test_files/copy/copy_to.test
@@ -55,6 +55,18 @@ Cannot open file: a/b/c/d.csv
 1.731000
 MAX(p.height)
 
+-CASE StringEscapeCopyTo
+
+-STATEMENT COPY (RETURN 100,'a string with "quotes"',5.6,'","') TO "${DATABASE_PATH}/string.csv"
+---- ok
+-STATEMENT CREATE NODE TABLE stringTable (number INT, str STRING, frac DOUBLE, str2 STRING, PRIMARY KEY (str))
+---- ok
+-STATEMENT COPY stringTable FROM '${DATABASE_PATH}/string.csv' (HEADER=TRUE)
+---- ok
+-STATEMENT MATCH (s:stringTable) RETURN s.number, s.str, s.frac, s.str2
+---- 1
+100|a string with "quotes"|5.600000|","
+
 -CASE CountResultCopyTo
 
 -STATEMENT COPY (MATCH (p:person) RETURN COUNT(*)) TO "${DATABASE_PATH}/count.csv"