chore(cubesql): ProtocolDetails - support session variables (#8587)

Author: ovr

rust/cubesql/cubesql/src/compile/engine/information_schema/mysql/variables.rs

@@ -1,6 +1,6 @@
 use std::{any::Any, sync::Arc};
 
-use crate::compile::engine::context::TableName;
+use crate::compile::{engine::context::TableName, DatabaseVariables};
 use async_trait::async_trait;
 use datafusion::{
     arrow::{
@@ -14,8 +14,6 @@ use datafusion::{
     physical_plan::{memory::MemoryExec, ExecutionPlan},
 };
 
-use crate::sql::database_variables::DatabaseVariables;
-
 pub struct PerfSchemaVariablesProvider {
     table_name: String,
     variables: DatabaseVariables,

rust/cubesql/cubesql/src/compile/engine/information_schema/postgres/pg_settings.rs

@@ -1,5 +1,6 @@
 use std::{any::Any, sync::Arc};
 
+use crate::compile::DatabaseVariables;
 use async_trait::async_trait;
 use datafusion::{
     arrow::{
@@ -13,8 +14,6 @@ use datafusion::{
     physical_plan::{memory::MemoryExec, ExecutionPlan},
 };
 
-use crate::sql::database_variables::DatabaseVariables;
-
 pub struct PgCatalogSettingsProvider {
     vars: DatabaseVariables,
 }

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

@@ -11,6 +11,7 @@ pub mod query_engine;
 pub mod rewrite;
 pub mod router;
 pub mod service;
+pub mod session;
 
 // Internal API
 pub mod test;
@@ -22,6 +23,7 @@ pub use protocol::*;
 pub use query_engine::*;
 pub use rewrite::rewriter::Rewriter;
 pub use router::*;
+pub use session::*;
 
 // Re-export base deps to minimise version maintenance for crate users such as cloud
 pub use datafusion::{self, arrow};

rust/cubesql/cubesql/src/compile/protocol.rs

@@ -1,5 +1,9 @@
-use crate::{compile::CubeContext, CubeError};
+use crate::{
+    compile::{CubeContext, DatabaseVariable, DatabaseVariables},
+    CubeError,
+};
 use datafusion::datasource;
+use log::error;
 use std::{
     fmt::Debug,
     hash::{Hash, Hasher},
@@ -13,11 +17,17 @@ pub trait DatabaseProtocolDetails: Debug + Send + Sync {
 
     fn support_transactions(&self) -> bool;
 
+    /// Get default state for session variables
+    fn get_session_default_variables(&self) -> DatabaseVariables;
+
+    /// Get default value for specific session variable
+    fn get_session_variable_default(&self, name: &str) -> Option<DatabaseVariable>;
+
     fn get_provider(
         &self,
         context: &CubeContext,
         tr: datafusion::catalog::TableReference,
-    ) -> Option<std::sync::Arc<dyn datasource::TableProvider>>;
+    ) -> Option<Arc<dyn datasource::TableProvider>>;
 
     fn table_name_by_table_provider(
         &self,
@@ -70,6 +80,28 @@ impl DatabaseProtocolDetails for DatabaseProtocol {
         }
     }
 
+    fn get_session_default_variables(&self) -> DatabaseVariables {
+        match &self {
+            DatabaseProtocol::MySQL => {
+                // TODO(ovr): Should we move it from session?
+                error!("get_session_default_variables was called on MySQL protocol");
+
+                DatabaseVariables::default()
+            }
+            DatabaseProtocol::PostgreSQL => {
+                // TODO(ovr): Should we move it from session?
+                error!("get_session_default_variables was called on PostgreSQL protocol");
+
+                DatabaseVariables::default()
+            }
+            DatabaseProtocol::Extension(ext) => ext.get_session_default_variables(),
+        }
+    }
+
+    fn get_session_variable_default(&self, name: &str) -> Option<DatabaseVariable> {
+        self.get_session_default_variables().get(name).cloned()
+    }
+
     fn get_provider(
         &self,
         context: &CubeContext,

rust/cubesql/cubesql/src/compile/router.rs

@@ -10,9 +10,9 @@ use crate::{
         engine::{df::planner::CubeQueryPlanner, udf::*, VariablesProvider},
         error::{CompilationError, CompilationResult},
         parser::parse_sql_to_statement,
+        DatabaseVariable, DatabaseVariablesToUpdate,
     },
     sql::{
-        database_variables::{DatabaseVariable, DatabaseVariablesToUpdate},
         dataframe,
         statement::{
             ApproximateCountDistinctVisitor, CastReplacer, RedshiftDatePartReplacer,

rust/cubesql/cubesql/src/compile/session.rs

@@ -0,0 +1,45 @@
+use datafusion::{scalar::ScalarValue, variable::VarType};
+use std::collections::HashMap;
+
+pub type DatabaseVariablesToUpdate = Vec<DatabaseVariable>;
+pub type DatabaseVariables = HashMap<String, DatabaseVariable>;
+
+#[derive(Debug, Clone)]
+pub struct DatabaseVariable {
+    pub name: String,
+    pub value: ScalarValue,
+    pub var_type: VarType,
+    pub readonly: bool,
+    // Postgres schema includes a range of additional parameters
+    pub additional_params: Option<HashMap<String, ScalarValue>>,
+}
+
+impl DatabaseVariable {
+    pub fn system(
+        name: String,
+        value: ScalarValue,
+        additional_params: Option<HashMap<String, ScalarValue>>,
+    ) -> Self {
+        Self {
+            name: name,
+            value: value,
+            var_type: VarType::System,
+            readonly: false,
+            additional_params,
+        }
+    }
+
+    pub fn user_defined(
+        name: String,
+        value: ScalarValue,
+        additional_params: Option<HashMap<String, ScalarValue>>,
+    ) -> Self {
+        Self {
+            name: name,
+            value: value,
+            var_type: VarType::UserDefined,
+            readonly: false,
+            additional_params,
+        }
+    }
+}

rust/cubesql/cubesql/src/sql/database_variables/mod.rs

@@ -1,53 +1,8 @@
-use std::collections::HashMap;
-
-use datafusion::{scalar::ScalarValue, variable::VarType};
+use crate::compile::DatabaseVariables;
 
 pub mod mysql;
 pub mod postgres;
 
-pub type DatabaseVariablesToUpdate = Vec<DatabaseVariable>;
-pub type DatabaseVariables = HashMap<String, DatabaseVariable>;
-
-#[derive(Debug, Clone)]
-pub struct DatabaseVariable {
-    pub name: String,
-    pub value: ScalarValue,
-    pub var_type: VarType,
-    pub readonly: bool,
-    // Postgres schema includes a range of additional parameters
-    pub additional_params: Option<HashMap<String, ScalarValue>>,
-}
-
-impl DatabaseVariable {
-    pub fn system(
-        name: String,
-        value: ScalarValue,
-        additional_params: Option<HashMap<String, ScalarValue>>,
-    ) -> Self {
-        Self {
-            name: name,
-            value: value,
-            var_type: VarType::System,
-            readonly: false,
-            additional_params,
-        }
-    }
-
-    pub fn user_defined(
-        name: String,
-        value: ScalarValue,
-        additional_params: Option<HashMap<String, ScalarValue>>,
-    ) -> Self {
-        Self {
-            name: name,
-            value: value,
-            var_type: VarType::UserDefined,
-            readonly: false,
-            additional_params,
-        }
-    }
-}
-
 pub fn mysql_default_session_variables() -> DatabaseVariables {
     mysql::session_vars::defaults()
 }

rust/cubesql/cubesql/src/sql/database_variables/mysql/global_vars.rs

@@ -1,9 +1,8 @@
 use std::collections::HashMap;
 
+use crate::compile::{DatabaseVariable, DatabaseVariables};
 use datafusion::scalar::ScalarValue;
 
-use crate::sql::database_variables::{DatabaseVariable, DatabaseVariables};
-
 pub fn defaults() -> DatabaseVariables {
     let mut variables: DatabaseVariables = HashMap::new();
 

rust/cubesql/cubesql/src/sql/database_variables/mysql/session_vars.rs

@@ -1,9 +1,8 @@
 use std::collections::HashMap;
 
+use crate::compile::{DatabaseVariable, DatabaseVariables};
 use datafusion::scalar::ScalarValue;
 
-use crate::sql::database_variables::{DatabaseVariable, DatabaseVariables};
-
 pub fn defaults() -> DatabaseVariables {
     let mut variables: DatabaseVariables = HashMap::new();
 

rust/cubesql/cubesql/src/sql/database_variables/postgres/global_vars.rs

@@ -2,7 +2,7 @@ use std::collections::HashMap;
 
 use datafusion::scalar::ScalarValue;
 
-use crate::sql::database_variables::{DatabaseVariable, DatabaseVariables};
+use crate::compile::{DatabaseVariable, DatabaseVariables};
 
 pub fn defaults() -> DatabaseVariables {
     let mut variables: DatabaseVariables = HashMap::new();