Fix IndexNotFound Exception after datasource creation

Signed-off-by: Vamsi Manohar <reddyvam@amazon.com>

datasources/src/main/java/org/opensearch/sql/datasources/service/DataSourceMetadataStorage.java

@@ -53,4 +53,13 @@ public interface DataSourceMetadataStorage {
    * @param datasourceName name of the {@link DataSource}.
    */
   void deleteDataSourceMetadata(String datasourceName);
+
+  /**
+   * Creates query result index for the given datasource. We have introduced this method to handle
+   * the issue: https://github.com/opensearch-project/sql/issues/2268 Once a long term solution is
+   * finalized, we can make a decision to remove or have it.
+   *
+   * @param dataSourceMetadata {@link DataSourceMetadata}.
+   */
+  void createDataSourceResultIndex(DataSourceMetadata dataSourceMetadata);
 }

datasources/src/main/java/org/opensearch/sql/datasources/service/DataSourceServiceImpl.java

@@ -81,6 +81,7 @@ public void createDataSource(DataSourceMetadata metadata) {
     if (!metadata.getName().equals(DEFAULT_DATASOURCE_NAME)) {
       this.dataSourceLoaderCache.getOrLoadDataSource(metadata);
       this.dataSourceMetadataStorage.createDataSourceMetadata(metadata);
+      this.dataSourceMetadataStorage.createDataSourceResultIndex(metadata);
     }
   }
 
