Runner implementation

async-query/src/main/java/org/opensearch/sql/spark/scheduler/job/AsyncQueryScheduledQueryJob.java

@@ -13,6 +13,10 @@
 import org.opensearch.jobscheduler.spi.ScheduledJobParameter;
 import org.opensearch.jobscheduler.spi.ScheduledJobRunner;
 import org.opensearch.plugins.Plugin;
+import org.opensearch.sql.spark.asyncquery.AsyncQueryExecutorService;
+import org.opensearch.sql.spark.asyncquery.model.NullAsyncQueryRequestContext;
+import org.opensearch.sql.spark.rest.model.CreateAsyncQueryRequest;
+import org.opensearch.sql.spark.rest.model.CreateAsyncQueryResponse;
 import org.opensearch.sql.spark.scheduler.model.OpenSearchScheduleQueryJobRequest;
 import org.opensearch.threadpool.ThreadPool;
 
@@ -30,7 +34,6 @@
  * plugin.
  */
 public class AsyncQueryScheduledQueryJob implements ScheduledJobRunner {
-
   private static final Logger log = LogManager.getLogger(AsyncQueryScheduledQueryJob.class);
 
   public static AsyncQueryScheduledQueryJob INSTANCE = new AsyncQueryScheduledQueryJob();
@@ -42,6 +45,7 @@ public static AsyncQueryScheduledQueryJob getJobRunnerInstance() {
   private ClusterService clusterService;
   private ThreadPool threadPool;
   private Client client;
+  private AsyncQueryExecutorService asyncQueryExecutorService;
 
   private AsyncQueryScheduledQueryJob() {
     // Singleton class, use getJobRunnerInstance method instead of constructor
@@ -59,8 +63,13 @@ public void setClient(Client client) {
     this.client = client;
   }
 
+  public void setAsyncQueryExecutorService(AsyncQueryExecutorService asyncQueryExecutorService) {
+    this.asyncQueryExecutorService = asyncQueryExecutorService;
+  }
+
   @Override
   public void runJob(ScheduledJobParameter jobParameter, JobExecutionContext context) {
+    // Parser will convert jobParameter to OpenSearchScheduleQueryJobRequest
     if (!(jobParameter instanceof OpenSearchScheduleQueryJobRequest)) {
       throw new IllegalStateException(
           "Job parameter is not instance of OpenSearchRefreshIndexJobRequest, type: "
@@ -79,15 +88,26 @@ public void runJob(ScheduledJobParameter jobParameter, JobExecutionContext conte
       throw new IllegalStateException("Client is not initialized.");
     }
 
+    if (this.asyncQueryExecutorService == null) {
+      throw new IllegalStateException("AsyncQueryExecutorService is not initialized.");
+    }
+
     Runnable runnable =
         () -> {
-          doRefresh(jobParameter.getName());
+          doRefresh((OpenSearchScheduleQueryJobRequest) jobParameter);
         };
     threadPool.generic().submit(runnable);
   }
 
-  void doRefresh(String refreshIndex) {
-    // TODO: add logic to refresh index
-    log.info("Scheduled refresh index job on : " + refreshIndex);
+  void doRefresh(OpenSearchScheduleQueryJobRequest request) {
+    // TODO: use internal logic to create refresh index query?
+    log.info("Scheduled refresh index job on : " + request.getName());
+    CreateAsyncQueryRequest createAsyncQueryRequest =
+        new CreateAsyncQueryRequest(
+            request.getScheduledQuery(), request.getDataSource(), request.getQueryLang());
+    CreateAsyncQueryResponse createAsyncQueryResponse =
+        asyncQueryExecutorService.createAsyncQuery(
+            createAsyncQueryRequest, new NullAsyncQueryRequestContext());
+    log.info("Created async query with queryId: " + createAsyncQueryResponse.getQueryId());
   }
 }

async-query/src/test/java/org/opensearch/sql/spark/scheduler/job/AsyncQueryScheduledQueryJobTest.java

@@ -12,6 +12,7 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
+import static org.mockito.ArgumentMatchers.any;
 
 import java.time.Instant;
 import org.junit.jupiter.api.BeforeEach;
@@ -26,6 +27,10 @@
 import org.opensearch.jobscheduler.spi.ScheduledJobParameter;
 import org.opensearch.sql.spark.scheduler.model.OpenSearchScheduleQueryJobRequest;
 import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.sql.spark.asyncquery.AsyncQueryExecutorService;
+import org.opensearch.sql.spark.asyncquery.model.NullAsyncQueryRequestContext;
+import org.opensearch.sql.spark.rest.model.CreateAsyncQueryRequest;
+import org.opensearch.sql.spark.rest.model.LangType;
 
 public class AsyncQueryScheduledQueryJobTest {
 
@@ -40,6 +45,8 @@ public class AsyncQueryScheduledQueryJobTest {
 
   @Mock private JobExecutionContext context;
 
+  @Mock private AsyncQueryExecutorService asyncQueryExecutorService;
+
   private AsyncQueryScheduledQueryJob jobRunner;
 
   private AsyncQueryScheduledQueryJob spyJobRunner;
@@ -51,6 +58,7 @@ public void setup() {
     jobRunner.setClient(null);
     jobRunner.setClusterService(null);
     jobRunner.setThreadPool(null);
+    jobRunner.setAsyncQueryExecutorService(null);
   }
 
   @Test
@@ -59,23 +67,28 @@ public void testRunJobWithCorrectParameter() {
     spyJobRunner.setClusterService(clusterService);
     spyJobRunner.setThreadPool(threadPool);
     spyJobRunner.setClient(client);
+    spyJobRunner.setAsyncQueryExecutorService(asyncQueryExecutorService);
 
-    OpenSearchScheduleQueryJobRequest jobParameter =
+    OpenSearchScheduleQueryJobRequest request =
         OpenSearchScheduleQueryJobRequest.builder()
             .jobId("testJob")
             .lastUpdateTime(Instant.now())
             .lockDurationSeconds(10L)
+            .scheduledQuery("REFRESH INDEX testIndex")
+            .dataSource("testDataSource")
+            .queryLang(LangType.SQL)
             .build();
 
-    spyJobRunner.runJob(jobParameter, context);
+    spyJobRunner.runJob(request, context);
 
     ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
     verify(threadPool.generic()).submit(captor.capture());
 
     Runnable runnable = captor.getValue();
     runnable.run();
 
-    verify(spyJobRunner).doRefresh(eq(jobParameter.getName()));
+    verify(spyJobRunner).doRefresh(eq(request));
+    verify(asyncQueryExecutorService).createAsyncQuery(any(CreateAsyncQueryRequest.class), any(NullAsyncQueryRequestContext.class));
   }
 
   @Test
@@ -131,6 +144,15 @@ public void testRunJobWithUninitializedServices() {
             () -> jobRunner.runJob(jobParameter, context),
             "Expected IllegalStateException but no exception was thrown");
     assertEquals("Client is not initialized.", exception.getMessage());
+
+    jobRunner.setAsyncQueryExecutorService(asyncQueryExecutorService);
+
+    exception =
+        assertThrows(
+            IllegalStateException.class,
+            () -> jobRunner.runJob(jobParameter, context),
+            "Expected IllegalStateException but no exception was thrown");
+    assertEquals("AsyncQueryExecutorService is not initialized.", exception.getMessage());
   }
 
   @Test
@@ -142,4 +164,4 @@ public void testGetJobRunnerInstanceMultipleCalls() {
     assertSame(instance1, instance2);
     assertSame(instance2, instance3);
   }
-}
+}

plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java

@@ -97,7 +97,7 @@
 import org.opensearch.sql.spark.rest.RestAsyncQueryManagementAction;
 import org.opensearch.sql.spark.scheduler.OpenSearchAsyncQueryScheduler;
 import org.opensearch.sql.spark.scheduler.OpenSearchScheduleQueryJobRequestParser;
-import org.opensearch.sql.spark.scheduler.job.ScheduledQueryJob;
+import org.opensearch.sql.spark.scheduler.job.AsyncQueryScheduledQueryJob;
 import org.opensearch.sql.spark.storage.SparkStorageFactory;
 import org.opensearch.sql.spark.transport.TransportCancelAsyncQueryRequestAction;
 import org.opensearch.sql.spark.transport.TransportCreateAsyncQueryRequestAction;
@@ -266,7 +266,7 @@ public String getJobIndex() {
 
   @Override
   public ScheduledJobRunner getJobRunner() {
-    return ScheduledQueryJob.getJobRunnerInstance();
+    return AsyncQueryScheduledQueryJob.getJobRunnerInstance();
   }
 
   @Override