Create new session if session is invalid

Signed-off-by: Peng Huo <penghuo@gmail.com>
opensearch-project · Oct 25, 2023 · 0699d58 · 0699d58
1 parent a5512f5
commit 0699d58
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 14 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/.../test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceImplSpecTest.java b/.../test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceImplSpecTest.java
@@ -460,24 +460,21 @@ 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() {