From 09f5e9a90028f26a2c4497f3f328ab3bce2bbdb4 Mon Sep 17 00:00:00 2001
From: ZENOTME <43447882+ZENOTME@users.noreply.github.com>
Date: Wed, 28 Feb 2024 21:30:07 +0900
Subject: [PATCH] fix(frontend): add subquery rewrite for parameter rewriter
 (#15095)

Co-authored-by: ZENOTME <st810918843@gmail.com>
---
 src/frontend/src/binder/bind_param.rs   | 18 ++++++++++++++++++
 src/tests/e2e_extended_mode/src/test.rs | 14 ++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/src/frontend/src/binder/bind_param.rs b/src/frontend/src/binder/bind_param.rs
index 9d51947b35063..df22fada66559 100644
--- a/src/frontend/src/binder/bind_param.rs
+++ b/src/frontend/src/binder/bind_param.rs
@@ -65,6 +65,11 @@ impl ExprRewriter for ParamRewriter {
         }
     }
 
+    fn rewrite_subquery(&mut self, mut subquery: crate::expr::Subquery) -> ExprImpl {
+        subquery.query.rewrite_exprs_recursive(self);
+        subquery.into()
+    }
+
     fn rewrite_parameter(&mut self, parameter: crate::expr::Parameter) -> ExprImpl {
         let data_type = parameter.return_type();
 
@@ -224,4 +229,17 @@ mod test {
             ),
         );
     }
+
+    #[tokio::test]
+    async fn subquery() {
+        expect_actual_eq(
+            create_expect_bound("select (select '1')"),
+            create_actual_bound(
+                "select (select $1)",
+                vec![],
+                vec![Some("1".into())],
+                vec![Format::Text],
+            ),
+        );
+    }
 }
diff --git a/src/tests/e2e_extended_mode/src/test.rs b/src/tests/e2e_extended_mode/src/test.rs
index 3e822bbd04d31..26c9b9d4c2c96 100644
--- a/src/tests/e2e_extended_mode/src/test.rs
+++ b/src/tests/e2e_extended_mode/src/test.rs
@@ -75,6 +75,7 @@ impl TestSuite {
         self.simple_cancel(true).await?;
         self.complex_cancel(false).await?;
         self.complex_cancel(true).await?;
+        self.subquery_with_param().await?;
         Ok(())
     }
 
@@ -541,4 +542,17 @@ impl TestSuite {
         new_client.execute("drop table t3", &[]).await?;
         Ok(())
     }
+
+    async fn subquery_with_param(&self) -> anyhow::Result<()> {
+        let client = self.create_client(false).await?;
+
+        let res = client
+            .query("select (select $1::SMALLINT)", &[&1024_i16])
+            .await
+            .unwrap();
+
+        assert_eq!(res[0].get::<usize, i16>(0), 1024_i16);
+
+        Ok(())
+    }
 }