Working on catalog and schema registration within ExecutionContext

Author: matthewmturner

datafusion/src/catalog/mod.rs

@@ -105,7 +105,21 @@ impl<'a> TableReference<'a> {
 
 impl<'a> From<&'a str> for TableReference<'a> {
     fn from(s: &'a str) -> Self {
-        Self::Bare { table: s }
+        let parts: Vec<&str> = s.split(".").collect();
+
+        match parts.len() {
+            1 => Self::Bare { table: s },
+            2 => Self::Partial {
+                schema: parts[0],
+                table: parts[1],
+            },
+            3 => Self::Full {
+                catalog: parts[0],
+                schema: parts[1],
+                table: parts[2],
+            },
+            _ => Self::Bare { table: s },
+        }
     }
 }
 

datafusion/src/execution/context.rs

@@ -249,7 +249,6 @@ impl ExecutionContext {
                 } else {
                     Some(Arc::new(schema.as_ref().to_owned().into()))
                 };
-
                 self.register_listing_table(name, location, options, provided_schema)
                     .await?;
                 let plan = LogicalPlanBuilder::empty(false).build()?;
@@ -289,25 +288,52 @@ impl ExecutionContext {
                 ..
             }) => {
                 // sqlparser doesnt accept database / catalog as parameter to CREATE SCHEMA
-                // so for now, we default to "public" schema
-                let catalog = self.catalog("public");
-                match catalog {
-                    Some(c) => {
-                        if if_not_exists {}
-                        if let Some(s) = c.schema(&schema_name) {
-                            Err(DataFusionError::Execution(format!(
-                                "Schema {:?} already exists",
-                                schema_name
-                            )))
-                        };
+                // so for now, we default to "datafusion" catalog
+                let default_catalog = "datafusion";
+                let catalog = self.catalog(default_catalog).ok_or_else(|| {
+                    DataFusionError::Execution(String::from(
+                        "Missing 'datafusion' catalog",
+                    ))
+                })?;
+
+                let schema = catalog.schema(&schema_name);
 
+                match (if_not_exists, schema) {
+                    //
+                    (true, Some(_)) => {
+                        println!("Schema '{:?}' already exists", &schema_name);
+                        let plan = LogicalPlanBuilder::empty(false).build()?;
+                        Ok(Arc::new(DataFrameImpl::new(self.state.clone(), &plan)))
+                    }
+                    (true, None) | (false, None) => {
+                        println!("Creating schema {:?}", schema_name);
                         let schema = Arc::new(MemorySchemaProvider::new());
-                        c.register_schema(&schema_name, schema);
+                        let plan = LogicalPlanBuilder::empty(false).build()?;
+                        schema.register_table(
+                            "test".into(),
+                            Arc::new(DataFrameImpl::new(self.state.clone(), &plan)),
+                        )?;
+                        let schem_reg_res = catalog.register_schema(&schema_name, schema);
+                        match schem_reg_res {
+                            Some(_) => {
+                                println!("Existing schema with name")
+                            }
+                            None => {
+                                println!("Succesfully registerd")
+                            }
+                        };
+                        // println!("Schemas pre reg: {:?}", catalog.schema_names);
+                        self.register_catalog(default_catalog, catalog);
+                        println!(
+                            "Schema names: {:?}",
+                            self.catalog(default_catalog).unwrap().schema_names()
+                        );
                         let plan = LogicalPlanBuilder::empty(false).build()?;
                         Ok(Arc::new(DataFrameImpl::new(self.state.clone(), &plan)))
                     }
-                    None => Err(DataFusionError::Execution(String::from(
-                        "'public' catalog does not exist",
+                    (false, Some(_)) => Err(DataFusionError::Execution(format!(
+                        "Schema '{:?}' already exists",
+                        schema_name
                     ))),
                 }
             }
@@ -563,10 +589,18 @@ impl ExecutionContext {
 
         let state = self.state.lock();
         let catalog = if state.config.information_schema {
-            Arc::new(CatalogWithInformationSchema::new(
-                Arc::downgrade(&state.catalog_list),
-                catalog,
-            ))
+            let is = state
+                .catalog_list
+                .catalog("datafusion")
+                .unwrap()
+                .schema("information_schema");
+            match is {
+                Some(_) => catalog,
+                None => Arc::new(CatalogWithInformationSchema::new(
+                    Arc::downgrade(&state.catalog_list),
+                    catalog,
+                )),
+            }
         } else {
             catalog
         };
@@ -1183,6 +1217,10 @@ impl ExecutionContextState {
         table_ref: impl Into<TableReference<'a>>,
     ) -> Result<Arc<dyn SchemaProvider>> {
         let resolved_ref = self.resolve_table_ref(table_ref.into());
+        println!(
+            "Resolved ref: {:?}:{:?}:{:?}",
+            resolved_ref.catalog, resolved_ref.schema, resolved_ref.table
+        );
 
         self.catalog_list
             .catalog(resolved_ref.catalog)

datafusion/tests/sql/information_schema.rs

@@ -18,7 +18,7 @@
 use async_trait::async_trait;
 use datafusion::{
     catalog::{
-        catalog::MemoryCatalogProvider,
+        catalog::{CatalogProvider, MemoryCatalogProvider},
         schema::{MemorySchemaProvider, SchemaProvider},
     },
     datasource::{TableProvider, TableType},