Parse ES|QL response body header more leniently (#903)

Co-authored-by: Laura Trotta <[email protected]>
elastic · Nov 14, 2024 · d130ce5 · d130ce5
1 parent 7001d64
commit d130ce5
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 14 deletions.
diff --git a/...-client/src/main/java/co/elastic/clients/elasticsearch/_helpers/esql/EsqlAdapterBase.java b/...-client/src/main/java/co/elastic/clients/elasticsearch/_helpers/esql/EsqlAdapterBase.java
@@ -34,26 +34,40 @@ public abstract class EsqlAdapterBase<T> implements EsqlAdapter<T> {
      * The caller can then read row arrays until finding an end array that closes the top-level array.
      */
     public static EsqlMetadata readHeader(JsonParser parser, JsonpMapper mapper) {
+        EsqlMetadata result = new EsqlMetadata();
+
         JsonpUtils.expectNextEvent(parser, JsonParser.Event.START_OBJECT);
-        JsonpUtils.expectNextEvent(parser, JsonParser.Event.KEY_NAME);
 
-        if (!"columns".equals(parser.getString())) {
-            throw new JsonpMappingException("Expecting a 'columns' property, but found '" + parser.getString() + "'", parser.getLocation());
+        parse: while (JsonpUtils.expectNextEvent(parser, JsonParser.Event.KEY_NAME) != null) {
+            switch (parser.getString()) {
+                case "values": {
+                    // We're done parsing header information
+                    break parse;
+                }
+                case "columns": {
+                    result.columns = JsonpDeserializer
+                        .arrayDeserializer(EsqlMetadata.EsqlColumn._DESERIALIZER)
+                        .deserialize(parser, mapper);
+                    break;
+                }
+                case "took": {
+                    JsonpUtils.expectNextEvent(parser, JsonParser.Event.VALUE_NUMBER);
+                    result.took = parser.getLong();
+                    break;
+                }
+                default: {
+                    // Ignore everything else
+                    JsonpUtils.skipValue(parser);
+                    break;
+                }
+            }
         }
 
-        List<EsqlMetadata.EsqlColumn> columns = JsonpDeserializer
-            .arrayDeserializer(EsqlMetadata.EsqlColumn._DESERIALIZER)
-            .deserialize(parser, mapper);
-
-        EsqlMetadata result = new EsqlMetadata();
-        result.columns = columns;
-
-        JsonpUtils.expectNextEvent(parser, JsonParser.Event.KEY_NAME);
-
-        if (!"values".equals(parser.getString())) {
-            throw new JsonpMappingException("Expecting a 'values' property, but found '" + parser.getString() + "'", parser.getLocation());
+        if (result.columns == null) {
+            throw new JsonpMappingException("Expecting a 'columns' property before 'values'.", parser.getLocation());
         }
 
+        // Beginning of the `values` property
         JsonpUtils.expectNextEvent(parser, JsonParser.Event.START_ARRAY);
 
         return result;

diff --git a/java-client/src/main/java/co/elastic/clients/elasticsearch/_helpers/esql/EsqlMetadata.java b/java-client/src/main/java/co/elastic/clients/elasticsearch/_helpers/esql/EsqlMetadata.java
@@ -26,6 +26,7 @@
 import co.elastic.clients.util.ObjectBuilder;
 import co.elastic.clients.util.ObjectBuilderBase;
 
+import javax.annotation.Nullable;
 import java.util.List;
 
 public class EsqlMetadata {
@@ -74,4 +75,7 @@ protected static void setupEsqlColumnDeserializer(ObjectDeserializer<EsqlColumn.
     }
 
     public List<EsqlColumn> columns;
+
+    @Nullable
+    public Long took;
 }
diff --git a/...-client/src/test/java/co/elastic/clients/elasticsearch/_helpers/esql/EsqlAdapterTest.java b/...-client/src/test/java/co/elastic/clients/elasticsearch/_helpers/esql/EsqlAdapterTest.java
@@ -23,6 +23,7 @@
 import co.elastic.clients.elasticsearch._helpers.esql.jdbc.ResultSetEsqlAdapter;
 import co.elastic.clients.elasticsearch._helpers.esql.objects.ObjectsEsqlAdapter;
 import co.elastic.clients.elasticsearch.esql.query.EsqlFormat;
+import co.elastic.clients.json.JsonpMappingException;
 import co.elastic.clients.json.jackson.JacksonJsonpMapper;
 import co.elastic.clients.testkit.MockHttpClient;
 import co.elastic.clients.transport.endpoints.BinaryResponse;
@@ -36,6 +37,7 @@
 public class EsqlAdapterTest extends Assertions {
 
     String json = "{\n" +
+        "  \"took\": 10," +
         "  \"columns\": [\n" +
         "\t{\"name\": \"avg_salary\", \"type\": \"double\"},\n" +
         "\t{\"name\": \"lang\",   \t\"type\": \"keyword\"}\n" +
@@ -56,6 +58,30 @@ public static class Data {
         public String lang;
     }
 
+    @Test
+    public void testMissingColumns() throws IOException {
+        String badJson = "{\n" +
+            "  \"took\": 10," +
+            "  \"values\": [\n" +
+            "\t[43760.0, \"Spanish\"],\n" +
+            "\t[48644.0, \"French\"],\n" +
+            "\t[48832.0, \"German\"]\n" +
+            "  ]\n" +
+            "}\n";
+
+        ElasticsearchClient esClient = new MockHttpClient()
+            .add("/_query", "application/json", badJson)
+            .client(new JacksonJsonpMapper());
+
+        JsonpMappingException jsonMappingException = assertThrows(JsonpMappingException.class, () -> {
+            esClient.esql().query(
+                ResultSetEsqlAdapter.INSTANCE,
+                "FROM employees | STATS avg_salary = AVG(salary) by country"
+            );
+        });
+        assertTrue(jsonMappingException.getMessage().contains("Expecting a 'columns' property"));
+    }
+
     @Test
     public void testObjectDeserializer() throws IOException {