Add _routing to SQL includes list

Signed-off-by: Andrew Carbonetto <[email protected]>
Bit-Quill · Jun 16, 2023 · 1fd3a57 · 1fd3a57
1 parent 94d5479
commit 1fd3a57
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 6 deletions.
diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/IdentifierIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/IdentifierIT.java
@@ -14,6 +14,8 @@
 import static org.opensearch.sql.util.TestUtils.performRequest;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import org.json.JSONArray;
 import org.json.JSONObject;
 import org.junit.jupiter.api.Test;
 import org.opensearch.client.Request;
@@ -99,6 +101,31 @@ public void testMetafieldIdentifierTest() throws IOException {
     verifyDataRows(result, rows(30, id, index, 1.0, 1.0, -2));
   }
 
+  @Test
+  public void testMetafieldIdentifierRoutingTest() throws IOException {
+    // create an index, but the contents doesn't matter
+    String id = "12345";
+    String index = "test.metafields";
+    new Index(index).addDoc("{\"age\": 30}", id);
+
+    // Execute using field metadata values
+    final JSONObject result = new JSONObject(executeQuery(
+        "SELECT _id, _index, _routing "
+            + "FROM " + index,
+        "jdbc"));
+
+    // Verify that the metadata values are returned when requested
+    verifySchema(result,
+        schema("_id", null, "keyword"),
+        schema("_index", null, "keyword"),
+        schema("_routing", null, "keyword"));
+    assertTrue(result.getJSONArray("schema").length() == 3);
+
+    // routing has the format: [thread_id][index][node] - where thread_id and node may be variable
+    // per run
+    assertTrue(result.getJSONArray("datarows").get(0).toString().contains("[" + index + "]"));
+  }
+
   @Test
   public void testMetafieldIdentifierWithAliasTest() throws IOException {
     // create an index, but the contents doesn't matter

diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/response/OpenSearchResponse.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/response/OpenSearchResponse.java
@@ -10,16 +10,15 @@
 import static org.opensearch.sql.opensearch.storage.OpenSearchIndex.METADATA_FIELD_ID;
 import static org.opensearch.sql.opensearch.storage.OpenSearchIndex.METADATA_FIELD_INDEX;
 import static org.opensearch.sql.opensearch.storage.OpenSearchIndex.METADATA_FIELD_MAXSCORE;
+import static org.opensearch.sql.opensearch.storage.OpenSearchIndex.METADATA_FIELD_ROUTING;
 import static org.opensearch.sql.opensearch.storage.OpenSearchIndex.METADATA_FIELD_SCORE;
 import static org.opensearch.sql.opensearch.storage.OpenSearchIndex.METADATA_FIELD_SORT;
 
 import com.google.common.collect.ImmutableMap;
 import java.util.Arrays;
-import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.stream.Collectors;
 import lombok.EqualsAndHashCode;
 import lombok.ToString;
@@ -146,8 +145,10 @@ public Iterator<ExprValue> iterator() {
                 if (maxScore != null) {
                   builder.put(METADATA_FIELD_MAXSCORE, maxScore);
                 }
-              } else { // if (metaDataField.equals(METADATA_FIELD_SORT)) {
+              } else if (metaDataField.equals(METADATA_FIELD_SORT)) {
                 builder.put(METADATA_FIELD_SORT, new ExprLongValue(hit.getSeqNo()));
+              } else { // if (metaDataField.equals(METADATA_FIELD_ROUTING)){
+                builder.put(METADATA_FIELD_ROUTING, new ExprStringValue(hit.getShard().toString()));
               }
             });
 

diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java
@@ -45,12 +45,15 @@ public class OpenSearchIndex implements Table {
   public static final String METADATA_FIELD_MAXSCORE = "_maxscore";
   public static final String METADATA_FIELD_SORT = "_sort";
 
+  public static final String METADATA_FIELD_ROUTING = "_routing";
+
   public static final java.util.Map<String, ExprType> METADATAFIELD_TYPE_MAP = Map.of(
       METADATA_FIELD_ID, ExprCoreType.STRING,
       METADATA_FIELD_INDEX, ExprCoreType.STRING,
       METADATA_FIELD_SCORE, ExprCoreType.FLOAT,
       METADATA_FIELD_MAXSCORE, ExprCoreType.FLOAT,
-      METADATA_FIELD_SORT, ExprCoreType.LONG
+      METADATA_FIELD_SORT, ExprCoreType.LONG,
+      METADATA_FIELD_ROUTING, ExprCoreType.STRING
   );
 
   /** OpenSearch client connection. */

diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/response/OpenSearchResponseTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/response/OpenSearchResponseTest.java
@@ -28,8 +28,10 @@
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.common.bytes.BytesArray;
 import org.opensearch.common.text.Text;
+import org.opensearch.index.shard.ShardId;
 import org.opensearch.search.SearchHit;
 import org.opensearch.search.SearchHits;
+import org.opensearch.search.SearchShardTarget;
 import org.opensearch.search.aggregations.Aggregations;
 import org.opensearch.search.fetch.subphase.highlight.HighlightField;
 import org.opensearch.sql.data.model.ExprFloatValue;
@@ -143,9 +145,13 @@ void iterator_metafields() {
                 new TotalHits(1L, TotalHits.Relation.EQUAL_TO),
                 3.75F));
 
+    ShardId shardId = new ShardId("index", "indexUUID", 42);
+    SearchShardTarget shardTarget = new SearchShardTarget("node", shardId, null, null);
+
     when(searchHit1.getSourceAsString()).thenReturn("{\"id1\", 1}");
     when(searchHit1.getId()).thenReturn("testId");
     when(searchHit1.getIndex()).thenReturn("testIndex");
+    when(searchHit1.getShard()).thenReturn(shardTarget);
     when(searchHit1.getScore()).thenReturn(3.75F);
     when(searchHit1.getSeqNo()).thenReturn(123456L);
 
@@ -155,11 +161,12 @@ void iterator_metafields() {
         "id1", new ExprIntegerValue(1),
         "_index", new ExprStringValue("testIndex"),
         "_id", new ExprStringValue("testId"),
+        "_routing", new ExprStringValue(shardTarget.toString()),
         "_sort", new ExprLongValue(123456L),
         "_score", new ExprFloatValue(3.75F),
         "_maxscore", new ExprFloatValue(3.75F)
     ));
-    List includes = List.of("id1", "_index", "_id", "_sort", "_score", "_maxscore");
+    List includes = List.of("id1", "_index", "_id", "_routing", "_sort", "_score", "_maxscore");
     int i = 0;
     for (ExprValue hit : new OpenSearchResponse(searchResponse, factory, includes)) {
       if (i == 0) {

diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java
@@ -187,9 +187,10 @@ void getReservedFieldTypes() {
     assertThat(
         fieldTypes,
         allOf(
-            aMapWithSize(5),
+            aMapWithSize(6),
             hasEntry("_id", ExprCoreType.STRING),
             hasEntry("_index", ExprCoreType.STRING),
+            hasEntry("_routing", ExprCoreType.STRING),
             hasEntry("_sort", ExprCoreType.LONG),
             hasEntry("_score", ExprCoreType.FLOAT),
             hasEntry("_maxscore", ExprCoreType.FLOAT)