branch-3.0: [BugFix](PreparedStatement) fix stmtId overflow #47950 (#…

…48116) Cherry-picked from #47950 Co-authored-by: lihangyu <[email protected]>
apache · Feb 22, 2025 · 3f15ddb · 3f15ddb
1 parent 3b85f31
commit 3f15ddb
Show file tree

Hide file tree

Showing 7 changed files with 101 additions and 7 deletions.
diff --git a/fe/fe-common/src/main/java/org/apache/doris/common/Config.java b/fe/fe-common/src/main/java/org/apache/doris/common/Config.java
@@ -140,6 +140,11 @@ public class Config extends ConfigBase {
             description = {"是否检查table锁泄漏", "Whether to check table lock leaky"})
     public static boolean check_table_lock_leaky = false;
 
+    @ConfField(mutable = true, masterOnly = false,
+            description = {"PreparedStatement stmtId 起始位置，仅用于测试",
+                    "PreparedStatement stmtId starting position, used for testing onl"})
+    public static long prepared_stmt_start_id = -1;
+
     @ConfField(description = {"插件的安装目录", "The installation directory of the plugin"})
     public static String plugin_dir = System.getenv("DORIS_HOME") + "/plugins";
 

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/PrepareCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/PrepareCommand.java
@@ -111,7 +111,7 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception {
         ctx.addPreparedStatementContext(name,
                 new PreparedStatementContext(this, ctx, ctx.getStatementContext(), name));
         if (ctx.getCommand() == MysqlCommand.COM_STMT_PREPARE) {
-            executor.sendStmtPrepareOK((int) ctx.getStmtId(), labels);
+            executor.sendStmtPrepareOK(Integer.parseInt(name), labels);
         }
     }
 

diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java
@@ -116,6 +116,8 @@ public enum ConnectType {
     // set for http_stream
     protected volatile TUniqueId loadId;
     protected volatile long backendId;
+    // range [Integer.MIN_VALUE, Integer.MAX_VALUE]
+    protected int preparedStmtId = Integer.MIN_VALUE;
     protected volatile LoadTaskInfo streamLoadInfo;
 
     protected volatile TUniqueId queryId = null;
@@ -415,10 +417,11 @@ public void addPreparedStatementContext(String stmtName, PreparedStatementContex
         }
         if (this.preparedStatementContextMap.size() > sessionVariable.maxPreparedStmtCount) {
             throw new UserException("Failed to create a server prepared statement"
-                    + "possibly because there are too many active prepared statements on server already."
+                    + " possibly because there are too many active prepared statements on server already."
                     + "set max_prepared_stmt_count with larger number than " + sessionVariable.maxPreparedStmtCount);
         }
         this.preparedStatementContextMap.put(stmtName, ctx);
+        incPreparedStmtId();
     }
 
     public void removePrepareStmt(String stmtName) {
@@ -445,6 +448,14 @@ public long getStmtId() {
         return stmtId;
     }
 
+    public long getPreparedStmtId() {
+        return preparedStmtId;
+    }
+
+    public void incPreparedStmtId() {
+        ++preparedStmtId;
+    }
+
     public long getBackendId() {
         return backendId;
     }

diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/MysqlConnectProcessor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/MysqlConnectProcessor.java
@@ -186,9 +186,7 @@ private void handleExecute() {
         // nererids
         PreparedStatementContext preparedStatementContext = ctx.getPreparedStementContext(String.valueOf(stmtId));
         if (preparedStatementContext == null) {
-            if (LOG.isDebugEnabled()) {
-                LOG.debug("No such statement in context, stmtId:{}", stmtId);
-            }
+            LOG.warn("No such statement in context, stmtId:{}", stmtId);
             ctx.getState().setError(ErrorCode.ERR_UNKNOWN_COM_ERROR,
                     "msg: Not supported such prepared statement");
             return;

diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java
@@ -709,9 +709,10 @@ private void executeByNereids(TUniqueId queryId) throws Exception {
             if (logicalPlan instanceof UnsupportedCommand || logicalPlan instanceof CreatePolicyCommand) {
                 throw new MustFallbackException("cannot prepare command " + logicalPlan.getClass().getSimpleName());
             }
-            logicalPlan = new PrepareCommand(String.valueOf(context.getStmtId()),
+            long stmtId = Config.prepared_stmt_start_id > 0
+                    ? Config.prepared_stmt_start_id : context.getPreparedStmtId();
+            logicalPlan = new PrepareCommand(String.valueOf(stmtId),
                     logicalPlan, statementContext.getPlaceholders(), originStmt);
-
         }
         // when we in transaction mode, we only support insert into command and transaction command
         if (context.isTxnModel()) {

diff --git a/regression-test/data/prepared_stmt_p0/prepared_stmt.out b/regression-test/data/prepared_stmt_p0/prepared_stmt.out
@@ -135,3 +135,36 @@ a
 -- !select24 --
 1	\N	[{"id":"1", "name":"doris"}, {"id":"2", "name":"apache"}, null]	\N
 
+-- !overflow_2 --
+2
+
+-- !overflow_2 --
+2
+
+-- !overflow_3 --
+3
+
+-- !overflow_3 --
+3
+
+-- !overflow_4 --
+4
+
+-- !overflow_4 --
+4
+
+-- !overflow_5 --
+5
+
+-- !overflow_5 --
+5
+
+-- !overflow_6 --
+6
+
+-- !overflow_6 --
+6
+
+-- !overflow_6 --
+6
+
diff --git a/regression-test/suites/prepared_stmt_p0/prepared_stmt.groovy b/regression-test/suites/prepared_stmt_p0/prepared_stmt.groovy
@@ -302,4 +302,50 @@ suite("test_prepared_stmt", "nonConcurrent") {
         assertEquals(com.mysql.cj.jdbc.ServerPreparedStatement, stmt_read.class)
         qe_select24 stmt_read
     }
+
+    // test stmtId overflow
+    def result2 = connect(user, password, url) {
+        // def stmt_read1 = prepareStatement "select 1"
+        // assertEquals(com.mysql.cj.jdbc.ServerPreparedStatement, stmt_read1.class)
+        // qe_overflow_1 stmt_read1
+        // stmt_read1.close()
+        // int max
+        sql """admin set frontend config("prepared_stmt_start_id" = "2147483647");"""
+        def stmt_read2 = prepareStatement "select 2"
+        assertEquals(com.mysql.cj.jdbc.ServerPreparedStatement, stmt_read2.class)
+        qe_overflow_2 stmt_read2
+        qe_overflow_2 stmt_read2
+        stmt_read2.close()
+        // int max + 1
+        sql """admin set frontend config("prepared_stmt_start_id" = "2147483648");"""
+        def stmt_read3 = prepareStatement "select 3"
+        // overflow throw NumberFormatExceptio and fallback to ClientPreparedStatement
+        assertEquals(com.mysql.cj.jdbc.ClientPreparedStatement, stmt_read3.class)
+        qe_overflow_3 stmt_read3
+        qe_overflow_3 stmt_read3
+        stmt_read3.close()
+        // int min 
+        sql """admin set frontend config("prepared_stmt_start_id" = "2147483646");"""
+        def stmt_read4 = prepareStatement "select 4"
+        assertEquals(com.mysql.cj.jdbc.ServerPreparedStatement, stmt_read4.class)
+        qe_overflow_4 stmt_read4
+        qe_overflow_4 stmt_read4
+        stmt_read4.close()
+
+        sql """admin set frontend config("prepared_stmt_start_id" = "123");"""
+        def stmt_read5 = prepareStatement "select 5"
+        assertEquals(com.mysql.cj.jdbc.ServerPreparedStatement, stmt_read5.class)
+        qe_overflow_5 stmt_read5
+        qe_overflow_5 stmt_read5
+        stmt_read5.close()
+
+        // set back
+        sql """admin set frontend config("prepared_stmt_start_id" = "-1");"""
+        def stmt_read6 = prepareStatement "select 6"
+        assertEquals(com.mysql.cj.jdbc.ServerPreparedStatement, stmt_read6.class)
+        qe_overflow_6 stmt_read6
+        qe_overflow_6 stmt_read6
+        qe_overflow_6 stmt_read6
+        stmt_read6.close()
+    }
 }