455.execute script short api

query/src/main/java/tech/ydb/query/script/ScriptClient.java

@@ -0,0 +1,92 @@
+package tech.ydb.query.script;
+
+import java.util.concurrent.CompletableFuture;
+
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+import tech.ydb.core.Result;
+import tech.ydb.core.Status;
+import tech.ydb.core.operation.Operation;
+import tech.ydb.core.operation.OperationTray;
+import tech.ydb.query.script.result.ScriptResultPart;
+import tech.ydb.query.script.settings.ExecuteScriptSettings;
+import tech.ydb.query.script.settings.FetchScriptSettings;
+import tech.ydb.query.script.settings.FindScriptSettings;
+import tech.ydb.table.query.Params;
+
+/**
+ * High-level API for executing YQL scripts and retrieving their results.
+ * <p>
+ * Provides convenience methods for starting script execution, tracking operation status,
+ * and fetching result sets with pagination support.
+ * <p>
+ * How to use
+ * <ul>
+ *     <li>startQueryScript - starting script execution or findQueryScript if script had already started</li>
+ *     <li>fetchQueryScriptStatus - wait for script execution</li>
+ *     <li>fetchQueryScriptResult - fetch script result if necessary</li>
+ * </ul>
+ * <p>Example with fetch
+ * <pre>{@code
+ *      Operation<Status> operation = scriptClient.startQueryScript("select...",Params.of(...), executeScriptSettings).join())
+ *      Status status = scriptClient.fetchQueryScriptStatus(operation, 1).join()
+ *      Result< ScriptResultPart> resultPartResult = scriptClient.fetchQueryScriptResult(operation, null, fetchScriptSettings).join()
+ *      ResultSetReader reader = scriptResultPart.getResultSetReader()
+ *      reader.next()
+ * }</pre>
+ * <p>Example without fetch
+ * <pre>{@code
+ * Status status = scriptClient.startQueryScript("select...",Params.of(...), executeScriptSettings)
+ *                             .thenCompose(p -> scriptClient.fetchQueryScriptStatus(p, 1))
+ *                             .join()
+ * }</pre>
+ * <p>Author: Evgeny Kuvardin
+ */
+public interface ScriptClient {
+
+    /**
+     * Returns operation metadata for a previously started script execution.
+     *
+     * @param operationId operation identifier
+     * @param settings    request settings
+     * @return future resolving to operation status
+     */
+    CompletableFuture<Operation<Status>> findQueryScript(String operationId, FindScriptSettings settings);
+
+    /**
+     * Starts execution of the given YQL script with optional parameters.
+     *
+     * @param query    YQL script text
+     * @param params   query parameters
+     * @param settings execution settings (TTL, resource pool, exec mode)
+     * @return future resolving to a long-running operation
+     */
+    CompletableFuture<Operation<Status>> startQueryScript(String query,
+                                                          Params params,
+                                                          ExecuteScriptSettings settings);
+
+    /**
+     * Wait for script execution and return status
+     *
+     * @param operation operation object returned when script started
+     * @param fetchRateSeconds How often should we check if the operation has finished
+     * @return future with result of script execution
+     */
+    default CompletableFuture<Status> fetchQueryScriptStatus(Operation<Status> operation, int fetchRateSeconds) {
+        return OperationTray.fetchOperation(operation, fetchRateSeconds);
+    }
+
+    /**
+     * Fetches script results incrementally.
+     *
+     * @param operation operation object returned when script started
+     * @param previous  previous result part, or {@code null} if fetching from start
+     * @param settings  fetch configuration
+     * @return future resolving to result part containing a result set fragment
+     */
+    CompletableFuture<Result<ScriptResultPart>> fetchQueryScriptResult(@Nonnull Operation<Status> operation,
+                                                                       @Nullable ScriptResultPart previous,
+                                                                       FetchScriptSettings settings);
+
+}

