fix cubesql after introducing cacheMode

KSDaemon · KSDaemon · commit 9e10d55a83a5 · 2025-09-18T12:48:43.000+03:00
diff --git a/rust/cubesql/cubesql/src/compile/builder.rs b/rust/cubesql/cubesql/src/compile/builder.rs
@@ -150,7 +150,6 @@ impl QueryBuilder {
                 } else {
                     None
                 },
-                cache: None,
                 ungrouped: None,
                 subquery_joins: None,
                 join_hints: None,
diff --git a/rust/cubesql/cubesql/src/compile/engine/df/scan.rs b/rust/cubesql/cubesql/src/compile/engine/df/scan.rs
@@ -1,4 +1,16 @@
+use crate::compile::date_parser::parse_date_str;
+use crate::{
+    compile::{
+        engine::df::wrapper::{CubeScanWrappedSqlNode, CubeScanWrapperNode, SqlQuery},
+        test::find_cube_scans_deep_search,
+    },
+    config::ConfigObj,
+    sql::AuthContextRef,
+    transport::{CubeStreamReceiver, LoadRequestMeta, SpanId, TransportService},
+    CubeError,
+};
 use async_trait::async_trait;
+use chrono::{Datelike, NaiveDate};
 use cubeclient::models::{V1LoadRequestQuery, V1LoadResponse};
 pub use datafusion::{
     arrow::{
@@ -18,28 +30,6 @@ pub use datafusion::{
         Partitioning, PhysicalPlanner, RecordBatchStream, SendableRecordBatchStream, Statistics,
     },
 };
-use futures::Stream;
-use log::warn;
-use std::{
-    any::Any,
-    borrow::Cow,
-    fmt,
-    sync::Arc,
-    task::{Context, Poll},
-};
-
-use crate::compile::date_parser::parse_date_str;
-use crate::{
-    compile::{
-        engine::df::wrapper::{CubeScanWrappedSqlNode, CubeScanWrapperNode, SqlQuery},
-        test::find_cube_scans_deep_search,
-    },
-    config::ConfigObj,
-    sql::AuthContextRef,
-    transport::{CubeStreamReceiver, LoadRequestMeta, SpanId, TransportService},
-    CubeError,
-};
-use chrono::{Datelike, NaiveDate};
 use datafusion::{
     arrow::{
         array::{
@@ -51,7 +41,17 @@ use datafusion::{
     execution::context::TaskContext,
     scalar::ScalarValue,
 };
+use futures::Stream;
+use log::warn;
 use serde_json::Value;
+use std::str::FromStr;
+use std::{
+    any::Any,
+    borrow::Cow,
+    fmt,
+    sync::Arc,
+    task::{Context, Poll},
+};
 
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct RegularMember {
@@ -79,10 +79,33 @@ impl MemberField {
     }
 }
 
+#[derive(Debug, Clone)]
+pub enum CacheMode {
+    StaleIfSlow,
+    StaleWhileRevalidate,
+    MustRevalidate,
+    NoCache,
+}
+
+impl FromStr for CacheMode {
+    type Err = String;
+
+    fn from_str(s: &str) -> std::result::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, Clone)]
 pub struct CubeScanOptions {
     pub change_user: Option<String>,
     pub max_records: Option<usize>,
+    pub cache_mode: Option<CacheMode>,
 }
 
 #[derive(Debug, Clone)]
@@ -679,6 +702,7 @@ async fn load_data(
                 meta,
                 schema,
                 member_fields,
+                options.cache_mode,
             )
             .await
             .map_err(|err| ArrowError::ComputeError(err.to_string()))?;
@@ -1189,6 +1213,7 @@ mod tests {
                 _meta_fields: LoadRequestMeta,
                 schema: SchemaRef,
                 member_fields: Vec<MemberField>,
+                _cache_mode: Option<CacheMode>,
             ) -> Result<Vec<RecordBatch>, CubeError> {
                 let response = r#"
                 {
@@ -1314,6 +1339,7 @@ mod tests {
             options: CubeScanOptions {
                 change_user: None,
                 max_records: None,
+                cache_mode: None,
             },
             transport: get_test_transport(),
             meta: get_test_load_meta(DatabaseProtocol::PostgreSQL),
diff --git a/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs b/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs
@@ -3453,7 +3453,6 @@ 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)
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/converter.rs b/rust/cubesql/cubesql/src/compile/rewrite/converter.rs
@@ -1,5 +1,4 @@
 pub use super::rewriter::CubeRunner;
-use crate::sql::session::CacheMode;
 use crate::{
     compile::{
         engine::df::{
@@ -39,8 +38,7 @@ use crate::{
     CubeError,
 };
 use cubeclient::models::{
-    Cache as V1LoadRequestCache, V1LoadRequestQuery, V1LoadRequestQueryFilterItem,
-    V1LoadRequestQueryTimeDimension,
+    V1LoadRequestQuery, V1LoadRequestQueryFilterItem, V1LoadRequestQueryTimeDimension,
 };
 use datafusion::{
     arrow::datatypes::{DataType, TimeUnit},
@@ -1584,26 +1582,6 @@ 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 {
@@ -2077,6 +2055,13 @@ impl LanguageToLogicalPlanConverter {
 
                 let member_fields = fields.iter().map(|(_, m)| m.clone()).collect();
 
+                let cache_mode = &*self
+                    .cube_context
+                    .session_state
+                    .cache_mode
+                    .read()
+                    .expect("failed to read lock for session cache_mode");
+
                 let node = Arc::new(CubeScanNode::new(
                     Arc::new(DFSchema::new_with_metadata(
                         fields.into_iter().map(|(f, _)| f).collect(),
@@ -2088,6 +2073,7 @@ impl LanguageToLogicalPlanConverter {
                     CubeScanOptions {
                         change_user,
                         max_records,
+                        cache_mode: cache_mode.clone(),
                     },
                     alias_to_cube.into_iter().map(|(_, c)| c).unique().collect(),
                     self.span_id.clone(),
diff --git a/rust/cubesql/cubesql/src/compile/test/mod.rs b/rust/cubesql/cubesql/src/compile/test/mod.rs
@@ -47,6 +47,7 @@ pub mod test_user_change;
 #[cfg(test)]
 pub mod test_wrapper;
 pub mod utils;
+use crate::compile::engine::df::scan::CacheMode;
 use crate::compile::{
     arrow::record_batch::RecordBatch, engine::df::scan::convert_transport_response,
 };
@@ -887,6 +888,7 @@ impl TransportService for TestConnectionTransport {
         meta: LoadRequestMeta,
         schema: SchemaRef,
         member_fields: Vec<MemberField>,
+        _cache_mode: Option<CacheMode>,
     ) -> Result<Vec<RecordBatch>, CubeError> {
         {
             let mut calls = self.load_calls.lock().await;
diff --git a/rust/cubesql/cubesql/src/sql/session.rs b/rust/cubesql/cubesql/src/sql/session.rs
@@ -1,7 +1,6 @@
 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},
@@ -10,6 +9,7 @@ use std::{
 use tokio_util::sync::CancellationToken;
 
 use super::{server_manager::ServerManager, session_manager::SessionManager, AuthContextRef};
+use crate::compile::engine::df::scan::CacheMode;
 use crate::{
     compile::{
         DatabaseProtocol, DatabaseProtocolDetails, DatabaseVariable, DatabaseVariables,
@@ -57,28 +57,6 @@ 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
diff --git a/rust/cubesql/cubesql/src/transport/mod.rs b/rust/cubesql/cubesql/src/transport/mod.rs
@@ -17,6 +17,7 @@ pub type CubeMetaHierarchy = cubeclient::models::V1CubeMetaHierarchy;
 pub type TransportLoadResponse = cubeclient::models::V1LoadResponse;
 pub type TransportLoadRequestQuery = cubeclient::models::V1LoadRequestQuery;
 pub type TransportLoadRequest = cubeclient::models::V1LoadRequest;
+pub type TransportLoadRequestCacheMode = cubeclient::models::Cache;
 pub type TransportMetaResponse = cubeclient::models::V1MetaResponse;
 pub type TransportError = cubeclient::models::V1Error;
 
diff --git a/rust/cubesql/cubesql/src/transport/service.rs b/rust/cubesql/cubesql/src/transport/service.rs
@@ -25,6 +25,8 @@ use tokio::{
 };
 use uuid::Uuid;
 
+use crate::compile::engine::df::scan::CacheMode;
+use crate::transport::TransportLoadRequestCacheMode;
 use crate::{
     compile::{
         engine::df::{
@@ -142,6 +144,7 @@ pub trait TransportService: Send + Sync + Debug {
         meta_fields: LoadRequestMeta,
         schema: SchemaRef,
         member_fields: Vec<MemberField>,
+        cache_mode: Option<CacheMode>,
     ) -> Result<Vec<RecordBatch>, CubeError>;
 
     async fn load_stream(
@@ -282,6 +285,7 @@ impl TransportService for HttpTransport {
         meta: LoadRequestMeta,
         schema: SchemaRef,
         member_fields: Vec<MemberField>,
+        cache_mode: Option<CacheMode>,
     ) -> Result<Vec<RecordBatch>, CubeError> {
         if meta.change_user().is_some() {
             return Err(CubeError::internal(
@@ -290,10 +294,23 @@ impl TransportService for HttpTransport {
             ));
         }
 
+        let v1_cache_mode = match cache_mode {
+            None => None,
+            Some(m) => match m {
+                CacheMode::StaleIfSlow => Some(TransportLoadRequestCacheMode::StaleIfSlow),
+                CacheMode::StaleWhileRevalidate => {
+                    Some(TransportLoadRequestCacheMode::StaleWhileRevalidate)
+                }
+                CacheMode::MustRevalidate => Some(TransportLoadRequestCacheMode::MustRevalidate),
+                CacheMode::NoCache => Some(TransportLoadRequestCacheMode::NoCache),
+            },
+        };
+
         // TODO: support meta_fields for HTTP
         let request = TransportLoadRequest {
             query: Some(query),
             query_type: Some("multi".to_string()),
+            cache: v1_cache_mode,
         };
         let response =
             cube_api::load_v1(&self.get_client_config_for_ctx(ctx), Some(request)).await?;
