From 1beb3d2a53a1abc2e4b3c7ef40787bc961387585 Mon Sep 17 00:00:00 2001 From: John Busby Date: Mon, 17 Aug 2026 14:53:36 -0700 Subject: [PATCH] feat(core): support ordered cube queries --- core/wren-core-py/tests/test_cube.py | 116 +++++++++- core/wren-core/core/src/mdl/cube.rs | 225 ++++++++++++++++++- core/wren-core/core/src/mdl/mod.rs | 2 +- core/wren-core/core/tests/cube_public_api.rs | 44 ++++ 4 files changed, 378 insertions(+), 9 deletions(-) create mode 100644 core/wren-core/core/tests/cube_public_api.rs diff --git a/core/wren-core-py/tests/test_cube.py b/core/wren-core-py/tests/test_cube.py index 74991d8040..f738b943b7 100644 --- a/core/wren-core-py/tests/test_cube.py +++ b/core/wren-core-py/tests/test_cube.py @@ -3,7 +3,6 @@ import json import pytest - from wren_core import cube_query_to_sql MANIFEST = json.dumps( @@ -56,6 +55,43 @@ } ) +ORDERING_MANIFEST = json.dumps( + { + "catalog": "test", + "schema": "public", + "models": [ + { + "name": "orders", + "tableReference": {"schema": "main", "table": "orders"}, + "columns": [ + {"name": "net_spend", "type": "double"}, + {"name": "merchant_name", "type": "varchar"}, + ], + } + ], + "cubes": [ + { + "name": "order_metrics", + "baseObject": "orders", + "measures": [ + { + "name": "net_spend", + "expression": "SUM(net_spend)", + "type": "DOUBLE", + } + ], + "dimensions": [ + { + "name": "merchant_name", + "expression": "merchant_name", + "type": "VARCHAR", + } + ], + } + ], + } +) + def test_basic_cube_query(): query = json.dumps( @@ -140,3 +176,81 @@ def test_invalid_manifest_json(): query = json.dumps({"cube": "order_metrics", "measures": ["revenue"]}) with pytest.raises(ValueError, match="Invalid manifest JSON"): cube_query_to_sql(query, "not json") + + +def test_order_by_uses_native_selected_ordinals_without_a_wrapper(): + query = json.dumps( + { + "cube": "order_metrics", + "measures": ["net_spend"], + "dimensions": ["merchant_name"], + "orderBy": [ + {"member": "net_spend", "direction": "desc"}, + {"member": "merchant_name", "direction": "asc"}, + ], + "limit": 5, + } + ) + + assert cube_query_to_sql(query, ORDERING_MANIFEST) == ( + "SELECT merchant_name AS merchant_name, SUM(net_spend) AS net_spend " + "FROM orders GROUP BY 1 ORDER BY 2 DESC, 1 ASC LIMIT 5" + ) + + +@pytest.mark.parametrize( + ("order_by", "message"), + [ + ( + [{"member": "not_selected", "direction": "desc"}], + "Cannot order by member 'not_selected': member is not selected by the query", + ), + ( + [ + {"member": "net_spend", "direction": "desc"}, + {"member": "net_spend", "direction": "asc"}, + ], + "Cannot order by member 'net_spend' more than once", + ), + ], +) +def test_order_by_rejects_unknown_and_duplicate_members(order_by, message): + query = json.dumps( + { + "cube": "order_metrics", + "measures": ["net_spend"], + "orderBy": order_by, + } + ) + + with pytest.raises(ValueError, match=message): + cube_query_to_sql(query, ORDERING_MANIFEST) + + +def test_order_by_rejects_non_lowercase_direction(): + query = json.dumps( + { + "cube": "order_metrics", + "measures": ["net_spend"], + "orderBy": [{"member": "net_spend", "direction": "DESC"}], + } + ) + + with pytest.raises(ValueError, match="Invalid CubeQuery JSON"): + cube_query_to_sql(query, ORDERING_MANIFEST) + + +def test_order_by_without_selected_members_is_rejected(): + query = json.dumps( + { + "cube": "order_metrics", + "measures": [], + "orderBy": [{"member": "net_spend", "direction": "desc"}], + } + ) + + with pytest.raises( + ValueError, + match="must include at least one measure, dimension, or time dimension", + ): + cube_query_to_sql(query, ORDERING_MANIFEST) diff --git a/core/wren-core/core/src/mdl/cube.rs b/core/wren-core/core/src/mdl/cube.rs index febf2bced5..2787d199d5 100644 --- a/core/wren-core/core/src/mdl/cube.rs +++ b/core/wren-core/core/src/mdl/cube.rs @@ -25,11 +25,36 @@ pub struct CubeQuery { #[serde(default)] pub filters: Vec, #[serde(default)] + pub order_by: Vec, + #[serde(default)] pub limit: Option, #[serde(default)] pub offset: Option, } +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct CubeOrderBy { + pub member: String, + pub direction: SortDirection, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum SortDirection { + Asc, + Desc, +} + +impl std::fmt::Display for SortDirection { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + SortDirection::Asc => "ASC", + SortDirection::Desc => "DESC", + }) + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct TimeDimensionFilter { @@ -206,6 +231,25 @@ fn validate_query( ); } } + let selected_members: HashSet<&str> = query + .measures + .iter() + .chain(&query.dimensions) + .map(String::as_str) + .chain(query.time_dimensions.iter().map(|td| td.dimension.as_str())) + .collect(); + let mut ordered_members = HashSet::new(); + for item in &query.order_by { + if !selected_members.contains(item.member.as_str()) { + return plan_err!( + "Cannot order by member '{}': member is not selected by the query", + item.member + ); + } + if !ordered_members.insert(item.member.as_str()) { + return plan_err!("Cannot order by member '{}' more than once", item.member); + } + } Ok(()) } @@ -320,19 +364,21 @@ fn build_sql( ) -> Result { let mut select_parts: Vec = Vec::new(); let mut group_ordinals: Vec = Vec::new(); - let mut order_ordinals: Vec = Vec::new(); + let mut selected_ordinals: HashMap = HashMap::new(); + let mut default_time_order: Option = None; let mut pos: usize = 1; - // 1. Time dimensions (appear first; drive ORDER BY) + // 1. Time dimensions (appear first; drive default ORDER BY) for td_filter in &query.time_dimensions { let td = time_dim_map[td_filter.dimension.as_str()]; let alias = format!("{}__{}", td_filter.dimension, td_filter.granularity); let expr = format!("DATE_TRUNC('{}', {})", td_filter.granularity, td.expression); select_parts.push(format!("{expr} AS {alias}")); group_ordinals.push(pos.to_string()); - if order_ordinals.is_empty() { - order_ordinals.push(pos.to_string()); - } + selected_ordinals + .entry(td_filter.dimension.clone()) + .or_insert(pos); + default_time_order.get_or_insert(pos); pos += 1; } @@ -341,6 +387,7 @@ fn build_sql( let expr = resolve_dim_expr(dim_name, dimension_map, time_dim_map); select_parts.push(format!("{expr} AS {dim_name}")); group_ordinals.push(pos.to_string()); + selected_ordinals.entry(dim_name.clone()).or_insert(pos); pos += 1; } @@ -348,6 +395,8 @@ fn build_sql( for measure_name in &query.measures { let expr = &resolved_exprs[measure_name]; select_parts.push(format!("{expr} AS {measure_name}")); + selected_ordinals.entry(measure_name.clone()).or_insert(pos); + pos += 1; } let mut sql = format!( @@ -374,11 +423,23 @@ fn build_sql( } // 5. GROUP BY / ORDER BY + let order_parts: Vec = if query.order_by.is_empty() { + default_time_order + .map(|ordinal| ordinal.to_string()) + .into_iter() + .collect() + } else { + query + .order_by + .iter() + .map(|item| format!("{} {}", selected_ordinals[&item.member], item.direction)) + .collect() + }; if !group_ordinals.is_empty() { sql.push_str(&format!(" GROUP BY {}", group_ordinals.join(", "))); } - if !order_ordinals.is_empty() { - sql.push_str(&format!(" ORDER BY {}", order_ordinals.join(", "))); + if !order_parts.is_empty() { + sql.push_str(&format!(" ORDER BY {}", order_parts.join(", "))); } // 6. LIMIT / OFFSET @@ -590,6 +651,7 @@ mod tests { dimensions: vec![], time_dimensions: vec![], filters: vec![], + order_by: vec![], limit: None, offset: None, } @@ -632,6 +694,89 @@ mod tests { ); } + #[test] + fn test_cube_query_order_by_json_uses_selected_output_ordinals() { + let query: CubeQuery = serde_json::from_str( + r#"{ + "cube": "OrdersCube", + "measures": ["revenue"], + "dimensions": ["status"], + "orderBy": [ + {"member": "revenue", "direction": "desc"}, + {"member": "status", "direction": "asc"} + ], + "limit": 5 + }"#, + ) + .unwrap(); + + assert_eq!( + cube_query_to_sql(&query, &orders_manifest()).unwrap(), + "SELECT status AS status, SUM(amount) AS revenue FROM orders GROUP BY 1 ORDER BY 2 DESC, 1 ASC LIMIT 5" + ); + } + + #[test] + fn test_cube_query_order_by_rejects_unselected_member() { + let query: CubeQuery = serde_json::from_str( + r#"{ + "cube": "OrdersCube", + "measures": ["revenue"], + "orderBy": [{"member": "order_count", "direction": "desc"}] + }"#, + ) + .unwrap(); + + let err = cube_query_to_sql(&query, &orders_manifest()) + .unwrap_err() + .to_string(); + assert!( + err.contains( + "Cannot order by member 'order_count': member is not selected by the query" + ), + "err={err}" + ); + } + + #[test] + fn test_cube_query_order_by_rejects_duplicate_members() { + let query: CubeQuery = serde_json::from_str( + r#"{ + "cube": "OrdersCube", + "measures": ["revenue"], + "orderBy": [ + {"member": "revenue", "direction": "desc"}, + {"member": "revenue", "direction": "asc"} + ] + }"#, + ) + .unwrap(); + + let err = cube_query_to_sql(&query, &orders_manifest()) + .unwrap_err() + .to_string(); + assert!( + err.contains("Cannot order by member 'revenue' more than once"), + "err={err}" + ); + } + + #[test] + fn test_cube_query_order_by_json_rejects_invalid_shapes_and_directions() { + for json in [ + r#"{"cube":"OrdersCube","measures":["revenue"],"orderBy":null}"#, + r#"{"cube":"OrdersCube","measures":["revenue"],"orderBy":{}}"#, + r#"{"cube":"OrdersCube","measures":["revenue"],"orderBy":[{"member":"revenue","direction":"ASC"}]}"#, + r#"{"cube":"OrdersCube","measures":["revenue"],"orderBy":[{"member":"revenue","direction":"sideways"}]}"#, + r#"{"cube":"OrdersCube","measures":["revenue"],"orderBy":[{"member":"revenue","direction":"desc","expression":"1; DROP TABLE orders"}]}"#, + ] { + assert!( + serde_json::from_str::(json).is_err(), + "json={json}" + ); + } + } + #[test] fn test_dimension_expression_used_not_name() { let mut q = query("OrdersCube"); @@ -664,6 +809,72 @@ mod tests { assert!(sql.contains("ORDER BY 1"), "sql={sql}"); } + #[test] + fn test_cube_query_explicit_order_replaces_default_time_order() { + let mut q = query("OrdersCube"); + q.measures = vec!["revenue".to_string()]; + q.time_dimensions = vec![TimeDimensionFilter { + dimension: "created_at".to_string(), + granularity: Granularity::Month, + date_range: None, + }]; + q.order_by = vec![CubeOrderBy { + member: "revenue".to_string(), + direction: SortDirection::Desc, + }]; + + assert_eq!( + cube_query_to_sql(&q, &orders_manifest()).unwrap(), + "SELECT DATE_TRUNC('month', created_at) AS created_at__month, SUM(amount) AS revenue FROM orders GROUP BY 1 ORDER BY 2 DESC" + ); + } + + #[test] + fn test_cube_query_explicit_order_of_time_dimension_uses_its_output_ordinal() { + let mut q = query("OrdersCube"); + q.measures = vec!["revenue".to_string()]; + q.time_dimensions = vec![TimeDimensionFilter { + dimension: "created_at".to_string(), + granularity: Granularity::Month, + date_range: None, + }]; + q.order_by = vec![CubeOrderBy { + member: "created_at".to_string(), + direction: SortDirection::Desc, + }]; + + assert_eq!( + cube_query_to_sql(&q, &orders_manifest()).unwrap(), + "SELECT DATE_TRUNC('month', created_at) AS created_at__month, SUM(amount) AS revenue FROM orders GROUP BY 1 ORDER BY 1 DESC" + ); + } + + #[test] + fn test_cube_query_missing_and_empty_order_by_preserve_identical_legacy_sql() { + let absent: CubeQuery = serde_json::from_str( + r#"{ + "cube":"OrdersCube","measures":["revenue"], + "timeDimensions":[{"dimension":"created_at","granularity":"month"}] + }"#, + ) + .unwrap(); + let empty: CubeQuery = serde_json::from_str( + r#"{ + "cube":"OrdersCube","measures":["revenue"], + "timeDimensions":[{"dimension":"created_at","granularity":"month"}], + "orderBy":[] + }"#, + ) + .unwrap(); + let expected = "SELECT DATE_TRUNC('month', created_at) AS created_at__month, SUM(amount) AS revenue FROM orders GROUP BY 1 ORDER BY 1"; + + let absent_sql = cube_query_to_sql(&absent, &orders_manifest()).unwrap(); + let empty_sql = cube_query_to_sql(&empty, &orders_manifest()).unwrap(); + assert_eq!(absent_sql, expected); + assert_eq!(empty_sql, expected); + assert_eq!(absent_sql, empty_sql); + } + #[test] fn test_time_dimension_with_date_range() { let mut q = query("OrdersCube"); diff --git a/core/wren-core/core/src/mdl/mod.rs b/core/wren-core/core/src/mdl/mod.rs index ca0b0e3827..030daf5914 100644 --- a/core/wren-core/core/src/mdl/mod.rs +++ b/core/wren-core/core/src/mdl/mod.rs @@ -39,7 +39,7 @@ pub mod builder { } pub mod context; pub(crate) mod cube; -pub use cube::{cube_query_to_sql, CubeQuery}; +pub use cube::{cube_query_to_sql, CubeOrderBy, CubeQuery, SortDirection}; pub(crate) mod dataset; mod dialect; pub mod function; diff --git a/core/wren-core/core/tests/cube_public_api.rs b/core/wren-core/core/tests/cube_public_api.rs new file mode 100644 index 0000000000..2e5adc0388 --- /dev/null +++ b/core/wren-core/core/tests/cube_public_api.rs @@ -0,0 +1,44 @@ +use wren_core::mdl::builder::{ + ColumnBuilder, CubeBuilder, CubeDimensionBuilder, ManifestBuilder, MeasureBuilder, + ModelBuilder, +}; +use wren_core::mdl::{cube_query_to_sql, CubeOrderBy, CubeQuery, SortDirection}; + +#[test] +fn cube_query_public_order_by_types_generate_sql_without_a_limit() { + let manifest = ManifestBuilder::new() + .model( + ModelBuilder::new("orders") + .table_reference("orders") + .column(ColumnBuilder::new("amount", "double").build()) + .column(ColumnBuilder::new("status", "varchar").build()) + .build(), + ) + .cube( + CubeBuilder::new("OrdersCube", "orders") + .measure(MeasureBuilder::new("revenue", "SUM(amount)", "number").build()) + .dimension( + CubeDimensionBuilder::new("status", "status", "string").build(), + ) + .build(), + ) + .build(); + let query = CubeQuery { + cube: "OrdersCube".to_string(), + measures: vec!["revenue".to_string()], + dimensions: vec!["status".to_string()], + time_dimensions: vec![], + filters: vec![], + order_by: vec![CubeOrderBy { + member: "revenue".to_string(), + direction: SortDirection::Desc, + }], + limit: None, + offset: None, + }; + + assert_eq!( + cube_query_to_sql(&query, &manifest).unwrap(), + "SELECT status AS status, SUM(amount) AS revenue FROM orders GROUP BY 1 ORDER BY 2 DESC" + ); +}