feat(cubesql): Initial support for INFORMATION_SCHEMA

This statement is required for DataGrip
cube-js · Nov 14, 2021 · d1fac9e · d1fac9e
1 parent 8a848aa
commit d1fac9e
Show file tree

Hide file tree

Showing 5 changed files with 52 additions and 7 deletions.
diff --git a/packages/cubejs-backend-native/Cargo.lock b/packages/cubejs-backend-native/Cargo.lock
diff --git a/rust/Cargo.lock b/rust/Cargo.lock
diff --git a/rust/cubesql/Cargo.toml b/rust/cubesql/Cargo.toml
@@ -9,7 +9,7 @@ documentation = "https://cube.dev/docs"
 homepage = "https://cube.dev"
 
 [dependencies]
-datafusion = { git = 'https://github.com/cube-js/arrow-datafusion.git', rev = "2ff38903e3fe3d66b75ead15e5bbb40fa6d36899", default-features = false }
+datafusion = { git = 'https://github.com/cube-js/arrow-datafusion.git', rev = "e6144de873addaa0b6442b12817875acc667cc11", default-features = false }
 anyhow = "1.0"
 thiserror = "1.0"
 cubeclient = { path = "../cubeclient" }

diff --git a/rust/cubesql/src/compile/mod.rs b/rust/cubesql/src/compile/mod.rs
@@ -3,6 +3,11 @@ use std::{backtrace::Backtrace, fmt};
 
 use chrono::{DateTime, Duration, TimeZone, Utc};
 
+use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit};
+use datafusion::catalog::catalog::MemoryCatalogProvider;
+use datafusion::catalog::schema::{MemorySchemaProvider, SchemaProvider};
+
+use datafusion::datasource::MemTable;
 use datafusion::sql::parser::Statement as DFStatement;
 use datafusion::sql::planner::SqlToRel;
 use datafusion::variable::VarType;
@@ -20,6 +25,7 @@ use cubeclient::models::{
 use crate::compile::parser::MySqlDialectWithBackTicks;
 use crate::mysql::dataframe;
 pub use crate::schema::ctx::*;
+use crate::schema::V1CubeMetaExt;
 use crate::CubeError;
 use crate::{
     compile::builder::QueryBuilder,
@@ -1167,6 +1173,10 @@ impl QueryPlanner {
             }
         };
 
+        if schema_name.to_lowercase() == "information_schema" {
+            return self.create_df_logical_plan(stmt.clone(), props);
+        }
+
         if schema_name.to_lowercase() != "db" {
             return Err(CompilationError::Unsupported(format!(
                 "Unable to access schema {}",
@@ -1276,7 +1286,8 @@ impl QueryPlanner {
         stmt: ast::Statement,
         props: &QueryPlannerExecutionProps,
     ) -> CompilationResult<QueryPlan> {
-        let mut ctx = ExecutionContext::new();
+        let mut ctx =
+            ExecutionContext::with_config(ExecutionConfig::new().with_information_schema(true));
 
         let variable_provider = SystemVar::new();
         ctx.register_variable(VarType::System, Arc::new(variable_provider));
@@ -1287,6 +1298,40 @@ impl QueryPlanner {
         ctx.register_udf(create_user_udf(props));
         ctx.register_udf(create_current_user_udf(props));
 
+        {
+            let schema_provider = MemorySchemaProvider::new();
+
+            for cube in &self.context.cubes {
+                let mut schema_fields = vec![];
+
+                for column in cube.get_columns() {
+                    let data_type = match column.mysql_type_as_str().as_str() {
+                        "int" => DataType::Int64,
+                        "time" => DataType::Timestamp(TimeUnit::Nanosecond, None),
+                        _ => DataType::Utf8,
+                    };
+
+                    schema_fields.push(Field::new(
+                        column.get_name(),
+                        data_type,
+                        column.mysql_can_be_null(),
+                    ));
+                }
+
+                let schema = Arc::new(Schema::new(schema_fields));
+                let provider = MemTable::try_new(schema.clone(), vec![vec![]]).unwrap();
+
+                schema_provider
+                    .register_table(cube.name.clone(), Arc::new(provider))
+                    .map_err(|err| CubeError::internal(err.to_string()));
+            }
+
+            let catalog_provider = MemoryCatalogProvider::new();
+            catalog_provider.register_schema("db", Arc::new(schema_provider));
+
+            ctx.register_catalog("db", Arc::new(catalog_provider));
+        }
+
         let state = ctx.state.lock().unwrap().clone();
         let df_query_planner = SqlToRel::new(&state);
 

diff --git a/rust/cubesql/src/config/injection.rs b/rust/cubesql/src/config/injection.rs
@@ -103,7 +103,7 @@ impl Injector {
             .read()
             .await
             .get(name)
-            .expect(&format!("Service is not found: {}", name))
+            .unwrap_or_else(|| panic!("Service is not found: {}", name))
             .clone();
         // println!("Locking service: {}", name);
         // TODO cycle depends lead to dead lock here
@@ -116,14 +116,14 @@ impl Injector {
         let factories = self.factories.read().await;
         let factory = factories
             .get(name)
-            .expect(&format!("Service not found: {}", name));
+            .unwrap_or_else(|| panic!("Service not found: {}", name));
         let service = factory(self.this.upgrade().unwrap()).await;
         // println!("Setting service: {}", name);
         self.services
             .write()
             .await
             .insert(name.to_string(), service.clone());
-        return service.clone().downcast(service).unwrap();
+        service.clone().downcast(service).unwrap()
     }
 
     pub async fn try_get_service<T: ?Sized + Send + Sync + 'static>(