From a309cd018a5b476f52249c7aca8c79e0c013f915 Mon Sep 17 00:00:00 2001
From: Weny Xu <wenymedia@gmail.com>
Date: Fri, 8 Mar 2024 15:31:20 +0800
Subject: [PATCH] fix: fix incorrect `COM_STMT_PREPARE` reply (#3463)

* fix: fix incorrect `COM_STMT_PREPARE` reply

* chore: use column name instead of index
---
 src/servers/src/mysql/handler.rs | 16 +++++++++++++++-
 tests-integration/tests/sql.rs   | 24 ++++++++++++------------
 2 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/src/servers/src/mysql/handler.rs b/src/servers/src/mysql/handler.rs
index 0f5fce59f369..8c1814580fc7 100644
--- a/src/servers/src/mysql/handler.rs
+++ b/src/servers/src/mysql/handler.rs
@@ -239,13 +239,27 @@ impl<W: AsyncWrite + Send + Sync + Unpin> AsyncMysqlShim<W> for MysqlInstanceShi
 
         debug_assert_eq!(params.len(), param_num - 1);
 
+        let columns = schema
+            .as_ref()
+            .map(|schema| {
+                schema
+                    .column_schemas()
+                    .iter()
+                    .map(|column_schema| {
+                        create_mysql_column(&column_schema.data_type, &column_schema.name)
+                    })
+                    .collect::<Result<Vec<_>>>()
+            })
+            .transpose()?
+            .unwrap_or_default();
+
         let stmt_id = self.save_plan(SqlPlan {
             query: query.to_string(),
             plan,
             schema,
         });
 
-        w.reply(stmt_id, &params, &[]).await?;
+        w.reply(stmt_id, &params, &columns).await?;
         crate::metrics::METRIC_MYSQL_PREPARED_COUNT
             .with_label_values(&[query_ctx.get_db_string().as_str()])
             .inc();
diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs
index 4d4a49d0d246..5b9e1e2f35f3 100644
--- a/tests-integration/tests/sql.rs
+++ b/tests-integration/tests/sql.rs
@@ -167,10 +167,10 @@ pub async fn test_mysql_crud(store_type: StorageType) {
     assert_eq!(rows.len(), 10);
 
     for (i, row) in rows.iter().enumerate() {
-        let ret: i64 = row.get(0);
-        let d: NaiveDate = row.get(1);
-        let dt: DateTime<Utc> = row.get(2);
-        let bytes: Vec<u8> = row.get(3);
+        let ret: i64 = row.get("i");
+        let d: NaiveDate = row.get("d");
+        let dt: DateTime<Utc> = row.get("dt");
+        let bytes: Vec<u8> = row.get("b");
         assert_eq!(ret, i as i64);
         let expected_d = NaiveDate::from_yo_opt(2015, 100).unwrap();
         assert_eq!(expected_d, d);
@@ -193,7 +193,7 @@ pub async fn test_mysql_crud(store_type: StorageType) {
     assert_eq!(rows.len(), 1);
 
     for row in rows {
-        let ret: i64 = row.get(0);
+        let ret: i64 = row.get("i");
         assert_eq!(ret, 6);
     }
 
@@ -358,9 +358,9 @@ pub async fn test_postgres_crud(store_type: StorageType) {
     assert_eq!(rows.len(), 10);
 
     for (i, row) in rows.iter().enumerate() {
-        let ret: i64 = row.get(0);
-        let d: NaiveDate = row.get(1);
-        let dt: NaiveDateTime = row.get(2);
+        let ret: i64 = row.get("i");
+        let d: NaiveDate = row.get("d");
+        let dt: NaiveDateTime = row.get("dt");
 
         assert_eq!(ret, i as i64);
 
@@ -381,7 +381,7 @@ pub async fn test_postgres_crud(store_type: StorageType) {
     assert_eq!(rows.len(), 1);
 
     for row in rows {
-        let ret: i64 = row.get(0);
+        let ret: i64 = row.get("i");
         assert_eq!(ret, 6);
     }
 
@@ -709,13 +709,13 @@ pub async fn test_mysql_prepare_stmt_insert_timestamp(store_type: StorageType) {
         .unwrap();
     assert_eq!(rows.len(), 3);
 
-    let x: DateTime<Utc> = rows[0].get(1);
+    let x: DateTime<Utc> = rows[0].get("ts");
     assert_eq!(x.to_string(), "2023-12-19 00:00:00 UTC");
 
-    let x: DateTime<Utc> = rows[1].get(1);
+    let x: DateTime<Utc> = rows[1].get("ts");
     assert_eq!(x.to_string(), "2023-12-19 13:19:01 UTC");
 
-    let x: DateTime<Utc> = rows[2].get(1);
+    let x: DateTime<Utc> = rows[2].get("ts");
     assert_eq!(x.to_string(), "2023-12-19 13:20:01.123 UTC");
 
     let _ = server.shutdown().await;