Remove iceberg catalog from executor for drop statements (#17)

* Merge

* Remove iceberg catalog from drop statement

* Fix

* Fix

Author: osipovartem

crates/catalog/src/catalog.rs

@@ -1,13 +1,20 @@
+use crate::df_error;
 use crate::schema::CachingSchema;
 use chrono::NaiveDateTime;
 use dashmap::DashMap;
 use datafusion::catalog::{CatalogProvider, SchemaProvider};
+use datafusion_common::DataFusionError;
+use futures::executor::block_on;
+use iceberg_rust::catalog::Catalog;
+use iceberg_rust_spec::namespace::Namespace;
+use snafu::futures::TryFutureExt;
 use std::fmt::{Display, Formatter};
 use std::{any::Any, sync::Arc};
 
 #[derive(Clone)]
 pub struct CachingCatalog {
     pub catalog: Arc<dyn CatalogProvider>,
+    pub iceberg_catalog: Option<Arc<dyn Catalog>>,
     pub catalog_type: CatalogType,
     pub schemas_cache: DashMap<String, Arc<CachingSchema>>,
     pub should_refresh: bool,
@@ -50,9 +57,14 @@ impl Display for CatalogType {
 }
 
 impl CachingCatalog {
-    pub fn new(catalog: Arc<dyn CatalogProvider>, name: String) -> Self {
+    pub fn new(
+        catalog_provider: Arc<dyn CatalogProvider>,
+        name: String,
+        iceberg_catalog: Option<Arc<dyn Catalog>>,
+    ) -> Self {
         Self {
-            catalog,
+            catalog: catalog_provider,
+            iceberg_catalog,
             schemas_cache: DashMap::new(),
             should_refresh: false,
             enable_information_schema: true,
@@ -128,6 +140,7 @@ impl CatalogProvider for CachingCatalog {
                         name: name.clone(),
                         schema,
                         tables_cache: DashMap::new(),
+                        iceberg_catalog: self.iceberg_catalog.clone(),
                     }),
                 );
             }
@@ -148,6 +161,7 @@ impl CatalogProvider for CachingCatalog {
                 name: name.to_string(),
                 schema: Arc::clone(&schema),
                 tables_cache: DashMap::new(),
+                iceberg_catalog: self.iceberg_catalog.clone(),
             });
 
             self.schemas_cache
@@ -173,6 +187,7 @@ impl CatalogProvider for CachingCatalog {
             name: name.to_string(),
             schema: Arc::clone(&schema),
             tables_cache: DashMap::new(),
+            iceberg_catalog: self.iceberg_catalog.clone(),
         });
         self.schemas_cache
             .insert(name.to_string(), Arc::clone(&caching_schema));
@@ -190,7 +205,22 @@ impl CatalogProvider for CachingCatalog {
         name: &str,
         cascade: bool,
     ) -> datafusion_common::Result<Option<Arc<dyn SchemaProvider>>> {
-        self.schemas_cache.remove(name);
-        self.catalog.deregister_schema(name, cascade)
+        let schema = self.schemas_cache.remove(name);
+
+        if let Some(catalog) = &self.iceberg_catalog {
+            let namespace = Namespace::try_new(std::slice::from_ref(&name.to_string()))
+                .map_err(|err| DataFusionError::External(Box::new(err)))?;
+            block_on(
+                catalog
+                    .drop_namespace(&namespace)
+                    .context(df_error::IcebergSnafu),
+            )?;
+        } else {
+            return self.catalog.deregister_schema(name, cascade);
+        }
+        if let Some((_, caching_schema)) = schema {
+            return Ok(Some(caching_schema));
+        }
+        Ok(None)
     }
 }

crates/catalog/src/catalog_list.rs

@@ -22,6 +22,7 @@ use datafusion::{
 };
 use datafusion_iceberg::catalog::catalog::IcebergCatalog as DataFusionIcebergCatalog;
 use futures::future::join_all;
