feat(cubesql): Support for metabase literal queries

rust/cubesql/.gitignore

@@ -6,3 +6,4 @@ upstream
 dist
 node_modules
 .vscode
+/cubesql/egraph-debug

rust/cubesql/cubesql/src/compile/mod.rs

@@ -62,8 +62,11 @@ use self::{
 use crate::{
     compile::{
         builder::QueryBuilder,
-        engine::udf::{
-            create_pg_is_other_temp_schema, create_pg_my_temp_schema, create_session_user_udf,
+        engine::{
+            df::scan::MemberField,
+            udf::{
+                create_pg_is_other_temp_schema, create_pg_my_temp_schema, create_session_user_udf,
+            },
         },
         rewrite::converter::LogicalPlanToLanguageConverter,
     },
@@ -1447,13 +1450,6 @@ struct QueryPlanner {
     session_manager: Arc<SessionManager>,
 }
 
-lazy_static! {
-    static ref METABASE_WORKAROUND: regex::Regex = regex::Regex::new(
-        r#"SELECT true AS "_" FROM "public"\."(?P<tblname>.*)" WHERE 1 <> 1 LIMIT 0"#
-    )
-    .unwrap();
-}
-
 impl QueryPlanner {
     pub fn new(
         state: Arc<SessionState>,
@@ -1475,24 +1471,6 @@ impl QueryPlanner {
         stmt: &ast::Statement,
         q: &Box<ast::Query>,
     ) -> CompilationResult<QueryPlan> {
-        // Metabase
-        if let Some(c) = METABASE_WORKAROUND.captures(&stmt.to_string()) {
-            let tblname = c.name("tblname").unwrap().as_str();
-            if self.meta.find_cube_with_name(tblname).is_some() {
-                return Ok(QueryPlan::MetaTabular(
-                    StatusFlags::empty(),
-                    Box::new(dataframe::DataFrame::new(
-                        vec![dataframe::Column::new(
-                            "_".to_string(),
-                            ColumnType::Int8,
-                            ColumnFlags::empty(),
-                        )],
-                        vec![],
-                    )),
-                ));
-            }
-        }
-
         // TODO move CUBESQL_REWRITE_ENGINE env to config
         let rewrite_engine = env::var("CUBESQL_REWRITE_ENGINE")
             .ok()
@@ -1691,7 +1669,7 @@ impl QueryPlanner {
                     schema
                         .fields()
                         .iter()
-                        .map(|f| Some(f.name().to_string()))
+                        .map(|f| MemberField::Member(f.name().to_string()))
                         .collect(),
                     query.request,
                     // @todo Remove after split!
@@ -3986,7 +3964,9 @@ ORDER BY \"COUNT(count)\" DESC"
         );
         assert_eq!(
             &cube_scan.member_fields,
-            &vec![Some("KibanaSampleDataEcommerce.count".to_string())]
+            &vec![MemberField::Member(
+                "KibanaSampleDataEcommerce.count".to_string()
+            )]
         );
     }
 
@@ -5617,7 +5597,7 @@ ORDER BY \"COUNT(count)\" DESC"
             QueryPlan::DataFusionSelect(flags, plan, ctx) => {
                 let df = DFDataFrame::new(ctx.state, &plan);
                 let batches = df.collect().await?;
-                let frame = batch_to_dataframe(&batches)?;
+                let frame = batch_to_dataframe(&df.schema().into(), &batches)?;
 
                 return Ok((frame.print(), flags));
             }
@@ -5698,14 +5678,14 @@ ORDER BY \"COUNT(count)\" DESC"
     #[tokio::test]
     async fn test_information_schema_stats_for_columns() -> Result<(), CubeError> {
         // This query is used by metabase for introspection
-        assert_eq!(
+        insta::assert_snapshot!(
+            "test_information_schema_stats_for_columns",
             execute_query("
             SELECT
                 A.TABLE_SCHEMA TABLE_CAT, NULL TABLE_SCHEM, A.TABLE_NAME, A.COLUMN_NAME, B.SEQ_IN_INDEX KEY_SEQ, B.INDEX_NAME PK_NAME
             FROM INFORMATION_SCHEMA.COLUMNS A, INFORMATION_SCHEMA.STATISTICS B
             WHERE A.COLUMN_KEY in ('PRI','pri') AND B.INDEX_NAME='PRIMARY'  AND (ISNULL(database()) OR (A.TABLE_SCHEMA = database())) AND (ISNULL(database()) OR (B.TABLE_SCHEMA = database())) AND A.TABLE_NAME = 'OutlierFingerprints'  AND B.TABLE_NAME = 'OutlierFingerprints'  AND A.TABLE_SCHEMA = B.TABLE_SCHEMA AND A.TABLE_NAME = B.TABLE_NAME AND A.COLUMN_NAME = B.COLUMN_NAME
-            ORDER BY A.COLUMN_NAME".to_string(), DatabaseProtocol::MySQL).await?,
-            "++\n++\n++"
+            ORDER BY A.COLUMN_NAME".to_string(), DatabaseProtocol::MySQL).await?
         );
 
         Ok(())
@@ -5986,13 +5966,13 @@ ORDER BY \"COUNT(count)\" DESC"
         );
 
         // Negative test, we dont define this variable
-        assert_eq!(
+        insta::assert_snapshot!(
+            "show_variables_like_aurora",
             execute_query(
                 "show variables like 'aurora_version';".to_string(),
                 DatabaseProtocol::MySQL
             )
-            .await?,
-            "++\n++\n++"
+            .await?
         );
 
         // All variables
@@ -6134,16 +6114,16 @@ ORDER BY \"COUNT(count)\" DESC"
 
     #[tokio::test]
     async fn test_tableau() -> Result<(), CubeError> {
-        assert_eq!(
+        insta::assert_snapshot!(
+            "tableau_table_name_column_name_query",
             execute_query(
                 "SELECT `table_name`, `column_name`
                 FROM `information_schema`.`columns`
                 WHERE `data_type`='enum' AND `table_schema`='db'"
                     .to_string(),
                 DatabaseProtocol::MySQL
             )
-            .await?,
-            "++\n++\n++"
+            .await?
         );
 
         insta::assert_snapshot!(
@@ -9522,4 +9502,28 @@ ORDER BY \"COUNT(count)\" DESC"
             }
         )
     }
+
+    #[test]
+    fn metabase_limit_0() {
+        init_logger();
+
+        let logical_plan = convert_select_to_query_plan(
+            "SELECT true AS \"_\" FROM \"public\".\"KibanaSampleDataEcommerce\" WHERE 1 <> 1 LIMIT 0".to_string(),
+            DatabaseProtocol::PostgreSQL
+        ).as_logical_plan();
+
+        assert_eq!(
+            logical_plan.find_cube_scan().request,
+            V1LoadRequestQuery {
+                measures: Some(vec![]),
+                dimensions: Some(vec![]),
+                segments: Some(vec![]),
+                time_dimensions: None,
+                order: None,
+                limit: Some(1),
+                offset: None,
+                filters: None
+            }
+        )
+    }
 }

rust/cubesql/cubesql/src/compile/rewrite/converter.rs

@@ -1,19 +1,22 @@
 use crate::{
     compile::{
-        engine::{df::scan::CubeScanNode, provider::CubeContext},
+        engine::{
+            df::scan::{CubeScanNode, MemberField},
+            provider::CubeContext,
+        },
         rewrite::{
             analysis::LogicalPlanAnalysis, rewriter::Rewriter, AggregateFunctionExprDistinct,
             AggregateFunctionExprFun, AggregateUDFExprFun, AliasExprAlias, AnyExprOp,
             BetweenExprNegated, BinaryExprOp, CastExprDataType, ColumnExprColumn, CubeScanAliases,
             CubeScanLimit, CubeScanTableName, DimensionName, EmptyRelationProduceOneRow,
             FilterMemberMember, FilterMemberOp, FilterMemberValues, FilterOpOp, InListExprNegated,
             JoinJoinConstraint, JoinJoinType, JoinLeftOn, JoinRightOn, LimitN, LiteralExprValue,
-            LogicalPlanLanguage, MeasureName, MemberErrorError, OrderAsc, OrderMember,
-            OuterColumnExprColumn, OuterColumnExprDataType, ProjectionAlias, ScalarFunctionExprFun,
-            ScalarUDFExprFun, ScalarVariableExprDataType, ScalarVariableExprVariable,
-            SegmentMemberMember, SortExprAsc, SortExprNullsFirst, TableScanLimit,
-            TableScanProjection, TableScanSourceTableName, TableScanTableName, TableUDFExprFun,
-            TimeDimensionDateRange, TimeDimensionGranularity, TimeDimensionName,
+            LiteralMemberValue, LogicalPlanLanguage, MeasureName, MemberErrorError, OrderAsc,
+            OrderMember, OuterColumnExprColumn, OuterColumnExprDataType, ProjectionAlias,
+            ScalarFunctionExprFun, ScalarUDFExprFun, ScalarVariableExprDataType,
+            ScalarVariableExprVariable, SegmentMemberMember, SortExprAsc, SortExprNullsFirst,
+            TableScanLimit, TableScanProjection, TableScanSourceTableName, TableScanTableName,
+            TableUDFExprFun, TimeDimensionDateRange, TimeDimensionGranularity, TimeDimensionName,
             TryCastExprDataType, UnionAlias, WindowFunctionExprFun, WindowFunctionExprWindowFrame,
         },
     },
@@ -33,6 +36,7 @@ use datafusion::{
         LogicalPlanBuilder, TableScan, Union,
     },
     physical_plan::planner::DefaultPhysicalPlanner,
+    scalar::ScalarValue,
     sql::planner::ContextProvider,
 };
 use egg::{EGraph, Id, RecExpr};
@@ -1081,7 +1085,7 @@ impl LanguageToLogicalPlanConverter {
                                             // TODO actually nullable. Just to fit tests
                                             false,
                                         ),
-                                        Some(measure.to_string()),
+                                        MemberField::Member(measure.to_string()),
                                     ));
                                 }
                                 LogicalPlanLanguage::TimeDimension(params) => {
@@ -1120,7 +1124,10 @@ impl LanguageToLogicalPlanConverter {
                                                 // TODO actually nullable. Just to fit tests
                                                 false,
                                             ),
-                                            Some(format!("{}.{}", dimension, granularity)),
+                                            MemberField::Member(format!(
+                                                "{}.{}",
+                                                dimension, granularity
+                                            )),
                                         ));
                                     }
                                 }
@@ -1146,7 +1153,7 @@ impl LanguageToLogicalPlanConverter {
                                             // TODO actually nullable. Just to fit tests
                                             false,
                                         ),
-                                        Some(dimension),
+                                        MemberField::Member(dimension),
                                     ));
                                 }
                                 LogicalPlanLanguage::Segment(params) => {
@@ -1160,7 +1167,22 @@ impl LanguageToLogicalPlanConverter {
                                             // TODO actually nullable. Just to fit tests
                                             false,
                                         ),
-                                        None,
+                                        MemberField::Literal(ScalarValue::Boolean(None)),
+                                    ));
+                                }
+                                LogicalPlanLanguage::LiteralMember(params) => {
+                                    let value =
+                                        match_data_node!(node_by_id, params[0], LiteralMemberValue);
+                                    let expr = self.to_expr(params[1])?;
+                                    fields.push((
+                                        DFField::new(
+                                            Some(&table_name),
+                                            &expr_name(&expr)?,
+                                            value.get_datatype(),
+                                            // TODO actually nullable. Just to fit tests
+                                            false,
+                                        ),
+                                        MemberField::Literal(value),
                                     ));
                                 }
                                 LogicalPlanLanguage::MemberError(params) => {
@@ -1321,10 +1343,7 @@ impl LanguageToLogicalPlanConverter {
                             ])
                         }
 
