[Feature] Integrate SQL/PPL with query-insights plugin

legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java

@@ -90,14 +90,26 @@ public class RestSqlAction extends BaseRestHandler {
    */
   private final BiFunction<SQLQueryRequest, RestChannel, Boolean> analyticsRouter;
 
+  /** Runs the SQL execution under a coordinator task so DSL searches link back to the SQL query. */
+  private final SqlCoordinatorTaskDispatcher coordinatorTaskDispatcher;
+
   public RestSqlAction(
       Settings settings,
       Injector injector,
       BiFunction<SQLQueryRequest, RestChannel, Boolean> analyticsRouter) {
+    this(settings, injector, analyticsRouter, SqlCoordinatorTaskDispatcher.PASSTHROUGH);
+  }
+
+  public RestSqlAction(
+      Settings settings,
+      Injector injector,
+      BiFunction<SQLQueryRequest, RestChannel, Boolean> analyticsRouter,
+      SqlCoordinatorTaskDispatcher coordinatorTaskDispatcher) {
     super();
     this.allowExplicitIndex = MULTI_ALLOW_EXPLICIT_INDEX.get(settings);
     this.newSqlQueryHandler = new RestSQLQueryAction(injector);
     this.analyticsRouter = analyticsRouter;
+    this.coordinatorTaskDispatcher = coordinatorTaskDispatcher;
   }
 
   @Override
@@ -156,11 +168,18 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
       // Route to analytics engine for non-Lucene (e.g., Parquet-backed) indices.
       // The router returns true and sends the response directly if it handled the request.
       final SQLQueryRequest finalRequest = newSqlRequest;
-      return channel -> {
-        if (!analyticsRouter.apply(finalRequest, channel)) {
-          delegateToV2Engine(request, client, sqlRequest, finalRequest, format, channel);
-        }
-      };
+      // Run under a coordinator task so every DSL search the SQL engine issues carries a parent
+      // reference back to this SQL query (used by query-insights to correlate and recover source).
+      return channel ->
+          coordinatorTaskDispatcher.dispatch(
+              client,
+              sqlRequest.getSql(),
+              channel,
+              ch -> {
+                if (!analyticsRouter.apply(finalRequest, ch)) {
+                  delegateToV2Engine(request, client, sqlRequest, finalRequest, format, ch);
+                }
+              });
     } catch (Exception e) {
       return channel -> handleException(channel, e);
     }

legacy/src/main/java/org/opensearch/sql/legacy/plugin/SqlCoordinatorTaskDispatcher.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.legacy.plugin;
+
+import java.util.function.Consumer;
+import org.opensearch.rest.RestChannel;
+import org.opensearch.transport.client.node.NodeClient;
+
+/**
+ * Runs a SQL execution under a coordinator task so the DSL search tasks it spawns carry a parent
+ * reference back to the originating SQL query.
+ *
+ * <p>The concrete implementation (wired in the plugin module) dispatches through a local transport
+ * action that registers the coordinator task; this seam exists because {@code RestSqlAction} lives
+ * in the {@code legacy} module and cannot depend on the plugin-module transport action directly.
+ * The default implementation simply runs the work with no coordinator task, preserving behavior for
+ * callers/tests that do not supply one.
+ */
+@FunctionalInterface
+public interface SqlCoordinatorTaskDispatcher {
+
+  /** Default: run the work directly without establishing a coordinator task. */
+  SqlCoordinatorTaskDispatcher PASSTHROUGH = (client, query, channel, work) -> work.accept(channel);
+
+  /**
+   * @param client node client used to dispatch the local coordinator-task action
+   * @param query original SQL query text (recorded on the coordinator task)
+   * @param channel REST channel the execution writes its response to
+   * @param work the SQL execution to run under the coordinator task
+   */
+  void dispatch(NodeClient client, String query, RestChannel channel, Consumer<RestChannel> work);
+}

legacy/src/test/java/org/opensearch/sql/legacy/plugin/SqlCoordinatorTaskDispatcherTest.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.legacy.plugin;
+
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
+import org.junit.Test;
+import org.opensearch.rest.RestChannel;
+import org.opensearch.transport.client.node.NodeClient;
+
+public class SqlCoordinatorTaskDispatcherTest {
+
+  @Test
+  public void passthroughRunsWorkWithGivenChannel() {
+    NodeClient client = mock(NodeClient.class);
+    RestChannel channel = mock(RestChannel.class);
+    AtomicBoolean ran = new AtomicBoolean(false);
+    AtomicReference<RestChannel> received = new AtomicReference<>();
+    Consumer<RestChannel> work =
+        ch -> {
+          ran.set(true);
+          received.set(ch);
+        };
+
+    SqlCoordinatorTaskDispatcher.PASSTHROUGH.dispatch(client, "SELECT 1", channel, work);
+
+    assertTrue(ran.get());
+    assertSame(channel, received.get());
+  }
+}

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/BackgroundSearchScanner.java

