feat(cubesql): Query rewrite cache (#7647)

* feat(cubesql): Introduce query rewrite cache

* Bind cache to protocol and fix tests

* Fix tests

* Remove mock calls check as it isn't predictable

* Do not cache LogicalPlan as it provides security checks

* Fix queries cache key

* Missing default hashing for language nodes introduces cache key ambiguity for ScalarVariableExprVariable

* Do not eval stable function such as timestamps before putting rewrite to cache
Add ConstantFolding to aliases and aliases push down for FilterSimplifyReplacer so it can find Alias wrapped stable functions and fold them to constant

* Disable cache by default but enable in tests

Author: paveltiunov

.github/workflows/rust-cubesql.yml

@@ -81,6 +81,7 @@ jobs:
           CUBESQL_TESTING_CUBE_URL: ${{ secrets.CUBESQL_TESTING_CUBE_URL }}
           CUBESQL_REWRITE_ENGINE: true
           CUBESQL_SQL_PUSH_DOWN: true
+          CUBESQL_REWRITE_CACHE: true
           CUBESQL_REWRITE_TIMEOUT: 60
         run: cd rust/cubesql && cargo tarpaulin --workspace --no-fail-fast --avoid-cfg-tarpaulin --out Xml
       - name: Upload code coverage
@@ -119,6 +120,7 @@ jobs:
           CUBESQL_TESTING_CUBE_TOKEN: ${{ secrets.CUBESQL_TESTING_CUBE_TOKEN }}
           CUBESQL_TESTING_CUBE_URL: ${{ secrets.CUBESQL_TESTING_CUBE_URL }}
           CUBESQL_SQL_PUSH_DOWN: true
+          CUBESQL_REWRITE_CACHE: true
         run: cd rust/cubesql && cargo test
 
   native_linux:

packages/cubejs-backend-native/Cargo.lock

@@ -240,7 +240,6 @@ describe('SQLInterface', () => {
         password: null,
       });
 
-      expect(meta.mock.calls.length).toEqual(3);
       expect(meta.mock.calls[0][0]).toEqual({
         request: {
           id: expect.any(String),

packages/cubejs-backend-native/test/sql.test.ts

@@ -61,6 +61,7 @@ sha1_smol = "1.0.0"
 tera = { version = "1", default-features = false }
 minijinja = { version = "1", features = ["json", "loader"] }
 lru = "0.12.1"
+sha2 = "0.10.8"
 
 
 [dev-dependencies]

rust/cubesql/Cargo.lock

@@ -4805,7 +4805,20 @@ pub fn register_fun_stubs(mut ctx: SessionContext) -> SessionContext {
         vol = Stable
     );
 
-    register_fun_stub!(udf, "eval_now", argc = 0, rettyp = Timestamp, vol = Stable);
+    register_fun_stub!(
+        udf,
+        "eval_now",
+        argc = 0,
+        rettyp = TimestampTz,
+        vol = Stable
+    );
+    register_fun_stub!(
+        udf,
+        "eval_utc_timestamp",
+        argc = 0,
+        rettyp = Timestamp,
+        vol = Stable
+    );
 
     ctx
 }

rust/cubesql/cubesql/Cargo.toml

@@ -15,7 +15,7 @@ use datafusion::{
         projection_drop_out::ProjectionDropOut,
         utils::from_plan,
     },
-    physical_plan::ExecutionPlan,
+    physical_plan::{planner::DefaultPhysicalPlanner, ExecutionPlan},
     prelude::*,
     scalar::ScalarValue,
     sql::{parser::Statement as DFStatement, planner::SqlToRel},
@@ -105,7 +105,10 @@ pub mod test;
 
 pub use crate::transport::ctx::*;
 use crate::{
-    compile::engine::df::wrapper::CubeScanWrapperNode,
+    compile::{
+        engine::df::wrapper::CubeScanWrapperNode,
+        rewrite::{analysis::LogicalPlanAnalysis, rewriter::Rewriter},
+    },
     transport::{LoadRequestMeta, SpanId, TransportService},
 };
 pub use error::{CompilationError, CompilationResult};
@@ -1387,6 +1390,8 @@ WHERE `TABLE_SCHEMA` = '{}'",
         qtrace: &mut Option<Qtrace>,
         span_id: Option<Arc<SpanId>>,
     ) -> CompilationResult<QueryPlan> {
+        self.reauthenticate_if_needed().await?;
+
         match &stmt {
             ast::Statement::Query(query) => match &query.body {
                 ast::SetExpr::Select(select) if select.into.is_some() => {
@@ -1451,15 +1456,64 @@ WHERE `TABLE_SCHEMA` = '{}'",
             qtrace.set_optimized_plan(&optimized_plan);
         }
 
-        let mut converter = LogicalPlanToLanguageConverter::new(Arc::new(cube_ctx));
+        let cube_ctx = Arc::new(cube_ctx);
+        let mut converter = LogicalPlanToLanguageConverter::new(cube_ctx.clone());
+        let mut query_params = Some(HashMap::new());
         let root = converter
-            .add_logical_plan(&optimized_plan)
+            .add_logical_plan_replace_params(&optimized_plan, &mut query_params)
             .map_err(|e| CompilationError::internal(e.to_string()))?;
 
-        self.reauthenticate_if_needed().await?;
+        let mut finalized_graph = self
+            .session_manager
+            .server
+            .compiler_cache
+            .rewrite(
+                self.state.auth_context().unwrap(),
+                cube_ctx.clone(),
+                converter.take_egraph(),
+                &query_params.unwrap(),
+                qtrace,
+            )
+            .await
+            .map_err(|e| match e.cause {
+                CubeErrorCauseType::Internal(_) => CompilationError::Internal(
+                    format!(
+                        "Error during rewrite: {}. Please check logs for additional information.",
+                        e.message
+                    ),
+                    e.to_backtrace().unwrap_or_else(|| Backtrace::capture()),
+                    Some(HashMap::from([
+                        ("query".to_string(), stmt.to_string()),
+                        (
+                            "sanitizedQuery".to_string(),
+                            SensitiveDataSanitizer::new().replace(&stmt).to_string(),
+                        ),
+                    ])),
+                ),
+                CubeErrorCauseType::User(_) => CompilationError::User(
+                    format!(
+                        "Error during rewrite: {}. Please check logs for additional information.",
+                        e.message
+                    ),
+                    Some(HashMap::from([
+                        ("query".to_string(), stmt.to_string()),
+                        (
+                            "sanitizedQuery".to_string(),
+                            SensitiveDataSanitizer::new().replace(&stmt).to_string(),
+                        ),
+                    ])),
+                ),
+            })?;
 
-        let result = converter
-            .take_rewriter()
+        // Replace Analysis as at least time has changed but it might be also context may affect rewriting in some other ways
+        finalized_graph.analysis = LogicalPlanAnalysis::new(
+            cube_ctx.clone(),
+            Arc::new(DefaultPhysicalPlanner::default()),
+        );
+
+        let mut rewriter = Rewriter::new(finalized_graph, cube_ctx.clone());
+
+        let result = rewriter
             .find_best_plan(
                 root,
                 self.state.auth_context().unwrap(),
@@ -3244,15 +3298,9 @@ mod tests {
 
         let logical_plan = &query_plan.print(true).unwrap();
 
-        let re = Regex::new(r"TimestampNanosecond\(\d+, None\)").unwrap();
-        let logical_plan = re
-            .replace_all(logical_plan, "TimestampNanosecond(0, None)")
-            .as_ref()
-            .to_string();
-
         assert_eq!(
             logical_plan,
-            "Projection: TimestampNanosecond(0, None) AS COL\
+            "Projection: CAST(utctimestamp() AS current_timestamp() AS Timestamp(Nanosecond, None)) AS COL\
             \n  EmptyRelation",
         );
     }
@@ -15679,7 +15727,10 @@ limit
             // Quarter of Year
             ["EXTRACT(QUARTER FROM t.\"order_date\")", "quarter"],
             // Year
-            ["CAST(EXTRACT(YEAR FROM t.\"order_date\") AS varchar)", "year"],
+            [
+                "CAST(EXTRACT(YEAR FROM t.\"order_date\") AS varchar)",
+                "year",
+            ],
         ];
 
         for [expr, expected_granularity] in &supported_granularities {
@@ -18248,7 +18299,7 @@ limit
                 segments: Some(vec![]),
                 time_dimensions: Some(vec![V1LoadRequestQueryTimeDimension {
                     dimension: "KibanaSampleDataEcommerce.order_date".to_string(),
-                    granularity: Some("quarter".to_string()),
+                    granularity: Some("month".to_string()),
                     date_range: None
                 }]),
                 order: None,

rust/cubesql/cubesql/src/compile/engine/udf.rs

@@ -833,6 +833,7 @@ impl LogicalPlanAnalysis {
                     _ => panic!("Expected Literal but got: {:?}", expr),
                 }
             }
+            LogicalPlanLanguage::AliasExpr(params) => constant_node(params[0]),
             LogicalPlanLanguage::ScalarUDFExpr(_) => {
                 let expr = node_to_expr(
                     enode,
@@ -858,6 +859,14 @@ impl LogicalPlanAnalysis {
                                 args: vec![],
                             },
                         )
+                    } else if &fun.name == "eval_utc_timestamp" {
+                        Self::eval_constant_expr(
+                            &egraph,
+                            &Expr::ScalarFunction {
+                                fun: BuiltinScalarFunction::UtcTimestamp,
+                                args: vec![],
+                            },
+                        )
                     } else if &fun.name == "str_to_date"
                         || &fun.name == "date_add"
                         || &fun.name == "date_sub"
@@ -883,13 +892,10 @@ impl LogicalPlanAnalysis {
                 .ok()?;
 
                 if let Expr::ScalarFunction { fun, .. } = &expr {
-                    if (fun.volatility() == Volatility::Immutable
-                        || fun.volatility() == Volatility::Stable)
-                        && !matches!(
-                            fun,
-                            BuiltinScalarFunction::CurrentDate | BuiltinScalarFunction::Now
-                        )
-                    {
+                    // Removed stable evaluation as it affects caching and SQL push down.
+                    // Whatever stable function should be evaluated it should be addressed as a special rewrite rule
+                    // as it seems LogicalPlanAnalysis can't change it's state.
+                    if fun.volatility() == Volatility::Immutable {
                         Self::eval_constant_expr(&egraph, &expr)
                     } else {
                         None
@@ -1185,18 +1191,29 @@ impl Analysis<LogicalPlanLanguage> for LogicalPlanAnalysis {
 
     fn modify(egraph: &mut EGraph<LogicalPlanLanguage, Self>, id: Id) {
         if let Some(ConstantFolding::Scalar(c)) = &egraph[id].data.constant {
+            // As ConstantFolding goes through Alias we can't add LiteralExpr at this level otherwise it gets dropped.
+            // In case there's wrapping node on top of Alias that can be evaluated to LiteralExpr further it gets replaced instead.
+            if let Some(Expr::Alias(_, _)) = egraph[id].data.original_expr.as_ref() {
+                return;
+            }
             // TODO: ideally all constants should be aliased, but this requires
             // rewrites to extract `.data.constant` instead of `literal_expr`.
-            let alias_name =
-                if c.is_null() || matches!(c, ScalarValue::Date32(_) | ScalarValue::Date64(_)) {
-                    egraph[id]
-                        .data
-                        .original_expr
-                        .as_ref()
-                        .map(|expr| expr.name(&DFSchema::empty()).unwrap())
-                } else {
-                    None
-                };
+            let alias_name = if c.is_null()
+                || matches!(
+                    c,
+                    ScalarValue::Date32(_)
+                        | ScalarValue::Date64(_)
+                        | ScalarValue::Int64(_)
+                        | ScalarValue::Float64(_)
+                ) {
+                egraph[id]
+                    .data
+                    .original_expr
+                    .as_ref()
+                    .map(|expr| expr.name(&DFSchema::empty()).unwrap())
+            } else {
+                None
+            };
             let c = c.clone();
             let value = egraph.add(LogicalPlanLanguage::LiteralExprValue(LiteralExprValue(c)));
             let literal_expr = egraph.add(LogicalPlanLanguage::LiteralExpr([value]));