IGNITE-23441 Sql. Cancellation of script execution

modules/api/src/main/java/org/apache/ignite/sql/IgniteSql.java

@@ -379,6 +379,16 @@ <T> CompletableFuture<AsyncResultSet<T>> executeAsync(
      */
     void executeScript(String query, @Nullable Object... arguments);
 
+    /**
+     * Executes a multi-statement SQL query.
+     *
+     * @param cancellationToken Cancellation token or {@code null}.
+     * @param query SQL query template.
+     * @param arguments Arguments for the template (optional).
+     * @throws SqlException If failed.
+     */
+    void executeScript(@Nullable CancellationToken cancellationToken, String query, @Nullable Object... arguments);
+
     /**
      * Executes a multi-statement SQL query.
      *
@@ -389,4 +399,15 @@ <T> CompletableFuture<AsyncResultSet<T>> executeAsync(
      */
     CompletableFuture<Void> executeScriptAsync(String query, @Nullable Object... arguments);
 
+    /**
+     * Executes a multi-statement SQL query.
+     *
+     * @param cancellationToken Cancellation token or {@code null}.
+     * @param query SQL query template.
+     * @param arguments Arguments for the template (optional).
+     * @return Operation future.
+     * @throws SqlException If failed.
+     */
+    CompletableFuture<Void> executeScriptAsync(@Nullable CancellationToken cancellationToken, String query, @Nullable Object... arguments);
+
 }

modules/client-handler/src/main/java/org/apache/ignite/client/handler/requests/sql/ClientSqlExecuteScriptRequest.java

@@ -53,8 +53,9 @@ public static CompletableFuture<Void> process(
         HybridTimestamp clientTs = HybridTimestamp.nullableHybridTimestamp(in.unpackLong());
         transactions.updateObservableTimestamp(clientTs);
 
+        // TODO https://issues.apache.org/jira/browse/IGNITE-23646 Pass cancellation token to the query processor.
         return IgniteSqlImpl.executeScriptCore(
-                sql, transactions.observableTimestampTracker(), () -> true, () -> {}, script, arguments, props.toSqlProps()
+                sql, transactions.observableTimestampTracker(), () -> true, () -> {}, script, null, arguments, props.toSqlProps()
         );
     }
 }

modules/client/src/main/java/org/apache/ignite/internal/client/sql/ClientSql.java

