feat(cubesql): Introduce CUBESQL_DISABLE_STRICT_AGG_TYPE_MATCH to avoid aggregation type checking during querying (#7316)

paveltiunov · web-flow · commit 61089f056af4 · 2023-10-26T22:57:13.000-07:00
diff --git a/rust/cubesql/cubesql/src/compile/context.rs b/rust/cubesql/cubesql/src/compile/context.rs
@@ -528,7 +528,8 @@ impl QueryContext {
         match selection {
             Selection::Measure(measure) => {
                 if measure.agg_type.is_some()
-                    && !measure.is_same_agg_type(&call_agg_type)
+                    // TODO not used
+                    && !measure.is_same_agg_type(&call_agg_type, false)
                 {
                     return Err(CompilationError::user(format!(
                         "Measure aggregation type doesn't match. The aggregation type for '{}' is '{}()' but '{}()' was provided",
diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs
@@ -1660,8 +1660,11 @@ mod tests {
     use crate::{
         compile::{
             rewrite::rewriter::Rewriter,
-            test::{get_string_cube_meta, get_test_tenant_ctx_with_meta},
+            test::{
+                get_string_cube_meta, get_test_session_with_config, get_test_tenant_ctx_with_meta,
+            },
         },
+        config::{ConfigObj, ConfigObjImpl},
         sql::{dataframe::batch_to_dataframe, types::StatusFlags},
     };
     use datafusion::{logical_plan::PlanVisitor, physical_plan::displayable};
@@ -1701,6 +1704,23 @@ mod tests {
         query.unwrap()
     }
 
+    async fn convert_select_to_query_plan_with_config(
+        query: String,
+        db: DatabaseProtocol,
+        config_obj: Arc<dyn ConfigObj>,
+    ) -> QueryPlan {
+        env::set_var("TZ", "UTC");
+
+        let query = convert_sql_to_cube_query(
+            &query,
+            get_test_tenant_ctx(),
+            get_test_session_with_config(db, config_obj).await,
+        )
+        .await;
+
+        query.unwrap()
+    }
+
     async fn convert_select_to_query_plan_with_meta(
         query: String,
         meta: Vec<V1CubeMeta>,
@@ -18628,6 +18648,40 @@ ORDER BY \"COUNT(count)\" DESC"
         );
     }
 
+    #[tokio::test]
+    async fn test_case_wrapper_non_strict_match() {
+        if !Rewriter::sql_push_down_enabled() {
+            return;
+        }
+        init_logger();
+
+        let mut config = ConfigObjImpl::default();
+
+        config.disable_strict_agg_type_match = true;
+
+        let query_plan = convert_select_to_query_plan_with_config(
+            "SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END, SUM(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1"
+                .to_string(),
+            DatabaseProtocol::PostgreSQL,
+            Arc::new(config)
+        )
+            .await;
+
+        let logical_plan = query_plan.as_logical_plan();
+        assert!(logical_plan
+            .find_cube_scan_wrapper()
+            .wrapped_sql
+            .unwrap()
+            .sql
+            .contains("CASE WHEN"));
+
+        let physical_plan = query_plan.as_physical_plan().await.unwrap();
+        println!(
+            "Physical plan: {}",
+            displayable(physical_plan.as_ref()).indent()
+        );
+    }
+
     #[tokio::test]
     async fn test_case_wrapper_ungrouped_sorted() {
         if !Rewriter::sql_push_down_enabled() {
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/members.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/members.rs
@@ -1975,6 +1975,12 @@ impl MemberRules {
         let cast_data_type_var = cast_data_type_var.map(|var| var!(var));
         let measure_out_var = var!(measure_out_var);
         let meta_context = self.cube_context.meta.clone();
+        let disable_strict_agg_type_match = self
+            .cube_context
+            .sessions
+            .server
+            .config_obj
+            .disable_strict_agg_type_match();
         move |egraph, subst| {
             if let Some(alias) = original_expr_name(egraph, subst[original_expr_var]) {
                 for measure_name in var_iter!(egraph[subst[measure_name_var]], MeasureName)
@@ -2040,6 +2046,7 @@ impl MemberRules {
                                                 cube_alias.to_string(),
                                                 subst[original_expr_var],
                                                 alias_to_cube.clone(),
+                                                disable_strict_agg_type_match,
                                             );
                                             return true;
                                         }
@@ -2088,6 +2095,12 @@ impl MemberRules {
         let cast_data_type_var = cast_data_type_var.map(|var| var!(var));
         let measure_out_var = measure_out_var.parse().unwrap();
         let meta_context = self.cube_context.meta.clone();
+        let disable_strict_agg_type_match = self
+            .cube_context
+            .sessions
+            .server
+            .config_obj
+            .disable_strict_agg_type_match();
         move |egraph, subst| {
             for column in measure_var
                 .map(|measure_var| {
@@ -2153,6 +2166,7 @@ impl MemberRules {
                                                 cube_alias,
                                                 subst[aggr_expr_var],
                                                 alias_to_cube,
+                                                disable_strict_agg_type_match,
                                             );
 
                                             return true;
@@ -2192,8 +2206,14 @@ impl MemberRules {
         cube_alias: String,
         expr: Id,
         alias_to_cube: Vec<((String, String), String)>,
+        disable_strict_agg_type_match: bool,
     ) {
-        if call_agg_type.is_some() && !measure.is_same_agg_type(call_agg_type.as_ref().unwrap()) {
+        if call_agg_type.is_some()
+            && !measure.is_same_agg_type(
+                call_agg_type.as_ref().unwrap(),
+                disable_strict_agg_type_match,
+            )
+        {
             let mut agg_type = measure
                 .agg_type
                 .as_ref()
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/split.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/split.rs
@@ -5461,6 +5461,12 @@ impl SplitRules {
         let distinct_var = distinct_var.map(|v| var!(v));
         let out_expr_var = out_expr_var.map(|v| var!(v));
         let meta = self.cube_context.meta.clone();
+        let disable_strict_agg_type_match = self
+            .cube_context
+            .sessions
+            .server
+            .config_obj
+            .disable_strict_agg_type_match();
         move |egraph, subst| {
             for alias_to_cube in var_iter!(
                 egraph[subst[cube_expr_var]],
@@ -5501,7 +5507,10 @@ impl SplitRules {
                                                     Some(&fun),
                                                     *distinct,
                                                 );
-                                                if !measure.is_same_agg_type(&agg_type.unwrap()) {
+                                                if !measure.is_same_agg_type(
+                                                    &agg_type.unwrap(),
+                                                    disable_strict_agg_type_match,
+                                                ) {
                                                     if let Some(expr_name) = original_expr_name(
                                                         egraph,
                                                         subst[aggr_expr_var],
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/aggregate.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/aggregate.rs
@@ -180,6 +180,12 @@ impl WrapperRules {
         let cube_members_var = var!(cube_members_var);
         let measure_out_var = var!(measure_out_var);
         let cube_context = self.cube_context.clone();
+        let disable_strict_agg_type_match = self
+            .cube_context
+            .sessions
+            .server
+            .config_obj
+            .disable_strict_agg_type_match();
         move |egraph, subst| {
             if let Some(alias) = original_expr_name(egraph, subst[original_expr_var]) {
                 for fun in fun_name_var
@@ -224,8 +230,10 @@ impl WrapperRules {
                                         cube_context.meta.find_measure_with_name(member.to_string())
                                     {
                                         if call_agg_type.is_none()
-                                            || measure
-                                                .is_same_agg_type(call_agg_type.as_ref().unwrap())
+                                            || measure.is_same_agg_type(
+                                                call_agg_type.as_ref().unwrap(),
+                                                disable_strict_agg_type_match,
+                                            )
                                         {
                                             let column_expr_column =
                                                 egraph.add(LogicalPlanLanguage::ColumnExprColumn(
diff --git a/rust/cubesql/cubesql/src/compile/test/mod.rs b/rust/cubesql/cubesql/src/compile/test/mod.rs
@@ -9,6 +9,7 @@ use datafusion::arrow::datatypes::SchemaRef;
 
 use crate::{
     compile::engine::df::{scan::MemberField, wrapper::SqlQuery},
+    config::{ConfigObj, ConfigObjImpl},
     sql::{
         session::DatabaseProtocol, AuthContextRef, AuthenticateResponse, HttpAuthContext,
         ServerManager, Session, SessionManager, SqlAuthService,
@@ -265,10 +266,18 @@ pub fn get_test_tenant_ctx_with_meta(meta: Vec<V1CubeMeta>) -> Arc<MetaContext>
 }
 
 pub async fn get_test_session(protocol: DatabaseProtocol) -> Arc<Session> {
+    get_test_session_with_config(protocol, Arc::new(ConfigObjImpl::default())).await
+}
+
+pub async fn get_test_session_with_config(
+    protocol: DatabaseProtocol,
+    config_obj: Arc<dyn ConfigObj>,
+) -> Arc<Session> {
     let server = Arc::new(ServerManager::new(
         get_test_auth(),
         get_test_transport(),
         None,
+        config_obj,
     ));
 
     let db_name = match &protocol {
diff --git a/rust/cubesql/cubesql/src/config/mod.rs b/rust/cubesql/cubesql/src/config/mod.rs
@@ -16,9 +16,11 @@ use crate::{
 use futures::future::join_all;
 use log::error;
 
-use mockall::automock;
-
-use std::env;
+use std::{
+    env,
+    fmt::{Debug, Display},
+    str::FromStr,
+};
 
 use std::sync::Arc;
 
@@ -108,15 +110,16 @@ pub struct Config {
     injector: Arc<Injector>,
 }
 
-#[automock]
-pub trait ConfigObj: DIService {
+pub trait ConfigObj: DIService + Debug {
     fn bind_address(&self) -> &Option<String>;
 
     fn postgres_bind_address(&self) -> &Option<String>;
 
     fn query_timeout(&self) -> u64;
 
     fn nonce(&self) -> &Option<Vec<u8>>;
+
+    fn disable_strict_agg_type_match(&self) -> bool;
 }
 
 #[derive(Debug, Clone)]
@@ -126,10 +129,36 @@ pub struct ConfigObjImpl {
     pub nonce: Option<Vec<u8>>,
     pub query_timeout: u64,
     pub timezone: Option<String>,
+    pub disable_strict_agg_type_match: bool,
+}
+
+impl ConfigObjImpl {
+    pub fn default() -> Self {
+        let query_timeout = env::var("CUBESQL_QUERY_TIMEOUT")
+            .ok()
+            .map(|v| v.parse::<u64>().unwrap())
+            .unwrap_or(120);
+        Self {
+            bind_address: env::var("CUBESQL_BIND_ADDR").ok().or_else(|| {
+                env::var("CUBESQL_PORT")
+                    .ok()
+                    .map(|v| format!("0.0.0.0:{}", v.parse::<u16>().unwrap()))
+            }),
+            postgres_bind_address: env::var("CUBESQL_PG_PORT")
+                .ok()
+                .map(|port| format!("0.0.0.0:{}", port.parse::<u16>().unwrap())),
+            nonce: None,
+            query_timeout,
+            timezone: Some("UTC".to_string()),
+            disable_strict_agg_type_match: env_parse(
+                "CUBESQL_DISABLE_STRICT_AGG_TYPE_MATCH",
+                false,
+            ),
+        }
+    }
 }
 
 crate::di_service!(ConfigObjImpl, [ConfigObj]);
-crate::di_service!(MockConfigObj, [ConfigObj]);
 
 impl ConfigObj for ConfigObjImpl {
     fn bind_address(&self) -> &Option<String> {
@@ -147,6 +176,10 @@ impl ConfigObj for ConfigObjImpl {
     fn query_timeout(&self) -> u64 {
         self.query_timeout
     }
+
+    fn disable_strict_agg_type_match(&self) -> bool {
+        self.disable_strict_agg_type_match
+    }
 }
 
 lazy_static! {
@@ -156,25 +189,9 @@ lazy_static! {
 
 impl Config {
     pub fn default() -> Config {
-        let query_timeout = env::var("CUBESQL_QUERY_TIMEOUT")
-            .ok()
-            .map(|v| v.parse::<u64>().unwrap())
-            .unwrap_or(120);
         Config {
             injector: Injector::new(),
-            config_obj: Arc::new(ConfigObjImpl {
-                bind_address: env::var("CUBESQL_BIND_ADDR").ok().or_else(|| {
-                    env::var("CUBESQL_PORT")
-                        .ok()
-                        .map(|v| format!("0.0.0.0:{}", v.parse::<u16>().unwrap()))
-                }),
-                postgres_bind_address: env::var("CUBESQL_PG_PORT")
-                    .ok()
-                    .map(|port| format!("0.0.0.0:{}", port.parse::<u16>().unwrap())),
-                nonce: None,
-                query_timeout,
-                timezone: Some("UTC".to_string()),
-            }),
+            config_obj: Arc::new(ConfigObjImpl::default()),
         }
     }
 
@@ -189,6 +206,7 @@ impl Config {
                 nonce: None,
                 query_timeout,
                 timezone,
+                disable_strict_agg_type_match: false,
             }),
         }
     }
@@ -231,6 +249,7 @@ impl Config {
                     i.get_service_typed().await,
                     i.get_service_typed().await,
                     config.nonce().clone(),
+                    config.clone(),
                 ))
             })
             .await;
@@ -287,4 +306,26 @@ impl Config {
     }
 }
 
+pub fn env_parse<T>(name: &str, default: T) -> T
+where
+    T: FromStr,
+    T::Err: Display,
+{
+    env_optparse(name).unwrap_or(default)
+}
+
+fn env_optparse<T>(name: &str) -> Option<T>
+where
+    T: FromStr,
+    T::Err: Display,
+{
+    env::var(name).ok().map(|x| match x.parse::<T>() {
+        Ok(v) => v,
+        Err(e) => panic!(
+            "Could not parse environment variable '{}' with '{}' value: {}",
+            name, x, e
+        ),
+    })
+}
+
 type LoopHandle = JoinHandle<Result<(), CubeError>>;
diff --git a/rust/cubesql/cubesql/src/sql/server_manager.rs b/rust/cubesql/cubesql/src/sql/server_manager.rs
diff --git a/rust/cubesql/cubesql/src/transport/ext.rs b/rust/cubesql/cubesql/src/transport/ext.rs