+use iceberg_rust::catalog::Catalog;
 use iceberg_rust::object_store::ObjectStoreBuilder;
 use iceberg_s3tables_catalog::S3TablesCatalog;
 use object_store::ObjectStore;
@@ -155,9 +156,8 @@ impl EmbucketCatalogList {
         err
     )]
     pub async fn register_catalogs(self: &Arc<Self>) -> Result<()> {
-        let mut all_catalogs = Vec::new();
         // Add metastore databases as catalogs
-        all_catalogs.extend(self.internal_catalogs().await?);
+        let all_catalogs = self.metastore_catalogs().await?;
         for catalog in all_catalogs {
             self.catalogs
                 .insert(catalog.name.clone(), Arc::new(catalog));
@@ -171,7 +171,7 @@ impl EmbucketCatalogList {
         skip(self),
         err
     )]
-    pub async fn internal_catalogs(&self) -> Result<Vec<CachingCatalog>> {
+    pub async fn metastore_catalogs(&self) -> Result<Vec<CachingCatalog>> {
         let mut catalogs = Vec::new();
         let databases = self
             .metastore
@@ -198,19 +198,23 @@ impl EmbucketCatalogList {
     }
 
     fn get_embucket_catalog(&self, db: &RwObject<Database>) -> Result<CachingCatalog> {
-        let iceberg_catalog = EmbucketIcebergCatalog::new(self.metastore.clone(), db.ident.clone())
-            .context(MetastoreSnafu)?;
-        let catalog: Arc<dyn CatalogProvider> = Arc::new(EmbucketCatalog::new(
+        let iceberg_catalog: Arc<dyn Catalog> = Arc::new(
+            EmbucketIcebergCatalog::new(self.metastore.clone(), db.ident.clone())
+                .context(MetastoreSnafu)?,
+        );
+        let catalog_provider: Arc<dyn CatalogProvider> = Arc::new(EmbucketCatalog::new(
             db.ident.clone(),
             self.metastore.clone(),
-            Arc::new(iceberg_catalog),
+            iceberg_catalog.clone(),
         ));
-        Ok(CachingCatalog::new(catalog, db.ident.clone())
-            .with_refresh(true)
-            .with_properties(Properties {
-                created_at: db.created_at,
-                updated_at: db.created_at,
-            }))
+        Ok(
+            CachingCatalog::new(catalog_provider, db.ident.clone(), Some(iceberg_catalog))
+                .with_refresh(true)
+                .with_properties(Properties {
+                    created_at: db.created_at,
+                    updated_at: db.created_at,
+                }),
+        )
     }
 
     #[tracing::instrument(
@@ -238,19 +242,23 @@ impl EmbucketCatalogList {
             .credentials_provider(SharedCredentialsProvider::new(creds))
             .region(Region::new(volume.region()))
             .build();
-        let catalog = S3TablesCatalog::new(
-            &config,
-            volume.arn.as_str(),
-            ObjectStoreBuilder::S3(Box::new(volume.s3_builder())),
-        )
-        .context(catalog_error::S3TablesSnafu)?;
+        let iceberg_catalog: Arc<dyn Catalog> = Arc::new(
+            S3TablesCatalog::new(
+                &config,
+                volume.arn.as_str(),
+                ObjectStoreBuilder::S3(Box::new(volume.s3_builder())),
+            )
+            .context(catalog_error::S3TablesSnafu)?,
+        );
 
-        let catalog = DataFusionIcebergCatalog::new(Arc::new(catalog), None)
+        let catalog = DataFusionIcebergCatalog::new(iceberg_catalog.clone(), None)
             .await
             .context(catalog_error::DataFusionSnafu)?;
-        Ok(CachingCatalog::new(Arc::new(catalog), name.to_string())
-            .with_refresh(true)
-            .with_catalog_type(CatalogType::S3tables))
+        Ok(
+            CachingCatalog::new(Arc::new(catalog), name.to_string(), Some(iceberg_catalog))
+                .with_refresh(false)
+                .with_catalog_type(CatalogType::S3tables),
+        )
     }
 
     /// Do not keep returned references to avoid deadlocks
@@ -302,6 +310,7 @@ impl EmbucketCatalogList {
                             schema: schema_provider,
                             tables_cache: DashMap::default(),
                             name: schema.clone(),
+                            iceberg_catalog: catalog_ref.iceberg_catalog.clone(),
                         };
                         catalog_ref
                             .schemas_cache
@@ -386,6 +395,7 @@ impl EmbucketCatalogList {
                             schema: schema_provider,
                             tables_cache: DashMap::default(),
                             name: schema.clone(),
+                            iceberg_catalog: catalog.iceberg_catalog.clone(),
                         };
                         let tables = schema.schema.table_names();
 
@@ -457,7 +467,7 @@ impl EmbucketCatalogList {
             let mut interval = interval(Duration::from_secs(interval_secs));
             loop {
                 interval.tick().await;
-                match self.internal_catalogs().await {
+                match self.metastore_catalogs().await {
                     Ok(catalogs) => {
                         for catalog in catalogs {
                             if self.catalogs.contains_key(&catalog.name) {
@@ -539,7 +549,7 @@ impl CatalogProviderList for EmbucketCatalogList {
         name: String,
         catalog: Arc<dyn CatalogProvider>,
     ) -> Option<Arc<dyn CatalogProvider>> {
-        let catalog = CachingCatalog::new(catalog, name);
+        let catalog = CachingCatalog::new(catalog, name, None);
         self.catalogs
             .insert(catalog.name.clone(), Arc::new(catalog))
             .map(|arc| {

crates/catalog/src/df_error.rs

@@ -1,4 +1,5 @@
 use error_stack_trace;
+use iceberg_rust::error::Error as IcebergError;
 use snafu::Location;
 use snafu::prelude::*;
 
@@ -45,6 +46,14 @@ pub enum DFExternalError {
         #[snafu(implicit)]
         location: Location,
     },
+
+    #[snafu(display("Iceberg error: {error}"))]
+    Iceberg {
+        #[snafu(source(from(IcebergError, Box::new)))]
+        error: Box<IcebergError>,
+        #[snafu(implicit)]
+        location: Location,
+    },
 }
 
 impl From<DFExternalError> for datafusion_common::DataFusionError {

crates/catalog/src/schema.rs

@@ -1,14 +1,20 @@
+use crate::df_error;
 use crate::table::CachingTable;
 use async_trait::async_trait;
 use dashmap::DashMap;
 use datafusion::catalog::{SchemaProvider, TableProvider};
 use datafusion_common::DataFusionError;
 use datafusion_expr::TableType;
+use futures::executor::block_on;
+use iceberg_rust::catalog::Catalog;
+use iceberg_rust_spec::identifier::Identifier;
+use snafu::futures::TryFutureExt;
 use std::any::Any;
 use std::sync::Arc;
 
 pub struct CachingSchema {
     pub schema: Arc<dyn SchemaProvider>,
+    pub iceberg_catalog: Option<Arc<dyn Catalog>>,
     pub name: String,
     pub tables_cache: DashMap<String, Arc<CachingTable>>,
 }
@@ -89,9 +95,15 @@ impl SchemaProvider for CachingSchema {
         name: &str,
     ) -> datafusion_common::Result<Option<Arc<dyn TableProvider>>> {
         let table = self.tables_cache.remove(name);
+
         if let Some((_, caching_table)) = table {
             if caching_table.table_type() != TableType::View {
-                return self.schema.deregister_table(name);
+                if let Some(catalog) = &self.iceberg_catalog {
+                    let ident = Identifier::new(std::slice::from_ref(&self.name), name);
+                    block_on(catalog.drop_table(&ident).context(df_error::IcebergSnafu))?;
+                } else {
+                    return self.schema.deregister_table(name);
+                }
             }
             return Ok(Some(caching_table as Arc<dyn TableProvider>));
         }