diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java index 4064b73d4a4..ae92a93b071 100644 --- a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java +++ b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java @@ -90,14 +90,26 @@ public class RestSqlAction extends BaseRestHandler { */ private final BiFunction 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 analyticsRouter) { + this(settings, injector, analyticsRouter, SqlCoordinatorTaskDispatcher.PASSTHROUGH); + } + + public RestSqlAction( + Settings settings, + Injector injector, + BiFunction 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); } diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/SqlCoordinatorTaskDispatcher.java b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/SqlCoordinatorTaskDispatcher.java new file mode 100644 index 00000000000..d2db1814de1 --- /dev/null +++ b/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. + * + *

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 work); +} diff --git a/legacy/src/test/java/org/opensearch/sql/legacy/plugin/SqlCoordinatorTaskDispatcherTest.java b/legacy/src/test/java/org/opensearch/sql/legacy/plugin/SqlCoordinatorTaskDispatcherTest.java new file mode 100644 index 00000000000..a1f4ac1662b --- /dev/null +++ b/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 received = new AtomicReference<>(); + Consumer work = + ch -> { + ran.set(true); + received.set(ch); + }; + + SqlCoordinatorTaskDispatcher.PASSTHROUGH.dispatch(client, "SELECT 1", channel, work); + + assertTrue(ran.get()); + assertSame(channel, received.get()); + } +} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/BackgroundSearchScanner.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/BackgroundSearchScanner.java index 3aa347b70fa..dbc38f4590a 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/BackgroundSearchScanner.java +++ b/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(); diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java b/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java index 31f14e3411b..f23d3e61617 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java +++ b/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 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,6 +329,30 @@ 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() { + @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> getActions() { @@ -330,6 +360,9 @@ public void onFailure(Exception e) { 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), diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/SqlQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/SqlQueryAction.java new file mode 100644 index 00000000000..9a7dc816f20 --- /dev/null +++ b/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 { + public static final String NAME = "cluster:admin/opensearch/sql"; + public static final SqlQueryAction INSTANCE = new SqlQueryAction(); + + private SqlQueryAction() { + super(NAME, TransportSqlQueryResponse::new); + } +} diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/SqlQueryTask.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/SqlQueryTask.java new file mode 100644 index 00000000000..608742040f1 --- /dev/null +++ b/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 headers) { + super(id, type, action, description, parentTaskId, headers); + } + + @Override + public boolean shouldCancelChildrenOnCancellation() { + return true; + } +} diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java index 772f1ec123f..3edd9a40908 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportPPLQueryAction.java +++ b/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 clearingListener = wrapWithProfilingClear(listener); diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryAction.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryAction.java new file mode 100644 index 00000000000..73d470cab7e --- /dev/null +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryAction.java @@ -0,0 +1,145 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +import java.io.IOException; +import java.util.concurrent.atomic.AtomicBoolean; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.HandledTransportAction; +import org.opensearch.common.inject.Inject; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.xcontent.MediaType; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.rest.RestChannel; +import org.opensearch.rest.RestRequest; +import org.opensearch.rest.RestResponse; +import org.opensearch.sql.opensearch.executor.OpenSearchQueryManager; +import org.opensearch.tasks.CancellableTask; +import org.opensearch.tasks.Task; +import org.opensearch.transport.TransportService; + +/** + * Establishes a coordinator {@link SqlQueryTask} for the duration of a SQL query. The transport + * framework registers the task (via {@link TransportSqlQueryRequest#createTask}) before {@link + * #doExecute} runs and unregisters it when the response listener completes. We bind the task to the + * executing thread's {@code OpenSearchQueryManager} ThreadLocal so the existing {@code + * applyParentTask} logic in {@code OpenSearchNodeClient} stamps it as the parent on every DSL + * search the SQL engine issues. The listener is completed once the wrapped REST channel emits its + * response, keeping the task alive across the (possibly asynchronous) execution. + */ +public class TransportSqlQueryAction + extends HandledTransportAction { + + @Inject + public TransportSqlQueryAction(TransportService transportService, ActionFilters actionFilters) { + super(SqlQueryAction.NAME, transportService, actionFilters, TransportSqlQueryRequest::new); + } + + @Override + protected void doExecute( + Task task, + TransportSqlQueryRequest request, + ActionListener listener) { + CancellableTask previous = OpenSearchQueryManager.getCancellableTask(); + if (task instanceof CancellableTask cancellableTask) { + OpenSearchQueryManager.setCancellableTask(cancellableTask); + } + + AtomicBoolean completed = new AtomicBoolean(false); + RestChannel channel = new CompletionSignalingChannel(request.getChannel(), listener, completed); + try { + request.getWork().accept(channel); + } catch (Exception e) { + // Only surface here if the channel hasn't already reported completion (success or error). + if (completed.compareAndSet(false, true)) { + listener.onFailure(e); + } + } finally { + // Restore the initiating thread's prior task so this pooled transport worker does not retain + // a stale reference. Any async workers continuing the execution have already captured the + // task synchronously during work.accept above, so clearing here does not race them. + if (previous != null) { + OpenSearchQueryManager.setCancellableTask(previous); + } else { + OpenSearchQueryManager.clearCancellableTask(); + } + } + } + + /** + * Delegates all channel behavior to the real channel, and completes the transport listener the + * first time a response is sent — which is what triggers unregistration of the coordinator task. + */ + private static final class CompletionSignalingChannel implements RestChannel { + private final RestChannel delegate; + private final ActionListener listener; + private final AtomicBoolean completed; + + CompletionSignalingChannel( + RestChannel delegate, + ActionListener listener, + AtomicBoolean completed) { + this.delegate = delegate; + this.listener = listener; + this.completed = completed; + } + + @Override + public void sendResponse(RestResponse response) { + try { + delegate.sendResponse(response); + } finally { + if (completed.compareAndSet(false, true)) { + listener.onResponse(new TransportSqlQueryResponse()); + } + } + } + + @Override + public XContentBuilder newBuilder() throws IOException { + return delegate.newBuilder(); + } + + @Override + public XContentBuilder newErrorBuilder() throws IOException { + return delegate.newErrorBuilder(); + } + + @Override + public XContentBuilder newBuilder(MediaType mediaType, boolean useFiltering) + throws IOException { + return delegate.newBuilder(mediaType, useFiltering); + } + + @Override + public XContentBuilder newBuilder( + MediaType mediaType, MediaType responseContentType, boolean useFiltering) + throws IOException { + return delegate.newBuilder(mediaType, responseContentType, useFiltering); + } + + @Override + public BytesStreamOutput bytesOutput() { + return delegate.bytesOutput(); + } + + @Override + public RestRequest request() { + return delegate.request(); + } + + @Override + public boolean detailedErrorsEnabled() { + return delegate.detailedErrorsEnabled(); + } + + @Override + public boolean detailedErrorStackTraceEnabled() { + return delegate.detailedErrorStackTraceEnabled(); + } + } +} diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryRequest.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryRequest.java new file mode 100644 index 00000000000..1d48b6e0319 --- /dev/null +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryRequest.java @@ -0,0 +1,83 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +import java.io.IOException; +import java.util.Map; +import java.util.function.Consumer; +import org.opensearch.action.ActionRequest; +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.rest.RestChannel; +import org.opensearch.tasks.Task; + +/** + * Request for {@link SqlQueryAction}. Dispatched only locally via {@code + * NodeClient.executeLocally}, so the {@link #work} and {@link #channel} references are carried + * in-process (never serialized). The {@link #query} text becomes the coordinator task description + * so consumers can recover the original SQL from the task. + */ +public class TransportSqlQueryRequest extends ActionRequest { + + /** Truncate very long queries so the task description stays bounded. */ + static final int MAX_DESCRIPTION_LENGTH = 4096; + + private final String query; + + /** The SQL execution to run under the coordinator task. Transient — local dispatch only. */ + private final transient Consumer work; + + /** The REST channel the execution writes its response to. Transient — local dispatch only. */ + private final transient RestChannel channel; + + public TransportSqlQueryRequest(String query, Consumer work, RestChannel channel) { + this.query = query == null ? "" : query; + this.work = work; + this.channel = channel; + } + + public TransportSqlQueryRequest(StreamInput in) throws IOException { + super(in); + this.query = in.readString(); + this.work = null; + this.channel = null; + } + + public Consumer getWork() { + return work; + } + + public RestChannel getChannel() { + return channel; + } + + @Override + public Task createTask( + long id, String type, String action, TaskId parentTaskId, Map headers) { + return new SqlQueryTask(id, type, action, getDescription(), parentTaskId, headers); + } + + @Override + public String getDescription() { + if (query.length() > MAX_DESCRIPTION_LENGTH) { + return query.substring(0, MAX_DESCRIPTION_LENGTH); + } + return query; + } + + @Override + public ActionRequestValidationException validate() { + return null; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeString(query); + } +} diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryResponse.java b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryResponse.java new file mode 100644 index 00000000000..4d1a891122f --- /dev/null +++ b/plugin/src/main/java/org/opensearch/sql/plugin/transport/TransportSqlQueryResponse.java @@ -0,0 +1,30 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +import java.io.IOException; +import org.opensearch.core.action.ActionResponse; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; + +/** + * Completion signal for {@link SqlQueryAction}. The actual SQL response is written directly to the + * REST channel by the execution path; this response only marks that the coordinator task can be + * unregistered, so it carries no payload. + */ +public class TransportSqlQueryResponse extends ActionResponse { + + public TransportSqlQueryResponse() {} + + public TransportSqlQueryResponse(StreamInput in) throws IOException { + super(in); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + // No payload: the SQL result is delivered via the REST channel, not this transport response. + } +} diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/SqlQueryTaskTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/SqlQueryTaskTest.java new file mode 100644 index 00000000000..a63af969dda --- /dev/null +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/SqlQueryTaskTest.java @@ -0,0 +1,34 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import java.util.Map; +import org.junit.Test; +import org.opensearch.core.tasks.TaskId; + +public class SqlQueryTaskTest { + + private SqlQueryTask newTask() { + return new SqlQueryTask( + 1, "transport", SqlQueryAction.NAME, "SELECT 1", TaskId.EMPTY_TASK_ID, Map.of()); + } + + @Test + public void testShouldCancelChildrenReturnsTrue() { + assertTrue(newTask().shouldCancelChildrenOnCancellation()); + } + + @Test + public void testCancellation() { + SqlQueryTask task = newTask(); + assertFalse(task.isCancelled()); + task.cancel("Test"); + assertTrue(task.isCancelled()); + } +} diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportSqlQueryActionTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportSqlQueryActionTest.java new file mode 100644 index 00000000000..baa69193b6d --- /dev/null +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportSqlQueryActionTest.java @@ -0,0 +1,228 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.io.IOException; +import java.util.HashSet; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.core.xcontent.MediaType; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.rest.RestChannel; +import org.opensearch.rest.RestRequest; +import org.opensearch.rest.RestResponse; +import org.opensearch.sql.opensearch.executor.OpenSearchQueryManager; +import org.opensearch.tasks.Task; +import org.opensearch.transport.TransportService; + +public class TransportSqlQueryActionTest { + + private TransportSqlQueryAction action; + + @Before + public void setUp() { + TransportService transportService = mock(TransportService.class); + action = new TransportSqlQueryAction(transportService, new ActionFilters(new HashSet<>())); + } + + @After + public void tearDown() { + // doExecute binds the task to a ThreadLocal and does not clear it on the caller thread. + OpenSearchQueryManager.clearCancellableTask(); + } + + private SqlQueryTask newTask() { + return new SqlQueryTask( + 1, "transport", SqlQueryAction.NAME, "SELECT 1", TaskId.EMPTY_TASK_ID, Map.of()); + } + + @SuppressWarnings("unchecked") + private ActionListener mockListener() { + return mock(ActionListener.class); + } + + @Test + public void completesListenerOnceWhenWorkSendsResponse() { + RestChannel delegate = mock(RestChannel.class); + RestResponse response = mock(RestResponse.class); + Consumer work = ch -> ch.sendResponse(response); + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", work, delegate); + ActionListener listener = mockListener(); + + action.doExecute(newTask(), request, listener); + + verify(delegate, times(1)).sendResponse(response); + verify(listener, times(1)).onResponse(any(TransportSqlQueryResponse.class)); + verify(listener, never()).onFailure(any()); + } + + @Test + public void bindsCancellableTaskDuringExecution() { + SqlQueryTask task = newTask(); + RestChannel delegate = mock(RestChannel.class); + RestResponse response = mock(RestResponse.class); + AtomicReference seen = new AtomicReference<>(); + Consumer work = + ch -> { + seen.set(OpenSearchQueryManager.getCancellableTask()); + ch.sendResponse(response); + }; + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", work, delegate); + + action.doExecute(task, request, mockListener()); + + assertSame(task, seen.get()); + } + + @Test + public void reportsFailureWhenWorkThrowsBeforeResponse() { + RestChannel delegate = mock(RestChannel.class); + RuntimeException boom = new RuntimeException("boom"); + Consumer work = + ch -> { + throw boom; + }; + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", work, delegate); + ActionListener listener = mockListener(); + + action.doExecute(newTask(), request, listener); + + verify(listener, times(1)).onFailure(boom); + verify(listener, never()).onResponse(any()); + } + + @Test + public void doesNotDoubleCompleteWhenWorkThrowsAfterResponse() { + RestChannel delegate = mock(RestChannel.class); + RestResponse response = mock(RestResponse.class); + Consumer work = + ch -> { + ch.sendResponse(response); + throw new RuntimeException("late failure after response already sent"); + }; + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", work, delegate); + ActionListener listener = mockListener(); + + action.doExecute(newTask(), request, listener); + + verify(listener, times(1)).onResponse(any(TransportSqlQueryResponse.class)); + verify(listener, never()).onFailure(any()); + } + + @Test + public void clearsCancellableTaskAfterExecution() { + RestChannel delegate = mock(RestChannel.class); + RestResponse response = mock(RestResponse.class); + Consumer work = ch -> ch.sendResponse(response); + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", work, delegate); + + action.doExecute(newTask(), request, mockListener()); + + // The pooled transport thread must not retain the task after doExecute returns. + assertNull(OpenSearchQueryManager.getCancellableTask()); + } + + @Test + public void clearsCancellableTaskEvenWhenWorkThrows() { + RestChannel delegate = mock(RestChannel.class); + Consumer work = + ch -> { + throw new RuntimeException("boom"); + }; + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", work, delegate); + + action.doExecute(newTask(), request, mockListener()); + + assertNull(OpenSearchQueryManager.getCancellableTask()); + } + + @Test + public void wrappedChannelDelegatesAllMethods() throws IOException { + RestChannel delegate = mock(RestChannel.class); + XContentBuilder builder = mock(XContentBuilder.class); + XContentBuilder errorBuilder = mock(XContentBuilder.class); + BytesStreamOutput bytesOutput = new BytesStreamOutput(); + RestRequest restRequest = mock(RestRequest.class); + MediaType mediaType = mock(MediaType.class); + when(delegate.newBuilder()).thenReturn(builder); + when(delegate.newErrorBuilder()).thenReturn(errorBuilder); + when(delegate.newBuilder(any(MediaType.class), anyBoolean())).thenReturn(builder); + when(delegate.newBuilder(any(MediaType.class), any(MediaType.class), anyBoolean())) + .thenReturn(builder); + when(delegate.bytesOutput()).thenReturn(bytesOutput); + when(delegate.request()).thenReturn(restRequest); + when(delegate.detailedErrorsEnabled()).thenReturn(true); + when(delegate.detailedErrorStackTraceEnabled()).thenReturn(false); + + AtomicReference failure = new AtomicReference<>(); + Consumer work = + ch -> { + try { + assertSame(builder, ch.newBuilder()); + assertSame(errorBuilder, ch.newErrorBuilder()); + assertSame(builder, ch.newBuilder(mediaType, true)); + assertSame(builder, ch.newBuilder(mediaType, mediaType, true)); + assertSame(bytesOutput, ch.bytesOutput()); + assertSame(restRequest, ch.request()); + assertTrue(ch.detailedErrorsEnabled()); + assertFalse(ch.detailedErrorStackTraceEnabled()); + } catch (AssertionError e) { + failure.set(e); + } catch (IOException e) { + failure.set(new AssertionError(e)); + } + }; + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", work, delegate); + + action.doExecute(newTask(), request, mockListener()); + + if (failure.get() != null) { + throw failure.get(); + } + verify(delegate, times(1)).newBuilder(); + verify(delegate, times(1)).newErrorBuilder(); + verify(delegate, times(1)).bytesOutput(); + verify(delegate, times(1)).request(); + } + + @Test + public void secondSendResponseDoesNotCompleteListenerTwice() { + RestChannel delegate = mock(RestChannel.class); + RestResponse first = mock(RestResponse.class); + RestResponse second = mock(RestResponse.class); + Consumer work = + ch -> { + ch.sendResponse(first); + ch.sendResponse(second); + }; + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", work, delegate); + ActionListener listener = mockListener(); + + action.doExecute(newTask(), request, listener); + + verify(listener, times(1)).onResponse(any(TransportSqlQueryResponse.class)); + } +} diff --git a/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportSqlQueryRequestTest.java b/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportSqlQueryRequestTest.java new file mode 100644 index 00000000000..c4e62830fe8 --- /dev/null +++ b/plugin/src/test/java/org/opensearch/sql/plugin/transport/TransportSqlQueryRequestTest.java @@ -0,0 +1,66 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.plugin.transport; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; +import static org.mockito.Mockito.mock; + +import java.util.Map; +import java.util.function.Consumer; +import org.junit.Test; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.rest.RestChannel; +import org.opensearch.tasks.Task; + +public class TransportSqlQueryRequestTest { + + @Test + public void testCreateTaskReturnsSqlQueryTask() { + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", ch -> {}, null); + Task task = + request.createTask(1, "transport", SqlQueryAction.NAME, TaskId.EMPTY_TASK_ID, Map.of()); + assertNotNull(task); + assertEquals(SqlQueryTask.class, task.getClass()); + } + + @Test + public void testDescriptionIsQueryText() { + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", ch -> {}, null); + assertEquals("SELECT 1", request.getDescription()); + } + + @Test + public void testDescriptionTruncatedToMaxLength() { + String longQuery = "x".repeat(TransportSqlQueryRequest.MAX_DESCRIPTION_LENGTH + 100); + TransportSqlQueryRequest request = new TransportSqlQueryRequest(longQuery, ch -> {}, null); + assertEquals( + TransportSqlQueryRequest.MAX_DESCRIPTION_LENGTH, request.getDescription().length()); + } + + @Test + public void testNullQueryBecomesEmptyDescription() { + TransportSqlQueryRequest request = new TransportSqlQueryRequest(null, ch -> {}, null); + assertEquals("", request.getDescription()); + } + + @Test + public void testCarriesWorkAndChannel() { + Consumer work = ch -> {}; + RestChannel channel = mock(RestChannel.class); + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", work, channel); + assertSame(work, request.getWork()); + assertSame(channel, request.getChannel()); + } + + @Test + public void testValidateReturnsNull() { + TransportSqlQueryRequest request = new TransportSqlQueryRequest("SELECT 1", ch -> {}, null); + assertNull(request.validate()); + } +}