feat(sql-parser): support two-layer nested aggregation queries

asap-common/dependencies/rs/asap_types/src/query_config.rs

@@ -22,6 +22,23 @@ pub enum QueryTimeAggregationOperator {
     Topk,
 }
 
+impl QueryTimeAggregationOperator {
+    /// The operator for an aggregation function name in any case, e.g.
+    /// PromQL's `sum` or SQL's `SUM`.
+    pub fn from_name(name: &str) -> Option<Self> {
+        match name.to_ascii_lowercase().as_str() {
+            "sum" => Some(Self::Sum),
+            "count" => Some(Self::Count),
+            "avg" => Some(Self::Avg),
+            "min" => Some(Self::Min),
+            "max" => Some(Self::Max),
+            "quantile" => Some(Self::Quantile),
+            "topk" => Some(Self::Topk),
+            _ => None,
+        }
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "snake_case")]
 pub enum QueryTimeGroupingMode {
@@ -193,6 +210,31 @@ mod tests {
         .is_err());
     }
 
+    #[test]
+    fn operator_names_parse_case_insensitively() {
+        assert!(matches!(
+            QueryTimeAggregationOperator::from_name("quantile"),
+            Some(QueryTimeAggregationOperator::Quantile)
+        ));
+        assert!(matches!(
+            QueryTimeAggregationOperator::from_name("QUANTILE"),
+            Some(QueryTimeAggregationOperator::Quantile)
+        ));
+        assert!(QueryTimeAggregationOperator::from_name("stddev").is_none());
+    }
+
+    /// The SQL matcher accepts these outer aggregations on its own list; the
+    /// planner turns them into operators, so every one must map.
+    #[test]
+    fn every_sql_outer_aggregation_has_an_operator() {
+        for name in sql_utilities::ast_matching::QUERY_TIME_AGGREGATIONS {
+            assert!(
+                QueryTimeAggregationOperator::from_name(name).is_some(),
+                "{name}"
+            );
+        }
+    }
+
     #[test]
     fn without_empty_labels_is_a_valid_promql_grouping() {
         assert!(QueryTimeAggregation {

asap-common/dependencies/rs/sql_utilities/src/ast_matching/sqlparser_test.rs

@@ -1403,4 +1403,199 @@ mod tests {
         )
         .is_none());
     }
+
+    // ── Two-layer nested aggregation ─────────────────────────────────────────
+
+    const NESTED_INNER: &str = "SELECT L1, L2, SUM(value) AS v FROM cpu_usage \
+         WHERE time BETWEEN DATEADD(s, -1, NOW()) AND NOW() GROUP BY L1, L2";
+
+    fn nested_pattern(sql: &str) -> crate::sqlpattern_matcher::SQLQuery {
+        let query_data =
+            parse_sql_query(sql).unwrap_or_else(|| panic!("Failed to parse query: {}", sql));
+        SQLPatternMatcher::new(create_test_schema(), 1.0).nested_query_to_pattern(&query_data)
+    }
+
+    fn assert_invalid_nested(sql: &str) {
+        let result = nested_pattern(sql);
+        assert_eq!(
+            result.error,
+            Some(QueryError::InvalidNestedQuery),
+            "{sql}: {:?}",
+            result.msg
+        );
+    }
+
+    #[test]
+    fn test_nested_query_matches_its_inner_layer() {
+        let result = nested_pattern(&format!(
+            "SELECT L1, SUM(v) FROM ({NESTED_INNER}) GROUP BY L1"
+        ));
+
+        assert!(result.is_valid(), "{:?}", result.msg);
+        assert_eq!(result.query_type, vec![QueryType::SpatioTemporal]);
+        let inner = result.outer_data().expect("inner layer");
+        assert_eq!(inner.aggregation_info.get_name(), "SUM");
+        assert_eq!(inner.aggregation_info.get_value_column_name(), "value");
+        assert_eq!(
+            inner.labels,
+            HashSet::from(["L1".to_string(), "L2".to_string()])
+        );
+    }
+
+    #[test]
+    fn test_nested_query_accepts_every_query_time_aggregation() {
+        for outer in [
+            "SUM(v)",
+            "COUNT(v)",
+            "AVG(v)",
+            "MIN(v)",
+            "MAX(v)",
+            "QUANTILE(0.9, v)",
+        ] {
+            let sql = format!("SELECT L1, {outer} FROM ({NESTED_INNER}) GROUP BY L1");
+            let result = nested_pattern(&sql);
+            assert!(result.is_valid(), "{sql}: {:?}", result.msg);
+        }
+    }
+
+    #[test]
+    fn test_nested_query_accepts_a_global_outer_aggregation() {
+        let result = nested_pattern(&format!("SELECT SUM(v) FROM ({NESTED_INNER})"));
+        assert!(result.is_valid(), "{:?}", result.msg);
+    }
+
+    #[test]
+    fn test_nested_query_rejects_an_outer_aggregate_not_over_the_inner_alias() {
+        assert_invalid_nested(&format!(
+            "SELECT L1, SUM(value) FROM ({NESTED_INNER}) GROUP BY L1"
+        ));
+        // Without an inner alias there is no column for the outer layer to aggregate.
+        assert_invalid_nested(
+            "SELECT L1, SUM(value) FROM (SELECT L1, SUM(value) FROM cpu_usage \
+             WHERE time BETWEEN DATEADD(s, -1, NOW()) AND NOW() GROUP BY L1) GROUP BY L1",
+        );
+    }
+
+    #[test]
+    fn test_nested_query_rejects_an_outer_label_missing_from_the_inner_layer() {
+        assert_invalid_nested(&format!(
+            "SELECT L3, SUM(v) FROM ({NESTED_INNER}) GROUP BY L3"
+        ));
+    }
+
+    #[test]
+    fn test_nested_query_rejects_count_distinct_outer_aggregation() {
+        assert_invalid_nested(&format!(
+            "SELECT L1, COUNT(DISTINCT v) FROM ({NESTED_INNER}) GROUP BY L1"
+        ));
+    }
+
+    #[test]
+    fn test_nested_query_rejects_outer_order_by_and_limit() {
+        assert_invalid_nested(&format!(
+            "SELECT L1, SUM(v) AS s FROM ({NESTED_INNER}) GROUP BY L1 ORDER BY s DESC LIMIT 3"
+        ));
+        assert_invalid_nested(&format!(
+            "SELECT L1, SUM(v) FROM ({NESTED_INNER}) GROUP BY L1 LIMIT 3"
+        ));
+    }
+
+    #[test]
+    fn test_nested_query_rejects_a_flat_query() {
+        assert_invalid_nested(NESTED_INNER);
+    }
+
+    #[test]
+    fn test_nested_query_propagates_inner_layer_errors() {
+        let result = nested_pattern(
+            "SELECT L1, SUM(v) FROM (SELECT L1, SUM(mb) AS v FROM cpu_usage \
+             WHERE time BETWEEN DATEADD(s, -1, NOW()) AND NOW() GROUP BY L1) GROUP BY L1",
+        );
+        assert_eq!(result.error, Some(QueryError::InvalidValueCol));
+    }
+
+    /// Each of these clauses would otherwise be silently dropped from the
+    /// nested query's data, so it must fail to parse instead.
+    #[test]
+    fn test_nested_query_rejects_clauses_it_cannot_represent() {
+        let rejected = [
+            // Outer WHERE (a filter, or a second time range).
+            format!("SELECT L1, SUM(v) FROM ({NESTED_INNER}) WHERE L1 = 'a' GROUP BY L1"),
+            format!(
+                "SELECT L1, SUM(v) FROM ({NESTED_INNER}) \
+                 WHERE time BETWEEN DATEADD(s, -1, NOW()) AND NOW() GROUP BY L1"
+            ),
+            // Outer HAVING.
+            format!("SELECT L1, SUM(v) FROM ({NESTED_INNER}) GROUP BY L1 HAVING SUM(v) > 1"),
+            // Inner ORDER BY / LIMIT (per-group or plain).
+            format!("SELECT L1, SUM(v) FROM ({NESTED_INNER} ORDER BY v DESC LIMIT 3 BY L1) GROUP BY L1"),
+            format!("SELECT L1, SUM(v) FROM ({NESTED_INNER} ORDER BY v DESC LIMIT 3) GROUP BY L1"),
+            format!("SELECT L1, SUM(v) FROM ({NESTED_INNER} LIMIT 3) GROUP BY L1"),
+            // Inner HAVING.
+            "SELECT L1, SUM(v) FROM (SELECT L1, L2, SUM(value) AS v FROM cpu_usage \
+             WHERE time BETWEEN DATEADD(s, -1, NOW()) AND NOW() GROUP BY L1, L2 HAVING SUM(value) > 1) \
+             GROUP BY L1"
+                .to_string(),
+            // Three layers.
+            format!(
+                "SELECT L1, SUM(m) FROM (SELECT L1, SUM(v) AS m FROM ({NESTED_INNER}) GROUP BY L1) GROUP BY L1"
+            ),
+            // Joins and extra FROM entries on the outer layer.
+            format!(
+                "SELECT L1, SUM(v) FROM ({NESTED_INNER}) t JOIN mem_usage m ON t.L1 = m.L1 GROUP BY L1"
+            ),
+            format!("SELECT L1, SUM(v) FROM ({NESTED_INNER}) t, mem_usage GROUP BY L1"),
+            // A join on the inner layer.
+            "SELECT L1, SUM(v) FROM (SELECT L1, L2, SUM(value) AS v FROM cpu_usage \
+             JOIN mem_usage ON cpu_usage.L1 = mem_usage.L1 \
+             WHERE time BETWEEN DATEADD(s, -1, NOW()) AND NOW() GROUP BY L1, L2) GROUP BY L1"
+                .to_string(),
+            // Inner QUALIFY.
+            "SELECT L1, SUM(v) FROM (SELECT L1, L2, SUM(value) AS v FROM cpu_usage \
+             WHERE time BETWEEN DATEADD(s, -1, NOW()) AND NOW() GROUP BY L1, L2 QUALIFY v > 1) \
+             GROUP BY L1"
+                .to_string(),
+        ];
+        for sql in rejected {
+            assert!(parse_sql_query(&sql).is_none(), "should not parse: {sql}");
+        }
+    }
+
+    /// A CTE would be served from whichever CTE comes first, whatever name
+    /// FROM references, so CTEs are rejected outright.
+    #[test]
+    fn test_ctes_are_rejected() {
+        let rejected = [
+            format!("WITH t AS ({NESTED_INNER}) SELECT L1, SUM(v) FROM t GROUP BY L1"),
+            format!(
+                "WITH a AS ({NESTED_INNER}), b AS (SELECT L1, L2, MAX(value) AS v FROM cpu_usage \
+                 WHERE time BETWEEN DATEADD(s, -1, NOW()) AND NOW() GROUP BY L1, L2) \
+                 SELECT L1, SUM(v) FROM b GROUP BY L1"
+            ),
+            format!(
+                "SELECT L1, SUM(v) FROM (WITH t AS ({NESTED_INNER}) SELECT * FROM t) GROUP BY L1"
+            ),
+        ];
+        for sql in rejected {
+            assert!(parse_sql_query(&sql).is_none(), "should not parse: {sql}");
+        }
+    }
+
+    /// Single-level queries reject the same unrepresentable clauses: each
+    /// would otherwise be dropped from the query's data.
+    #[test]
+    fn test_flat_query_rejects_joins_and_qualify() {
+        let window = "WHERE time BETWEEN DATEADD(s, -1, NOW()) AND NOW()";
+        let rejected = [
+            format!(
+                "SELECT L1, SUM(value) FROM cpu_usage JOIN mem_usage ON cpu_usage.L1 = mem_usage.L1 \
+                 {window} GROUP BY L1"
+            ),
+            format!("SELECT L1, SUM(value) FROM cpu_usage, mem_usage {window} GROUP BY L1"),
+            format!("SELECT L1, SUM(value) AS s FROM cpu_usage {window} GROUP BY L1 QUALIFY s > 1"),
+        ];
+        for sql in rejected {
+            assert!(parse_sql_query(&sql).is_none(), "should not parse: {sql}");
+        }
+    }
 }

