query-generator: Allow where in + aggregates

Lifts most of the restriction of
`ff819dc0d9cede2af458e17a30830f2a1e843821`, which prevented WHERE IN
and aggregates from being generated for the same query. After
`8c9bc345a6a42f071bf1b621047f840eb9b31379`, we can generate most
aggregations alongside `IN` clauses, but we still need to disallow
`DISTINCT`, either as a plain modifier on a column or as a modifier on
certain aggregation function (`sum`, `count`).

Change-Id: I76ac3573d864e39a8d3b91a6155fc32914332dce
Reviewed-on: https://gerrit.readyset.name/c/readyset/+/7479
Tested-by: Buildkite CI
Reviewed-by: Ethan Donowitz <[email protected]>

query-generator/src/lib.rs

@@ -1070,6 +1070,14 @@ impl AggregateType {
             AggregateType::Min { column_type } => column_type.clone(),
         }
     }
+
+    pub fn is_distinct(&self) -> bool {
+        match self {
+            AggregateType::Count { distinct, .. } => *distinct,
+            AggregateType::Sum { distinct, .. } => *distinct,
+            _ => false,
+        }
+    }
 }
 
 /// Parameters for generating an arbitrary FilterRhs
@@ -2243,9 +2251,10 @@ impl Arbitrary for Operations {
     fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
         any_with::<Vec<QueryOperation>>(args)
             .prop_map(|mut ops| {
-                // Don't generate an aggregate or distinct in the same query as a WHERE IN clause,
+                // Don't generate an aggregate with distinct keyword or a plain distinct
+                // in the same query as a WHERE IN clause,
                 // since we don't support those queries (ENG-2942)
-                let mut agg_or_distinct_found = false;
+                let mut distinct_found = false;
                 let mut in_parameter_found = false;
 
                 // Don't generate an OR filter in the same query as a parameter of any kind, since
@@ -2254,16 +2263,24 @@ impl Arbitrary for Operations {
                 let mut or_filter_found = false;
 
                 ops.retain(|op| match op {
-                    QueryOperation::ColumnAggregate(_) | QueryOperation::Distinct => {
+                    QueryOperation::ColumnAggregate(agg) if agg.is_distinct() => {
+                        if in_parameter_found {
+                            false
+                        } else {
+                            distinct_found = true;
+                            true
+                        }
+                    }
+                    QueryOperation::Distinct => {
                         if in_parameter_found {
                             false
                         } else {
-                            agg_or_distinct_found = true;
+                            distinct_found = true;
                             true
                         }
                     }
                     QueryOperation::InParameter { .. } => {
-                        if agg_or_distinct_found || or_filter_found {
+                        if distinct_found | or_filter_found {
                             false
                         } else {
                             in_parameter_found = true;