From d4b56aefdb38573b4bf66a8626dacbb3134d01ba Mon Sep 17 00:00:00 2001
From: Taishi Naka <62321668+lemonadern@users.noreply.github.com>
Date: Fri, 23 Aug 2024 13:50:41 +0900
Subject: [PATCH] Support arithmetic expressions in limit clause (#82)

---
 Cargo.lock                                    |  2 +-
 crates/uroborosql-fmt/src/cst/expr.rs         |  1 +
 .../uroborosql-fmt/src/cst/expr/expr_seq.rs   | 11 +++++++++-
 .../src/visitor/clause/limit.rs               | 17 ++++----------
 .../testfiles/dst/select/limit_offset.sql     | 18 +++++++++++++++
 .../testfiles/src/select/limit_offset.sql     | 22 +++++++++++++++++++
 6 files changed, 56 insertions(+), 15 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 3e429ff..2583911 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -333,7 +333,7 @@ dependencies = [
 [[package]]
 name = "tree-sitter-sql"
 version = "0.0.2"
-source = "git+https://github.com/future-architect/tree-sitter-sql#155225ebc4ec40dbfaf9f4d3743a89bd6945234e"
+source = "git+https://github.com/future-architect/tree-sitter-sql#f410342a6364ea5972cee463169e0a1db934a56c"
 dependencies = [
  "cc",
  "tree-sitter",
diff --git a/crates/uroborosql-fmt/src/cst/expr.rs b/crates/uroborosql-fmt/src/cst/expr.rs
index 0dbc06f..32a818d 100644
--- a/crates/uroborosql-fmt/src/cst/expr.rs
+++ b/crates/uroborosql-fmt/src/cst/expr.rs
@@ -187,6 +187,7 @@ impl Expr {
             Expr::Boolean(boolean) => boolean.set_head_comment(comment),
             Expr::ColumnList(col_list) => col_list.set_head_comment(comment),
             // primary, aligned, boolean以外の式は現状、バインドパラメータがつくことはない
+            Expr::ExprSeq(expr_seq) => expr_seq.set_head_comment_to_first_child(comment),
             _ => unimplemented!(),
         }
     }
diff --git a/crates/uroborosql-fmt/src/cst/expr/expr_seq.rs b/crates/uroborosql-fmt/src/cst/expr/expr_seq.rs
index a851353..045fb4a 100644
--- a/crates/uroborosql-fmt/src/cst/expr/expr_seq.rs
+++ b/crates/uroborosql-fmt/src/cst/expr/expr_seq.rs
@@ -1,5 +1,5 @@
 use crate::{
-    cst::{Location, Position},
+    cst::{Comment, Location, Position},
     error::UroboroSQLFmtError,
     util::{tab_size, to_tab_num},
 };
@@ -34,6 +34,15 @@ impl ExprSeq {
         self.loc.clone()
     }
 
+    /// 先頭の Expr にバインドパラメータをセットする
+    pub(crate) fn set_head_comment_to_first_child(&mut self, comment: Comment) {
+        if let Some(first_expr) = self.exprs.first_mut() {
+            first_expr.set_head_comment(comment);
+        } else {
+            unreachable!()
+        }
+    }
+
     pub(crate) fn is_multi_line(&self) -> bool {
         self.exprs.iter().any(|e| e.is_multi_line())
     }
diff --git a/crates/uroborosql-fmt/src/visitor/clause/limit.rs b/crates/uroborosql-fmt/src/visitor/clause/limit.rs
index 72ac16b..bf02669 100644
--- a/crates/uroborosql-fmt/src/visitor/clause/limit.rs
+++ b/crates/uroborosql-fmt/src/visitor/clause/limit.rs
@@ -3,7 +3,7 @@ use tree_sitter::TreeCursor;
 use crate::{
     cst::*,
     error::UroboroSQLFmtError,
-    visitor::{ensure_kind, error_annotation_from_cursor, Visitor, COMMENT},
+    visitor::{ensure_kind, Visitor, COMMENT},
 };
 
 impl Visitor {
@@ -28,13 +28,6 @@ impl Visitor {
         }
 
         match cursor.node().kind() {
-            "number" => {
-                // numberをExprに格納
-                let number = self.visit_expr(cursor, src)?;
-                // numberからBody::SingleLineを作成
-                let body = Body::SingleLine(Box::new(SingleLine::new(number)));
-                limit_clause.set_body(body);
-            }
             "ALL" => {
                 // "LIMIT ALL"というキーワードと捉えて構造体に格納
                 let all_kw = PrimaryExpr::with_node(cursor.node(), src, PrimaryExprKind::Keyword);
@@ -43,11 +36,9 @@ impl Visitor {
                 limit_clause.set_body(body);
             }
             _ => {
-                return Err(UroboroSQLFmtError::UnexpectedSyntax(format!(
-                    r#"visit_limit_clause(): expected node is number or ALL, but actual {}\n{}"#,
-                    cursor.node().kind(),
-                    error_annotation_from_cursor(cursor, src)
-                )));
+                let expr = self.visit_expr(cursor, src)?;
+                let body = Body::SingleLine(Box::new(SingleLine::new(expr)));
+                limit_clause.set_body(body);
             }
         }
 
diff --git a/crates/uroborosql-fmt/testfiles/dst/select/limit_offset.sql b/crates/uroborosql-fmt/testfiles/dst/select/limit_offset.sql
index cbae197..aac82b1 100644
--- a/crates/uroborosql-fmt/testfiles/dst/select/limit_offset.sql
+++ b/crates/uroborosql-fmt/testfiles/dst/select/limit_offset.sql
@@ -34,3 +34,21 @@ order by
 limit	/*$hoge*/all
 offset	5
 ;
+select
+	tbl1.column1	as	column1
+from
+	table1	tbl1
+order by
+	tbl1.column2	desc
+limit	1	+	2
+offset	5
+;
+select
+	tbl1.column1	as	column1
+from
+	table1	tbl1
+order by
+	tbl1.column2	desc
+limit	/*$hoge*/100	+	1
+offset	5
+;
diff --git a/crates/uroborosql-fmt/testfiles/src/select/limit_offset.sql b/crates/uroborosql-fmt/testfiles/src/select/limit_offset.sql
index e71ed10..5347944 100644
--- a/crates/uroborosql-fmt/testfiles/src/select/limit_offset.sql
+++ b/crates/uroborosql-fmt/testfiles/src/select/limit_offset.sql
@@ -42,3 +42,25 @@ ORDER BY
     DESC
 limit /*$hoge*/all
 OFFSET 5;
+
+SELECT
+    TBL1.COLUMN1    
+    AS  COLUMN1
+FROM
+    TABLE1  TBL1
+ORDER BY
+    TBL1.COLUMN2    
+    DESC
+limit 1 + 2
+OFFSET 5;
+
+SELECT
+    TBL1.COLUMN1    
+    AS  COLUMN1
+FROM
+    TABLE1  TBL1
+ORDER BY
+    TBL1.COLUMN2    
+    DESC
+limit /*$hoge*/100 + 1
+OFFSET 5;