chore(cubesql): Initial support for custom protocols (#8572)

Author: ovr

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

@@ -1,7 +1,8 @@
+use cubesql::compile::DatabaseProtocol;
 use cubesql::compile::{convert_sql_to_cube_query, get_df_batches};
 use cubesql::config::processing_loop::ShutdownMode;
 use cubesql::config::ConfigObj;
-use cubesql::sql::{DatabaseProtocol, SessionManager};
+use cubesql::sql::SessionManager;
 use cubesql::transport::TransportService;
 use futures::StreamExt;
 

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

@@ -1,101 +1,165 @@
-use std::sync::Arc;
+use std::{any::Any, collections::HashMap, sync::Arc};
 
+use async_trait::async_trait;
+use cubeclient::models::V1CubeMeta;
 use datafusion::{
-    arrow::datatypes::DataType,
-    error::Result,
-    scalar::ScalarValue,
-    variable::{VarProvider, VarType},
+    arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit},
+    datasource::{self, TableProvider},
+    error::DataFusionError,
+    execution::context::SessionState as DFSessionState,
+    logical_plan::Expr,
+    physical_plan::{udaf::AggregateUDF, udf::ScalarUDF, udtf::TableUDF, ExecutionPlan},
+    sql::planner::ContextProvider,
 };
-use log::warn;
 
-use crate::sql::{session::DatabaseProtocol, ServerManager, SessionState};
+use crate::{
+    compile::{DatabaseProtocolDetails, MetaContext},
+    sql::{ColumnType, SessionManager, SessionState},
+    transport::V1CubeMetaExt,
+    CubeError,
+};
 
