Add comma escape character (#104274)

astefan · web-flow · commit 1f44cc5df914 · 2024-01-16T16:50:36.000+02:00
* Add escape character for comma in CsvTestsDataLoader
* Handle escaped commas in .csv-spec values as well
* Change the way the indices are populated when running the loader from command line
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/CsvTestUtils.java b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/CsvTestUtils.java
@@ -61,6 +61,9 @@ public final class CsvTestUtils {
     private static final int MAX_WIDTH = 20;
     private static final CsvPreference CSV_SPEC_PREFERENCES = new CsvPreference.Builder('"', '|', "\r\n").build();
     private static final String NULL_VALUE = "null";
+    private static final char ESCAPE_CHAR = '\\';
+    public static final String COMMA_ESCAPING_REGEX = "(?<!\\" + ESCAPE_CHAR + "),";
+    public static final String ESCAPED_COMMA_SEQUENCE = ESCAPE_CHAR + ",";
 
     private CsvTestUtils() {}
 
@@ -229,6 +232,8 @@ public void close() {
      * Takes a csv String and converts it to a String array. Also, it recognizes an opening bracket "[" in one string and a closing "]"
      * in another string and it creates a single concatenated comma-separated String of all the values between the opening bracket entry
      * and the closing bracket entry. In other words, entries enclosed by "[]" are returned as a single element.
+     *
+     * Commas can be escaped with \ (backslash) character.
      */
     static String[] multiValuesAwareCsvToStringArray(String csvLine, int lineNumber) {
         var mvCompressedEntries = new ArrayList<String>();
@@ -237,14 +242,20 @@ static String[] multiValuesAwareCsvToStringArray(String csvLine, int lineNumber)
 
         int pos = 0;          // current position in the csv String
         int commaPos;         // current "," character position
+        int previousCommaPos = 0;
         while ((commaPos = csvLine.indexOf(",", pos)) != -1 || pos <= csvLine.length()) {
+            if (commaPos > 0 && csvLine.charAt(commaPos - 1) == ESCAPE_CHAR) {// skip the escaped comma
+                pos = commaPos + 1;// moving on to the next character after comma
+                continue;
+            }
+
             boolean isLastElement = commaPos == -1;
-            String entry = csvLine.substring(pos, isLastElement ? csvLine.length() : commaPos).trim();
+            String entry = csvLine.substring(previousCommaPos, isLastElement ? csvLine.length() : commaPos).trim();
             if (entry.startsWith("[")) {
                 if (previousMvValue != null || (isLastElement && entry.endsWith("]") == false)) {
                     String message = "Error line [{}:{}]: Unexpected start of a multi-value field value; current token [{}], "
                         + (isLastElement ? "no closing point" : "previous token [{}]");
-                    throw new IllegalArgumentException(format(message, lineNumber, pos, entry, previousMvValue));
+                    throw new IllegalArgumentException(format(message, lineNumber, previousCommaPos, entry, previousMvValue));
                 }
                 if (entry.endsWith("]")) {
                     if (entry.length() > 2) {// single-valued multivalue field :shrug:
@@ -263,7 +274,7 @@ static String[] multiValuesAwareCsvToStringArray(String csvLine, int lineNumber)
                         format(
                             "Error line [{}:{}]: Unexpected end of a multi-value field value (no previous starting point); found [{}]",
                             lineNumber,
-                            pos,
+                            previousCommaPos,
                             entry
                         )
                     );
@@ -279,8 +290,8 @@ static String[] multiValuesAwareCsvToStringArray(String csvLine, int lineNumber)
                             format(
                                 "Error line [{}:{}]: Unexpected missing value in a multi-value column; found [{}]",
                                 lineNumber,
-                                pos,
-                                csvLine.substring(pos - 1)
+                                previousCommaPos,
+                                csvLine.substring(previousCommaPos - 1)
                             )
                         );
                     }
@@ -290,12 +301,22 @@ static String[] multiValuesAwareCsvToStringArray(String csvLine, int lineNumber)
                 }
             }
             pos = 1 + (isLastElement ? csvLine.length() : commaPos);// break out of the loop if it reached its last element
+            previousCommaPos = pos;
         }
         return mvCompressedEntries.toArray(String[]::new);
     }
 
     public record ExpectedResults(List<String> columnNames, List<Type> columnTypes, List<List<Object>> values) {}
 
+    /**
+     * The method loads a section of a .csv-spec file representing the results of executing the query of that section.
+     * It reads both the schema (field names and their types) and the row values.
+     * Values starting with an opening square bracket and ending with a closing square bracket are considered multi-values. Inside
+     * these multi-values, commas separate the individual values and escaped commas are allowed with a prefixed \
+     * default \ (backslash) character.
+     * @param csv a string representing the header and row values of a single query execution result
+     * @return data structure with column names, their types and values
+     */
     public static ExpectedResults loadCsvSpecValues(String csv) {
         List<String> columnNames;
         List<Type> columnTypes;
@@ -338,13 +359,21 @@ public static ExpectedResults loadCsvSpecValues(String csv) {
                     if (value.startsWith("[") ^ value.endsWith("]")) {
                         throw new IllegalArgumentException("Incomplete multi-value (opening and closing square brackets) found " + value);
                     }
-                    if (value.contains(",") && value.startsWith("[")) {// commas outside a multi-value should be ok
-                        List<Object> listOfMvValues = new ArrayList<>();
-                        for (String mvValue : delimitedListToStringArray(value.substring(1, value.length() - 1), ",")) {
-                            listOfMvValues.add(columnTypes.get(i).convert(mvValue.trim()));
+                    if (value.contains(",") && value.startsWith("[")) {
+                        // split on commas but ignoring escaped commas
+                        String[] multiValues = value.substring(1, value.length() - 1).split(COMMA_ESCAPING_REGEX);
+                        if (multiValues.length > 0) {
+                            List<Object> listOfMvValues = new ArrayList<>();
+                            for (String mvValue : multiValues) {
+                                listOfMvValues.add(columnTypes.get(i).convert(mvValue.trim().replace(ESCAPED_COMMA_SEQUENCE, ",")));
+                            }
+                            rowValues.add(listOfMvValues);
+                        } else {
+                            rowValues.add(columnTypes.get(i).convert(value.replace(ESCAPED_COMMA_SEQUENCE, ",")));
                         }
-                        rowValues.add(listOfMvValues);
                     } else {
+                        // The value considered here is the one where any potential escaped comma is kept as is (with the escape char)
+                        // TODO if we'd want escaped commas outside multi-values fields, we'd have to adjust this value here as well
                         rowValues.add(columnTypes.get(i).convert(value));
                     }
                 }
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/CsvTestsDataLoader.java b/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/CsvTestsDataLoader.java
@@ -42,8 +42,9 @@
 import java.util.Map;
 import java.util.Set;
 
-import static org.elasticsearch.common.Strings.delimitedListToStringArray;
 import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
+import static org.elasticsearch.xpack.esql.CsvTestUtils.COMMA_ESCAPING_REGEX;
+import static org.elasticsearch.xpack.esql.CsvTestUtils.ESCAPED_COMMA_SEQUENCE;
 import static org.elasticsearch.xpack.esql.CsvTestUtils.multiValuesAwareCsvToStringArray;
 
 public class CsvTestsDataLoader {
@@ -137,17 +138,33 @@ public static void main(String[] args) throws IOException {
         }
 
         try (RestClient client = builder.build()) {
-            loadDataSetIntoEs(client);
+            loadDataSetIntoEs(client, (restClient, indexName, indexMapping) -> {
+                Request request = new Request("PUT", "/" + indexName);
+                request.setJsonEntity("{\"mappings\":" + indexMapping + "}");
+                restClient.performRequest(request);
+            });
         }
     }
 
+    private static void loadDataSetIntoEs(RestClient client, IndexCreator indexCreator) throws IOException {
+        loadDataSetIntoEs(client, LogManager.getLogger(CsvTestsDataLoader.class), indexCreator);
+    }
+
     public static void loadDataSetIntoEs(RestClient client) throws IOException {
-        loadDataSetIntoEs(client, LogManager.getLogger(CsvTestsDataLoader.class));
+        loadDataSetIntoEs(client, (restClient, indexName, indexMapping) -> {
+            ESRestTestCase.createIndex(restClient, indexName, null, indexMapping, null);
+        });
     }
 
     public static void loadDataSetIntoEs(RestClient client, Logger logger) throws IOException {
+        loadDataSetIntoEs(client, logger, (restClient, indexName, indexMapping) -> {
+            ESRestTestCase.createIndex(restClient, indexName, null, indexMapping, null);
+        });
+    }
+
+    private static void loadDataSetIntoEs(RestClient client, Logger logger, IndexCreator indexCreator) throws IOException {
         for (var dataSet : CSV_DATASET_MAP.values()) {
-            load(client, dataSet.indexName, "/" + dataSet.mappingFileName, "/" + dataSet.dataFileName, logger);
+            load(client, dataSet.indexName, "/" + dataSet.mappingFileName, "/" + dataSet.dataFileName, logger, indexCreator);
         }
         forceMerge(client, CSV_DATASET_MAP.keySet(), logger);
         for (var policy : ENRICH_POLICIES) {
@@ -169,7 +186,14 @@ private static void loadEnrichPolicy(RestClient client, String policyName, Strin
         client.performRequest(request);
     }
 
-    private static void load(RestClient client, String indexName, String mappingName, String dataName, Logger logger) throws IOException {
+    private static void load(
+        RestClient client,
+        String indexName,
+        String mappingName,
+        String dataName,
+        Logger logger,
+        IndexCreator indexCreator
+    ) throws IOException {
         URL mapping = CsvTestsDataLoader.class.getResource(mappingName);
         if (mapping == null) {
             throw new IllegalArgumentException("Cannot find resource " + mappingName);
@@ -178,14 +202,10 @@ private static void load(RestClient client, String indexName, String mappingName
         if (data == null) {
             throw new IllegalArgumentException("Cannot find resource " + dataName);
         }
-        createTestIndex(client, indexName, readTextFile(mapping));
+        indexCreator.createIndex(client, indexName, readTextFile(mapping));
         loadCsvData(client, indexName, data, CsvTestsDataLoader::createParser, logger);
     }
 
-    private static void createTestIndex(RestClient client, String indexName, String mapping) throws IOException {
-        ESRestTestCase.createIndex(client, indexName, null, mapping, null);
-    }
-
     public static String readTextFile(URL resource) throws IOException {
         try (BufferedReader reader = TestUtils.reader(resource)) {
             StringBuilder b = new StringBuilder();
@@ -198,6 +218,20 @@ public static String readTextFile(URL resource) throws IOException {
     }
 
     @SuppressWarnings("unchecked")
+    /**
+     * Loads a classic csv file in an ES cluster using a RestClient.
+     * The structure of the file is as follows:
+     * - commented lines should start with "//"
+     * - the first non-comment line from the file is the schema line (comma separated field_name:ES_data_type elements)
+     *   - sub-fields should be placed after the root field using a dot notation for the name:
+     *       root_field:long,root_field.sub_field:integer
+     *   - a special _id field can be used in the schema and the values of this field will be used in the bulk request as actual doc ids
+     * - all subsequent non-comment lines represent the values that will be used to build the _bulk request
+     * - an empty string "" refers to a null value
+     * - a value starting with an opening square bracket "[" and ending with a closing square bracket "]" refers to a multi-value field
+     *   - multi-values are comma separated
+     *   - commas inside multivalue fields can be escaped with \ (backslash) character
+     */
     private static void loadCsvData(
         RestClient client,
         String indexName,
@@ -278,9 +312,11 @@ private static void loadCsvData(
                                     if (i > 0 && row.length() > 0) {
                                         row.append(",");
                                     }
-                                    if (entries[i].contains(",")) {// multi-value
+                                    // split on comma ignoring escaped commas
+                                    String[] multiValues = entries[i].split(COMMA_ESCAPING_REGEX);
+                                    if (multiValues.length > 0) {// multi-value
                                         StringBuilder rowStringValue = new StringBuilder("[");
-                                        for (String s : delimitedListToStringArray(entries[i], ",")) {
+                                        for (String s : multiValues) {
                                             rowStringValue.append("\"" + s + "\",");
                                         }
                                         // remove the last comma and put a closing bracket instead
@@ -289,6 +325,8 @@ private static void loadCsvData(
                                     } else {
                                         entries[i] = "\"" + entries[i] + "\"";
                                     }
+                                    // replace any escaped commas with single comma
+                                    entries[i].replace(ESCAPED_COMMA_SEQUENCE, ",");
                                     row.append("\"" + columns[i] + "\":" + entries[i]);
                                 } catch (Exception e) {
                                     throw new IllegalArgumentException(
@@ -356,4 +394,8 @@ private static XContentParser createParser(XContent xContent, InputStream data)
     public record TestsDataset(String indexName, String mappingFileName, String dataFileName) {}
 
     public record EnrichConfig(String policyName, String policyFileName) {}
+
+    private interface IndexCreator {
+        void createIndex(RestClient client, String indexName, String mapping) throws IOException;
+    }
 }
