|
1 | 1 | use datafusion::physical_plan::displayable; |
2 | 2 | use std::sync::Arc; |
3 | | - |
| 3 | +use regex::Regex; |
4 | 4 | use crate::{ |
5 | 5 | compile::{ |
6 | 6 | rewrite::rewriter::Rewriter, |
@@ -965,6 +965,96 @@ async fn test_case_wrapper_escaping() { |
965 | 965 | .contains("\\\\\\\\\\\\`")); |
966 | 966 | } |
967 | 967 |
|
| 968 | +/// Test aliases for grouped CubeScan in wrapper |
| 969 | +/// qualifiers from join should get remapped to single from alias |
| 970 | +/// long generated aliases from Datafusion should get shortened |
| 971 | +#[tokio::test] |
| 972 | +async fn test_join_wrapper_cubescan_aliasing() { |
| 973 | + if !Rewriter::sql_push_down_enabled() { |
| 974 | + return; |
| 975 | + } |
| 976 | + init_testing_logger(); |
| 977 | + |
| 978 | + let query_plan = convert_select_to_query_plan( |
| 979 | + // language=PostgreSQL |
| 980 | + r#" |
| 981 | +WITH |
| 982 | +-- This subquery should be represented as CubeScan(ungrouped=false) inside CubeScanWrapper |
| 983 | +cube_scan_subq AS ( |
| 984 | + SELECT |
| 985 | + logs_alias.content logs_content, |
| 986 | + DATE_TRUNC('month', kibana_alias.last_mod) last_mod_month, |
| 987 | + CASE |
| 988 | + WHEN sum(kibana_alias."sumPrice") IS NOT NULL |
| 989 | + THEN sum(kibana_alias."sumPrice") |
| 990 | + ELSE 0 |
| 991 | + END sum_price |
| 992 | + FROM KibanaSampleDataEcommerce kibana_alias |
| 993 | + JOIN Logs logs_alias |
| 994 | + ON kibana_alias.__cubeJoinField = logs_alias.__cubeJoinField |
| 995 | + GROUP BY |
| 996 | + logs_content, |
| 997 | + last_mod_month |
| 998 | +), |
| 999 | +filter_subq AS ( |
| 1000 | + SELECT |
| 1001 | + Logs.content logs_content_filter |
| 1002 | + FROM Logs |
| 1003 | + GROUP BY |
| 1004 | + logs_content_filter |
| 1005 | +) |
| 1006 | +SELECT |
| 1007 | + logs_content |
| 1008 | +FROM cube_scan_subq |
| 1009 | +WHERE |
| 1010 | + -- This subquery filter should trigger wrapping of whole query |
| 1011 | + logs_content IN ( |
| 1012 | + SELECT |
| 1013 | + logs_content_filter |
| 1014 | + FROM filter_subq |
| 1015 | + ) |
| 1016 | +GROUP BY |
| 1017 | + logs_content |
| 1018 | +ORDER BY |
| 1019 | + logs_content ASC |
| 1020 | +; |
| 1021 | +"# |
| 1022 | + .to_string(), |
| 1023 | + DatabaseProtocol::PostgreSQL, |
| 1024 | + ) |
| 1025 | + .await; |
| 1026 | + |
| 1027 | + let physical_plan = query_plan.as_physical_plan().await.unwrap(); |
| 1028 | + println!( |
| 1029 | + "Physical plan: {}", |
| 1030 | + displayable(physical_plan.as_ref()).indent() |
| 1031 | + ); |
| 1032 | + |
| 1033 | + let logical_plan = query_plan.as_logical_plan(); |
| 1034 | + let sql = logical_plan |
| 1035 | + .find_cube_scan_wrapper() |
| 1036 | + .wrapped_sql |
| 1037 | + .unwrap() |
| 1038 | + .sql; |
| 1039 | + |
| 1040 | + // Check that all aliases from different tables have same qualifier, and that names are simple and short |
| 1041 | + // logs_content => logs_alias.content |
| 1042 | + // last_mod_month => DATE_TRUNC('month', kibana_alias.last_mod), |
| 1043 | + // sum_price => CASE WHEN sum(kibana_alias."sumPrice") ... END |
| 1044 | + assert!(Regex::new(r#""logs_alias"."[a-zA-Z0-9_]+" "logs_content""#) |
| 1045 | + .unwrap() |
| 1046 | + .is_match(&sql) |
| 1047 | + ); |
| 1048 | + assert!(Regex::new(r#""logs_alias"."[a-zA-Z0-9_]{1,16}" "last_mod_month""#) |
| 1049 | + .unwrap() |
| 1050 | + .is_match(&sql) |
| 1051 | + ); |
| 1052 | + assert!(Regex::new(r#"CASE WHEN "logs_alias"."[a-zA-Z0-9_]+" IS NOT NULL THEN "logs_alias"."[a-zA-Z0-9_]+" ELSE 0 END "sum_price""#) |
| 1053 | + .unwrap() |
| 1054 | + .is_match(&sql) |
| 1055 | + ); |
| 1056 | +} |
| 1057 | + |
968 | 1058 | /// Test that WrappedSelect(... limit=Some(0) ...) will render it correctly |
969 | 1059 | #[tokio::test] |
970 | 1060 | async fn test_wrapper_limit_zero() { |
|
0 commit comments