From a86dc09e81bd0ddf79246758367848346f67bb57 Mon Sep 17 00:00:00 2001 From: Dmitry Patsura Date: Fri, 16 Aug 2024 10:47:11 +0200 Subject: [PATCH] chore(cubesql): ProtocolDetails - support session variables (#8587) --- .../information_schema/mysql/variables.rs | 4 +- .../postgres/pg_settings.rs | 3 +- rust/cubesql/cubesql/src/compile/mod.rs | 2 + rust/cubesql/cubesql/src/compile/protocol.rs | 36 +++++++++++++- rust/cubesql/cubesql/src/compile/router.rs | 2 +- rust/cubesql/cubesql/src/compile/session.rs | 45 ++++++++++++++++++ .../cubesql/src/sql/database_variables/mod.rs | 47 +------------------ .../database_variables/mysql/global_vars.rs | 3 +- .../database_variables/mysql/session_vars.rs | 3 +- .../postgres/global_vars.rs | 2 +- .../postgres/session_vars.rs | 2 +- .../cubesql/cubesql/src/sql/server_manager.rs | 9 +--- rust/cubesql/cubesql/src/sql/session.rs | 25 ++++------ 13 files changed, 99 insertions(+), 84 deletions(-) create mode 100644 rust/cubesql/cubesql/src/compile/session.rs diff --git a/rust/cubesql/cubesql/src/compile/engine/information_schema/mysql/variables.rs b/rust/cubesql/cubesql/src/compile/engine/information_schema/mysql/variables.rs index c2394ed34f303..aaa32935c982e 100644 --- a/rust/cubesql/cubesql/src/compile/engine/information_schema/mysql/variables.rs +++ b/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, diff --git a/rust/cubesql/cubesql/src/compile/engine/information_schema/postgres/pg_settings.rs b/rust/cubesql/cubesql/src/compile/engine/information_schema/postgres/pg_settings.rs index 7f9b6cb9bbd6b..22400eaa3833b 100644 --- a/rust/cubesql/cubesql/src/compile/engine/information_schema/postgres/pg_settings.rs +++ b/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, } diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs index bc0652845d4c0..a3788ef5ce948 100644 --- a/rust/cubesql/cubesql/src/compile/mod.rs +++ b/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}; diff --git a/rust/cubesql/cubesql/src/compile/protocol.rs b/rust/cubesql/cubesql/src/compile/protocol.rs index f7fc4741a39a5..936a2f07e832b 100644 --- a/rust/cubesql/cubesql/src/compile/protocol.rs +++ b/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; + fn get_provider( &self, context: &CubeContext, tr: datafusion::catalog::TableReference, - ) -> Option>; + ) -> Option>; 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 { + self.get_session_default_variables().get(name).cloned() + } + fn get_provider( &self, context: &CubeContext, diff --git a/rust/cubesql/cubesql/src/compile/router.rs b/rust/cubesql/cubesql/src/compile/router.rs index 405d8972c8359..afaedf68ea4a5 100644 --- a/rust/cubesql/cubesql/src/compile/router.rs +++ b/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, diff --git a/rust/cubesql/cubesql/src/compile/session.rs b/rust/cubesql/cubesql/src/compile/session.rs new file mode 100644 index 0000000000000..d07bf0b3a7c9f --- /dev/null +++ b/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; +pub type DatabaseVariables = HashMap; + +#[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>, +} + +impl DatabaseVariable { + pub fn system( + name: String, + value: ScalarValue, + additional_params: Option>, + ) -> 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>, + ) -> Self { + Self { + name: name, + value: value, + var_type: VarType::UserDefined, + readonly: false, + additional_params, + } + } +} diff --git a/rust/cubesql/cubesql/src/sql/database_variables/mod.rs b/rust/cubesql/cubesql/src/sql/database_variables/mod.rs index bd42bcb17358f..58bc7464f8bc1 100644 --- a/rust/cubesql/cubesql/src/sql/database_variables/mod.rs +++ b/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; -pub type DatabaseVariables = HashMap; - -#[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>, -} - -impl DatabaseVariable { - pub fn system( - name: String, - value: ScalarValue, - additional_params: Option>, - ) -> 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>, - ) -> 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() } diff --git a/rust/cubesql/cubesql/src/sql/database_variables/mysql/global_vars.rs b/rust/cubesql/cubesql/src/sql/database_variables/mysql/global_vars.rs index 10d65beef2b1f..b7b010a8ab25a 100644 --- a/rust/cubesql/cubesql/src/sql/database_variables/mysql/global_vars.rs +++ b/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(); diff --git a/rust/cubesql/cubesql/src/sql/database_variables/mysql/session_vars.rs b/rust/cubesql/cubesql/src/sql/database_variables/mysql/session_vars.rs index 10d65beef2b1f..b7b010a8ab25a 100644 --- a/rust/cubesql/cubesql/src/sql/database_variables/mysql/session_vars.rs +++ b/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(); diff --git a/rust/cubesql/cubesql/src/sql/database_variables/postgres/global_vars.rs b/rust/cubesql/cubesql/src/sql/database_variables/postgres/global_vars.rs index 92af8161a0fe8..5068393135045 100644 --- a/rust/cubesql/cubesql/src/sql/database_variables/postgres/global_vars.rs +++ b/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(); diff --git a/rust/cubesql/cubesql/src/sql/database_variables/postgres/session_vars.rs b/rust/cubesql/cubesql/src/sql/database_variables/postgres/session_vars.rs index 857c3e799b437..4470d2bde4005 100644 --- a/rust/cubesql/cubesql/src/sql/database_variables/postgres/session_vars.rs +++ b/rust/cubesql/cubesql/src/sql/database_variables/postgres/session_vars.rs @@ -1,7 +1,7 @@ use datafusion::scalar::ScalarValue; use std::collections::HashMap; -use crate::sql::database_variables::{DatabaseVariable, DatabaseVariables}; +use crate::compile::{DatabaseVariable, DatabaseVariables}; pub fn defaults() -> DatabaseVariables { let mut variables: DatabaseVariables = HashMap::new(); diff --git a/rust/cubesql/cubesql/src/sql/server_manager.rs b/rust/cubesql/cubesql/src/sql/server_manager.rs index 5625b3eeccbc5..6bdad074ae8b4 100644 --- a/rust/cubesql/cubesql/src/sql/server_manager.rs +++ b/rust/cubesql/cubesql/src/sql/server_manager.rs @@ -1,12 +1,9 @@ use crate::{ - compile::DatabaseProtocol, + compile::{DatabaseProtocol, DatabaseVariables, DatabaseVariablesToUpdate}, config::ConfigObj, sql::{ compiler_cache::CompilerCache, - database_variables::{ - mysql_default_global_variables, postgres_default_global_variables, - DatabaseVariablesToUpdate, - }, + database_variables::{mysql_default_global_variables, postgres_default_global_variables}, SqlAuthService, }, transport::TransportService, @@ -14,8 +11,6 @@ use crate::{ }; use std::sync::{Arc, RwLock as RwLockSync, RwLockReadGuard, RwLockWriteGuard}; -use super::database_variables::DatabaseVariables; - #[derive(Debug)] pub struct ServerConfiguration { /// Max number of prepared statements which can be allocated per connection diff --git a/rust/cubesql/cubesql/src/sql/session.rs b/rust/cubesql/cubesql/src/sql/session.rs index 9b352c013c104..1be8bd42765a1 100644 --- a/rust/cubesql/cubesql/src/sql/session.rs +++ b/rust/cubesql/cubesql/src/sql/session.rs @@ -8,17 +8,14 @@ use std::{ }; use tokio_util::sync::CancellationToken; -use super::{ - database_variables::DatabaseVariables, server_manager::ServerManager, - session_manager::SessionManager, AuthContextRef, -}; +use super::{server_manager::ServerManager, session_manager::SessionManager, AuthContextRef}; use crate::{ - compile::{DatabaseProtocol, DatabaseProtocolDetails}, + compile::{ + DatabaseProtocol, DatabaseProtocolDetails, DatabaseVariable, DatabaseVariables, + DatabaseVariablesToUpdate, + }, sql::{ - database_variables::{ - mysql_default_session_variables, postgres_default_session_variables, DatabaseVariable, - DatabaseVariablesToUpdate, - }, + database_variables::{mysql_default_session_variables, postgres_default_session_variables}, extended::PreparedStatement, temp_tables::TempTableManager, }, @@ -328,10 +325,7 @@ impl SessionState { _ => match &self.protocol { DatabaseProtocol::MySQL => return MYSQL_DEFAULT_VARIABLES.clone(), DatabaseProtocol::PostgreSQL => return POSTGRES_DEFAULT_VARIABLES.clone(), - DatabaseProtocol::Extension(ext) => unimplemented!( - "Session.all_variables is not implemented for custom protocol: {:?}", - ext - ), + DatabaseProtocol::Extension(ext) => ext.get_session_default_variables(), }, } } @@ -349,10 +343,7 @@ impl SessionState { DatabaseProtocol::PostgreSQL => { POSTGRES_DEFAULT_VARIABLES.get(name).map(|v| v.clone()) } - DatabaseProtocol::Extension(ext) => unimplemented!( - "Session.get_variable is not implemented for custom protocol: {:?}", - ext - ), + DatabaseProtocol::Extension(ext) => ext.get_session_variable_default(name), }, } }