New Changes to handle bug:#866

Signed-off-by: vamsi-amazon <[email protected]>

Author: vamsimanohar

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

@@ -98,6 +98,57 @@ public void testSourceFieldQuery() throws IOException {
         actual);
   }
 
+  @Test
+  public void testErrorsAcrossMultipleQueries() throws IOException {
+    Request request1 = new Request("PUT", "/test.one/_doc/1?refresh=true");
+    request1.setJsonEntity("{\"name\": \"hello\", \"age\": 20}");
+    client().performRequest(request1);
+    Request request2 = new Request("PUT", "/test.one/_doc/2?refresh=true");
+    request2.setJsonEntity("{\"name\": \"world\", \"age\": 30}");
+    client().performRequest(request2);
+
+    String actual = executeByStandaloneQueryEngine("source=test.one | fields name");
+    assertEquals(
+        "{\n"
+            + "  \"schema\": [\n"
+            + "    {\n"
+            + "      \"name\": \"name\",\n"
+            + "      \"type\": \"string\"\n"
+            + "    }\n"
+            + "  ],\n"
+            + "  \"datarows\": [\n"
+            + "    [\n"
+            + "      \"hello\"\n"
+            + "    ],\n"
+            + "    [\n"
+            + "      \"world\"\n"
+            + "    ]\n"
+            + "  ],\n"
+            + "  \"total\": 2,\n"
+            + "  \"size\": 2\n"
+            + "}",
+        actual);
+
+    String actual2 = executeByStandaloneQueryEngine("source=test.one | where age=30 | fields name");
+    assertEquals(
+        "{\n"
+            + "  \"schema\": [\n"
+            + "    {\n"
+            + "      \"name\": \"name\",\n"
+            + "      \"type\": \"string\"\n"
+            + "    }\n"
+            + "  ],\n"
+            + "  \"datarows\": [\n"
+            + "    [\n"
+            + "      \"world\"\n"
+            + "    ]\n"
+            + "  ],\n"
+            + "  \"total\": 1,\n"
+            + "  \"size\": 1\n"
+            + "}",
+        actual2);
+  }
+
   private String executeByStandaloneQueryEngine(String query) {
     AtomicReference<String> actual = new AtomicReference<>();
     pplService.execute(

plugin/src/main/java/org/opensearch/sql/plugin/catalog/CatalogServiceImpl.java

@@ -8,10 +8,12 @@
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.DeserializationFeature;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.ImmutableSet;
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URISyntaxException;
 import java.security.PrivilegedExceptionAction;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -37,7 +39,7 @@ public class CatalogServiceImpl implements CatalogService {
 
   private static final Logger LOG = LogManager.getLogger();
 
-  public static final String OPEN_SEARCH = "opensearch";
+  public static StorageEngine defaultOpenSearchStorageEngine;
 
   private Map<String, StorageEngine> storageEngineMap = new HashMap<>();
 
@@ -80,21 +82,22 @@ public void loadConnectors(Settings settings) {
   @Override
   public StorageEngine getStorageEngine(String catalog) {
     if (catalog == null || !storageEngineMap.containsKey(catalog)) {
-      return storageEngineMap.get(OPEN_SEARCH);
+      return defaultOpenSearchStorageEngine;
     }
     return storageEngineMap.get(catalog);
   }
 
   @Override
   public Set<String> getCatalogs() {
-    return storageEngineMap.keySet()
-        .stream().filter(catalog -> !catalog.equals(OPEN_SEARCH))
-        .collect(Collectors.toSet());
+    return Collections.unmodifiableSet(storageEngineMap.keySet());
   }
 
   @Override
   public void registerOpenSearchStorageEngine(StorageEngine storageEngine) {
-    storageEngineMap.put(OPEN_SEARCH, storageEngine);
+    if(storageEngine == null) {
+      throw new IllegalArgumentException("Default storage engine can't be null");
+    }
+    defaultOpenSearchStorageEngine = storageEngine;
   }
 
   private <T> T doPrivileged(PrivilegedExceptionAction<T> action) {

plugin/src/test/java/org/opensearch/sql/plugin/catalog/CatalogServiceImplTest.java

@@ -15,15 +15,21 @@
 import lombok.SneakyThrows;
 import org.junit.Assert;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnitRunner;
 import org.opensearch.common.settings.MockSecureSettings;
 import org.opensearch.common.settings.Settings;
+import org.opensearch.sql.storage.StorageEngine;
 
-
+@RunWith(MockitoJUnitRunner.class)
 public class CatalogServiceImplTest {
 
   public static final String CATALOG_SETTING_METADATA_KEY =
       "plugins.query.federation.catalog.config";
 
+  @Mock
+  private StorageEngine storageEngine;
 
   @SneakyThrows
   @Test
@@ -73,6 +79,19 @@ public void testLoadConnectorsWithMalformedJson() {
         () -> CatalogServiceImpl.getInstance().loadConnectors(settings));
   }
 
+  @SneakyThrows
+  @Test
+  public void testGetStorageEngineAfterGetCatalogs() {
+    Settings settings = getCatalogSettings("empty_catalog.json");
+    CatalogServiceImpl.getInstance().registerOpenSearchStorageEngine(storageEngine);
+    CatalogServiceImpl.getInstance().loadConnectors(settings);
+    Set<String> expected = new HashSet<>();
+    Assert.assertEquals(expected, CatalogServiceImpl.getInstance().getCatalogs());
+    Assert.assertEquals(storageEngine, CatalogServiceImpl.getInstance().getStorageEngine(null));
+    Assert.assertEquals(expected, CatalogServiceImpl.getInstance().getCatalogs());
+    Assert.assertEquals(storageEngine, CatalogServiceImpl.getInstance().getStorageEngine(null));
+  }
+
 
   private Settings getCatalogSettings(String filename) throws URISyntaxException, IOException {
     MockSecureSettings mockSecureSettings = new MockSecureSettings();

plugin/src/test/resources/empty_catalog.json

@@ -0,0 +1 @@
+[]