fix and add SELECT 1 case

eldenmoon · Apr 29, 2024 · 4cdb798 · 4cdb798
1 parent 3118c26
commit 4cdb798
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 7 deletions.
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
@@ -132,6 +132,7 @@
 import org.apache.doris.mysql.MysqlChannel;
 import org.apache.doris.mysql.MysqlCommand;
 import org.apache.doris.mysql.MysqlEofPacket;
+import org.apache.doris.mysql.MysqlOkPacket;
 import org.apache.doris.mysql.MysqlSerializer;
 import org.apache.doris.mysql.ProxyMysqlChannel;
 import org.apache.doris.mysql.privilege.PrivPredicate;
@@ -1760,7 +1761,8 @@ private void handleQueryStmt() throws Exception {
 
         // handle selects that fe can do without be, so we can make sql tools happy, especially the setup step.
         // TODO FE not support doris field type conversion to arrow field type.
-        if (!context.getConnectType().equals(ConnectType.ARROW_FLIGHT_SQL)) {
+        if (!context.getConnectType().equals(ConnectType.ARROW_FLIGHT_SQL)
+                    && context.getCommand() != MysqlCommand.COM_STMT_EXECUTE) {
             Optional<ResultSet> resultSet = planner.handleQueryInFe(parsedStmt);
             if (resultSet.isPresent()) {
                 sendResultSet(resultSet.get());
@@ -2699,13 +2701,18 @@ private void sendStmtPrepareOK() throws IOException {
                 serializer.writeField(colNames.get(i), Type.fromPrimitiveType(types.get(i)));
                 context.getMysqlChannel().sendOnePacket(serializer.toByteBuffer());
             }
+            serializer.reset();
+            if (!context.getMysqlChannel().clientDeprecatedEOF()) {
+                MysqlEofPacket eofPacket = new MysqlEofPacket(context.getState());
+                eofPacket.writeTo(serializer);
+            } else {
+                MysqlOkPacket okPacket = new MysqlOkPacket(context.getState());
+                okPacket.writeTo(serializer);
+            }
+            context.getMysqlChannel().sendOnePacket(serializer.toByteBuffer());
         }
-        // send EOF if nessessary
-        if (!context.getMysqlChannel().clientDeprecatedEOF()) {
-            context.getState().setEof();
-        } else {
-            context.getState().setOk();
-        }
+        context.getMysqlChannel().flush();
+        context.getState().setNoop();
     }
 
     private void sendFields(List<String> colNames, List<Type> types) throws IOException {

diff --git a/regression-test/data/prepared_stmt_p0/prepared_stmt.out b/regression-test/data/prepared_stmt_p0/prepared_stmt.out
@@ -47,3 +47,9 @@
 -- !select3 --
 1	1	user1	30	1234	12345
 
+-- !select4 --
+10
+
+-- !select5 --
+1
+
diff --git a/regression-test/suites/prepared_stmt_p0/prepared_stmt.groovy b/regression-test/suites/prepared_stmt_p0/prepared_stmt.groovy
@@ -156,5 +156,12 @@ suite("test_prepared_stmt") {
      stmt_read.setInt(2, 1234)
      stmt_read.setInt(3, 1)
      qe_select3 stmt_read
+
+     stmt_read = prepareStatement "SELECT 10"
+     assertEquals(stmt_read.class, com.mysql.cj.jdbc.ServerPreparedStatement);
+     qe_select4 stmt_read
+     stmt_read = prepareStatement "SELECT 1"
+     assertEquals(stmt_read.class, com.mysql.cj.jdbc.ServerPreparedStatement);
+     qe_select5 stmt_read
     }
 }