Match AT warning scan to effective filters

Author: nicosuave

test/sql/measures.test

@@ -298,6 +298,18 @@ GROUP BY region
 );
 ----
 
+statement error
+CREATE TABLE ctas_warning_parenthesized_cte_result AS (
+WITH cte_warning_parenthesized AS (
+    SELECT region, 100.0 * AGGREGATE(revenue) / AGGREGATE(revenue) AT (ALL region) AS pct_of_total
+    FROM sales_v
+    WHERE year = 2023
+    GROUP BY region
+)
+SELECT * FROM cte_warning_parenthesized
+);
+----
+
 statement error
 CREATE TABLE ctas_warning_empty_result AS
 SELECT region, 100.0 * AGGREGATE(revenue) / AGGREGATE(revenue) AT (ALL region) AS pct_of_total

yardstick-rs/src/sql/measures.rs

@@ -2056,7 +2056,7 @@ fn expand_cte_queries(sql: &str) -> CteExpansion {
     }
 }
 
-fn top_level_parenthesized_query_body(sql: &str) -> Option<&str> {
+fn top_level_parenthesized_query_body_range(sql: &str) -> Option<(usize, usize)> {
     let upper = sql.to_uppercase();
     let mut idx = 0;
     while let Some(as_pos) = find_top_level_keyword(sql, "AS", idx) {
@@ -2067,7 +2067,8 @@ fn top_level_parenthesized_query_body(sql: &str) -> Option<&str> {
                 || matches_keyword_at(&upper, body_start, "WITH")
             {
                 if let Ok((_, body)) = balanced_parens(&sql[open_pos + 1..]) {
-                    return Some(body);
+                    let start = open_pos + 1;
+                    return Some((start, start + body.len()));
                 }
             }
         }
@@ -2076,6 +2077,10 @@ fn top_level_parenthesized_query_body(sql: &str) -> Option<&str> {
     None
 }
 
+fn top_level_parenthesized_query_body(sql: &str) -> Option<&str> {
+    top_level_parenthesized_query_body_range(sql).map(|(start, end)| &sql[start..end])
+}
+
 fn extract_base_relation_sql(view_query: &str) -> Option<String> {
     let query = view_query.trim().trim_end_matches(';').trim();
     if query.is_empty() {
@@ -5955,6 +5960,26 @@ fn dimension_key_in_group_by(ident: &str, group_by_cols: &[String]) -> bool {
     })
 }
 
+fn effective_at_where_filter_identifiers(modifiers: &[ContextModifier]) -> Vec<String> {
+    if modifiers
+        .iter()
+        .any(|modifier| matches!(modifier, ContextModifier::AllGlobal))
+    {
+        return Vec::new();
+    }
+
+    modifiers
+        .iter()
+        .find_map(|modifier| {
+            if let ContextModifier::Where(condition) = modifier {
+                Some(extract_where_filter_identifiers(condition))
+            } else {
+                None
+            }
+        })
+        .unwrap_or_default()
+}
+
 fn warning_for_at_all_ungrouped_where(
     measure_name: &str,
     modifiers: &[ContextModifier],
@@ -6001,19 +6026,7 @@ fn warning_for_at_all_ungrouped_where(
         .iter()
         .enumerate()
         .filter_map(|(idx, modifier)| {
-            if let ContextModifier::Where(condition) = modifier {
-                let cleared_by_global_all = modifiers
-                    .iter()
-                    .enumerate()
-                    .any(|(all_idx, other)| {
-                        all_idx != idx && matches!(other, ContextModifier::AllGlobal)
-                    });
-                if cleared_by_global_all {
-                    None
-                } else {
-                    Some(extract_where_filter_identifiers(condition))
-                }
-            } else if let ContextModifier::Set(dim, _) = modifier {
+            if let ContextModifier::Set(dim, _) = modifier {
                 let dim_key = normalize_dimension_key(dim.split('.').next_back().unwrap_or(dim).trim());
                 let removed_by_all =
                     modifiers.iter().enumerate().any(|(all_idx, other)| {
@@ -6041,6 +6054,10 @@ fn warning_for_at_all_ungrouped_where(
         })
         .flatten()
         .collect();
+    let encoded_filter_dims: HashSet<String> = encoded_filter_dims
+        .into_iter()
+        .chain(effective_at_where_filter_identifiers(modifiers))
+        .collect();
     let mut ungrouped_filters: Vec<String> = extract_where_filter_identifiers(where_clause)
         .into_iter()
         .filter(|ident| source_dims.is_empty() || source_dims.contains(ident))
@@ -6067,7 +6084,30 @@ fn warning_for_at_all_ungrouped_where(
 
 /// Expand AGGREGATE() with AT modifiers in SQL
 pub fn expand_aggregate_with_at(sql: &str) -> AggregateExpandResult {
-    let cte_expansion = expand_cte_queries(sql);
+    let cte_expansion = if let Some((body_start, body_end)) =
+        top_level_parenthesized_query_body_range(sql)
+    {
+        let body_sql = &sql[body_start..body_end];
+        let body_expansion = expand_cte_queries(body_sql);
+        if body_expansion.had_aggregate
+            || body_expansion.sql != body_sql
+            || !body_expansion.warnings.is_empty()
+        {
+            let mut expanded_sql = sql.to_string();
+            if body_expansion.sql != body_sql {
+                expanded_sql.replace_range(body_start..body_end, &body_expansion.sql);
+            }
+            CteExpansion {
+                sql: expanded_sql,
+                had_aggregate: body_expansion.had_aggregate,
+                warnings: body_expansion.warnings,
+            }
+        } else {
+            expand_cte_queries(sql)
+        }
+    } else {
+        expand_cte_queries(sql)
+    };
     let mut sql = cte_expansion.sql;
     let mut had_aggregate = cte_expansion.had_aggregate;
     let mut warnings = cte_expansion.warnings;
@@ -7520,6 +7560,49 @@ FROM orders"#;
         assert!(warning.is_none());
     }
 
+    #[test]
+    fn test_at_all_warns_when_filter_only_in_overwritten_at_where() {
+        let source_dims = HashSet::from([
+            normalize_identifier_name("year"),
+            normalize_identifier_name("region"),
+            normalize_identifier_name("category"),
+        ]);
+        let warning = warning_for_at_all_ungrouped_where(
+            "revenue",
+            &[
+                ContextModifier::All("region".to_string()),
+                ContextModifier::Where("category = 'A'".to_string()),
+                ContextModifier::Where("year = 2024".to_string()),
+            ],
+            Some("year = 2024"),
+            &["region".to_string()],
+            &source_dims,
+        )
+        .unwrap();
+        assert!(warning.contains("year"));
+    }
+
+    #[test]
+    fn test_at_all_does_not_warn_when_filter_in_effective_at_where() {
+        let source_dims = HashSet::from([
+            normalize_identifier_name("year"),
+            normalize_identifier_name("region"),
+            normalize_identifier_name("category"),
+        ]);
+        let warning = warning_for_at_all_ungrouped_where(
+            "revenue",
+            &[
+                ContextModifier::All("region".to_string()),
+                ContextModifier::Where("year = 2024".to_string()),
+                ContextModifier::Where("category = 'A'".to_string()),
+            ],
+            Some("year = 2024"),
+            &["region".to_string()],
+            &source_dims,
+        );
+        assert!(warning.is_none());
+    }
+
     #[test]
     fn test_at_all_does_not_warn_when_set_encodes_filter_dimension() {
         let source_dims = HashSet::from([