@@ -90,6 +91,8 @@ public void updateDataSource(DataSourceMetadata dataSourceMetadata) {
     if (!dataSourceMetadata.getName().equals(DEFAULT_DATASOURCE_NAME)) {
       this.dataSourceLoaderCache.getOrLoadDataSource(dataSourceMetadata);
       this.dataSourceMetadataStorage.updateDataSourceMetadata(dataSourceMetadata);
+      // If resultIndex is changed in update call
+      this.dataSourceMetadataStorage.createDataSourceResultIndex(dataSourceMetadata);
     } else {
       throw new UnsupportedOperationException(
           "Not allowed to update default datasource :" + DEFAULT_DATASOURCE_NAME);

datasources/src/main/java/org/opensearch/sql/datasources/storage/OpenSearchDataSourceMetadataStorage.java

@@ -56,6 +56,11 @@ public class OpenSearchDataSourceMetadataStorage implements DataSourceMetadataSt
   private static final Integer DATASOURCE_QUERY_RESULT_SIZE = 10000;
   private static final String DATASOURCE_INDEX_SETTINGS_FILE_NAME =
       "datasources-index-settings.yml";
+  private static final String DATASOURCE_QUERY_RESULT_INDEX_MAPPING_FILE_NAME =
+      "query_execution_result_mapping.yml";
+  private static final String DATASOURCE_QUERY_RESULT_INDEX_SETTING_FILE_NAME =
+      "query_execution_result_settings.yml";
+
   private static final Logger LOG = LogManager.getLogger();
   private final Client client;
   private final ClusterService clusterService;
@@ -183,6 +188,44 @@ public void deleteDataSourceMetadata(String datasourceName) {
     }
   }
 
+  @Override
+  public void createDataSourceResultIndex(DataSourceMetadata dataSourceMetadata) {
+    if (!this.clusterService.state().routingTable().hasIndex(dataSourceMetadata.getResultIndex())) {
+      try {
+        InputStream mappingFileStream =
+            OpenSearchDataSourceMetadataStorage.class
+                .getClassLoader()
+                .getResourceAsStream(DATASOURCE_QUERY_RESULT_INDEX_MAPPING_FILE_NAME);
+        InputStream settingsFileStream =
+            OpenSearchDataSourceMetadataStorage.class
+                .getClassLoader()
+                .getResourceAsStream(DATASOURCE_QUERY_RESULT_INDEX_SETTING_FILE_NAME);
+        CreateIndexRequest createIndexRequest =
+            new CreateIndexRequest(dataSourceMetadata.getResultIndex());
+        createIndexRequest
+            .mapping(IOUtils.toString(mappingFileStream, StandardCharsets.UTF_8), XContentType.YAML)
+            .settings(
+                IOUtils.toString(settingsFileStream, StandardCharsets.UTF_8), XContentType.YAML);
+        ActionFuture<CreateIndexResponse> createIndexResponseActionFuture;
+        try (ThreadContext.StoredContext ignored =
+            client.threadPool().getThreadContext().stashContext()) {
+          createIndexResponseActionFuture = client.admin().indices().create(createIndexRequest);
+        }
+        CreateIndexResponse createIndexResponse = createIndexResponseActionFuture.actionGet();
+        if (createIndexResponse.isAcknowledged()) {
+          LOG.info(
+              "Index: {} creation Acknowledged", dataSourceMetadata.fromNameToCustomResultIndex());
+        } else {
+          throw new RuntimeException("Index creation is not acknowledged.");
+        }
+      } catch (Throwable e) {
+        LOG.error("Datasource Query Result Index Creation failed due to :{}", e.getMessage());
+        // Suppressing all exceptions as this is not an important step in datasource creation
+        // workflow.
+      }
+    }
+  }
+
   private void createDataSourcesIndex() {
     try {
       InputStream mappingFileStream =

datasources/src/main/resources/query_execution_result_mapping.yml

@@ -0,0 +1,37 @@
+---
+##
+# Copyright OpenSearch Contributors
+# SPDX-License-Identifier: Apache-2.0
+##
+
+# Schema file for the query_execution_request index
+# Also "dynamic" is set to "false" so that other fields can be added.,
+dynamic: false
+properties:
+  result:
+    type: object
+    enabled: false
+  schema:
+    type: object
+    enabled: false
+  jobRunId:
+    type: keyword
+  applicationId:
+    type: keyword
+  dataSourceName:
+    type: keyword
+  status:
+    type: keyword
+  queryId:
+    type: keyword
+  queryText:
+    type: text
+  sessionId:
+    type: keyword
+  updateTime:
+    type: date
+    format: strict_date_time||epoch_millis
+  error:
+    type: text
+  queryRunTime:
+    type: long

datasources/src/main/resources/query_execution_result_settings.yml

@@ -0,0 +1,11 @@
+---
+##
+# Copyright OpenSearch Contributors
+# SPDX-License-Identifier: Apache-2.0
+##
+
+# Settings file for the query_execution_result index
+index:
+  number_of_shards: "1"
+  auto_expand_replicas: "0-2"
+  number_of_replicas: "0"

datasources/src/test/java/org/opensearch/sql/datasources/storage/OpenSearchDataSourceMetadataStorageTest.java

@@ -598,6 +598,56 @@ public void testDeleteDataSourceMetadataWithUnexpectedResult() {
     Mockito.verify(client.threadPool().getThreadContext(), Mockito.times(1)).stashContext();
   }
 
+  @Test
+  public void testCreateDataSourceResultIndexWithOutCreatingIndex() {
+    DataSourceMetadata dataSourceMetadata = getDataSourceMetadata();
+    Mockito.when(
+            clusterService.state().routingTable().hasIndex(dataSourceMetadata.getResultIndex()))
+        .thenReturn(Boolean.TRUE);
+    this.openSearchDataSourceMetadataStorage.createDataSourceResultIndex(dataSourceMetadata);
+
+    Mockito.verify(client.admin().indices(), Mockito.times(0)).create(ArgumentMatchers.any());
+    Mockito.verify(client, Mockito.times(0)).index(ArgumentMatchers.any());
+    Mockito.verify(client.threadPool().getThreadContext(), Mockito.times(0)).stashContext();
+  }
+
+  @Test
+  public void testCreateDataSourceResultIndexWhenIndexIsNotAvailable() {
+    DataSourceMetadata dataSourceMetadata = getDataSourceMetadata();
+    Mockito.when(
+            clusterService.state().routingTable().hasIndex(dataSourceMetadata.getResultIndex()))
+        .thenReturn(Boolean.FALSE);
+    Mockito.when(client.admin().indices().create(ArgumentMatchers.any()))
+        .thenReturn(createIndexResponseActionFuture);
+    Mockito.when(createIndexResponseActionFuture.actionGet())
+        .thenReturn(new CreateIndexResponse(true, true, dataSourceMetadata.getResultIndex()));
+
+    this.openSearchDataSourceMetadataStorage.createDataSourceResultIndex(dataSourceMetadata);
+
+    Mockito.verify(client.admin().indices(), Mockito.times(1)).create(ArgumentMatchers.any());
+    Mockito.verify(client, Mockito.times(0)).index(ArgumentMatchers.any());
+    Mockito.verify(client.threadPool().getThreadContext(), Mockito.times(1)).stashContext();
+  }
+
+  @Test
+  public void testCreateDataSourceResultIndexWithException() {
+
+    DataSourceMetadata dataSourceMetadata = getDataSourceMetadata();
+    Mockito.when(
+            clusterService.state().routingTable().hasIndex(dataSourceMetadata.getResultIndex()))
+        .thenReturn(Boolean.FALSE);
+    Mockito.when(client.admin().indices().create(ArgumentMatchers.any()))
+        .thenReturn(createIndexResponseActionFuture);
+    Mockito.when(createIndexResponseActionFuture.actionGet())
+        .thenReturn(new CreateIndexResponse(false, false, dataSourceMetadata.getResultIndex()));
+
+    this.openSearchDataSourceMetadataStorage.createDataSourceResultIndex(dataSourceMetadata);
+
+    Mockito.verify(client.admin().indices(), Mockito.times(1)).create(ArgumentMatchers.any());
+    Mockito.verify(client, Mockito.times(0)).index(ArgumentMatchers.any());
+    Mockito.verify(client.threadPool().getThreadContext(), Mockito.times(1)).stashContext();
+  }
+
   private String getBasicDataSourceMetadataString() throws JsonProcessingException {
     DataSourceMetadata dataSourceMetadata = new DataSourceMetadata();
     dataSourceMetadata.setName("testDS");

integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java

@@ -264,6 +264,9 @@ public void updateDataSourceMetadata(DataSourceMetadata dataSourceMetadata) {}
 
       @Override
       public void deleteDataSourceMetadata(String datasourceName) {}
+
+      @Override
+      public void createDataSourceResultIndex(DataSourceMetadata dataSourceMetadata) {}
     };
   }
 

spark/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceSpec.java

@@ -145,8 +145,6 @@ public void setup() {
             null);
     dataSourceService.createDataSource(otherDm);
     stateStore = new StateStore(client, clusterService);
-    createIndexWithMappings(dm.getResultIndex(), loadResultIndexMappings());
-    createIndexWithMappings(otherDm.getResultIndex(), loadResultIndexMappings());
   }
 
   @After