Add an extension point to allow SQL based registration of external catalogs

datafusion-cli/src/catalog.rs

@@ -57,6 +57,10 @@ impl CatalogProviderList for DynamicObjectStoreCatalog {
         self.inner.register_catalog(name, catalog)
     }
 
+    fn deregister_catalog(&self, name: &str) -> Result<Option<Arc<dyn CatalogProvider>>> {
+        self.inner.deregister_catalog(name)
+    }
+
     fn catalog_names(&self) -> Vec<String> {
         self.inner.catalog_names()
     }

datafusion-examples/examples/data_io/catalog.rs

@@ -279,6 +279,11 @@ impl CatalogProviderList for CustomCatalogProviderList {
         Some(catalog)
     }
 
+    fn deregister_catalog(&self, name: &str) -> Result<Option<Arc<dyn CatalogProvider>>> {
+        let mut cats = self.catalogs.write().unwrap();
+        Ok(cats.remove(name))
+    }
+
     /// Retrieves the list of available catalog names
     fn catalog_names(&self) -> Vec<String> {
         let cats = self.catalogs.read().unwrap();

datafusion/catalog/src/catalog.rs

@@ -16,7 +16,9 @@
 // under the License.
 
 // Re-export from this module for backwards compatibility.
-pub use datafusion_session::{CatalogProvider, CatalogProviderList};
+pub use datafusion_session::{
+    CatalogProvider, CatalogProviderFactory, CatalogProviderList,
+};
 // Re-export so users can access this type through `datafusion_catalog` and
 // `datafusion::catalog` without depending directly on `datafusion_session`.
 pub use datafusion_session::EmptyCatalogProviderList;

datafusion/catalog/src/dynamic_file/catalog.rs

@@ -49,6 +49,13 @@ impl CatalogProviderList for DynamicFileCatalog {
         self.inner.register_catalog(name, catalog)
     }
 
+    fn deregister_catalog(
+        &self,
+        name: &str,
+    ) -> datafusion_common::Result<Option<Arc<dyn CatalogProvider>>> {
+        self.inner.deregister_catalog(name)
+    }
+
     fn catalog_names(&self) -> Vec<String> {
         self.inner.catalog_names()
     }

datafusion/catalog/src/memory/catalog.rs

@@ -20,7 +20,7 @@
 
 use crate::{CatalogProvider, CatalogProviderList, SchemaProvider};
 use dashmap::DashMap;
-use datafusion_common::exec_err;
+use datafusion_common::{Result, exec_err};
 use std::sync::Arc;
 
 /// Simple in-memory list of catalogs
@@ -54,6 +54,10 @@ impl CatalogProviderList for MemoryCatalogProviderList {
         self.catalogs.insert(name, catalog)
     }
 
+    fn deregister_catalog(&self, name: &str) -> Result<Option<Arc<dyn CatalogProvider>>> {
+        Ok(self.catalogs.remove(name).map(|(_, catalog)| catalog))
+    }
+
     fn catalog_names(&self) -> Vec<String> {
         self.catalogs.iter().map(|c| c.key().clone()).collect()
     }
@@ -97,15 +101,15 @@ impl CatalogProvider for MemoryCatalogProvider {
         &self,
         name: &str,
         schema: Arc<dyn SchemaProvider>,
-    ) -> datafusion_common::Result<Option<Arc<dyn SchemaProvider>>> {
+    ) -> Result<Option<Arc<dyn SchemaProvider>>> {
         Ok(self.schemas.insert(name.into(), schema))
     }
 
     fn deregister_schema(
         &self,
         name: &str,
         cascade: bool,
-    ) -> datafusion_common::Result<Option<Arc<dyn SchemaProvider>>> {
+    ) -> Result<Option<Arc<dyn SchemaProvider>>> {
         if let Some(schema) = self.schema(name) {
             let table_names = schema.table_names();
             match (table_names.is_empty(), cascade) {

datafusion/core/src/execution/context/mod.rs

@@ -28,7 +28,8 @@ use crate::execution::session_state::SessionStateBuilder;
 use crate::{
     catalog::listing_schema::ListingSchemaProvider,
     catalog::{
-        CatalogProvider, CatalogProviderList, TableProvider, TableProviderFactory,
+        CatalogProvider, CatalogProviderFactory, CatalogProviderList, TableProvider,
+        TableProviderFactory,
     },
     dataframe::DataFrame,
     datasource::listing::{
@@ -44,10 +45,10 @@ use crate::{
     logical_expr::AggregateUDF,
     logical_expr::ScalarUDF,
     logical_expr::{
-        CreateCatalog, CreateCatalogSchema, CreateExternalTable, CreateFunction,
-        CreateMemoryTable, CreateView, DropCatalogSchema, DropFunction, DropTable,
-        DropView, Execute, LogicalPlan, LogicalPlanBuilder, Prepare, ResetVariable,
-        SetVariable, TableType, UNNAMED_TABLE,
+        CreateCatalog, CreateCatalogSchema, CreateExternalCatalog, CreateExternalTable,
+        CreateFunction, CreateMemoryTable, CreateView, DropCatalog, DropCatalogSchema,
+        DropFunction, DropTable, DropView, Execute, LogicalPlan, LogicalPlanBuilder,
+        Prepare, ResetVariable, SetVariable, TableType, UNNAMED_TABLE,
     },
     physical_expr::PhysicalExpr,
     physical_plan::ExecutionPlan,
@@ -561,6 +562,19 @@ impl SessionContext {
         self.state.read().table_factories().get(file_type).cloned()
     }
 
+    /// Return the [`CatalogProviderFactory`] that is registered for the
+    /// specified catalog type, if any.
+    pub fn catalog_factory(
+        &self,
+        catalog_type: &str,
+    ) -> Option<Arc<dyn CatalogProviderFactory>> {
+        self.state
+            .read()
+            .catalog_factories()
+            .get(catalog_type)
+            .cloned()
+    }
+
     /// Return the `enable_ident_normalization` of this Session
     pub fn enable_ident_normalization(&self) -> bool {
         self.state
@@ -705,9 +719,15 @@ impl SessionContext {
                         self.create_catalog_schema(cmd)
                     }
                     DdlStatement::CreateCatalog(cmd) => self.create_catalog(cmd),
+                    DdlStatement::CreateExternalCatalog(cmd) => {
+                        (Box::pin(async move { self.create_external_catalog(&cmd).await })
+                            as std::pin::Pin<Box<dyn Future<Output = _> + Send>>)
+                            .await
+                    }
                     DdlStatement::DropTable(cmd) => Box::pin(self.drop_table(cmd)).await,
                     DdlStatement::DropView(cmd) => Box::pin(self.drop_view(cmd)).await,
                     DdlStatement::DropCatalogSchema(cmd) => self.drop_schema(cmd),
+                    DdlStatement::DropCatalog(cmd) => self.drop_catalog(cmd),
                     DdlStatement::CreateFunction(cmd) => {
                         Box::pin(self.create_function(*cmd)).await
                     }
@@ -1043,6 +1063,49 @@ impl SessionContext {
         }
     }
 
+    async fn create_external_catalog(
+        &self,
+        cmd: &CreateExternalCatalog,
+    ) -> Result<DataFrame> {
+        let exists = self.catalog(cmd.catalog_name.as_str()).is_some();
+
+        match (cmd.if_not_exists, cmd.or_replace, exists) {
+            (true, false, true) => self.return_empty_dataframe(),
+            (true, true, true) => {
+                exec_err!("'IF NOT EXISTS' cannot coexist with 'REPLACE'")
+            }
+            (false, false, true) => {
+                exec_err!("External catalog '{}' already exists", cmd.catalog_name)
+            }
+            (_, _, _) => {
+                let new_catalog = self.create_custom_catalog(cmd).await?;
+                self.state
+                    .write()
+                    .catalog_list()
+                    .register_catalog(cmd.catalog_name.clone(), new_catalog);
+                self.return_empty_dataframe()
+            }
+        }
+    }
+
+    async fn create_custom_catalog(
+        &self,
+        cmd: &CreateExternalCatalog,
+    ) -> Result<Arc<dyn CatalogProvider>> {
+        let state = self.state.read().clone();
+        let catalog_type = cmd.catalog_type.to_uppercase();
+        let factory = state
+            .catalog_factories()
+            .get(catalog_type.as_str())
+            .ok_or_else(|| {
+                exec_datafusion_err!(
+                    "Unable to find catalog factory for {}",
+                    cmd.catalog_type
+                )
+            })?;
+        factory.create(&state, cmd).await
+    }
+
     async fn drop_table(&self, cmd: DropTable) -> Result<DataFrame> {
         let DropTable {
             name, if_exists, ..
@@ -1106,6 +1169,22 @@ impl SessionContext {
         exec_err!("Schema '{schema_ref}' doesn't exist.")
     }
 
+    fn drop_catalog(&self, cmd: DropCatalog) -> Result<DataFrame> {
+        let DropCatalog {
+            name, if_exists, ..
+        } = cmd;
+        let dereg = self
+            .state
+            .write()
+            .catalog_list()
+            .deregister_catalog(&name)?;
+        match (dereg, if_exists) {
+            (Some(_), _) => self.return_empty_dataframe(),
+            (None, true) => self.return_empty_dataframe(),
+            (None, false) => exec_err!("Catalog '{name}' doesn't exist."),
+        }
+    }
+
     fn set_variable(&self, stmt: SetVariable) -> Result<()> {
         let SetVariable { variable, value } = stmt;
 

datafusion/core/src/execution/session_state.rs

@@ -23,7 +23,9 @@ use std::collections::{HashMap, HashSet};
 use std::fmt::Debug;
 use std::sync::Arc;
 
-use crate::catalog::{CatalogProviderList, SchemaProvider, TableProviderFactory};
+use crate::catalog::{
+    CatalogProviderFactory, CatalogProviderList, SchemaProvider, TableProviderFactory,
+};
 use crate::datasource::file_format::FileFormatFactory;
 #[cfg(feature = "sql")]
 use crate::datasource::provider_as_source;
@@ -202,6 +204,16 @@ struct SessionStateInner {
     ///
     /// [`TableProvider`]: crate::catalog::TableProvider
     table_factories: HashMap<String, Arc<dyn TableProviderFactory>>,
+    /// CatalogProviderFactories for different catalog implementations.
+    ///
+    /// Maps strings like "ICEBERG" to an instance of [`CatalogProviderFactory`]
+    ///
+    /// This is used to create [`CatalogProvider`] instances for the
+    /// `CREATE EXTERNAL CATALOG ... STORED AS <TYPE>` statement, for catalogs
+    /// backed by an external implementation.
+    ///
+    /// [`CatalogProvider`]: crate::catalog::CatalogProvider
+    catalog_factories: HashMap<String, Arc<dyn CatalogProviderFactory>>,
     /// Runtime environment
     runtime_env: Arc<RuntimeEnv>,
     /// [FunctionFactory] to support pluggable user defined function handler.
@@ -246,6 +258,7 @@ impl Debug for SessionState {
             .field("execution_props", &self.execution_props)
             .field("table_options", &self.inner.table_options)
             .field("table_factories", &self.inner.table_factories)
+            .field("catalog_factories", &self.inner.catalog_factories)
             .field("function_factory", &self.inner.function_factory)
             .field("cache_factory", &self.inner.cache_factory)
             .field("expr_planners", &self.inner.expr_planners);
@@ -494,6 +507,18 @@ impl SessionState {
         &mut Arc::make_mut(&mut self.inner).table_factories
     }
 
+    /// Get the catalog factories
+    pub fn catalog_factories(&self) -> &HashMap<String, Arc<dyn CatalogProviderFactory>> {
+        &self.inner.catalog_factories
+    }
+
+    /// Get the catalog factories
+    pub fn catalog_factories_mut(
+        &mut self,
+    ) -> &mut HashMap<String, Arc<dyn CatalogProviderFactory>> {
+        &mut Arc::make_mut(&mut self.inner).catalog_factories
+    }
+
     /// Parse an SQL string into an DataFusion specific AST
     /// [`Statement`]. See [`SessionContext::sql`] for running queries.
     ///
@@ -1134,6 +1159,7 @@ pub struct SessionStateBuilder {
     table_options: Option<TableOptions>,
     execution_props: Option<ExecutionProps>,
     table_factories: Option<HashMap<String, Arc<dyn TableProviderFactory>>>,
+    catalog_factories: Option<HashMap<String, Arc<dyn CatalogProviderFactory>>>,
     runtime_env: Option<Arc<RuntimeEnv>>,
     function_factory: Option<Arc<dyn FunctionFactory>>,
     cache_factory: Option<Arc<dyn CacheFactory>>,
@@ -1177,6 +1203,7 @@ impl SessionStateBuilder {
             config: None,
             execution_props: None,
             table_factories: None,
+            catalog_factories: None,
             runtime_env: None,
             function_factory: None,
             cache_factory: None,
@@ -1242,6 +1269,7 @@ impl SessionStateBuilder {
             table_options: Some(existing.table_options),
             execution_props: Some(execution_props),
             table_factories: Some(existing.table_factories),
+            catalog_factories: Some(existing.catalog_factories),
             runtime_env: Some(existing.runtime_env),
             function_factory: existing.function_factory,
             cache_factory: existing.cache_factory,
@@ -1556,6 +1584,27 @@ impl SessionStateBuilder {
         self
     }
 
+    /// Add a [`CatalogProviderFactory`] to the map of factories
+    pub fn with_catalog_factory(
+        mut self,
+        key: String,
+        catalog_factory: Arc<dyn CatalogProviderFactory>,
+    ) -> Self {
+        let mut catalog_factories = self.catalog_factories.unwrap_or_default();
+        catalog_factories.insert(key, catalog_factory);
+        self.catalog_factories = Some(catalog_factories);
+        self
+    }
+
+    /// Set the map of [`CatalogProviderFactory`]s
+    pub fn with_catalog_factories(
+        mut self,
+        catalog_factories: HashMap<String, Arc<dyn CatalogProviderFactory>>,
+    ) -> Self {
+        self.catalog_factories = Some(catalog_factories);
+        self
+    }
+
     /// Set the [`RuntimeEnv`]
     pub fn with_runtime_env(mut self, runtime_env: Arc<RuntimeEnv>) -> Self {
         self.runtime_env = Some(runtime_env);
@@ -1655,6 +1704,7 @@ impl SessionStateBuilder {
             config,
             execution_props,
             table_factories,
+            catalog_factories,
             runtime_env,
             function_factory,
             cache_factory,
@@ -1696,6 +1746,7 @@ impl SessionStateBuilder {
             }),
             config,
             table_factories: table_factories.unwrap_or_default(),
+            catalog_factories: catalog_factories.unwrap_or_default(),
             runtime_env,
             function_factory,
             cache_factory,
@@ -1937,6 +1988,13 @@ impl SessionStateBuilder {
         &mut self.table_factories
     }
 
+    /// Returns the current catalog_factories value
+    pub fn catalog_factories(
+        &mut self,
+    ) -> &mut Option<HashMap<String, Arc<dyn CatalogProviderFactory>>> {
+        &mut self.catalog_factories
+    }
+
     /// Returns the current runtime_env value
     pub fn runtime_env(&mut self) -> &mut Option<Arc<RuntimeEnv>> {
         &mut self.runtime_env
@@ -1989,6 +2047,7 @@ impl Debug for SessionStateBuilder {
             .field("execution_props", &self.execution_props)
             .field("table_options", &self.table_options)
             .field("table_factories", &self.table_factories)
+            .field("catalog_factories", &self.catalog_factories)
             .field("function_factory", &self.function_factory)
             .field("cache_factory", &self.cache_factory)
             .field("expr_planners", &self.expr_planners);