@@ -21,8 +21,10 @@
 import org.opensearch.sql.monitor.profile.ProfileContext;
 import org.opensearch.sql.monitor.profile.QueryProfiling;
 import org.opensearch.sql.opensearch.client.OpenSearchClient;
+import org.opensearch.sql.opensearch.executor.OpenSearchQueryManager;
 import org.opensearch.sql.opensearch.request.OpenSearchRequest;
 import org.opensearch.sql.opensearch.response.OpenSearchResponse;
+import org.opensearch.tasks.CancellableTask;
 
 /**
  * Utility class for asynchronously scanning an index. This lets us send background requests to the
@@ -106,13 +108,40 @@ public boolean isScanDone() {
   public void startScanning(OpenSearchRequest request) {
     if (isAsync()) {
       ProfileContext ctx = QueryProfiling.current();
+      // Capture the coordinator task on the calling thread so the background search thread can
+      // re-establish it. Without this, applyParentTask() in OpenSearchNodeClient sees a null task
+      // on the background pool and the prefetched DSL search task loses its parent (SQL/PPL) link.
+      CancellableTask task = OpenSearchQueryManager.getCancellableTask();
       nextBatchFuture =
           CompletableFuture.supplyAsync(
-              () -> QueryProfiling.withCurrentContext(ctx, () -> client.search(request)),
+              () -> QueryProfiling.withCurrentContext(ctx, () -> searchWithTask(request, task)),
               backgroundExecutor);
     }
   }
 
+  /**
+   * Runs {@code client.search} with the given coordinator task bound to this (pooled) thread's
+   * ThreadLocal, so parent-task linkage is applied to the outgoing DSL search request. Restores the
+   * thread's previous task afterward to keep the shared background pool clean.
+   */
+  private OpenSearchResponse searchWithTask(
+      OpenSearchRequest request, @Nullable CancellableTask task) {
+    if (task == null) {
+      return client.search(request);
+    }
+    CancellableTask previous = OpenSearchQueryManager.getCancellableTask();
+    OpenSearchQueryManager.setCancellableTask(task);
+    try {
+      return client.search(request);
+    } finally {
+      if (previous != null) {
+        OpenSearchQueryManager.setCancellableTask(previous);
+      } else {
+        OpenSearchQueryManager.clearCancellableTask();
+      }
+    }
+  }
+
   private OpenSearchResponse getCurrentResponse(OpenSearchRequest request) {
     if (isAsync()) {
       try {
@@ -176,8 +205,10 @@ public SearchBatchResult fetchNextBatch(OpenSearchRequest request) {
 
       // Pre-fetch next batch if needed
       if (!stopIteration && isAsync()) {
+        CancellableTask task = OpenSearchQueryManager.getCancellableTask();
         nextBatchFuture =
-            CompletableFuture.supplyAsync(() -> client.search(request), backgroundExecutor);
+            CompletableFuture.supplyAsync(
+                () -> searchWithTask(request, task), backgroundExecutor);
       }
     } else {
       iterator = Collections.emptyIterator();

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

@@ -102,6 +102,7 @@
 import org.opensearch.sql.legacy.metrics.Metrics;
 import org.opensearch.sql.legacy.plugin.RestSqlAction;
 import org.opensearch.sql.legacy.plugin.RestSqlStatsAction;
+import org.opensearch.sql.legacy.plugin.SqlCoordinatorTaskDispatcher;
 import org.opensearch.sql.opensearch.client.OpenSearchNodeClient;
 import org.opensearch.sql.opensearch.setting.OpenSearchSettings;
 import org.opensearch.sql.opensearch.storage.OpenSearchDataSourceFactory;
@@ -119,8 +120,12 @@
 import org.opensearch.sql.plugin.rest.RestQuerySettingsAction;
 import org.opensearch.sql.plugin.rest.RestUnifiedQueryAction;
 import org.opensearch.sql.plugin.transport.PPLQueryAction;
+import org.opensearch.sql.plugin.transport.SqlQueryAction;
 import org.opensearch.sql.plugin.transport.TransportPPLQueryAction;
 import org.opensearch.sql.plugin.transport.TransportPPLQueryResponse;
+import org.opensearch.sql.plugin.transport.TransportSqlQueryAction;
+import org.opensearch.sql.plugin.transport.TransportSqlQueryRequest;
+import org.opensearch.sql.plugin.transport.TransportSqlQueryResponse;
 import org.opensearch.sql.prometheus.storage.PrometheusStorageFactory;
 import org.opensearch.sql.protocol.response.format.JsonResponseFormatter;
 import org.opensearch.sql.protocol.response.format.JsonResponseFormatter.Style;
@@ -217,7 +222,8 @@ public List<RestHandler> getRestHandlers(
     return Arrays.asList(
         new RestPPLQueryAction(),
         new RestPPLGrammarAction(),
-        new RestSqlAction(settings, injector, createSqlAnalyticsRouter()),
+        new RestSqlAction(
+            settings, injector, createSqlAnalyticsRouter(), createSqlCoordinatorTaskDispatcher()),
         new RestSqlStatsAction(settings, restController),
         new RestPPLStatsAction(settings, restController),
         new RestQuerySettingsAction(settings, restController),
@@ -323,13 +329,40 @@ public void onFailure(Exception e) {
     };
   }
 
+  /**
+   * Dispatcher that runs each SQL execution under a coordinator {@link SqlQueryTask} via a local
+   * transport action, so the DSL search tasks the SQL engine spawns reference the SQL query as
+   * their parent.
+   */
+  private SqlCoordinatorTaskDispatcher createSqlCoordinatorTaskDispatcher() {
+    return (nodeClient, query, channel, work) ->
+        nodeClient.executeLocally(
+            SqlQueryAction.INSTANCE,
+            new TransportSqlQueryRequest(query, work, channel),
+            new ActionListener<TransportSqlQueryResponse>() {
+              @Override
+              public void onResponse(TransportSqlQueryResponse response) {
+                // The SQL result was already written to the REST channel by the execution path.
+              }
+
+              @Override
+              public void onFailure(Exception e) {
+                // Only reached when execution failed before sending any response.
+                RestSqlAction.handleException(channel, e);
+              }
+            });
+  }
+
   /** Register action and handler so that transportClient can find proxy for action. */
   @Override
   public List<ActionHandler<? extends ActionRequest, ? extends ActionResponse>> getActions() {
     return Arrays.asList(
         new ActionHandler<>(
             new ActionType<>(PPLQueryAction.NAME, TransportPPLQueryResponse::new),
             TransportPPLQueryAction.class),
+        new ActionHandler<>(
+            new ActionType<>(SqlQueryAction.NAME, TransportSqlQueryResponse::new),
+            TransportSqlQueryAction.class),
         new ActionHandler<>(
             new ActionType<>(
                 TransportCreateDataSourceAction.NAME, CreateDataSourceActionResponse::new),

plugin/src/main/java/org/opensearch/sql/plugin/transport/SqlQueryAction.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.plugin.transport;
+
+import org.opensearch.action.ActionType;
+
+/**
+ * Internal action used to create a coordinator task for a SQL query. Not exposed as a public REST
+ * API; it is dispatched locally from {@code RestSqlAction} via {@code NodeClient.executeLocally} so
+ * the transport framework registers a {@link SqlQueryTask} for the duration of the query.
+ */
+public class SqlQueryAction extends ActionType<TransportSqlQueryResponse> {
+  public static final String NAME = "cluster:admin/opensearch/sql";
+  public static final SqlQueryAction INSTANCE = new SqlQueryAction();
+
+  private SqlQueryAction() {
+    super(NAME, TransportSqlQueryResponse::new);
+  }
+}

plugin/src/main/java/org/opensearch/sql/plugin/transport/SqlQueryTask.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.plugin.transport;
+
+import java.util.Map;
+import org.opensearch.core.tasks.TaskId;
+import org.opensearch.tasks.CancellableTask;
+
+/**
+ * Coordinator task for a SQL query. Registering this task in the {@code TaskManager} gives every
+ * DSL search task spawned by the SQL engine a parent reference back to the originating SQL query,
+ * so downstream consumers (e.g. query-insights) can correlate the DSL searches with their SQL
+ * source and look up the original request from the task description.
+ */
+public class SqlQueryTask extends CancellableTask {
+  public SqlQueryTask(
+      long id,
+      String type,
+      String action,
+      String description,
+      TaskId parentTaskId,
+      Map<String, String> headers) {
+    super(id, type, action, description, parentTaskId, headers);
+  }
+
+  @Override
+  public boolean shouldCancelChildrenOnCancellation() {
+    return true;
+  }
+}

plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java

@@ -179,6 +179,7 @@ protected void doExecute(
 
     // in order to use PPL service, we need to convert TransportPPLQueryRequest to PPLQueryRequest
     PPLQueryRequest transformedRequest = transportRequest.toPPLQueryRequest();
+
     QueryContext.setProfile(transformedRequest.profile());
     ActionListener<TransportPPLQueryResponse> clearingListener = wrapWithProfilingClear(listener);