feat(sql-utilities,query-engine): add SQL ORDER BY and LIMIT support

akanksha-akkihal · akanksha-akkihal · commit 8140ea78dca7 · 2026-05-24T20:00:29.000Z
diff --git a/asap-common/dependencies/rs/sql_utilities/src/ast_matching/sqlhelper.rs b/asap-common/dependencies/rs/sql_utilities/src/ast_matching/sqlhelper.rs
@@ -89,10 +89,28 @@ impl SQLSchema {
 #[derive(Debug, Clone)]
 pub struct SQLQueryData {
     pub aggregation_info: AggregationInfo,
+    /// Alias of the aggregate function in SELECT, e.g. `agg(v) AS p99` → `Some("p99")`.
+    /// Captured separately from `aggregation_info` because it's presentational only:
+    /// two queries that differ solely in alias must still match the same template.
+    pub aggregation_alias: Option<String>,
     pub metric: String,
     pub labels: HashSet<String>,
     pub time_info: TimeInfo,
     pub subquery: Option<Box<SQLQueryData>>,
+    /// `ORDER BY` items in source order. Empty when no ORDER BY is present.
+    /// Excluded from `matches_sql_pattern` since ordering is post-aggregation.
+    pub order_by: Vec<OrderByItem>,
+    /// `LIMIT N`. None when no LIMIT is present. Excluded from `matches_sql_pattern`.
+    pub limit: Option<u64>,
+}
+
+/// Single `ORDER BY` clause item: a column reference plus sort direction.
+/// `column` is either a GROUP BY identifier or the aggregate alias.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct OrderByItem {
+    pub column: String,
+    /// `true` for ASC (the default when neither ASC nor DESC is specified), `false` for DESC.
+    pub ascending: bool,
 }
 
 #[derive(Debug, Clone)]
diff --git a/asap-common/dependencies/rs/sql_utilities/src/ast_matching/sqlparser_test.rs b/asap-common/dependencies/rs/sql_utilities/src/ast_matching/sqlparser_test.rs
@@ -700,10 +700,17 @@ mod tests {
     }
 
     #[test]
