pass cache mode to cubeScan

Author: KSDaemon

packages/cubejs-backend-native/src/node_export.rs

@@ -256,6 +256,17 @@ async fn handle_sql_query(
                 .await?;
         }
 
+        let cache_enum = cache_mode.parse().map_err(|e| CubeError::user(e))?;
+
+        {
+            let mut cm = session
+                .state
+                .cache_mode
+                .write()
+                .expect("failed to unlock session cache_mode for change");
+            *cm = Some(cache_enum);
+        }
+
         let session_clone = Arc::clone(&session);
         let span_id_clone = span_id.clone();
 

rust/cubesql/cubeclient/src/models/mod.rs

@@ -24,8 +24,8 @@ pub use self::v1_error::V1Error;
 pub mod v1_load_request;
 pub use self::v1_load_request::V1LoadRequest;
 pub mod v1_load_request_query;
-pub use self::v1_load_request_query::V1LoadRequestQuery;
 pub use self::v1_load_request_query::Cache;
+pub use self::v1_load_request_query::V1LoadRequestQuery;
 pub mod v1_load_request_query_filter_base;
 pub use self::v1_load_request_query_filter_base::V1LoadRequestQueryFilterBase;
 pub mod v1_load_request_query_filter_item;

rust/cubesql/cubesql/src/compile/builder.rs

@@ -150,6 +150,7 @@ impl QueryBuilder {
                 } else {
                     None
                 },
+                cache: None,
                 ungrouped: None,
                 subquery_joins: None,
                 join_hints: None,

rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs

@@ -3453,6 +3453,7 @@ impl WrappedSelectNode {
             } else {
                 None
             },
+            cache: None,
             // TODO is it okay to just override limit?
             limit: if let Some(limit) = self.limit {
                 Some(limit as i32)

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

@@ -1,4 +1,5 @@
 pub use super::rewriter::CubeRunner;
+use crate::sql::session::CacheMode;
 use crate::{
     compile::{
         engine::df::{
@@ -38,7 +39,8 @@ use crate::{
     CubeError,
 };
 use cubeclient::models::{
-    V1LoadRequestQuery, V1LoadRequestQueryFilterItem, V1LoadRequestQueryTimeDimension,
+    Cache as V1LoadRequestCache, V1LoadRequestQuery, V1LoadRequestQueryFilterItem,
+    V1LoadRequestQueryTimeDimension,
 };
 use datafusion::{
     arrow::datatypes::{DataType, TimeUnit},
@@ -1582,6 +1584,26 @@ impl LanguageToLogicalPlanConverter {
                 let mut query_time_dimensions = Vec::new();
                 let mut query_order = Vec::new();
                 let mut query_dimensions = Vec::new();
+                let cache_mode = &*self
+                    .cube_context
+                    .session_state
+                    .cache_mode
+                    .read()
+                    .expect("failed to read lock for session cache_mode");
+
+                let v1_cache_mode = match cache_mode {
+                    None => None,
+                    Some(m) => match m {
+                        CacheMode::StaleIfSlow => Some(V1LoadRequestCache::StaleIfSlow),
+                        CacheMode::StaleWhileRevalidate => {
+                            Some(V1LoadRequestCache::StaleWhileRevalidate)
+                        }
+                        CacheMode::MustRevalidate => Some(V1LoadRequestCache::MustRevalidate),
+                        CacheMode::NoCache => Some(V1LoadRequestCache::NoCache),
+                    },
+                };
+
+                query.cache = v1_cache_mode;
 
                 for m in members {
                     match m {

rust/cubesql/cubesql/src/sql/session.rs

@@ -1,6 +1,7 @@
 use datafusion::scalar::ScalarValue;
 use log::trace;
 use rand::Rng;
+use std::str::FromStr;
 use std::{
     collections::HashMap,
     sync::{Arc, LazyLock, RwLock as RwLockSync, Weak},
@@ -56,6 +57,28 @@ pub enum QueryState {
     },
 }
 
+#[derive(Debug)]
+pub enum CacheMode {
+    StaleIfSlow,
+    StaleWhileRevalidate,
+    MustRevalidate,
+    NoCache,
+}
+
+impl FromStr for CacheMode {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "stale-if-slow" => Ok(Self::StaleIfSlow),
+            "stale-while-revalidate" => Ok(Self::StaleWhileRevalidate),
+            "must-revalidate" => Ok(Self::MustRevalidate),
+            "no-cache" => Ok(Self::NoCache),
+            other => Err(format!("Unknown cache mode: {}", other)),
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct SessionState {
     // connection id, immutable
@@ -90,6 +113,8 @@ pub struct SessionState {
     pub statements: RWLockAsync<HashMap<String, PreparedStatement>>,
 
     auth_context_expiration: Duration,
+
+    pub cache_mode: RwLockSync<Option<CacheMode>>,
 }
 
 impl SessionState {
@@ -120,6 +145,7 @@ impl SessionState {
             query: RwLockSync::new(QueryState::None),
             statements: RWLockAsync::new(HashMap::new()),
             auth_context_expiration,
+            cache_mode: RwLockSync::new(None),
         }
     }