@@ -188,6 +188,13 @@ public void executeScript(String query, @Nullable Object... arguments) {
         }
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void executeScript(@Nullable CancellationToken cancellationToken, String query, @Nullable Object... arguments) {
+        // TODO https://issues.apache.org/jira/browse/IGNITE-23646 Support cancellation token.
+        throw new UnsupportedOperationException();
+    }
+
     /** {@inheritDoc} */
     @Override
     public CompletableFuture<AsyncResultSet<SqlRow>> executeAsync(
@@ -335,6 +342,14 @@ public CompletableFuture<Void> executeScriptAsync(String query, @Nullable Object
         return ch.serviceAsync(ClientOp.SQL_EXEC_SCRIPT, payloadWriter, null);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public CompletableFuture<Void> executeScriptAsync(@Nullable CancellationToken cancellationToken, String query,
+            @Nullable Object... arguments) {
+        // TODO https://issues.apache.org/jira/browse/IGNITE-23646 Support cancellation token.
+        throw new UnsupportedOperationException();
+    }
+
     private static void packProperties(
             PayloadOutputChannel w,
             @Nullable Map<String, Object> statementProps) {

modules/runner/src/main/java/org/apache/ignite/internal/restart/RestartProofIgniteSql.java

@@ -206,11 +206,22 @@ public void executeScript(String query, @Nullable Object... arguments) {
         attachmentLock.consumeAttached(ignite -> ignite.sql().executeScript(query, arguments));
     }
 
+    @Override
+    public void executeScript(@Nullable CancellationToken cancellationToken, String query, @Nullable Object... arguments) {
+        attachmentLock.consumeAttached(ignite -> ignite.sql().executeScript(cancellationToken, query, arguments));
+    }
+
     @Override
     public CompletableFuture<Void> executeScriptAsync(String query, @Nullable Object... arguments) {
         return attachmentLock.attachedAsync(ignite -> ignite.sql().executeScriptAsync(query, arguments));
     }
 
+    @Override
+    public CompletableFuture<Void> executeScriptAsync(@Nullable CancellationToken cancellationToken, String query,
+            @Nullable Object... arguments) {
+        return attachmentLock.attachedAsync(ignite -> ignite.sql().executeScriptAsync(cancellationToken, query, arguments));
+    }
+
     @Override
     public <T> T unwrap(Class<T> classToUnwrap) {
         return attachmentLock.attached(ignite -> Wrappers.unwrap(ignite.sql(), classToUnwrap));

modules/sql-engine/src/integrationTest/java/org/apache/ignite/internal/sql/api/ItSqlApiBaseTest.java

@@ -25,6 +25,9 @@
 import static org.apache.ignite.internal.sql.engine.util.SqlTestUtils.assertThrowsSqlException;
 import static org.apache.ignite.internal.testframework.IgniteTestUtils.assertThrowsWithCause;
 import static org.apache.ignite.internal.testframework.IgniteTestUtils.assertThrowsWithCode;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.greaterThan;
+import static org.hamcrest.Matchers.is;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
@@ -36,15 +39,20 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import org.apache.ignite.Ignite;
 import org.apache.ignite.internal.catalog.commands.CatalogUtils;
+import org.apache.ignite.internal.client.sql.ClientSql;
 import org.apache.ignite.internal.sql.BaseSqlIntegrationTest;
 import org.apache.ignite.internal.sql.ColumnMetadataImpl;
 import org.apache.ignite.internal.sql.ColumnMetadataImpl.ColumnOriginImpl;
+import org.apache.ignite.internal.testframework.IgniteTestUtils;
 import org.apache.ignite.internal.tx.TxManager;
+import org.apache.ignite.lang.CancelHandle;
+import org.apache.ignite.lang.CancellationToken;
 import org.apache.ignite.lang.CursorClosedException;
 import org.apache.ignite.lang.ErrorGroups.Common;
 import org.apache.ignite.lang.ErrorGroups.Sql;
@@ -65,10 +73,12 @@
 import org.apache.ignite.table.Table;
 import org.apache.ignite.tx.Transaction;
 import org.apache.ignite.tx.TransactionOptions;
+import org.awaitility.Awaitility;
 import org.hamcrest.Matcher;
 import org.jetbrains.annotations.Nullable;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.AssertionFailureBuilder;
+import org.junit.jupiter.api.Assumptions;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
@@ -929,6 +939,46 @@ public void testQueryTimeout() {
         });
     }
 
+    @Test
+    public void cancelScript() {
+        IgniteSql sql = igniteSql();
+
+        // TODO https://issues.apache.org/jira/browse/IGNITE-23646 Support cancellation token and remove assumption.
+        Assumptions.assumeFalse(sql instanceof ClientSql, "The thin client does not yet support canceling requests.");
+
+        sql("CREATE TABLE test (id INT PRIMARY KEY);");
+
+        String query = "INSERT INTO test SELECT x FROM system_range(0, 10000000000); SELECT 1;";
+
+        CancelHandle cancelHandle = CancelHandle.create();
+        CancellationToken token = cancelHandle.token();
+
+        CompletableFuture<Void> scriptFut = IgniteTestUtils.runAsync(() -> executeScript(sql, token, query));
+
+        // Wait until FIRST script statement is started to execute.
+        Awaitility.await().untilAsserted(() -> assertThat(queryProcessor().runningQueries(), greaterThan(1)));
+
+        assertThat(scriptFut.isDone(), is(false));
+
+        cancelHandle.cancel();
+
+        assertThrowsSqlException(
+                Sql.EXECUTION_CANCELLED_ERR,
+                "The query was cancelled while executing.",
+                () -> IgniteTestUtils.await(scriptFut)
+        );
+
+        assertThat(queryProcessor().runningQueries(), is(0));
+
+        // Checks the exception that is thrown if a query is canceled before a cursor is obtained.
+        assertThrowsSqlException(
+                Sql.EXECUTION_CANCELLED_ERR,
+                "The query was cancelled while executing.",
+                () -> executeScript(sql, token, "SELECT 1; SELECT 2;")
+        );
+        assertThat(queryProcessor().runningQueries(), is(0));
+    }
+
     @Test
     public void testQueryTimeoutIsPropagatedFromTheServer() throws Exception {
         Statement stmt = igniteSql().statementBuilder()
@@ -1041,6 +1091,8 @@ protected ResultProcessor execute(IgniteSql sql, String query, Object... args) {
 
     protected abstract void executeScript(IgniteSql sql, String query, Object... args);
 
+    protected abstract void executeScript(IgniteSql sql, CancellationToken token, String query, Object... args);
+
     protected abstract void rollback(Transaction outerTx);
 
     protected abstract void commit(Transaction outerTx);

modules/sql-engine/src/integrationTest/java/org/apache/ignite/internal/sql/api/ItSqlAsynchronousApiTest.java

@@ -17,6 +17,7 @@
 
 package org.apache.ignite.internal.sql.api;
 
+import static org.apache.ignite.internal.sql.engine.util.SqlTestUtils.assertThrowsSqlException;
 import static org.apache.ignite.internal.testframework.IgniteTestUtils.await;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -120,6 +121,17 @@ public void cancelQueryString() throws InterruptedException {
 
             return sql.executeAsync(transaction, token, query);
         });
+
+        // Checks the exception that is thrown if a query is canceled before a cursor is obtained.
+        CancelHandle cancelHandle = CancelHandle.create();
+        CancellationToken token = cancelHandle.token();
+        cancelHandle.cancel();
+
+        assertThrowsSqlException(
+                Sql.EXECUTION_CANCELLED_ERR,
+                "The query was cancelled while executing.",
+                () -> await(sql.executeAsync(null, token, "SELECT 1"))
+        );
     }
 
     @Test
@@ -221,6 +233,11 @@ protected void executeScript(IgniteSql sql, String query, Object... args) {
         await(sql.executeScriptAsync(query, args));
     }
 
+    @Override
+    protected void executeScript(IgniteSql sql, CancellationToken cancellationToken, String query, Object... args) {
+        await(sql.executeScriptAsync(cancellationToken, query, args));
+    }
+
     @Override
     protected void rollback(Transaction tx) {
         await(tx.rollbackAsync());

modules/sql-engine/src/integrationTest/java/org/apache/ignite/internal/sql/api/ItSqlSynchronousApiTest.java

@@ -17,6 +17,7 @@
 
 package org.apache.ignite.internal.sql.api;
 
+import static org.apache.ignite.internal.sql.engine.util.SqlTestUtils.assertThrowsSqlException;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertInstanceOf;
@@ -76,6 +77,18 @@ public void cancelQueryString() throws InterruptedException {
             Transaction transaction = igniteTx().begin();
             return sql.execute(transaction, token, statement);
         });
+
+        // Checks the exception that is thrown if a query is canceled before a cursor is obtained.
+        CancelHandle cancelHandle = CancelHandle.create();
+        CancellationToken token = cancelHandle.token();
+        cancelHandle.cancel();
+
+        //noinspection resource
+        assertThrowsSqlException(
+                Sql.EXECUTION_CANCELLED_ERR,
+                "The query was cancelled while executing.",
+                () -> sql.execute(null, token, "SELECT 1")
+        );
     }
 
     @Test
@@ -160,6 +173,11 @@ protected void executeScript(IgniteSql sql, String query, Object... args) {
         sql.executeScript(query, args);
     }
 
+    @Override
+    protected void executeScript(IgniteSql sql, CancellationToken cancellationToken, String query, Object... args) {
+        sql.executeScript(cancellationToken, query, args);
+    }
+
     @Override
     protected void rollback(Transaction tx) {
         tx.rollback();

modules/sql-engine/src/integrationTest/java/org/apache/ignite/internal/sql/engine/BaseSqlMultiStatementTest.java

@@ -38,6 +38,7 @@
 import org.apache.ignite.internal.sql.engine.util.CursorUtils;
 import org.apache.ignite.internal.tx.InternalTransaction;
 import org.apache.ignite.internal.util.AsyncCursor.BatchedResult;
+import org.apache.ignite.lang.CancellationToken;
 import org.jetbrains.annotations.Nullable;
 import org.junit.jupiter.api.AfterEach;
 
@@ -68,21 +69,27 @@ void executeScript(String query, Object ... params) {
 
     /** Fully executes multi-statements query without reading cursor data. */
     void executeScript(String query, @Nullable InternalTransaction tx, Object ... params) {
-        iterateThroughResultsAndCloseThem(runScript(query, tx, params));
+        iterateThroughResultsAndCloseThem(runScript(tx, null, query, params));
     }
 
     /** Initiates multi-statements query execution. */
-    AsyncSqlCursor<InternalSqlRow> runScript(String query) {
-        return runScript(query, null);
+    AsyncSqlCursor<InternalSqlRow> runScript(String query, Object... params) {
+        return runScript(null, null, query, params);
     }
 
-    AsyncSqlCursor<InternalSqlRow> runScript(String query, @Nullable InternalTransaction tx, Object ... params) {
+    /** Initiates multi-statements query execution. */
+    AsyncSqlCursor<InternalSqlRow> runScript(CancellationToken cancellationToken, String query, Object... params) {
+        return runScript(null, cancellationToken, query, params);
+    }
+
+    AsyncSqlCursor<InternalSqlRow> runScript(@Nullable InternalTransaction tx, @Nullable CancellationToken cancellationToken, String query,
+            Object... params) {
         SqlProperties properties = SqlPropertiesHelper.newBuilder()
                 .set(QueryProperty.ALLOWED_QUERY_TYPES, SqlQueryType.ALL)
                 .build();
 
         AsyncSqlCursor<InternalSqlRow> cursor = await(
-                queryProcessor().queryAsync(properties, observableTimeTracker(), tx, null, query, params)
+                queryProcessor().queryAsync(properties, observableTimeTracker(), tx, cancellationToken, query, params)
         );
 
         return Objects.requireNonNull(cursor);