-                        if query_measures.len() == 0
-                            && query_dimensions.len() == 0
-                            && query_time_dimensions.len() == 0
-                        {
+                        if fields.len() == 0 {
                             return Err(CubeError::internal(
                                 "Can't detect Cube query and it may be not supported yet"
                                     .to_string(),

rust/cubesql/cubesql/src/compile/rewrite/cost.rs

@@ -72,6 +72,7 @@ impl CostFunction<LogicalPlanLanguage> for BestCubePlan {
         let cube_members = match enode {
             LogicalPlanLanguage::Measure(_) => 1,
             LogicalPlanLanguage::Dimension(_) => 1,
+            LogicalPlanLanguage::LiteralMember(_) => 1,
             LogicalPlanLanguage::TimeDimensionGranularity(TimeDimensionGranularity(Some(_))) => 1,
             LogicalPlanLanguage::MemberError(_) => 1,
             _ => 0,

rust/cubesql/cubesql/src/compile/rewrite/mod.rs

@@ -228,6 +228,10 @@ crate::plan_to_language! {
             name: String,
             expr: Arc<Expr>,
         },
+        LiteralMember {
+            value: ScalarValue,
+            expr: Arc<Expr>,
+        },
         Order {
             member: String,
             asc: bool,
@@ -800,6 +804,10 @@ fn segment_expr(name: impl Display, expr: impl Display) -> String {
     format!("(Segment {} {})", name, expr)
 }
 
+fn literal_member(value: impl Display, expr: impl Display) -> String {
+    format!("(LiteralMember {} {})", value, expr)
+}
+
 fn time_dimension_expr(
     name: impl Display,
     granularity: impl Display,