query/src/main/java/tech/ydb/query/script/ScriptRpc.java

@@ -0,0 +1,48 @@
+package tech.ydb.query.script;
+
+import java.util.concurrent.CompletableFuture;
+
+import tech.ydb.core.Result;
+import tech.ydb.core.Status;
+import tech.ydb.core.grpc.GrpcRequestSettings;
+import tech.ydb.core.operation.Operation;
+import tech.ydb.proto.query.YdbQuery;
+
+/**
+ * Low-level RPC interface for executing YQL scripts and fetching their results using gRPC.
+ * <p>
+ * Provides direct bindings to the YDB QueryService API
+ * Used internally by {@link tech.ydb.query.script.ScriptClient} implementations.
+ *
+ * <p>Author: Evgeny Kuvardin
+ */
+public interface ScriptRpc {
+
+    /**
+     * Retrieves a previously created operation by its ID.
+     *
+     * @param operationId ID of the operation to fetch
+     * @return future resolving to the operation metadata and status
+     */
+    CompletableFuture<Operation<Status>> getOperation(String operationId);
+
+    /**
+     * Executes a script as a long-running operation.
+     *
+     * @param request  execution request describing the script and execution mode {@link YdbQuery.ExecuteScriptRequest}
+     * @param settings RPC request settings including timeout, trace ID, etc.
+     * @return future resolving to an {@link Operation} representing the script execution
+     */
+    CompletableFuture<Operation<Status>> executeScript(
+            YdbQuery.ExecuteScriptRequest request, GrpcRequestSettings settings);
+
+    /**
+     * Fetches partial results for a previously executed script.
+     *
+     * @param request  fetch request including token, result set index, etc. {@link YdbQuery.FetchScriptResultsRequest}
+     * @param settings RPC settings for this request
+     * @return future resolving to the result fetch response {@link Result} of {@link YdbQuery.FetchScriptResultsResponse}
+     */
+    CompletableFuture<Result<YdbQuery.FetchScriptResultsResponse>> fetchScriptResults(
+            YdbQuery.FetchScriptResultsRequest request, GrpcRequestSettings settings);
+}

query/src/main/java/tech/ydb/query/script/impl/ScriptClientImpl.java

