Support for ALL/DISTINCT/ORDER BY in aggregate functions

crates/uroborosql-fmt/src/cst/expr/function.rs

@@ -1,25 +1,193 @@
-use itertools::repeat_n;
+use itertools::{repeat_n, Itertools};
 
 use crate::{
-    cst::{Clause, Location},
+    cst::{AlignInfo, AlignedExpr, Clause, Comment, Location},
     error::UroboroSQLFmtError,
-    util::{convert_keyword_case, is_line_overflow, tab_size, to_tab_num},
+    util::{convert_keyword_case, is_line_overflow, tab_size, to_tab_num, trim_bind_param},
 };
 
-use super::column_list::ColumnList;
-
 /// FunctionCallがユーザ定義関数か組み込み関数か示すEnum
 #[derive(Debug, Clone)]
 pub(crate) enum FunctionCallKind {
     UserDefined,
     BuiltIn,
 }
 
+/// 関数呼び出しの引数を表す
+#[derive(Debug, Clone)]
+pub(crate) struct FunctionCallArgs {
+    all_distinct: Option<Clause>,
+    exprs: Vec<AlignedExpr>,
+    order_by: Option<Clause>,
+    loc: Location,
+    /// 複数行で出力するかを指定するフラグ。
+    /// デフォルトでは false (つまり、単一行で出力する) になっている。
+    force_multi_line: bool,
+}
+
+impl FunctionCallArgs {
+    pub(crate) fn new(exprs: Vec<AlignedExpr>, loc: Location) -> FunctionCallArgs {
+        Self {
+            all_distinct: None,
+            exprs,
+            order_by: None,
+            loc,
+            force_multi_line: false,
+        }
+    }
+
+    pub(crate) fn force_multi_line(&self) -> bool {
+        self.force_multi_line
+    }
+
+    pub(crate) fn add_expr(&mut self, cols: AlignedExpr) {
+        self.loc.append(cols.loc());
+        self.exprs.push(cols);
+    }
+
+    pub(crate) fn set_all_distinct(&mut self, all_distinct: Clause) {
+        self.all_distinct = Some(all_distinct)
+    }
+
+    pub(crate) fn set_order_by(&mut self, order_by: Clause) {
+        self.order_by = Some(order_by)
+    }
+
+    pub(crate) fn append_loc(&mut self, loc: Location) {
+        self.loc.append(loc)
+    }
+
+    pub(crate) fn set_trailing_comment(
+        &mut self,
+        comment: Comment,
+    ) -> Result<(), UroboroSQLFmtError> {
+        // exprs は必ず1つ以上要素を持っている
+        let last = self.exprs.last_mut().unwrap();
+        if last.loc().is_same_line(&comment.loc()) {
+            last.set_trailing_comment(comment)
+        } else {
+            Err(UroboroSQLFmtError::IllegalOperation(format!(
+                "set_trailing_comment:{:?} is not trailing comment!",
+                comment
+            )))
+        }
+    }
+
+    pub(crate) fn last_line_len(&self, acc: usize) -> usize {
+        if self.is_multi_line() {
+            ")".len()
+        } else {
+            let mut current_len = acc + "(".len();
+
+            self.exprs.iter().enumerate().for_each(|(i, col)| {
+                current_len += col.last_line_len_from_left(current_len);
+                if i != self.exprs.len() - 1 {
+                    current_len += ", ".len()
+                }
+            });
+            current_len + ")".len()
+        }
+    }
+
+    /// 列リストを複数行で描画するかを指定する。
+    /// true を与えたら必ず複数行で描画され、false を与えたらできるだけ単一行で描画する。
+    pub(crate) fn set_force_multi_line(&mut self, b: bool) {
+        self.force_multi_line = b
+    }
+
+    /// 複数行で描画するかどうかを bool 型の値で取得する。
+    /// 複数行で描画する場合は true を返す。
+    /// 自身の is_multi_line のオプションの値だけでなく、各列が単一行かどうか、末尾コメントを持つかどうかも考慮する。
+    pub(crate) fn is_multi_line(&self) -> bool {
+        self.force_multi_line
+            || self.all_distinct.is_some()
+            || self.order_by.is_some()
+            || self
+                .exprs
+                .iter()
+                .any(|a| a.is_multi_line() || a.has_trailing_comment())
+    }
+
+    /// カラムリストをrenderする。
+    /// 自身の is_multi_line() が true になる場合には複数行で描画し、false になる場合単一行で描画する。
+    pub(crate) fn render(&self, depth: usize) -> Result<String, UroboroSQLFmtError> {
+        // depth は開きかっこを描画する行のインデントの深さ
+        let mut result = String::new();
+
+        if self.is_multi_line() {
+            // 各列を複数行に出力する
+            result.push_str("(\n");
+
+            // ALL/DISTINCT
+            if let Some(all_distinct) = &self.all_distinct {
+                result.push_str(&all_distinct.render(depth + 1)?);
+            }
+
+            // 各引数の描画
+            {
+                // ALL/DISTINCT、ORDER BYがある場合はインデントを1つ深くする
+                let depth = if self.all_distinct.is_some() || self.order_by.is_some() {
+                    depth + 1
+                } else {
+                    depth
+                };
+
+                // 最初の行のインデント
+                result.extend(repeat_n('\t', depth + 1));
+
+                // 各要素間の改行、カンマ、インデント
+                let mut separator = "\n".to_string();
+                separator.extend(repeat_n('\t', depth));
+                separator.push_str(",\t");
+
+                // Vec<AlignedExpr> -> Vec<&AlignedExpr>
+                let aligned_exprs = self.exprs.iter().collect_vec();
+                let align_info = AlignInfo::from(aligned_exprs);
+
+                result.push_str(
+                    &self
+                        .exprs
+                        .iter()
+                        .map(|a| a.render_align(depth + 1, &align_info))
+                        .collect::<Result<Vec<_>, _>>()?
+                        .join(&separator),
+                );
+            }
+
+            // ORDER BY
+            if let Some(order_by) = &self.order_by {
+                result.push('\n');
+                result.push_str(&order_by.render(depth + 1)?);
+            } else {
+                result.push('\n');
+            }
+
+            result.extend(repeat_n('\t', depth));
+            result.push(')');
+        } else {
+            // 単一行で描画する
+            // ALL/DISTINCT、ORDER BYがある場合は複数行で描画するのでこの分岐には到達しない
+            result.push('(');
+            result.push_str(
+                &self
+                    .exprs
+                    .iter()
+                    .filter_map(|e| e.render(depth + 1).ok())
+                    .join(", "),
+            );
+            result.push(')');
+        }
+
+        // 閉じかっこの後の改行は呼び出し元が担当
+        Ok(result)
+    }
+}
+
 /// 関数呼び出しを表す
 #[derive(Debug, Clone)]
 pub(crate) struct FunctionCall {
     name: String,
-    args: ColumnList,
+    args: FunctionCallArgs,
     /// OVER句が持つ句 (PARTITION BY、ORDER BY)
     /// None であるならば OVER句自体がない
     over_window_definition: Option<Vec<Clause>>,
@@ -33,7 +201,7 @@ pub(crate) struct FunctionCall {
 impl FunctionCall {
     pub(crate) fn new(
         name: impl Into<String>,
-        args: ColumnList,
+        args: FunctionCallArgs,
         kind: FunctionCallKind,
         loc: Location,
     ) -> FunctionCall {

crates/uroborosql-fmt/src/visitor/expr/function.rs

@@ -6,7 +6,7 @@ use crate::{
     cst::*,
     error::UroboroSQLFmtError,
     util::convert_keyword_case,
-    visitor::{ensure_kind, Visitor},
+    visitor::{create_clause, ensure_kind, Visitor, COMMA, COMMENT},
 };
 
 impl Visitor {
@@ -25,7 +25,8 @@ impl Visitor {
         cursor.goto_next_sibling();
 
         ensure_kind(cursor, "(")?;
-        let args = self.visit_column_list(cursor, src)?;
+
+        let args = self.visit_function_call_args(cursor, src)?;
         cursor.goto_next_sibling();
 
         let mut func_call = FunctionCall::new(
@@ -112,4 +113,70 @@ impl Visitor {
 
         Ok(clauses)
     }
+
+    /// 関数の引数をFunctionCallArgsで返す
+    /// 引数は "(" [ ALL | DISTINCT ] expression [ , ... ] [ order_by_clause ] ")" という構造になっている
+    pub(crate) fn visit_function_call_args(
+        &mut self,
+        cursor: &mut TreeCursor,
+        src: &str,
+    ) -> Result<FunctionCallArgs, UroboroSQLFmtError> {
+        let mut function_call_args =
+            FunctionCallArgs::new(vec![], Location::new(cursor.node().range()));
+
+        ensure_kind(cursor, "(")?;
+
+        cursor.goto_next_sibling();
+
+        // 引数が空の場合
+        if cursor.node().kind() == ")" {
+            return Ok(function_call_args);
+        }
+
+        match cursor.node().kind() {
+            "ALL" | "DISTINCT" => {
+                let all_distinct_clause = create_clause(cursor, src, cursor.node().kind())?;
+
+                function_call_args.set_all_distinct(all_distinct_clause);
+
+                cursor.goto_next_sibling();
+            }
+            _ => {}
+        }
+
+        let first_expr = self.visit_expr(cursor, src)?.to_aligned();
+        function_call_args.add_expr(first_expr);
+
+        // [ , ... ] [ order_by_clause ] ")"
+        while cursor.goto_next_sibling() {
+            function_call_args.append_loc(Location::new(cursor.node().range()));
+
+            match cursor.node().kind() {
+                COMMA => {
+                    cursor.goto_next_sibling();
+                    let expr = self.visit_expr(cursor, src)?.to_aligned();
+                    function_call_args.add_expr(expr);
+                }
+                ")" => break,
+                COMMENT => {
+                    // 末尾コメントを想定する
+                    let comment = Comment::new(cursor.node(), src);
+                    function_call_args.set_trailing_comment(comment)?
+                }
+                "order_by_clause" => {
+                    let order_by = self.visit_order_by_clause(cursor, src)?;
+                    function_call_args.set_order_by(order_by);
+                }
+                _ => {
+                    return Err(UroboroSQLFmtError::Unimplemented(format!(
+                        "visit_function_call_args(): Unexpected node\nnode_kind: {}\n{:#?}",
+                        cursor.node().kind(),
+                        cursor.node().range(),
+                    )));
+                }
+            }
+        }
+
+        Ok(function_call_args)
+    }
 }

crates/uroborosql-fmt/src/visitor/expr/type_cast.rs

@@ -51,7 +51,7 @@ impl Visitor {
             aligned.add_rhs(Some(as_keyword), Expr::Primary(Box::new(type_name)));
             let loc = aligned.loc();
 
-            let args = ColumnList::new(vec![aligned], loc);
+            let args = FunctionCallArgs::new(vec![aligned], loc);
 
             let function =
                 FunctionCall::new(cast_keyword, args, FunctionCallKind::BuiltIn, cast_loc);
@@ -88,7 +88,7 @@ impl Visitor {
 
                 let function = FunctionCall::new(
                     cast_keyword,
-                    ColumnList::new(vec![aligned], expr.loc()),
+                    FunctionCallArgs::new(vec![aligned], expr.loc()),
                     FunctionCallKind::BuiltIn,
                     cast_loc,
                 );

crates/uroborosql-fmt/testfiles/dst/select/aggregate_function.sql

@@ -0,0 +1,22 @@
+SELECT
+	COUNT(
+		DISTINCT
+			TBL.COL1
+	)
+;
+SELECT
+	COUNT(
+		ALL
+			TBL.COL1
+	)
+;
+SELECT
+	STRING_AGG(
+		DISTINCT
+			TBL.COLUMN1
+		,	','
+		ORDER BY
+			TBL.COLUMN2
+		,	TBL.COLUMN3
+	)
+;

crates/uroborosql-fmt/testfiles/src/select/aggregate_function.sql

@@ -0,0 +1,8 @@
+select
+	count(distinct	tbl.col1);
+
+select
+	count(all	tbl.col1);
+
+select
+	string_agg(distinct	tbl.column1, ',' order by	tbl.column2, tbl.column3);