Add ut

Signed-off-by: wyb <[email protected]>
StarRocks · Sep 20, 2024 · 64d2634 · 64d2634
1 parent 17a48db
commit 64d2634
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 6 deletions.
diff --git a/fe/fe-core/src/main/java/com/starrocks/qe/ConnectContext.java b/fe/fe-core/src/main/java/com/starrocks/qe/ConnectContext.java
@@ -956,7 +956,7 @@ public void checkTimeout(long now) {
                 killFlag = true;
 
                 String suggestedMsg = String.format("please increase the '%s' session variable",
-                        isExecLoadType() ? "insert_timeout" : "query_timeout");
+                        isExecLoadType() ? SessionVariable.INSERT_TIMEOUT : SessionVariable.QUERY_TIMEOUT);
                 errMsg = ErrorCode.ERR_TIMEOUT.formatErrorMsg(getExecType(), timeoutSecond, suggestedMsg);
             }
         }

diff --git a/fe/fe-core/src/main/java/com/starrocks/qe/ResultReceiver.java b/fe/fe-core/src/main/java/com/starrocks/qe/ResultReceiver.java
@@ -141,16 +141,16 @@ public RowBatch getNext(Status status) throws TException {
             LOG.warn("fetch result execution exception, finstId={}", DebugUtil.printId(finstId), e);
             if (e.getMessage().contains("time out")) {
                 // if timeout, we set error code to TIMEOUT, and it will not retry querying.
-                status.setStatus(new Status(TStatusCode.TIMEOUT, ErrorCode.ERR_TIMEOUT.formatErrorMsg(
-                        "Query", timeoutMs / 1000, "please increase the 'query_timeout' session variable")));
+                status.setStatus(new Status(TStatusCode.TIMEOUT, ErrorCode.ERR_TIMEOUT.formatErrorMsg("Query", timeoutMs / 1000,
+                        String.format("please increase the '%s' session variable", SessionVariable.QUERY_TIMEOUT))));
             } else {
                 status.setRpcStatus(e.getMessage());
                 SimpleScheduler.addToBlocklist(backendId);
             }
         } catch (TimeoutException e) {
             LOG.warn("fetch result timeout, finstId={}", DebugUtil.printId(finstId), e);
-            status.setInternalErrorStatus(ErrorCode.ERR_TIMEOUT.formatErrorMsg(
-                    "Query", timeoutMs / 1000, "please increase the 'query_timeout' session variable"));
+            status.setInternalErrorStatus(ErrorCode.ERR_TIMEOUT.formatErrorMsg("Query", timeoutMs / 1000,
+                    String.format("please increase the '%s' session variable", SessionVariable.QUERY_TIMEOUT)));
             if (MetricRepo.hasInit) {
                 MetricRepo.COUNTER_QUERY_TIMEOUT.increase(1L);
             }

diff --git a/fe/fe-core/src/main/java/com/starrocks/qe/StmtExecutor.java b/fe/fe-core/src/main/java/com/starrocks/qe/StmtExecutor.java
@@ -2247,7 +2247,7 @@ public void handleDMLStmt(ExecPlan execPlan, DmlStmt stmt) throws Exception {
                                         "or set parallel_fragment_exec_instance_num to a lower value in session variable");
                     } else {
                         ErrorReport.reportDdlException(ErrorCode.ERR_TIMEOUT, getExecType(), timeout,
-                                "please increase the 'insert_timeout' session variable");
+                                String.format("please increase the '%s' session variable", SessionVariable.INSERT_TIMEOUT));
                     }
                 }
             }

diff --git a/fe/fe-core/src/test/java/com/starrocks/qe/StmtExecutorTest.java b/fe/fe-core/src/test/java/com/starrocks/qe/StmtExecutorTest.java
@@ -15,8 +15,10 @@
 package com.starrocks.qe;
 
 import com.starrocks.server.GlobalStateMgr;
+import com.starrocks.sql.ast.StatementBase;
 import com.starrocks.sql.parser.AstBuilder;
 import com.starrocks.sql.parser.SqlParser;
+import com.starrocks.utframe.UtFrameUtils;
 import mockit.Expectations;
 import mockit.Mocked;
 import org.junit.Assert;
@@ -49,4 +51,40 @@ public void testIsForwardToLeader(@Mocked GlobalStateMgr state) {
         Assert.assertFalse(new StmtExecutor(new ConnectContext(),
                 SqlParser.parseSingleStatement("show frontends", SqlModeHelper.MODE_DEFAULT)).isForwardToLeader());
     }
+
+    @Test
+    public void testExecType() {
+        ConnectContext ctx = UtFrameUtils.createDefaultCtx();
+        ConnectContext.threadLocalInfo.set(ctx);
+
+        StatementBase stmt = SqlParser.parseSingleStatement("select * from t1", SqlModeHelper.MODE_DEFAULT);
+        StmtExecutor executor = new StmtExecutor(new ConnectContext(), stmt);
+        Assert.assertEquals("Query", executor.getExecType());
+        Assert.assertFalse(executor.isExecLoadType());
+        Assert.assertEquals(ConnectContext.get().getSessionVariable().getQueryTimeoutS(), executor.getExecTimeout());
+
+        stmt = SqlParser.parseSingleStatement("insert into t1 select * from t2", SqlModeHelper.MODE_DEFAULT);
+        executor = new StmtExecutor(new ConnectContext(), stmt);
+        Assert.assertEquals("Insert", executor.getExecType());
+        Assert.assertTrue(executor.isExecLoadType());
+        Assert.assertEquals(ConnectContext.get().getSessionVariable().getInsertTimeoutS(), executor.getExecTimeout());
+
+        stmt = SqlParser.parseSingleStatement("create table t1 as select * from t2", SqlModeHelper.MODE_DEFAULT);
+        executor = new StmtExecutor(new ConnectContext(), stmt);
+        Assert.assertEquals("Insert", executor.getExecType());
+        Assert.assertTrue(executor.isExecLoadType());
+        Assert.assertEquals(ConnectContext.get().getSessionVariable().getInsertTimeoutS(), executor.getExecTimeout());
+
+        stmt = SqlParser.parseSingleStatement("update t1 set k1 = 1 where k2 = 1", SqlModeHelper.MODE_DEFAULT);
+        executor = new StmtExecutor(new ConnectContext(), stmt);
+        Assert.assertEquals("Update", executor.getExecType());
+        Assert.assertTrue(executor.isExecLoadType());
+        Assert.assertEquals(ConnectContext.get().getSessionVariable().getInsertTimeoutS(), executor.getExecTimeout());
+
+        stmt = SqlParser.parseSingleStatement("delete from t1 where k2 = 1", SqlModeHelper.MODE_DEFAULT);
+        executor = new StmtExecutor(new ConnectContext(), stmt);
+        Assert.assertEquals("Delete", executor.getExecType());
+        Assert.assertTrue(executor.isExecLoadType());
+        Assert.assertEquals(ConnectContext.get().getSessionVariable().getInsertTimeoutS(), executor.getExecTimeout());
+    }
 }