add count distinct

SeaQL · Sep 11, 2023 · 567df82 · 567df82
1 parent 36c03c3
commit 567df82
Show file tree

Hide file tree

Showing 3 changed files with 89 additions and 6 deletions.
diff --git a/src/backend/query_builder.rs b/src/backend/query_builder.rs
@@ -310,7 +310,12 @@ pub trait QueryBuilder:
             }
             SimpleExpr::FunctionCall(func) => {
                 self.prepare_function(&func.func, sql);
-                self.prepare_tuple(&func.args, sql);
+                write!(sql, "(").unwrap();
+                if func.distinct {
+                    write!(sql, "DISTINCT ").unwrap();
+                }
+                self.prepare_comma_seperated_sequence(&func.args, sql);
+                write!(sql, ")").unwrap();
             }
             SimpleExpr::Binary(left, op, right) => match (op, right.as_ref()) {
                 (BinOper::In, SimpleExpr::Tuple(t)) if t.is_empty() => {
@@ -510,8 +515,12 @@ pub trait QueryBuilder:
             }
             TableRef::FunctionCall(func, alias) => {
                 self.prepare_function(&func.func, sql);
-                self.prepare_tuple(&func.args, sql);
-                write!(sql, " AS ").unwrap();
+                write!(sql, "(").unwrap();
+                if func.distinct {
+                    write!(sql, "DISTINCT ").unwrap();
+                }
+                self.prepare_comma_seperated_sequence(&func.args, sql);
+                write!(sql, ") AS ").unwrap();
                 alias.prepare(sql.as_writer(), self.quote());
             }
             _ => self.prepare_table_ref_iden(table_ref, sql),
@@ -946,16 +955,21 @@ pub trait QueryBuilder:
         });
     }
 
-    /// Translate [`SimpleExpr::Tuple`] into SQL statement.
-    fn prepare_tuple(&self, exprs: &[SimpleExpr], sql: &mut dyn SqlWriter) {
-        write!(sql, "(").unwrap();
+    /// Write a comma seperated sequence of [`SimpleExpr`]s.
+    fn prepare_comma_seperated_sequence(&self, exprs: &[SimpleExpr], sql: &mut dyn SqlWriter) {
         exprs.iter().fold(true, |first, expr| {
             if !first {
                 write!(sql, ", ").unwrap();
             }
             self.prepare_simple_expr(expr, sql);
             false
         });
+    }
+
+    /// Translate [`SimpleExpr::Tuple`] into SQL statement.
+    fn prepare_tuple(&self, exprs: &[SimpleExpr], sql: &mut dyn SqlWriter) {
+        write!(sql, "(").unwrap();
+        self.prepare_comma_seperated_sequence(exprs, sql);
         write!(sql, ")").unwrap();
     }
 

diff --git a/src/expr.rs b/src/expr.rs
@@ -1430,6 +1430,35 @@ impl Expr {
         Func::count(self.left).into()
     }
 
+    /// Express a `COUNT` function with the DISTINCT modifier.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use sea_query::{tests_cfg::*, *};
+    ///
+    /// let query = Query::select()
+    ///     .expr(Expr::col((Char::Table, Char::SizeW)).count_distinct())
+    ///     .from(Char::Table)
+    ///     .to_owned();
+    ///
+    /// assert_eq!(
+    ///     query.to_string(MysqlQueryBuilder),
+    ///     r#"SELECT COUNT(DISTINCT `character`.`size_w`) FROM `character`"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(PostgresQueryBuilder),
+    ///     r#"SELECT COUNT(DISTINCT "character"."size_w") FROM "character""#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(SqliteQueryBuilder),
+    ///     r#"SELECT COUNT(DISTINCT "character"."size_w") FROM "character""#
+    /// );
+    /// ```
+    pub fn count_distinct(self) -> SimpleExpr {
+        Func::count_distinct(self.left).into()
+    }
+
     /// Express a `IF NULL` function.
     ///
     /// # Examples

diff --git a/src/func.rs b/src/func.rs
@@ -33,13 +33,15 @@ pub enum Function {
 #[derive(Debug, Clone, PartialEq)]
 pub struct FunctionCall {
     pub(crate) func: Function,
+    pub(crate) distinct: bool,
     pub(crate) args: Vec<SimpleExpr>,
 }
 
 impl FunctionCall {
     pub(crate) fn new(func: Function) -> Self {
         Self {
             func,
+            distinct: false,
             args: Vec::new(),
         }
     }
@@ -62,6 +64,12 @@ impl FunctionCall {
         self
     }
 
+    /// Add the `DISTINCT` modifier to the first argument
+    pub fn distinct(mut self) -> Self {
+        self.distinct = true;
+        self
+    }
+
     pub fn get_func(&self) -> &Function {
         &self.func
     }
@@ -307,6 +315,38 @@ impl Func {
         FunctionCall::new(Function::Count).arg(expr)
     }
 
+    /// Call `COUNT` function with the `DISTINCT` modifier.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use sea_query::{tests_cfg::*, *};
+    ///
+    /// let query = Query::select()
+    ///     .expr(Func::count_distinct(Expr::col((Char::Table, Char::Id))))
+    ///     .from(Char::Table)
+    ///     .to_owned();
+    ///
+    /// assert_eq!(
+    ///     query.to_string(MysqlQueryBuilder),
+    ///     r#"SELECT COUNT(DISTINCT `character`.`id`) FROM `character`"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(PostgresQueryBuilder),
+    ///     r#"SELECT COUNT(DISTINCT "character"."id") FROM "character""#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(SqliteQueryBuilder),
+    ///     r#"SELECT COUNT(DISTINCT "character"."id") FROM "character""#
+    /// );
+    /// ```
+    pub fn count_distinct<T>(expr: T) -> FunctionCall
+    where
+        T: Into<SimpleExpr>,
+    {
+        FunctionCall::new(Function::Count).arg(expr).distinct()
+    }
+
     /// Call `CHAR_LENGTH` function.
     ///
     /// # Examples