-    fn test_order_by_is_rejected() {
-        assert!(parse_sql_query(
-            "SELECT AVG(value) FROM cpu_usage WHERE time BETWEEN DATEADD(s, -10, NOW()) AND NOW() GROUP BY L1, L2 ORDER BY L1"
-        ).is_none());
+    fn test_order_by_groupby_column_default_ascending() {
+        // Bare `ORDER BY L1` (no ASC/DESC) defaults to ascending. The order_by item
+        // must reflect that the column is a GROUP BY key.
+        let q = parse_sql_query(
+            "SELECT AVG(value) FROM cpu_usage WHERE time BETWEEN DATEADD(s, -10, NOW()) AND NOW() GROUP BY L1, L2 ORDER BY L1",
+        )
+        .expect("ORDER BY group-by column should parse");
+        assert_eq!(q.order_by.len(), 1);
+        assert_eq!(q.order_by[0].column, "L1");
+        assert!(q.order_by[0].ascending);
+        assert_eq!(q.limit, None);
     }
 
     // ── scrape_interval > 1s regression tests (issue #201) ───────────────────
@@ -919,4 +926,112 @@ mod tests {
         assert!(query.labels.contains("L2"));
         assert_eq!(query.aggregation_info.get_name(), "SUM");
     }
+
+    // ── ORDER BY / LIMIT support ─────────────────────────────────────────────
+    //
+    // The parser must accept ORDER BY (possibly multi-column, with optional ASC/DESC)
+    // and LIMIT N, capturing them in SQLQueryData for the engine to apply post-aggregation.
+    // ORDER BY columns must reference either the aggregate alias or a GROUP BY column.
+    // The aggregate alias is captured separately so `ORDER BY <alias>` can resolve.
+
+    #[test]
+    fn test_order_by_aggregate_alias_desc_limit_n() {
+        // Top-N user case: ORDER BY <agg alias> DESC LIMIT 10.
+        let q = parse_sql_query(
+            "SELECT L1, L2, QUANTILE(0.99, value) AS p99 FROM cpu_usage \
+             WHERE time BETWEEN DATEADD(s, -11, NOW()) AND DATEADD(s, -10, NOW()) \
+             GROUP BY L1, L2 \
+             ORDER BY p99 DESC LIMIT 10",
+        )
+        .expect("ORDER BY <agg alias> DESC LIMIT N should parse");
+        assert_eq!(q.aggregation_alias.as_deref(), Some("p99"));
+        assert_eq!(q.order_by.len(), 1);
+        assert_eq!(q.order_by[0].column, "p99");
+        assert!(!q.order_by[0].ascending);
+        assert_eq!(q.limit, Some(10));
+    }
+
+    #[test]
+    fn test_order_by_multiple_columns_mixed_directions() {
+        let q = parse_sql_query(
+            "SELECT QUANTILE(0.99, value) AS p99 FROM cpu_usage \
+             WHERE time BETWEEN DATEADD(s, -1, NOW()) AND NOW() \
+             GROUP BY L1, L2 \
+             ORDER BY L1 ASC, p99 DESC",
+        )
+        .expect("multi-column ORDER BY with mixed directions should parse");
+        assert_eq!(q.order_by.len(), 2);
+        assert_eq!(q.order_by[0].column, "L1");
+        assert!(q.order_by[0].ascending);
+        assert_eq!(q.order_by[1].column, "p99");
+        assert!(!q.order_by[1].ascending);
+    }
+
+    #[test]
+    fn test_limit_only_no_orderby() {
+        let q = parse_sql_query(
+            "SELECT SUM(value) FROM cpu_usage \
+             WHERE time BETWEEN DATEADD(s, -10, NOW()) AND NOW() \
+             GROUP BY L1 \
+             LIMIT 5",
+        )
+        .expect("LIMIT without ORDER BY should parse");
+        assert!(q.order_by.is_empty());
+        assert_eq!(q.limit, Some(5));
+    }
+
+    #[test]
+    fn test_order_by_unknown_column_rejected() {
+        // mystery_col is neither the aggregate alias nor a GROUP BY column.
+        assert!(parse_sql_query(
+            "SELECT QUANTILE(0.99, value) AS p99 FROM cpu_usage \
+             WHERE time BETWEEN DATEADD(s, -10, NOW()) AND NOW() \
+             GROUP BY L1 \
+             ORDER BY mystery_col",
+        )
+        .is_none());
+    }
+
+    #[test]
+    fn test_order_by_expression_rejected() {
+        assert!(parse_sql_query(
+            "SELECT QUANTILE(0.99, value) AS p99 FROM cpu_usage \
+             WHERE time BETWEEN DATEADD(s, -10, NOW()) AND NOW() \
+             GROUP BY L1 \
+             ORDER BY p99 + 1",
+        )
+        .is_none());
+    }
+
+    #[test]
+    fn test_limit_with_offset_rejected() {
+        // OFFSET is not supported (no pagination semantics in the precompute model).
+        assert!(parse_sql_query(
+            "SELECT SUM(value) FROM cpu_usage \
+             WHERE time BETWEEN DATEADD(s, -10, NOW()) AND NOW() \
+             GROUP BY L1 \
+             LIMIT 5 OFFSET 3",
+        )
+        .is_none());
+    }
+
+    #[test]
+    fn test_matches_template_ignores_order_by_and_limit() {
+        // A registered template without ORDER BY / LIMIT must still match an incoming
+        // query that adds them — they're presentational, not structural.
+        let template = parse_sql_query(
+            "SELECT QUANTILE(0.99, value) AS p99 FROM cpu_usage \
+             WHERE time BETWEEN DATEADD(s, -10, NOW()) AND NOW() \
+             GROUP BY L1, L2",
+        )
+        .unwrap();
+        let incoming = parse_sql_query(
+            "SELECT QUANTILE(0.99, value) AS top FROM cpu_usage \
+             WHERE time BETWEEN DATEADD(s, -10, '2025-10-01 00:00:10') AND '2025-10-01 00:00:10' \
+             GROUP BY L1, L2 \
+             ORDER BY top DESC LIMIT 25",
+        )
+        .unwrap();
+        assert!(incoming.matches_sql_pattern(&template));
+    }
 }
