Merge branch 'main' into dropIndexFixPr

common/src/main/java/org/opensearch/sql/common/setting/Settings.java

@@ -34,13 +34,18 @@ public enum Key {
     QUERY_SIZE_LIMIT("plugins.query.size_limit"),
     ENCYRPTION_MASTER_KEY("plugins.query.datasources.encryption.masterkey"),
     DATASOURCES_URI_HOSTS_DENY_LIST("plugins.query.datasources.uri.hosts.denylist"),
+    DATASOURCES_LIMIT("plugins.query.datasources.limit"),
 
     METRICS_ROLLING_WINDOW("plugins.query.metrics.rolling_window"),
     METRICS_ROLLING_INTERVAL("plugins.query.metrics.rolling_interval"),
     SPARK_EXECUTION_ENGINE_CONFIG("plugins.query.executionengine.spark.config"),
     CLUSTER_NAME("cluster.name"),
     SPARK_EXECUTION_SESSION_ENABLED("plugins.query.executionengine.spark.session.enabled"),
-    SPARK_EXECUTION_SESSION_LIMIT("plugins.query.executionengine.spark.session.limit");
+    SPARK_EXECUTION_SESSION_LIMIT("plugins.query.executionengine.spark.session.limit"),
+    SESSION_INDEX_TTL("plugins.query.executionengine.spark.session.index.ttl"),
+    RESULT_INDEX_TTL("plugins.query.executionengine.spark.result.index.ttl"),
+    AUTO_INDEX_MANAGEMENT_ENABLED(
+        "plugins.query.executionengine.spark.auto_index_management.enabled");
 
     @Getter private final String keyValue;
 

core/src/main/java/org/opensearch/sql/datasource/model/DataSourceMetadata.java

@@ -16,7 +16,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import lombok.AllArgsConstructor;
+import java.util.function.Function;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
 import lombok.Setter;
@@ -25,11 +25,24 @@
 
 @Getter
 @Setter
-@AllArgsConstructor
 @EqualsAndHashCode
 @JsonIgnoreProperties(ignoreUnknown = true)
 public class DataSourceMetadata {
 
+  public static final String DEFAULT_RESULT_INDEX = "query_execution_result";
+  public static final int MAX_RESULT_INDEX_NAME_SIZE = 255;
+  // OS doesn’t allow uppercase: https://tinyurl.com/yse2xdbx
+  public static final String RESULT_INDEX_NAME_PATTERN = "[a-z0-9_-]+";
+  public static String INVALID_RESULT_INDEX_NAME_SIZE =
+      "Result index name size must contains less than "
+          + MAX_RESULT_INDEX_NAME_SIZE
+          + " characters";
+  public static String INVALID_CHAR_IN_RESULT_INDEX_NAME =
+      "Result index name has invalid character. Valid characters are a-z, 0-9, -(hyphen) and"
+          + " _(underscore)";
+  public static String INVALID_RESULT_INDEX_PREFIX =
+      "Result index must start with " + DEFAULT_RESULT_INDEX;
+
   @JsonProperty private String name;
 
   @JsonProperty private String description;
@@ -44,18 +57,31 @@ public class DataSourceMetadata {
 
   @JsonProperty private String resultIndex;
 
+  public static Function<String, String> DATASOURCE_TO_RESULT_INDEX =
+      datasourceName -> String.format("%s_%s", DEFAULT_RESULT_INDEX, datasourceName);
+
   public DataSourceMetadata(
       String name,
+      String description,
       DataSourceType connector,
       List<String> allowedRoles,
       Map<String, String> properties,
       String resultIndex) {
     this.name = name;
+    String errorMessage = validateCustomResultIndex(resultIndex);
+    if (errorMessage != null) {
+      throw new IllegalArgumentException(errorMessage);
+    }
+    if (resultIndex == null) {
+      this.resultIndex = fromNameToCustomResultIndex();
+    } else {
+      this.resultIndex = resultIndex;
+    }
+
     this.connector = connector;
-    this.description = StringUtils.EMPTY;
+    this.description = description;
     this.properties = properties;
     this.allowedRoles = allowedRoles;
-    this.resultIndex = resultIndex;
   }
 
   public DataSourceMetadata() {
@@ -71,9 +97,56 @@ public DataSourceMetadata() {
   public static DataSourceMetadata defaultOpenSearchDataSourceMetadata() {
     return new DataSourceMetadata(
         DEFAULT_DATASOURCE_NAME,
+        StringUtils.EMPTY,
         DataSourceType.OPENSEARCH,
         Collections.emptyList(),
         ImmutableMap.of(),
         null);
   }
+
+  public String validateCustomResultIndex(String resultIndex) {
+    if (resultIndex == null) {
+      return null;
+    }
+    if (resultIndex.length() > MAX_RESULT_INDEX_NAME_SIZE) {
+      return INVALID_RESULT_INDEX_NAME_SIZE;
+    }
+    if (!resultIndex.matches(RESULT_INDEX_NAME_PATTERN)) {
+      return INVALID_CHAR_IN_RESULT_INDEX_NAME;
+    }
+    if (resultIndex != null && !resultIndex.startsWith(DEFAULT_RESULT_INDEX)) {
+      return INVALID_RESULT_INDEX_PREFIX;
+    }
+    return null;
+  }
+
+  /**
+   * Since we are using datasource name to create result index, we need to make sure that the final
+   * name is valid
+   *
+   * @param resultIndex result index name
+   * @return valid result index name
+   */
+  private String convertToValidResultIndex(String resultIndex) {
+    // Limit Length
+    if (resultIndex.length() > MAX_RESULT_INDEX_NAME_SIZE) {
+      resultIndex = resultIndex.substring(0, MAX_RESULT_INDEX_NAME_SIZE);
+    }
+
+    // Pattern Matching: Remove characters that don't match the pattern
+    StringBuilder validChars = new StringBuilder();
+    for (char c : resultIndex.toCharArray()) {
+      if (String.valueOf(c).matches(RESULT_INDEX_NAME_PATTERN)) {
+        validChars.append(c);
+      }
+    }
+    return validChars.toString();
+  }
+
+  public String fromNameToCustomResultIndex() {
+    if (name == null) {
+      throw new IllegalArgumentException("Datasource name cannot be null");
+    }
+    return convertToValidResultIndex(DATASOURCE_TO_RESULT_INDEX.apply(name.toLowerCase()));
+  }
 }

core/src/test/java/org/opensearch/sql/analysis/AnalyzerTestBase.java

@@ -19,6 +19,7 @@
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.tuple.Pair;
 import org.opensearch.sql.DataSourceSchemaName;
 import org.opensearch.sql.analysis.symbol.Namespace;
@@ -197,6 +198,7 @@ public Set<DataSourceMetadata> getDataSourceMetadata(boolean isDefaultDataSource
               ds ->
                   new DataSourceMetadata(
                       ds.getName(),
+                      StringUtils.EMPTY,
                       ds.getConnectorType(),
                       Collections.emptyList(),
                       ImmutableMap.of(),

core/src/test/java/org/opensearch/sql/planner/physical/datasource/DataSourceTableScanTest.java

@@ -18,6 +18,7 @@
 import java.util.LinkedHashMap;
 import java.util.Set;
 import java.util.stream.Collectors;
+import org.apache.commons.lang3.StringUtils;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -62,6 +63,7 @@ void testIterator() {
                 dataSource ->
                     new DataSourceMetadata(
                         dataSource.getName(),
+                        StringUtils.EMPTY,
                         dataSource.getConnectorType(),
                         Collections.emptyList(),
                         ImmutableMap.of(),

datasources/build.gradle

@@ -16,6 +16,7 @@ repositories {
 dependencies {
     implementation project(':core')
     implementation project(':protocol')
+    implementation project(':opensearch')
     implementation group: 'org.opensearch', name: 'opensearch', version: "${opensearch_version}"
     implementation group: 'org.opensearch', name: 'opensearch-x-content', version: "${opensearch_version}"
     implementation group: 'org.opensearch', name: 'common-utils', version: "${opensearch_build}"
@@ -35,7 +36,7 @@ dependencies {
 test {
     useJUnitPlatform()
     testLogging {
-        events "passed", "skipped", "failed"
+        events "skipped", "failed"
         exceptionFormat "full"
     }
 }

datasources/src/main/java/org/opensearch/sql/datasources/transport/TransportCreateDataSourceAction.java

@@ -7,6 +7,7 @@
 
 package org.opensearch.sql.datasources.transport;
 
+import static org.opensearch.sql.common.setting.Settings.Key.DATASOURCES_LIMIT;
 import static org.opensearch.sql.protocol.response.format.JsonResponseFormatter.Style.PRETTY;
 
 import org.opensearch.action.ActionType;
@@ -30,6 +31,7 @@ public class TransportCreateDataSourceAction
       new ActionType<>(NAME, CreateDataSourceActionResponse::new);
 
   private DataSourceService dataSourceService;
+  private org.opensearch.sql.opensearch.setting.OpenSearchSettings settings;
 
   /**
    * TransportCreateDataSourceAction action for creating datasource.
@@ -42,33 +44,44 @@ public class TransportCreateDataSourceAction
   public TransportCreateDataSourceAction(
       TransportService transportService,
       ActionFilters actionFilters,
-      DataSourceServiceImpl dataSourceService) {
+      DataSourceServiceImpl dataSourceService,
+      org.opensearch.sql.opensearch.setting.OpenSearchSettings settings) {
     super(
         TransportCreateDataSourceAction.NAME,
         transportService,
         actionFilters,
         CreateDataSourceActionRequest::new);
     this.dataSourceService = dataSourceService;
+    this.settings = settings;
   }
 
   @Override
   protected void doExecute(
       Task task,
       CreateDataSourceActionRequest request,
       ActionListener<CreateDataSourceActionResponse> actionListener) {
-    try {
-      DataSourceMetadata dataSourceMetadata = request.getDataSourceMetadata();
-      dataSourceService.createDataSource(dataSourceMetadata);
-      String responseContent =
-          new JsonResponseFormatter<String>(PRETTY) {
-            @Override
-            protected Object buildJsonObject(String response) {
-              return response;
-            }
-          }.format("Created DataSource with name " + dataSourceMetadata.getName());
-      actionListener.onResponse(new CreateDataSourceActionResponse(responseContent));
-    } catch (Exception e) {
-      actionListener.onFailure(e);
+    int dataSourceLimit = settings.getSettingValue(DATASOURCES_LIMIT);
+    if (dataSourceService.getDataSourceMetadata(false).size() >= dataSourceLimit) {
+      actionListener.onFailure(
+          new IllegalStateException(
+              String.format(
+                  "domain concurrent datasources can not" + " exceed %d", dataSourceLimit)));
+    } else {
+      try {
+
+        DataSourceMetadata dataSourceMetadata = request.getDataSourceMetadata();
+        dataSourceService.createDataSource(dataSourceMetadata);
+        String responseContent =
+            new JsonResponseFormatter<String>(PRETTY) {
+              @Override
+              protected Object buildJsonObject(String response) {
+                return response;
+              }
+            }.format("Created DataSource with name " + dataSourceMetadata.getName());
+        actionListener.onResponse(new CreateDataSourceActionResponse(responseContent));
+      } catch (Exception e) {
+        actionListener.onFailure(e);
+      }
     }
   }
 }

datasources/src/test/java/org/opensearch/sql/datasources/service/DataSourceServiceImplTest.java

@@ -29,6 +29,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import org.apache.commons.lang3.StringUtils;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -382,6 +383,7 @@ void testRemovalOfAuthorizationInfo() {
     DataSourceMetadata dataSourceMetadata =
         new DataSourceMetadata(
             "testDS",
+            StringUtils.EMPTY,
             DataSourceType.PROMETHEUS,
             Collections.singletonList("prometheus_access"),
             properties,
@@ -407,6 +409,7 @@ void testRemovalOfAuthorizationInfoForAccessKeyAndSecretKye() {
     DataSourceMetadata dataSourceMetadata =
         new DataSourceMetadata(
             "testDS",
+            StringUtils.EMPTY,
             DataSourceType.PROMETHEUS,
             Collections.singletonList("prometheus_access"),
             properties,
@@ -434,6 +437,7 @@ void testRemovalOfAuthorizationInfoForGlueWithRoleARN() {
     DataSourceMetadata dataSourceMetadata =
         new DataSourceMetadata(
             "testGlue",
+            StringUtils.EMPTY,
             DataSourceType.S3GLUE,
             Collections.singletonList("glue_access"),
             properties,
@@ -498,6 +502,7 @@ void testGetRawDataSourceMetadata() {
     DataSourceMetadata dataSourceMetadata =
         new DataSourceMetadata(
             "testDS",
+            StringUtils.EMPTY,
             DataSourceType.PROMETHEUS,
             Collections.singletonList("prometheus_access"),
             properties,

datasources/src/test/java/org/opensearch/sql/datasources/transport/TransportCreateDataSourceActionTest.java

@@ -1,14 +1,19 @@
 package org.opensearch.sql.datasources.transport;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.fail;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.opensearch.sql.common.setting.Settings.Key.DATASOURCES_LIMIT;
 
 import java.util.HashSet;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Answers;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Captor;
 import org.mockito.Mock;
@@ -21,6 +26,7 @@
 import org.opensearch.sql.datasources.model.transport.CreateDataSourceActionRequest;
 import org.opensearch.sql.datasources.model.transport.CreateDataSourceActionResponse;
 import org.opensearch.sql.datasources.service.DataSourceServiceImpl;
+import org.opensearch.sql.opensearch.setting.OpenSearchSettings;
 import org.opensearch.tasks.Task;
 import org.opensearch.transport.TransportService;
 
@@ -29,9 +35,13 @@ public class TransportCreateDataSourceActionTest {
 
   @Mock private TransportService transportService;
   @Mock private TransportCreateDataSourceAction action;
-  @Mock private DataSourceServiceImpl dataSourceService;
+
+  @Mock(answer = Answers.RETURNS_DEEP_STUBS)
+  private DataSourceServiceImpl dataSourceService;
+
   @Mock private Task task;
   @Mock private ActionListener<CreateDataSourceActionResponse> actionListener;
+  @Mock private OpenSearchSettings settings;
 
   @Captor
   private ArgumentCaptor<CreateDataSourceActionResponse>
@@ -43,7 +53,9 @@ public class TransportCreateDataSourceActionTest {
   public void setUp() {
     action =
         new TransportCreateDataSourceAction(
-            transportService, new ActionFilters(new HashSet<>()), dataSourceService);
+            transportService, new ActionFilters(new HashSet<>()), dataSourceService, settings);
+    when(dataSourceService.getDataSourceMetadata(false).size()).thenReturn(1);
+    when(settings.getSettingValue(DATASOURCES_LIMIT)).thenReturn(20);
   }
 
   @Test
@@ -79,4 +91,30 @@ public void testDoExecuteWithException() {
     Assertions.assertTrue(exception instanceof RuntimeException);
     Assertions.assertEquals("Error", exception.getMessage());
   }
+
+  @Test
+  public void testDataSourcesLimit() {
+    DataSourceMetadata dataSourceMetadata = new DataSourceMetadata();
+    dataSourceMetadata.setName("test_datasource");
+    dataSourceMetadata.setConnector(DataSourceType.PROMETHEUS);
+    CreateDataSourceActionRequest request = new CreateDataSourceActionRequest(dataSourceMetadata);
+    when(dataSourceService.getDataSourceMetadata(false).size()).thenReturn(1);
+    when(settings.getSettingValue(DATASOURCES_LIMIT)).thenReturn(1);
+
+    action.doExecute(
+        task,
+        request,
+        new ActionListener<CreateDataSourceActionResponse>() {
+          @Override
+          public void onResponse(CreateDataSourceActionResponse createDataSourceActionResponse) {
+            fail();
+          }
+
+          @Override
+          public void onFailure(Exception e) {
+            assertEquals("domain concurrent datasources can not exceed 1", e.getMessage());
+          }
+        });
+    verify(dataSourceService, times(0)).createDataSource(dataSourceMetadata);
+  }
 }