Support multiple indices in PPL and SQL

Signed-off-by: penghuo <[email protected]>

core/src/main/java/org/opensearch/sql/ast/tree/Relation.java

@@ -7,7 +7,9 @@
 package org.opensearch.sql.ast.tree;
 
 import com.google.common.collect.ImmutableList;
+import java.util.Arrays;
 import java.util.List;
+import java.util.stream.Collectors;
 import lombok.AllArgsConstructor;
 import lombok.EqualsAndHashCode;
 import lombok.RequiredArgsConstructor;
@@ -23,7 +25,18 @@
 @EqualsAndHashCode(callSuper = false)
 @RequiredArgsConstructor
 public class Relation extends UnresolvedPlan {
-  private final UnresolvedExpression tableName;
+  private static final String COMMA = ",";
+
+  private final List<UnresolvedExpression> tableName;
+
+  public Relation(UnresolvedExpression tableName) {
+    this(tableName, null);
+  }
+
+  public Relation(UnresolvedExpression tableName, String alias) {
+    this.tableName = Arrays.asList(tableName);
+    this.alias = alias;
+  }
 
   /**
    * Optional alias name for the relation.
@@ -36,7 +49,9 @@ public class Relation extends UnresolvedPlan {
    * @return    table name
    */
   public String getTableName() {
-    return tableName.toString();
+    return tableName.stream()
+        .map(UnresolvedExpression::toString)
+        .collect(Collectors.joining(COMMA));
   }
 
   /**

core/src/test/java/org/opensearch/sql/ast/tree/RelationTest.java

@@ -3,12 +3,12 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-
 package org.opensearch.sql.ast.tree;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.opensearch.sql.ast.dsl.AstDSL.qualifiedName;
 
+import java.util.Arrays;
 import org.junit.jupiter.api.Test;
 
 class RelationTest {
@@ -26,4 +26,9 @@ void should_return_alias_if_aliased() {
     assertEquals("t", relation.getTableNameOrAlias());
   }
 
+  @Test
+  void comma_seperated_index_return_concat_table_names() {
+    Relation relation = new Relation(Arrays.asList(qualifiedName("test1"), qualifiedName("test2")));
+    assertEquals("test1,test2", relation.getTableNameOrAlias());
+  }
 }

docs/user/general/identifiers.rst

@@ -38,7 +38,7 @@ Examples
 
 Here are examples for using index pattern directly without quotes::
 
-    os> SELECT * FROM *cc*nt*;
+    os> SELECT * FROM *cc*nts;
     fetched rows / total rows = 4/4
     +------------------+-------------+----------------------+-----------+----------+--------+------------+---------+-------+-----------------------+------------+
     | account_number   | firstname   | address              | balance   | gender   | city   | employer   | state   | age   | email                 | lastname   |
@@ -140,3 +140,39 @@ The second example is to show a field name qualified by index alias specified. S
     +--------+-------+--------------------+
 
 Note that in both examples above, the qualifier is removed in response. This happens only when identifiers selected is a simple field name. In other cases, expressions rather than an atom field, the column name in response is exactly the same as the text in ``SELECT``clause.
+
+Multiple Indices
+================
+
+Description
+-----------
+
+To query multiple indices, you could
+
+1. Include ``*`` in index name, this is an index pattern for wildcard match.
+2. Delimited multiple indices and seperated them by ``,``. Note: no space allowed between each index.
+
+
+Examples
+---------
+
+Query wildcard indices::
+
+    os> SELECT count(*) as cnt FROM acc*;
+    fetched rows / total rows = 1/1
+    +-------+
+    | cnt   |
+    |-------|
+    | 5     |
+    +-------+
+
+
+Query delimited multiple indices seperated by ``,``::
+
+    os> SELECT count(*) as cnt FROM `accounts,account2`;
+    fetched rows / total rows = 1/1
+    +-------+
+    | cnt   |
+    |-------|
+    | 5     |
+    +-------+

docs/user/ppl/general/identifiers.rst

@@ -94,4 +94,51 @@ Identifiers are treated in case sensitive manner. So it must be exactly same as
 Examples
 --------
 
-For example, if you run ``source=Accounts``, it will end up with an index not found exception from our plugin because the actual index name is under lower case.
+For example, if you run ``source=Accounts``, it will end up with an index not found exception from our plugin because the actual index name is under lower case.
+
+Multiple Indices
+================
+
+Description
+-----------
+
+To query multiple indices, you could
+
+1. Include ``*`` in index name, this is an index pattern for wildcard match.
+2. Include multiple indices and seperated them by ``,``.
+3. Delimited multiple indices and seperated them by ``,``. Note: no space allowed between each index.
+
+
+Examples
+---------
+
+Query wildcard indices::
+
+    os> source=acc* | stats count();
+    fetched rows / total rows = 1/1
+    +-----------+
+    | count()   |
+    |-----------|
+    | 5         |
+    +-----------+
+
+Query multiple indices seperated by ``,``::
+
+    os> source=accounts, account2 | stats count();
+    fetched rows / total rows = 1/1
+    +-----------+
+    | count()   |
+    |-----------|
+    | 5         |
+    +-----------+
+
+Query delimited multiple indices seperated by ``,``::
+
+    os> source=`accounts,account2` | stats count();
+    fetched rows / total rows = 1/1
+    +-----------+
+    | count()   |
+    |-----------|
+    | 5         |
+    +-----------+
+

doctest/test_data/account2.json

@@ -0,0 +1 @@
+{"account_number":2}

doctest/test_docs.py

@@ -22,6 +22,7 @@
 ACCOUNTS = "accounts"
 EMPLOYEES = "employees"
 PEOPLE = "people"
+ACCOUNT2 = "account2"
 
 
 class DocTestConnection(OpenSearchConnection):
@@ -84,6 +85,7 @@ def set_up_test_indices(test):
     set_up(test)
     load_file("accounts.json", index_name=ACCOUNTS)
     load_file("people.json", index_name=PEOPLE)
+    load_file("account2.json", index_name=ACCOUNT2)
 
 
 def load_file(filename, index_name):
@@ -112,7 +114,7 @@ def set_up(test):
 
 def tear_down(test):
     # drop leftover tables after each test
-    test_data_client.indices.delete(index=[ACCOUNTS, EMPLOYEES, PEOPLE], ignore_unavailable=True)
+    test_data_client.indices.delete(index=[ACCOUNTS, EMPLOYEES, PEOPLE, ACCOUNT2], ignore_unavailable=True)
 
 
 docsuite = partial(doctest.DocFileSuite,

opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchClient.java

@@ -27,7 +27,7 @@ public interface OpenSearchClient {
    * @param indexExpression index expression
    * @return index mapping(s) from index name to its mapping
    */
-  Map<String, IndexMapping> getIndexMappings(String indexExpression);
+  Map<String, IndexMapping> getIndexMappings(String... indexExpression);
 
   /**
    * Perform search query in the search request.

opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClient.java

@@ -77,10 +77,10 @@ public OpenSearchNodeClient(ClusterService clusterService,
    *     thrown if no index matched.
    */
   @Override
-  public Map<String, IndexMapping> getIndexMappings(String indexExpression) {
+  public Map<String, IndexMapping> getIndexMappings(String... indexExpression) {
     try {
       ClusterState state = clusterService.state();
-      String[] concreteIndices = resolveIndexExpression(state, new String[] {indexExpression});
+      String[] concreteIndices = resolveIndexExpression(state, indexExpression);
 
       return populateIndexMappings(
           state.metadata().findMappings(concreteIndices, ALL_TYPES, ALL_FIELDS));

opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchRestClient.java

@@ -42,7 +42,7 @@ public class OpenSearchRestClient implements OpenSearchClient {
   private final RestHighLevelClient client;
 
   @Override
-  public Map<String, IndexMapping> getIndexMappings(String indexExpression) {
+  public Map<String, IndexMapping> getIndexMappings(String... indexExpression) {
     GetMappingsRequest request = new GetMappingsRequest().indices(indexExpression);
     try {
       GetMappingsResponse response = client.indices().getMapping(request, RequestOptions.DEFAULT);

opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java

@@ -38,9 +38,9 @@ public class OpenSearchQueryRequest implements OpenSearchRequest {
   public static final TimeValue DEFAULT_QUERY_TIMEOUT = TimeValue.timeValueMinutes(1L);
 
   /**
-   * Index name.
+   * {@link OpenSearchRequest.IndexName}.
    */
-  private final String indexName;
+  private final IndexName indexName;
 
   /**
    * Search request source builder.
@@ -65,6 +65,14 @@ public class OpenSearchQueryRequest implements OpenSearchRequest {
    */
   public OpenSearchQueryRequest(String indexName, int size,
                                 OpenSearchExprValueFactory factory) {
+    this(new IndexName(indexName), size, factory);
+  }
+
+  /**
+   * Constructor of ElasticsearchQueryRequest.
+   */
+  public OpenSearchQueryRequest(IndexName indexName, int size,
+      OpenSearchExprValueFactory factory) {
     this.indexName = indexName;
     this.sourceBuilder = new SearchSourceBuilder();
     sourceBuilder.from(0);
@@ -97,7 +105,7 @@ public void clean(Consumer<String> cleanAction) {
   @VisibleForTesting
   protected SearchRequest searchRequest() {
     return new SearchRequest()
-        .indices(indexName)
+        .indices(indexName.getIndexNames())
         .source(sourceBuilder);
   }
 }

opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java

@@ -8,6 +8,7 @@
 
 import java.util.function.Consumer;
 import java.util.function.Function;
+import lombok.EqualsAndHashCode;
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.action.search.SearchScrollRequest;
@@ -19,7 +20,6 @@
  * OpenSearch search request.
  */
 public interface OpenSearchRequest {
-
   /**
    * Apply the search action or scroll action on request based on context.
    *
@@ -49,4 +49,28 @@ OpenSearchResponse search(Function<SearchRequest, SearchResponse> searchAction,
    * @return ElasticsearchExprValueFactory.
    */
   OpenSearchExprValueFactory getExprValueFactory();
+
+  /**
+   * OpenSearch Index Name.
+   * Indices are seperated by ",".
+   */
+  @EqualsAndHashCode
+  class IndexName {
+    private static final String COMMA = ",";
+
+    private final String[] indexNames;
+
+    public IndexName(String indexName) {
+      this.indexNames = indexName.split(COMMA);
+    }
+
+    public String[] getIndexNames() {
+      return indexNames;
+    }
+
+    @Override
+    public String toString() {
+      return String.join(COMMA, indexNames);
+    }
+  }
 }

opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java

@@ -29,16 +29,17 @@
  * Maintain scroll ID between calls to client search method
  */
 @EqualsAndHashCode
-@RequiredArgsConstructor
 @Getter
 @ToString
 public class OpenSearchScrollRequest implements OpenSearchRequest {
 
   /** Default scroll context timeout in minutes. */
   public static final TimeValue DEFAULT_SCROLL_TIMEOUT = TimeValue.timeValueMinutes(1L);
 
-  /** Index name. */
-  private final String indexName;
+  /**
+   * {@link OpenSearchRequest.IndexName}.
+   */
+  private final IndexName indexName;
 
   /** Index name. */
   @EqualsAndHashCode.Exclude
@@ -49,11 +50,20 @@ public class OpenSearchScrollRequest implements OpenSearchRequest {
    * Scroll id which is set after first request issued. Because ElasticsearchClient is shared by
    * multi-thread so this state has to be maintained here.
    */
-  @Setter private String scrollId;
+  @Setter
+  private String scrollId;
 
   /** Search request source builder. */
   private final SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
 
+  public OpenSearchScrollRequest(IndexName indexName, OpenSearchExprValueFactory exprValueFactory) {
+    this.indexName = indexName;
+    this.exprValueFactory = exprValueFactory;
+  }
+
+  public OpenSearchScrollRequest(String indexName, OpenSearchExprValueFactory exprValueFactory) {
+    this(new IndexName(indexName), exprValueFactory);
+  }
 
   @Override
   public OpenSearchResponse search(Function<SearchRequest, SearchResponse> searchAction,
@@ -87,7 +97,7 @@ public void clean(Consumer<String> cleanAction) {
    */
   public SearchRequest searchRequest() {
     return new SearchRequest()
-        .indices(indexName)
+        .indices(indexName.getIndexNames())
         .scroll(DEFAULT_SCROLL_TIMEOUT)
         .source(sourceBuilder);
   }