diff --git a/asap-common/dependencies/rs/sql_utilities/src/ast_matching/sqlpattern_matcher.rs b/asap-common/dependencies/rs/sql_utilities/src/ast_matching/sqlpattern_matcher.rs
@@ -53,10 +53,13 @@ impl SQLQuery {
 
         let query_data = SQLQueryData {
             aggregation_info: aggregation,
+            aggregation_alias: None,
             metric,
             labels,
             time_info: time,
             subquery: None,
+            order_by: Vec::new(),
+            limit: None,
         };
 
         self.query_data.push(query_data);
diff --git a/asap-common/dependencies/rs/sql_utilities/src/ast_matching/sqlpattern_parser.rs b/asap-common/dependencies/rs/sql_utilities/src/ast_matching/sqlpattern_parser.rs
@@ -1,5 +1,5 @@
 use crate::sqlhelper::SQLSchema;
-use crate::sqlhelper::{AggregationInfo, SQLQueryData, TimeInfo};
+use crate::sqlhelper::{AggregationInfo, OrderByItem, SQLQueryData, TimeInfo};
 use sqlparser::ast::*;
 use std::collections::HashSet;
 
@@ -36,20 +36,106 @@ impl SQLPatternParser {
     }
 
     fn parse_query_node(&self, query: &Query) -> Option<SQLQueryData> {
-        if query.order_by.is_some() {
-            println!("ORDER BY is not supported");
-            return None;
-        }
+        // Parse ORDER BY / LIMIT before walking into the SELECT body. Both are properties
+        // of the outer Query, not of the inner Select. Any unsupported sub-shape (positional
+        // refs, expressions, OFFSET, NULLS FIRST/LAST, ClickHouse `ORDER BY ALL`, etc.)
+        // bails the whole query rather than silently dropping it.
+        let order_by_items = self.parse_order_by_items(query)?;
+        let limit = self.parse_limit_value(query)?;
 
         // Convert CTE to subquery if present
         let query = self.cte_to_subquery(query);
 
-        match &query.body.as_ref() {
-            SetExpr::Select(select) => self.parse_select(select),
+        let mut data = match &query.body.as_ref() {
+            SetExpr::Select(select) => self.parse_select(select)?,
             _ => {
                 println!("Not a SELECT statement");
-                None
+                return None;
+            }
+        };
+
+        // ORDER BY columns must reference either the aggregate alias or a group-by key.
+        // Anything else (e.g. `ORDER BY some_other_column`) is rejected to avoid the
+        // engine returning an arbitrary order in cases where the user assumed the
+        // column would resolve.
+        for item in &order_by_items {
+            let valid = data.aggregation_alias.as_deref() == Some(item.column.as_str())
+                || data.labels.contains(&item.column);
+            if !valid {
+                return None;
+            }
+        }
+
+        data.order_by = order_by_items;
+        data.limit = limit;
+        Some(data)
+    }
+
+    /// Convert `query.order_by` into a flat `Vec<OrderByItem>`.
+    /// Returns `Some(vec![])` when no ORDER BY is present.
+    /// Returns `None` for any unsupported shape (positional refs, expressions,
+    /// `WITH FILL`, `NULLS FIRST/LAST`, ClickHouse `ORDER BY ALL`, `INTERPOLATE`).
+    fn parse_order_by_items(&self, query: &Query) -> Option<Vec<OrderByItem>> {
+        let order_by = match &query.order_by {
+            None => return Some(Vec::new()),
+            Some(ob) => ob,
+        };
+        if order_by.interpolate.is_some() {
+            return None;
+        }
+        let exprs = match &order_by.kind {
+            OrderByKind::Expressions(e) => e,
+            // `ORDER BY ALL` (DuckDB / ClickHouse extension) is not supported.
+            OrderByKind::All(_) => return None,
+        };
+        let mut items = Vec::with_capacity(exprs.len());
+        for ob in exprs {
+            if ob.with_fill.is_some() || ob.options.nulls_first.is_some() {
+                return None;
             }
+            let column = match &ob.expr {
+                Expr::Identifier(ident) => ident.value.clone(),
+                _ => return None,
+            };
+            // Default direction is ASC when neither ASC nor DESC is written.
+            let ascending = ob.options.asc.unwrap_or(true);
+            items.push(OrderByItem { column, ascending });
+        }
+        Some(items)
+    }
+
+    /// Convert `query.limit_clause` into an `Option<u64>`.
+    /// Returns `Some(None)` when no LIMIT is present.
+    /// Returns `None` for any unsupported shape (OFFSET, `LIMIT BY`, MySQL `LIMIT a, b`,
+    /// non-literal expressions, `LIMIT ALL`).
+    fn parse_limit_value(&self, query: &Query) -> Option<Option<u64>> {
+        let clause = match &query.limit_clause {
+            None => return Some(None),
+            Some(c) => c,
+        };
+        let limit_expr = match clause {
+            // MySQL-style `LIMIT a, b` (offset-comma-limit) is not supported.
+            LimitClause::OffsetCommaLimit { .. } => return None,
+            LimitClause::LimitOffset {
+                limit,
+                offset,
+                limit_by,
+            } => {
+                if offset.is_some() || !limit_by.is_empty() {
+                    return None;
+                }
+                match limit {
+                    None => return Some(None), // `LIMIT ALL` or no LIMIT
+                    Some(e) => e,
+                }
+            }
+        };
+        match limit_expr {
+            Expr::Value(ValueWithSpan {
+                value: Value::Number(n, _),
+                ..
+            }) => n.parse::<u64>().ok().map(Some),
+            _ => None,
         }
     }
 
@@ -88,7 +174,7 @@ impl SQLPatternParser {
     fn parse_select(&self, select: &Select) -> Option<SQLQueryData> {
         let (metric, has_subquery) = self.get_metric(select)?;
 
-        let aggregation = self.get_aggregation(select)?;
+        let (aggregation, aggregation_alias) = self.get_aggregation(select)?;
 
         let group_bys = self.get_groupbys(select)?;
 
@@ -118,17 +204,20 @@ impl SQLPatternParser {
 
             Some(SQLQueryData {
                 aggregation_info: aggregation,
+                aggregation_alias,
                 metric,
                 labels: group_bys,
                 time_info,
                 subquery: None,
+                order_by: Vec::new(),
+                limit: None,
             })
         } else {
             // Parse subquery
             let subquery = match &select.from[0].relation {
                 TableFactor::Derived { subquery, .. } => match subquery.body.as_ref() {
                     SetExpr::Select(inner_select) => {
-                        let inner_aggregation = self.get_aggregation(inner_select)?;
+                        let (inner_aggregation, inner_alias) = self.get_aggregation(inner_select)?;
                         let inner_group_bys = self.get_groupbys(inner_select)?;
                         if !self.select_identifiers_subset_of(inner_select, &inner_group_bys) {
                             return None;
@@ -137,10 +226,13 @@ impl SQLPatternParser {
 
                         Some(Box::new(SQLQueryData {
                             aggregation_info: inner_aggregation,
+                            aggregation_alias: inner_alias,
                             metric: metric.clone(),
                             labels: inner_group_bys,
                             time_info,
                             subquery: None,
+                            order_by: Vec::new(),
+                            limit: None,
                         }))
                     }
                     _ => None,
@@ -150,10 +242,13 @@ impl SQLPatternParser {
 
             Some(SQLQueryData {
                 aggregation_info: aggregation,
+                aggregation_alias,
                 metric,
                 labels: group_bys,
                 time_info: TimeInfo::new("UNUSED".to_string(), -1.0, -1_f64),
                 subquery: Some(subquery),
+                order_by: Vec::new(),
+                limit: None,
             })
         }
     }
@@ -238,17 +333,20 @@ impl SQLPatternParser {
         true
     }
 
-    fn get_aggregation(&self, select: &Select) -> Option<AggregationInfo> {
+    fn get_aggregation(&self, select: &Select) -> Option<(AggregationInfo, Option<String>)> {
         // Find the (single) aggregate function in the projection list. Other
         // projection items must be plain column references — these are expected to
         // be group-by keys (e.g. `SELECT g1, g2, SUM(v) FROM t GROUP BY g1, g2`).
         // Anything else (multiple aggregates, computed expressions, literals, *)
         // is rejected since the structural pattern model only tracks one statistic.
+        // Also captures the aggregate's alias if the SELECT writes `agg(v) AS <alias>`,
+        // so `ORDER BY <alias>` can resolve later.
         let mut agg_func: Option<&Function> = None;
+        let mut agg_alias: Option<String> = None;
         for item in &select.projection {
-            let expr = match item {
-                SelectItem::UnnamedExpr(expr) => expr,
-                SelectItem::ExprWithAlias { expr, .. } => expr,
+            let (expr, alias) = match item {
+                SelectItem::UnnamedExpr(expr) => (expr, None),
+                SelectItem::ExprWithAlias { expr, alias } => (expr, Some(alias.value.clone())),
                 _ => return None,
             };
             match expr {
@@ -257,6 +355,7 @@ impl SQLPatternParser {
                         return None;
                     }
                     agg_func = Some(f);
+                    agg_alias = alias;
                 }
                 Expr::Identifier(_) | Expr::CompoundIdentifier(_) => {}
                 _ => return None,
@@ -332,7 +431,7 @@ impl SQLPatternParser {
             name
         };
 
-        Some(AggregationInfo::new(normalized_name, col, args))
+        Some((AggregationInfo::new(normalized_name, col, args), agg_alias))
     }
 
     fn get_metric(&self, select: &Select) -> Option<(String, bool)> {
diff --git a/asap-query-engine/src/engines/simple_engine/sql.rs b/asap-query-engine/src/engines/simple_engine/sql.rs
