Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for COUNT(DISTINCT a) #700

Merged
merged 1 commit into from
Dec 14, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 20 additions & 6 deletions src/backend/query_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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() => {
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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();
}

Expand Down
29 changes: 29 additions & 0 deletions src/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
40 changes: 40 additions & 0 deletions src/func.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
}
}
Expand All @@ -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
}
Expand Down Expand Up @@ -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
Expand Down