Skip to content

Commit

Permalink
fix(cubesql): Allow to pass measure as an argument in COUNT function (#…
Browse files Browse the repository at this point in the history 
…4063)
  • Loading branch information
@ovr
ovr committed Feb 10, 2022
1 parent a243837 commit c48c7ea
Show file tree
Hide file tree
Showing 2 changed files with 109 additions and 71 deletions.
120 changes: 63 additions & 57 deletions rust/cubesql/cubesql/src/compile/context.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -67,10 +67,10 @@ impl QueryContext {

pub fn find_dimension_for_identifier(
&self,
column_name: &String,
identifier: &String,
) -> Option<V1CubeMetaDimension> {
for dimension in self.meta.dimensions.iter() {
if dimension.get_real_name().eq_ignore_ascii_case(column_name) {
if dimension.get_real_name().eq_ignore_ascii_case(identifier) {
return Some(dimension.clone());
}
}
Expand Down Expand Up @@ -450,31 +450,33 @@ impl QueryContext {
};

let fn_name = f.name.to_string().to_ascii_lowercase();
if fn_name.eq(&"count".to_string()) && !f.distinct {
if measure_name != "*" {
return Err(CompilationError::User(format!(
"Unable to use '{}' as argument to aggregation function '{}()'",
measure_name,
f.name.to_string(),
)));
}
let (selection_opt, call_agg_type) = if fn_name.eq(&"count".to_string()) && !f.distinct {
if &measure_name == "*" {
let measure_for_argument = self.meta.measures.iter().find(|measure| {
if measure.agg_type.is_some() {
let agg_type = measure.agg_type.clone().unwrap();
agg_type.eq(&"count".to_string())
} else {
false
}
});

let measure_for_argument = self.meta.measures.iter().find(|measure| {
if measure.agg_type.is_some() {
let agg_type = measure.agg_type.clone().unwrap();
agg_type.eq(&"count".to_string())
if let Some(measure) = measure_for_argument {
(
Some(Selection::Measure(measure.clone())),
"count".to_string(),
)
} else {
false
return Err(CompilationError::User(format!(
"Unable to find measure with count type: {}",
f
)));
}
});

if let Some(measure) = measure_for_argument {
Ok(Selection::Measure(measure.clone()))
} else {
Err(CompilationError::User(format!(
"Unable to find measure with count type: {}",
f
)))
(
self.find_selection_for_identifier(&measure_name, true),
"count".to_string(),
)
}
} else {
let mut call_agg_type = fn_name;
Expand All @@ -491,44 +493,48 @@ impl QueryContext {
)));
}

let selection_opt = self.find_selection_for_identifier(&measure_name, true);
if let Some(selection) = selection_opt {
match selection {
Selection::Measure(measure) => {
if measure.agg_type.is_some()
&& !measure.is_same_agg_type(&call_agg_type)
{
return Err(CompilationError::User(format!(
"Measure aggregation type doesn't match. The aggregation type for '{}' is '{}()' but '{}()' was provided",
measure.get_real_name(),
measure.agg_type.unwrap_or("unknown".to_string()).to_uppercase(),
f.name.to_string(),
)));
} else {
// @todo Should we throw an exception?
};

Ok(Selection::Measure(measure))
}
Selection::Dimension(t) | Selection::TimeDimension(t, _) => {
Err(CompilationError::User(format!(
"Dimension '{}' was used with the aggregate function '{}()'. Please use a measure instead",
t.get_real_name(),
f.name.to_string(),
)))
}
Selection::Segment(s) => Err(CompilationError::User(format!(
"Unable to use segment '{}' as measure in aggregation function '{}()'",
s.get_real_name(),
(
self.find_selection_for_identifier(&measure_name, true),
call_agg_type,
)
};

let selection = selection_opt.ok_or_else(|| {
CompilationError::User(format!(
"Unable to find measure with name '{}' which is used as argument to aggregation function '{}()'",
measure_name,
f.name.to_string(),
))
})?;
match selection {
Selection::Measure(measure) => {
if measure.agg_type.is_some()
&& !measure.is_same_agg_type(&call_agg_type)
{
return Err(CompilationError::User(format!(
"Measure aggregation type doesn't match. The aggregation type for '{}' is '{}()' but '{}()' was provided",
measure.get_real_name(),
measure.agg_type.unwrap_or("unknown".to_string()).to_uppercase(),
f.name.to_string(),
))),
}
} else {
)));
} else {
// @todo Should we throw an exception?
};

Ok(Selection::Measure(measure))
}
Selection::Dimension(t) | Selection::TimeDimension(t, _) => {
Err(CompilationError::User(format!(
"Unable to find measure with name '{}' for {}",
measure_name, f
"Dimension '{}' was used with the aggregate function '{}()'. Please use a measure instead",
t.get_real_name(),
f.name.to_string(),
)))
}
Selection::Segment(s) => Err(CompilationError::User(format!(
"Unable to use segment '{}' as measure in aggregation function '{}()'",
s.get_real_name(),
f.name.to_string(),
))),
}
}

Expand Down
60 changes: 46 additions & 14 deletions rust/cubesql/cubesql/src/compile/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2828,6 +2828,28 @@ mod tests {
.unwrap(),
),
),
(
"SELECT COUNT(count) FROM KibanaSampleDataEcommerce".to_string(),
V1LoadRequestQuery {
measures: Some(vec!["KibanaSampleDataEcommerce.count".to_string()]),
dimensions: Some(vec![]),
segments: Some(vec![]),
time_dimensions: None,
order: None,
limit: None,
offset: None,
filters: None,
},
Arc::new(
DFSchema::new(vec![DFField::new(
None,
"KibanaSampleDataEcommerce.count",
DataType::Int64,
false,
)])
.unwrap(),
),
),
(
"SELECT COUNT(DISTINCT agentCount) FROM Logs".to_string(),
V1LoadRequestQuery {
Expand Down Expand Up @@ -2907,6 +2929,28 @@ mod tests {
#[test]
fn test_select_error() {
let variants = vec![
// Count agg fn
(
"SELECT COUNT(maxPrice) FROM KibanaSampleDataEcommerce".to_string(),
CompilationError::User("Measure aggregation type doesn't match. The aggregation type for 'maxPrice' is 'MAX()' but 'COUNT()' was provided".to_string()),
),
(
"SELECT COUNT(order_date) FROM KibanaSampleDataEcommerce".to_string(),
CompilationError::User("Dimension 'order_date' was used with the aggregate function 'COUNT()'. Please use a measure instead".to_string()),
),
(
"SELECT COUNT(2) FROM KibanaSampleDataEcommerce".to_string(),
CompilationError::User("Unable to use number '2' as argument to aggregation function".to_string()),
),
(
"SELECT COUNT(unknownIdentifier) FROM KibanaSampleDataEcommerce".to_string(),
CompilationError::User("Unable to find measure with name 'unknownIdentifier' which is used as argument to aggregation function 'COUNT()'".to_string()),
),
// Another aggregation functions
(
"SELECT COUNT(DISTINCT *) FROM KibanaSampleDataEcommerce".to_string(),
CompilationError::User("Unable to use '*' as argument to aggregation function 'COUNT()' (only COUNT() supported)".to_string()),
),
(
"SELECT MAX(*) FROM KibanaSampleDataEcommerce".to_string(),
CompilationError::User("Unable to use '*' as argument to aggregation function 'MAX()' (only COUNT() supported)".to_string()),
Expand All @@ -2921,7 +2965,7 @@ mod tests {
),
(
"SELECT MAX(unknownIdentifier) FROM KibanaSampleDataEcommerce".to_string(),
CompilationError::User("Unable to find measure with name 'unknownIdentifier' for MAX(unknownIdentifier)".to_string()),
CompilationError::User("Unable to find measure with name 'unknownIdentifier' which is used as argument to aggregation function 'MAX()'".to_string()),
),
// Check restrictions for segments usage
(
Expand All @@ -2936,18 +2980,6 @@ mod tests {
"SELECT COUNT(*) FROM KibanaSampleDataEcommerce ORDER BY is_male DESC".to_string(),
CompilationError::User("Unable to use segment 'is_male' in ORDER BY".to_string()),
),
(
"SELECT COUNT(2) FROM KibanaSampleDataEcommerce".to_string(),
CompilationError::User("Unable to use number '2' as argument to aggregation function".to_string()),
),
(
"SELECT COUNT(unknownIdentifier) FROM KibanaSampleDataEcommerce".to_string(),
CompilationError::User("Unable to use 'unknownIdentifier' as argument to aggregation function 'COUNT()'".to_string()),
),
(
"SELECT COUNT(DISTINCT *) FROM KibanaSampleDataEcommerce".to_string(),
CompilationError::User("Unable to use '*' as argument to aggregation function 'COUNT()' (only COUNT() supported)".to_string()),
),
];

for (input_query, expected_error) in variants.iter() {
Expand All @@ -2960,7 +2992,7 @@ mod tests {

match &query {
Ok(_) => panic!("Query ({}) should return error", input_query),
Err(e) => assert_eq!(e, expected_error),
Err(e) => assert_eq!(e, expected_error, "for {}", input_query),
}
}
}
Expand Down

0 comments on commit c48c7ea

Please sign in to comment.