@@ -0,0 +1,150 @@
+package tech.ydb.query.script.impl;
+
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+import javax.annotation.WillNotClose;
+
+import com.google.protobuf.Duration;
+
+import tech.ydb.core.Result;
+import tech.ydb.core.Status;
+import tech.ydb.core.grpc.GrpcRequestSettings;
+import tech.ydb.core.grpc.GrpcTransport;
+import tech.ydb.core.operation.Operation;
+import tech.ydb.core.settings.BaseRequestSettings;
+import tech.ydb.proto.query.YdbQuery;
+import tech.ydb.query.script.ScriptClient;
+import tech.ydb.query.script.ScriptRpc;
+import tech.ydb.query.script.result.ScriptResultPart;
+import tech.ydb.query.script.settings.ExecuteScriptSettings;
+import tech.ydb.query.script.settings.FetchScriptSettings;
+import tech.ydb.query.script.settings.FindScriptSettings;
+import tech.ydb.query.settings.QueryExecMode;
+import tech.ydb.query.settings.QueryStatsMode;
+import tech.ydb.table.query.Params;
+
+/**
+ * Default implementation of {@link ScriptClient} using {@link ScriptRpc} for RPC calls.
+ * <p>
+ * Handles script execution lifecycle: starting scripts, polling their status,
+ * and retrieving result sets in streaming fashion.
+ *
+ * <p>Author: Evgeny Kuvardin
+ */
+public class ScriptClientImpl implements ScriptClient {
+
+    private final ScriptRpc scriptRpc;
+
+    ScriptClientImpl(ScriptRpc scriptRpc) {
+        this.scriptRpc = scriptRpc;
+    }
+
+    public static ScriptClient newClient(@WillNotClose GrpcTransport transport) {
+        return new ScriptClientImpl(ScriptRpcImpl.useTransport(transport));
+    }
+
+    @Override
+    public CompletableFuture<Operation<Status>> findQueryScript(String operationId, FindScriptSettings settings) {
+        return scriptRpc.getOperation(operationId);
+    }
+
+    @Override
+    public CompletableFuture<Operation<Status>> startQueryScript(String query,
+                                                                 Params params,
+                                                                 ExecuteScriptSettings settings) {
+        YdbQuery.ExecuteScriptRequest.Builder request = YdbQuery.ExecuteScriptRequest.newBuilder()
+                .setExecMode(mapExecMode(settings.getExecMode()))
+                .setStatsMode(mapStatsMode(settings.getStatsMode()))
+                .setScriptContent(YdbQuery.QueryContent.newBuilder()
+                        .setSyntax(YdbQuery.Syntax.SYNTAX_YQL_V1)
+                        .setText(query)
+                        .build());
+
+        java.time.Duration ttl = settings.getTtl();
+        if (ttl != null) {
+            request.setResultsTtl(Duration.newBuilder().setNanos(settings.getTtl().getNano()));
+        }
+
+        String resourcePool = settings.getResourcePool();
+        if (resourcePool != null && !resourcePool.isEmpty()) {
+            request.setPoolId(resourcePool);
+        }
+
+        request.putAllParameters(params.toPb());
+
+        GrpcRequestSettings options = makeGrpcRequestSettings(settings);
+
+        return scriptRpc.executeScript(request.build(), options);
+    }
+
+    @Override
+    public CompletableFuture<Result<ScriptResultPart>> fetchQueryScriptResult(@Nonnull Operation<Status> operation,
+                                                                              @Nullable ScriptResultPart previous,
+                                                                              FetchScriptSettings settings) {
+        YdbQuery.FetchScriptResultsRequest.Builder requestBuilder = YdbQuery.FetchScriptResultsRequest.newBuilder();
+
+        if (previous != null && previous.getNextFetchToken() != null) {
+            requestBuilder.setFetchToken(previous.getNextFetchToken());
+        }
+
+        if (settings.getRowsLimit() > 0) {
+            requestBuilder.setRowsLimit(settings.getRowsLimit());
+        }
+
+        requestBuilder.setOperationId(operation.getId());
+
+        if (settings.getSetResultSetIndex() >= 0) {
+            requestBuilder.setResultSetIndex(settings.getSetResultSetIndex());
+        }
+
+        GrpcRequestSettings options = makeGrpcRequestSettings(settings);
+
+        return scriptRpc.fetchScriptResults(requestBuilder.build(), options)
+                .thenApply(p -> p.map(ScriptResultPart::new));
+    }
+
+    private GrpcRequestSettings makeGrpcRequestSettings(BaseRequestSettings settings) {
+        String traceId = settings.getTraceId() == null ? UUID.randomUUID().toString() : settings.getTraceId();
+        return GrpcRequestSettings.newBuilder()
+                .withDeadline(settings.getRequestTimeout())
+                .withTraceId(traceId)
+                .build();
+    }
+
+    private static YdbQuery.ExecMode mapExecMode(QueryExecMode mode) {
+        switch (mode) {
+            case EXECUTE:
+                return YdbQuery.ExecMode.EXEC_MODE_EXECUTE;
+            case EXPLAIN:
+                return YdbQuery.ExecMode.EXEC_MODE_EXPLAIN;
+            case PARSE:
+                return YdbQuery.ExecMode.EXEC_MODE_PARSE;
+            case VALIDATE:
+                return YdbQuery.ExecMode.EXEC_MODE_VALIDATE;
+
+            case UNSPECIFIED:
+            default:
+                return YdbQuery.ExecMode.EXEC_MODE_UNSPECIFIED;
+        }
+    }
+
+    private static YdbQuery.StatsMode mapStatsMode(QueryStatsMode mode) {
+        switch (mode) {
+            case NONE:
+                return YdbQuery.StatsMode.STATS_MODE_NONE;
+            case BASIC:
+                return YdbQuery.StatsMode.STATS_MODE_BASIC;
+            case FULL:
+                return YdbQuery.StatsMode.STATS_MODE_FULL;
+            case PROFILE:
+                return YdbQuery.StatsMode.STATS_MODE_PROFILE;
+
+            case UNSPECIFIED:
+            default:
+                return YdbQuery.StatsMode.STATS_MODE_UNSPECIFIED;
+        }
+    }
+}

