Include joined aliases in warning matching

Author: nicosuave

yardstick-rs/src/sql/measures.rs

@@ -4259,7 +4259,7 @@ fn extract_dimension_exprs_from_query(sql: &str) -> HashMap<String, String> {
 fn extract_from_clause_info_ffi(sql: &str) -> FromClauseInfo {
     let mut info = FromClauseInfo::default();
 
-    if let Ok(select_info) = parser_ffi::parse_select(sql) {
+    if let Ok(Ok(select_info)) = std::panic::catch_unwind(|| parser_ffi::parse_select(sql)) {
         for (i, table_ref) in select_info.tables.iter().enumerate() {
             let effective_name = table_ref
                 .alias
@@ -4283,6 +4283,66 @@ fn extract_from_clause_info_ffi(sql: &str) -> FromClauseInfo {
     info
 }
 
+fn extract_from_clause_info(sql: &str) -> FromClauseInfo {
+    #[cfg(test)]
+    let mut info = FromClauseInfo::default();
+    #[cfg(not(test))]
+    let mut info = extract_from_clause_info_ffi(sql);
+
+    supplement_from_clause_info_from_sql(&mut info, sql);
+    info
+}
+
+fn supplement_from_clause_info_from_sql(info: &mut FromClauseInfo, sql: &str) {
+    let mut search_pos = 0;
+    while let Some((keyword_pos, keyword)) = next_from_or_join_keyword(sql, search_pos) {
+        let table_start = skip_ws_and_comments(sql, keyword_pos + keyword.len());
+        if table_start >= sql.len() || sql.as_bytes()[table_start] == b'(' {
+            search_pos = table_start.saturating_add(1);
+            continue;
+        }
+
+        let Some(table_end) = parse_qualified_name_span(sql, table_start) else {
+            search_pos = table_start.saturating_add(1);
+            continue;
+        };
+        let raw_table = sql[table_start..table_end].trim();
+        let table_name =
+            extract_last_qualified_identifier(raw_table).unwrap_or_else(|| raw_table.to_string());
+
+        let alias_start = skip_ws_and_comments(sql, table_end);
+        let (effective_name, has_alias, next_pos) =
+            if let Some((alias_end, alias)) = parse_alias_span(sql, alias_start) {
+                (alias, true, alias_end)
+            } else {
+                (table_name.clone(), false, table_end)
+            };
+
+        let table_info = TableInfo {
+            name: table_name,
+            effective_name: effective_name.clone(),
+            has_alias,
+        };
+        if keyword.eq_ignore_ascii_case("FROM") && info.primary_table.is_none() {
+            info.primary_table = Some(table_info.clone());
+        }
+        info.tables.entry(effective_name).or_insert(table_info);
+        search_pos = next_pos;
+    }
+}
+
+fn next_from_or_join_keyword(sql: &str, search_pos: usize) -> Option<(usize, &'static str)> {
+    let from_pos = find_top_level_keyword(sql, "FROM", search_pos);
+    let join_pos = find_top_level_keyword(sql, "JOIN", search_pos);
+    match (from_pos, join_pos) {
+        (Some(from), Some(join)) if from <= join => Some((from, "FROM")),
+        (Some(_), Some(join)) => Some((join, "JOIN")),
+        (Some(from), None) => Some((from, "FROM")),
+        (None, Some(join)) => Some((join, "JOIN")),
+        (None, None) => None,
+    }
+}
+
 /// Information about a resolved measure
 #[derive(Debug, Clone)]
 pub struct ResolvedMeasure {
@@ -6501,7 +6561,7 @@ pub fn expand_aggregate_with_at(sql: &str) -> AggregateExpandResult {
     // Prefer parser-FFI FROM extraction (supports JOIN aliases when SQL parses there),
     // then fall back to string extraction for AGGREGATE/AT syntax that parser-FFI cannot parse.
     let context_sql = top_level_parenthesized_query_body(&sql).unwrap_or(sql.as_str());
-    let mut from_info = extract_from_clause_info_ffi(context_sql);
+    let mut from_info = extract_from_clause_info(context_sql);
     let (primary_table_name, existing_alias) = if let Some(pt) = from_info.primary_table.clone() {
         let alias = if pt.has_alias {
             Some(pt.effective_name.clone())
@@ -9075,6 +9135,17 @@ GROUP BY s.year";
         assert_eq!(find_alias_for_view(&info, "nonexistent"), None);
     }
 
+    #[test]
+    fn test_extract_from_clause_info_string_fallback_join_aliases() {
+        let sql = "SELECT r.region, AGGREGATE(refunds) AT (ALL region) \
+                   FROM fact_orders_v o JOIN fact_returns_v r ON o.id = r.order_id";
+        let info = extract_from_clause_info(sql);
+
+        assert_eq!(info.primary_table.as_ref().unwrap().name, "fact_orders_v");
+        assert_eq!(find_alias_for_view(&info, "fact_orders_v"), Some("o"));
+        assert_eq!(find_alias_for_view(&info, "fact_returns_v"), Some("r"));
+    }
+
     #[test]
     #[serial]
     fn test_expand_aggregate_with_join_measure_from_first_table() {
@@ -9128,6 +9199,33 @@ GROUP BY s.year";
         assert!(result.expanded_sql.contains("FROM orders_v"));
     }
 
+    #[test]
+    fn test_at_all_warns_for_joined_measure_alias_where_filter() {
+        let sql = "SELECT r.region, AGGREGATE(refunds) AT (ALL region) AS all_region_refunds \
+                   FROM fact_orders_v o JOIN fact_returns_v r ON o.id = r.order_id";
+        let from_info = extract_from_clause_info(sql);
+        let mut source_qualifiers = HashSet::from([normalize_identifier_name("fact_returns_v")]);
+        if let Some(alias) = find_alias_for_view(&from_info, "fact_returns_v") {
+            source_qualifiers.insert(normalize_identifier_name(alias));
+        }
+        let source_dims = HashSet::from([
+            normalize_identifier_name("year"),
+            normalize_identifier_name("region"),
+        ]);
+
+        let warning = warning_for_at_all_ungrouped_where_with_qualifiers(
+            "refunds",
+            &[ContextModifier::All("region".to_string())],
+            Some("r.year = 2023"),
+            &["r.region".to_string()],
+            &source_dims,
+            &source_qualifiers,
+        )
+        .unwrap();
+
+        assert!(warning.contains("year"));
+    }
+
     #[test]
     fn test_parse_at_modifier_ad_hoc_dimension() {
         // Ad hoc dimension with function expression