Create new session if client provided session is invalid (#2368) (#2371)

* Create new session if session is invalid * fix code style * fix UT * fix error response --------- (cherry picked from commit e16da37) Signed-off-by: Peng Huo <[email protected]> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
opensearch-project · Oct 26, 2023 · 5ab7858 · 5ab7858
1 parent b620a56
commit 5ab7858
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 35 deletions.
diff --git a/spark/src/main/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcher.java b/spark/src/main/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcher.java
@@ -219,10 +219,9 @@ private DispatchQueryResponse handleNonIndexQuery(DispatchQueryRequest dispatchQ
         // get session from request
         SessionId sessionId = new SessionId(dispatchQueryRequest.getSessionId());
         Optional<Session> createdSession = sessionManager.getSession(sessionId);
-        if (createdSession.isEmpty()) {
-          throw new IllegalArgumentException("no session found. " + sessionId);
+        if (createdSession.isPresent()) {
+          session = createdSession.get();
         }
-        session = createdSession.get();
       }
       if (session == null || !session.isReady()) {
         // create session if not exist or session dead/fail

diff --git a/spark/src/main/java/org/opensearch/sql/spark/execution/statement/StatementModel.java b/spark/src/main/java/org/opensearch/sql/spark/execution/statement/StatementModel.java
@@ -36,7 +36,7 @@ public class StatementModel extends StateModel {
   public static final String QUERY_ID = "queryId";
   public static final String SUBMIT_TIME = "submitTime";
   public static final String ERROR = "error";
-  public static final String UNKNOWN = "unknown";
+  public static final String UNKNOWN = "";
   public static final String STATEMENT_DOC_TYPE = "statement";
 
   private final String version;

diff --git a/.../test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceImplSpecTest.java b/.../test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceImplSpecTest.java
@@ -45,6 +45,7 @@
 import org.opensearch.common.settings.ClusterSettings;
 import org.opensearch.common.settings.Setting;
 import org.opensearch.common.settings.Settings;
+import org.opensearch.core.common.Strings;
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.plugins.Plugin;
@@ -227,6 +228,7 @@ public void withSessionCreateAsyncQueryThenGetResultThenCancel() {
     // 2. fetch async query result.
     AsyncQueryExecutionResponse asyncQueryResults =
         asyncQueryExecutorService.getAsyncQueryResults(response.getQueryId());
+    assertTrue(Strings.isEmpty(asyncQueryResults.getError()));
     assertEquals(StatementState.WAITING.getState(), asyncQueryResults.getStatus());
 
     // 3. cancel async query.
@@ -460,24 +462,22 @@ public void recreateSessionIfNotReady() {
   }
 
   @Test
-  public void submitQueryInInvalidSessionThrowException() {
+  public void submitQueryInInvalidSessionWillCreateNewSession() {
     LocalEMRSClient emrsClient = new LocalEMRSClient();
     AsyncQueryExecutorService asyncQueryExecutorService =
         createAsyncQueryExecutorService(emrsClient);
 
     // enable session
     enableSession(true);
 
-    // 1. create async query.
-    SessionId sessionId = SessionId.newSessionId(DATASOURCE);
-    IllegalArgumentException exception =
-        assertThrows(
-            IllegalArgumentException.class,
-            () ->
-                asyncQueryExecutorService.createAsyncQuery(
-                    new CreateAsyncQueryRequest(
-                        "select 1", DATASOURCE, LangType.SQL, sessionId.getSessionId())));
-    assertEquals("no session found. " + sessionId, exception.getMessage());
+    // 1. create async query with invalid sessionId
+    SessionId invalidSessionId = SessionId.newSessionId(DATASOURCE);
+    CreateAsyncQueryResponse asyncQuery =
+        asyncQueryExecutorService.createAsyncQuery(
+            new CreateAsyncQueryRequest(
+                "select 1", DATASOURCE, LangType.SQL, invalidSessionId.getSessionId()));
+    assertNotNull(asyncQuery.getSessionId());
+    assertNotEquals(invalidSessionId.getSessionId(), asyncQuery.getSessionId());
   }
 
   private DataSourceServiceImpl createDataSourceService() {

diff --git a/spark/src/test/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcherTest.java b/spark/src/test/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcherTest.java
@@ -327,26 +327,6 @@ void testDispatchSelectQueryReuseSession() {
     Assertions.assertEquals(MOCK_SESSION_ID, dispatchQueryResponse.getSessionId());
   }
 
-  @Test
-  void testDispatchSelectQueryInvalidSession() {
-    String query = "select * from my_glue.default.http_logs";
-    DispatchQueryRequest queryRequest = dispatchQueryRequestWithSessionId(query, "invalid");
-
-    doReturn(true).when(sessionManager).isEnabled();
-    doReturn(Optional.empty()).when(sessionManager).getSession(any());
-    DataSourceMetadata dataSourceMetadata = constructMyGlueDataSourceMetadata();
-    when(dataSourceService.getRawDataSourceMetadata("my_glue")).thenReturn(dataSourceMetadata);
-    doNothing().when(dataSourceUserAuthorizationHelper).authorizeDataSource(dataSourceMetadata);
-    IllegalArgumentException exception =
-        Assertions.assertThrows(
-            IllegalArgumentException.class, () -> sparkQueryDispatcher.dispatch(queryRequest));
-
-    verifyNoInteractions(emrServerlessClient);
-    verify(sessionManager, never()).createSession(any());
-    Assertions.assertEquals(
-        "no session found. " + new SessionId("invalid"), exception.getMessage());
-  }
-
   @Test
   void testDispatchSelectQueryFailedCreateSession() {
     String query = "select * from my_glue.default.http_logs";