From 9480b7fd2a550c32345ec1b71c54b8cefbbac1b0 Mon Sep 17 00:00:00 2001
From: Mike Swierczek <441523+Michael-S@users.noreply.github.com>
Date: Wed, 6 Nov 2024 17:03:54 -0500
Subject: [PATCH] Fix: CSV and Raw output, escape quotes (#3063)

Fixes #3050

Signed-off-by: Mike Swierczek <441523+Michael-S@users.noreply.github.com>
(cherry picked from commit cfe38d7d6678dfbd5597f1098a1a47e29f84a708)
Signed-off-by: Simeon Widdis <sawiddis@amazon.com>
---
 .../protocol/response/format/FlatResponseBase.java |  6 +++++-
 .../response/format/CsvResponseFormatterTest.java  |  8 ++++++++
 .../response/format/RawResponseFormatterTest.java  | 14 +++++++++++---
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseBase.java b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseBase.java
index 98e79a4048..be2517bf6c 100644
--- a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseBase.java
+++ b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseBase.java
@@ -84,6 +84,10 @@ protected List<List<String>> formatData(List<List<String>> lines) {
 
   protected String quoteIfRequired(String separator, String cell) {
     final String quote = "\"";
-    return cell.contains(separator) ? quote + cell.replaceAll("\"", "\"\"") + quote : cell;
+    if (cell != null && (cell.contains(separator) || cell.contains(quote))) {
+      return quote + cell.replaceAll(quote, quote + quote) + quote;
+    } else {
+      return cell;
+    }
   }
 }
diff --git a/protocol/src/test/java/org/opensearch/sql/protocol/response/format/CsvResponseFormatterTest.java b/protocol/src/test/java/org/opensearch/sql/protocol/response/format/CsvResponseFormatterTest.java
index 1436959a9d..e231bcf4f0 100644
--- a/protocol/src/test/java/org/opensearch/sql/protocol/response/format/CsvResponseFormatterTest.java
+++ b/protocol/src/test/java/org/opensearch/sql/protocol/response/format/CsvResponseFormatterTest.java
@@ -108,8 +108,16 @@ void quoteIfRequired() {
     QueryResult response =
         new QueryResult(
             schema,
+<<<<<<< HEAD
             Arrays.asList(tupleValue(ImmutableMap.of("na,me", "John,Smith", ",,age", "30,,,"))));
     String expected = "\"na,me\",\",,age\"%n" + "\"John,Smith\",\"30,,,\"";
+=======
+            Arrays.asList(
+                tupleValue(ImmutableMap.of("na,me", "John,Smith", ",,age", "30,,,")),
+                tupleValue(ImmutableMap.of("na,me", "\"Janice Jones", ",,age", "26\""))));
+    String expected =
+        "\"na,me\",\",,age\"%n\"John,Smith\",\"30,,,\"%n\"\"\"Janice Jones\",\"26\"\"\"";
+>>>>>>> cfe38d7d6 (Fix: CSV and Raw output, escape quotes (#3063))
     assertEquals(format(expected), formatter.format(response));
   }
 
diff --git a/protocol/src/test/java/org/opensearch/sql/protocol/response/format/RawResponseFormatterTest.java b/protocol/src/test/java/org/opensearch/sql/protocol/response/format/RawResponseFormatterTest.java
index fd057437a0..ebdadcd50b 100644
--- a/protocol/src/test/java/org/opensearch/sql/protocol/response/format/RawResponseFormatterTest.java
+++ b/protocol/src/test/java/org/opensearch/sql/protocol/response/format/RawResponseFormatterTest.java
@@ -120,10 +120,18 @@ void quoteIfRequired() {
     QueryResult response =
         new QueryResult(
             schema,
-            Arrays.asList(tupleValue(ImmutableMap.of("na|me", "John|Smith", "||age", "30|||"))));
-    String expected = "\"na|me\"|\"||age\"%n" + "\"John|Smith\"|\"30|||\"";
+            Arrays.asList(
+                tupleValue(ImmutableMap.of("na|me", "John|Smith", "||age", "30|||")),
+                tupleValue(ImmutableMap.of("na|me", "Ja\"ne J\"ones", "||age", "\"40\""))));
+    String expected =
+        "\"na|me\"|\"||age\"%n"
+            + "\"John|Smith\"|\"30|||\"%n"
+            + "\"Ja\"\"ne J\"\"ones\"|\"\"\"40\"\"\"";
     assertEquals(format(expected), getRawFormatter().format(response));
-    String expectedPretty = "\"na|me\"     |\"||age\"%n" + "\"John|Smith\"|\"30|||\"";
+    String expectedPretty =
+        "\"na|me\"         |\"||age\" %n"
+            + "\"John|Smith\"    |\"30|||\" %n"
+            + "\"Ja\"\"ne J\"\"ones\"|\"\"\"40\"\"\"";
     assertEquals(format(expectedPretty), getRawFormatterPretty().format(response));
   }