asap-common/dependencies/rs/sql_utilities/src/ast_matching/sqlpattern_matcher.rs

@@ -23,9 +23,15 @@ pub enum QueryError {
     IllegalAggregationFn,
     SpatialDurationSmall,
     NestedQueryUnsupported,
+    /// A two-layer query whose outer layer is not a query-time aggregation
+    /// over its inner layer's output.
+    InvalidNestedQuery,
     MissingTimeRange,
 }
 
+/// Aggregations an outer layer can apply to its inner layer's rows.
+pub const QUERY_TIME_AGGREGATIONS: [&str; 6] = ["SUM", "COUNT", "AVG", "MIN", "MAX", "QUANTILE"];
+
 #[derive(Debug)]
 pub struct SQLQuery {
     pub query_type: Vec<QueryType>,
@@ -294,4 +300,43 @@ impl SQLPatternMatcher {
 
         sql_query
     }
+
+    /// Matches a two-layer query whose outer layer aggregates its inner
+    /// layer's rows, e.g. `SELECT job, SUM(v) FROM (<inner>) GROUP BY job`.
+    /// Returns the inner layer's pattern; the outer layer runs at query time.
+    pub fn nested_query_to_pattern(&self, query: &SQLQueryData) -> SQLQuery {
+        let invalid = |msg: String| {
+            SQLQuery::new(Vec::new(), Some(QueryError::InvalidNestedQuery), Some(msg))
+        };
+        let Some(inner) = query.subquery.as_deref() else {
+            return invalid("query has no inner layer".to_string());
+        };
+
+        if !query.order_by.is_empty() || query.limit.is_some() {
+            return invalid("ORDER BY / LIMIT on the outer layer is not supported".to_string());
+        }
+
+        let outer_aggregation = query.aggregation_info.get_name();
+        if !QUERY_TIME_AGGREGATIONS.contains(&outer_aggregation) {
+            return invalid(format!(
+                "outer aggregation {outer_aggregation} cannot run over the inner layer"
+            ));
+        }
+
+        let outer_column = query.aggregation_info.get_value_column_name();
+        if inner.aggregation_alias.as_deref() != Some(outer_column) {
+            return invalid(format!(
+                "outer aggregation reads {outer_column}, which is not the inner aggregate's alias"
+            ));
+        }
+
+        if !query.labels.is_subset(&inner.labels) {
+            let missing: HashSet<_> = query.labels.difference(&inner.labels).collect();
+            return invalid(format!(
+                "outer layer groups by {missing:?}, which the inner layer does not output"
+            ));
+        }
+
+        self.query_info_to_pattern(inner)
+    }
 }