query/src/main/java/tech/ydb/query/script/impl/ScriptRpcImpl.java

@@ -0,0 +1,79 @@
+package tech.ydb.query.script.impl;
+
+import java.util.concurrent.CompletableFuture;
+
+import javax.annotation.WillNotClose;
+
+import tech.ydb.core.Result;
+import tech.ydb.core.Status;
+import tech.ydb.core.grpc.GrpcRequestSettings;
+import tech.ydb.core.grpc.GrpcTransport;
+import tech.ydb.core.operation.Operation;
+import tech.ydb.core.operation.OperationBinder;
+import tech.ydb.proto.OperationProtos;
+import tech.ydb.proto.operation.v1.OperationServiceGrpc;
+import tech.ydb.proto.query.YdbQuery;
+import tech.ydb.proto.query.v1.QueryServiceGrpc;
+import tech.ydb.query.script.ScriptRpc;
+
+/**
+ * Default gRPC-based implementation of {@link ScriptRpc}.
+ * <p>
+ * Uses {@link GrpcTransport} to communicate with YDB QueryService and OperationService.
+ * Provides async unary calls for executing scripts and retrieving results or operation metadata.
+ *
+ * <p>Author: Evgeny Kuvardin
+ */
+public class ScriptRpcImpl implements ScriptRpc {
+
+    private final GrpcTransport transport;
+
+    private ScriptRpcImpl(GrpcTransport grpcTransport) {
+        this.transport = grpcTransport;
+    }
+
+    /**
+     * Creates a new RPC instance bound to the given gRPC transport.
+     *
+     * @param grpcTransport transport instance (not closed by this class)
+     * @return new {@link ScriptRpcImpl} instance
+     */
+    public static ScriptRpcImpl useTransport(@WillNotClose GrpcTransport grpcTransport) {
+        return new ScriptRpcImpl(grpcTransport);
+    }
+
+    @Override
+    public CompletableFuture<Operation<Status>> getOperation(String operationId) {
+        OperationProtos.GetOperationRequest request = OperationProtos.GetOperationRequest.newBuilder()
+                .setId(operationId)
+                .build();
+
+        GrpcRequestSettings settings = GrpcRequestSettings.newBuilder().build();
+
+        return transport
+                .unaryCall(OperationServiceGrpc.getGetOperationMethod(), settings, request)
+                .thenApply(
+                        OperationBinder.bindAsync(transport,
+                                OperationProtos.GetOperationResponse::getOperation
+                        ));
+    }
+
+    @Override
+    public CompletableFuture<Operation<Status>> executeScript(
+            YdbQuery.ExecuteScriptRequest request, GrpcRequestSettings settings) {
+
+        return transport.unaryCall(QueryServiceGrpc.getExecuteScriptMethod(), settings, request)
+                .thenApply(
+                        OperationBinder.bindAsync(transport,
+                                op -> op
+                        ));
+    }
+
+    @Override
+    public CompletableFuture<Result<YdbQuery.FetchScriptResultsResponse>> fetchScriptResults(
+            YdbQuery.FetchScriptResultsRequest request, GrpcRequestSettings settings) {
+
+        return transport
+                .unaryCall(QueryServiceGrpc.getFetchScriptResultsMethod(), settings, request);
+    }
+}