-pub struct VariablesProvider {
-    session: Arc<SessionState>,
-    server: Arc<ServerManager>,
+#[derive(Clone)]
+pub struct CubeContext {
+    /// Internal state for the context (default)
+    pub state: Arc<DFSessionState>,
+    /// References
+    pub meta: Arc<MetaContext>,
+    pub sessions: Arc<SessionManager>,
+    pub session_state: Arc<SessionState>,
 }
 
-impl VariablesProvider {
-    pub fn new(session: Arc<SessionState>, server: Arc<ServerManager>) -> Self {
-        Self { session, server }
+impl CubeContext {
+    pub fn new(
+        state: Arc<DFSessionState>,
+        meta: Arc<MetaContext>,
+        sessions: Arc<SessionManager>,
+        session_state: Arc<SessionState>,
+    ) -> Self {
+        Self {
+            state,
+            meta,
+            sessions,
+            session_state,
+        }
+    }
+
+    pub fn table_name_by_table_provider(
+        &self,
+        table_provider: Arc<dyn datasource::TableProvider>,
+    ) -> Result<String, CubeError> {
+        self.session_state
+            .protocol
+            .table_name_by_table_provider(table_provider)
     }
 
-    fn get_session_value(&self, identifier: Vec<String>, var_type: VarType) -> Result<ScalarValue> {
-        let key = if identifier.len() > 1 {
-            let ignore_first = identifier[0].to_ascii_lowercase() == "@@session".to_owned();
-            if ignore_first {
-                identifier[1..].concat()
-            } else {
-                identifier.concat()[1..].to_string()
-            }
-        } else {
-            identifier.concat()[1..].to_string()
-        };
-
-        if let Some(var) = self.session.get_variable(&key) {
-            if var.var_type == var_type {
-                return Ok(var.value.clone());
-            }
+    pub fn get_function<T>(&self, name: &str, udfs: &HashMap<String, Arc<T>>) -> Option<Arc<T>> {
+        if name.starts_with("pg_catalog.") {
+            return udfs.get(&format!("{}", &name[11..name.len()])).cloned();
         }
 
-        warn!("Unknown session variable: {}", key);
+        udfs.get(name).cloned()
+    }
+}
 
-        Ok(ScalarValue::Utf8(None))
+impl ContextProvider for CubeContext {
+    fn get_table_provider(
+        &self,
+        tr: datafusion::catalog::TableReference,
+    ) -> Option<Arc<dyn TableProvider>> {
+        return self.session_state.protocol.get_provider(&self.clone(), tr);
     }
 
-    fn get_global_value(&self, identifier: Vec<String>) -> Result<ScalarValue> {
-        let key = if identifier.len() > 1 {
-            let ignore_first = identifier[0].to_ascii_lowercase() == "@@global".to_owned();
-
-            if ignore_first {
-                identifier[1..].concat()
-            } else {
-                identifier.concat()[2..].to_string()
-            }
-        } else {
-            identifier.concat()[2..].to_string()
-        };
-
-        if let Some(var) = self
-            .server
-            .read_variables(DatabaseProtocol::MySQL)
-            .get(&key)
-        {
-            if var.var_type == VarType::System {
-                return Ok(var.value.clone());
-            }
-        }
+    fn get_function_meta(&self, name: &str) -> Option<Arc<ScalarUDF>> {
+        self.get_function(name, &self.state.scalar_functions)
+    }
+
+    fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> {
+        self.get_function(name, &self.state.aggregate_functions)
+    }
 
-        warn!("Unknown system variable: {}", key);
+    fn get_table_function_meta(&self, name: &str) -> Option<Arc<TableUDF>> {
+        self.get_function(name, &self.state.table_functions)
+    }
 
-        Ok(ScalarValue::Utf8(None))
+    fn get_variable_type(&self, _variable_names: &[String]) -> Option<DataType> {
+        Some(DataType::Utf8)
     }
 }
 
-impl VarProvider for VariablesProvider {
-    /// get variable value
-    fn get_value(&self, identifier: Vec<String>) -> Result<ScalarValue> {
-        let first_word_vec: Vec<char> = identifier[0].chars().collect();
-        if first_word_vec.len() < 2 {
-            return Ok(ScalarValue::Utf8(None));
-        }
+pub trait TableName {
+    fn table_name(&self) -> &str;
+}
+
+pub struct CubeTableProvider {
+    cube: V1CubeMeta,
+}
 
-        match (&first_word_vec[0], &first_word_vec[1]) {
-            ('@', '@') => {
-                if identifier.len() > 1
-                    && identifier[0].to_ascii_lowercase() == "@@session".to_owned()
-                {
-                    return self.get_session_value(identifier, VarType::System);
-                }
-
-                return self.get_global_value(identifier);
-            }
-            ('@', _) => return self.get_session_value(identifier, VarType::UserDefined),
-            (_, _) => return Ok(ScalarValue::Utf8(None)),
-        };
+impl CubeTableProvider {
+    pub fn new(cube: V1CubeMeta) -> Self {
+        Self { cube }
     }
+}
 
-    fn get_type(&self, _var_names: &[String]) -> Option<DataType> {
-        Some(DataType::Utf8)
+impl TableName for CubeTableProvider {
+    fn table_name(&self) -> &str {
+        &self.cube.name
+    }
+}
+
+#[async_trait]
+impl TableProvider for CubeTableProvider {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn schema(&self) -> SchemaRef {
+        Arc::new(Schema::new(
+            self.cube
+                .get_columns()
+                .into_iter()
+                .map(|c| {
+                    Field::new(
+                        c.get_name(),
+                        match c.get_column_type() {
+                            ColumnType::Date(large) => {
+                                if large {
+                                    DataType::Date64
+                                } else {
+                                    DataType::Date32
+                                }
+                            }
+                            ColumnType::Interval(unit) => DataType::Interval(unit),
+                            ColumnType::String => DataType::Utf8,
+                            ColumnType::VarStr => DataType::Utf8,
+                            ColumnType::Boolean => DataType::Boolean,
+                            ColumnType::Double => DataType::Float64,
+                            ColumnType::Int8 => DataType::Int64,
+                            ColumnType::Int32 => DataType::Int64,
+                            ColumnType::Int64 => DataType::Int64,
+                            ColumnType::Blob => DataType::Utf8,
+                            ColumnType::Decimal(p, s) => DataType::Decimal(p, s),
+                            ColumnType::List(field) => DataType::List(field.clone()),
+                            ColumnType::Timestamp => {
+                                DataType::Timestamp(TimeUnit::Millisecond, None)
+                            }
+                        },
+                        true,
+                    )
+                })
+                .collect(),
+        ))
+    }
+
+    async fn scan(
+        &self,
+        _projection: &Option<Vec<usize>>,
+        _filters: &[Expr],
+        _limit: Option<usize>,
+    ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
+        Err(DataFusionError::Plan(format!(
+            "Not rewritten table scan node for '{}' cube",
+            self.cube.name
+        )))
     }
 }

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

@@ -0,0 +1,139 @@
+use std::sync::Arc;
+
+use datafusion::datasource::{self, TableProvider};
+
+use super::information_schema::mysql::{
+    collations::InfoSchemaCollationsProvider as MySqlSchemaCollationsProvider,
+    columns::InfoSchemaColumnsProvider as MySqlSchemaColumnsProvider,
+    key_column_usage::InfoSchemaKeyColumnUsageProvider as MySqlSchemaKeyColumnUsageProvider,
+    processlist::InfoSchemaProcesslistProvider as MySqlSchemaProcesslistProvider,
+    referential_constraints::InfoSchemaReferentialConstraintsProvider as MySqlSchemaReferentialConstraintsProvider,
+    schemata::InfoSchemaSchemataProvider as MySqlSchemaSchemataProvider,
+    statistics::InfoSchemaStatisticsProvider as MySqlSchemaStatisticsProvider,
+    tables::InfoSchemaTableProvider as MySqlSchemaTableProvider,
+    variables::PerfSchemaVariablesProvider as MySqlPerfSchemaVariablesProvider,
+};
+use crate::{
+    compile::{
+        engine::{CubeContext, CubeTableProvider, TableName},
+        DatabaseProtocol,
+    },
+    CubeError,
+};
+
+impl DatabaseProtocol {
+    pub fn get_mysql_table_name(
+        &self,
+        table_provider: Arc<dyn TableProvider>,
+    ) -> Result<String, CubeError> {
+        let any = table_provider.as_any();
+        Ok(if let Some(t) = any.downcast_ref::<CubeTableProvider>() {
+            t.table_name().to_string()
+        } else if let Some(t) = any.downcast_ref::<MySqlSchemaTableProvider>() {
+            t.table_name().to_string()
+        } else if let Some(t) = any.downcast_ref::<MySqlSchemaColumnsProvider>() {
+            t.table_name().to_string()
+        } else if let Some(t) = any.downcast_ref::<MySqlSchemaStatisticsProvider>() {
+            t.table_name().to_string()
+        } else if let Some(t) = any.downcast_ref::<MySqlSchemaKeyColumnUsageProvider>() {
+            t.table_name().to_string()
+        } else if let Some(t) = any.downcast_ref::<MySqlSchemaSchemataProvider>() {
+            t.table_name().to_string()
+        } else if let Some(t) = any.downcast_ref::<MySqlSchemaReferentialConstraintsProvider>() {
+            t.table_name().to_string()
+        } else if let Some(t) = any.downcast_ref::<MySqlSchemaCollationsProvider>() {
+            t.table_name().to_string()
+        } else if let Some(t) = any.downcast_ref::<MySqlPerfSchemaVariablesProvider>() {
+            t.table_name().to_string()
+        } else if let Some(t) = any.downcast_ref::<MySqlSchemaProcesslistProvider>() {
+            t.table_name().to_string()
+        } else {
+            return Err(CubeError::internal(format!(
+                "Unknown table provider with schema: {:?}",
+                table_provider.schema()
+            )));
+        })
+    }
+
+    pub(crate) fn get_mysql_provider(
+        &self,
+        context: &CubeContext,
+        tr: datafusion::catalog::TableReference,
+    ) -> Option<std::sync::Arc<dyn datasource::TableProvider>> {
+        let (db, table) = match tr {
+            datafusion::catalog::TableReference::Partial { schema, table, .. } => {
+                (schema.to_ascii_lowercase(), table.to_ascii_lowercase())
+            }
+            datafusion::catalog::TableReference::Full {
+                catalog: _,
+                schema,
+                table,
+            } => (schema.to_ascii_lowercase(), table.to_ascii_lowercase()),
+            datafusion::catalog::TableReference::Bare { table } => {
+                ("db".to_string(), table.to_ascii_lowercase())
+            }
+        };
+
+        match db.as_str() {
+            "db" => {
+                if let Some(cube) = context
+                    .meta
+                    .cubes
+                    .iter()
+                    .find(|c| c.name.eq_ignore_ascii_case(&table))
+                {
+                    // TODO .clone()
+                    return Some(Arc::new(CubeTableProvider::new(cube.clone())));
+                } else {
+                    return None;
+                }
+            }
+            "information_schema" => match table.as_str() {
+                "tables" => {
+                    return Some(Arc::new(MySqlSchemaTableProvider::new(
+                        context.meta.clone(),
+                    )))
+                }
+                "columns" => {
+                    return Some(Arc::new(MySqlSchemaColumnsProvider::new(
+                        context.meta.clone(),
+                    )))
+                }
+                "statistics" => return Some(Arc::new(MySqlSchemaStatisticsProvider::new())),
+                "key_column_usage" => {
+                    return Some(Arc::new(MySqlSchemaKeyColumnUsageProvider::new()))
+                }
+                "schemata" => return Some(Arc::new(MySqlSchemaSchemataProvider::new())),
+                "processlist" => {
+                    return Some(Arc::new(MySqlSchemaProcesslistProvider::new(
+                        context.sessions.clone(),
+                    )))
+                }
+                "referential_constraints" => {
+                    return Some(Arc::new(MySqlSchemaReferentialConstraintsProvider::new()))
+                }
+                "collations" => return Some(Arc::new(MySqlSchemaCollationsProvider::new())),
+                _ => return None,
+            },
+            "performance_schema" => match table.as_str() {
+                "global_variables" => {
+                    return Some(Arc::new(MySqlPerfSchemaVariablesProvider::new(
+                        "performance_schema.global_variables".to_string(),
+                        context
+                            .sessions
+                            .server
+                            .all_variables(context.session_state.protocol.clone()),
+                    )))
+                }
+                "session_variables" => {
+                    return Some(Arc::new(MySqlPerfSchemaVariablesProvider::new(
+                        "performance_schema.session_variables".to_string(),
+                        context.session_state.all_variables(),
+                    )))
+                }
+                _ => return None,
+            },
+            _ => return None,
+        }
+    }
+}