From 1a8714d159f6f7d2ba3c8f33a1935770e1844750 Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sat, 15 Aug 2026 14:41:53 +0200 Subject: [PATCH 01/11] refactor(datafusion): rename catalog provider modules --- .../datafusion/src/{catalog.rs => catalog_provider.rs} | 2 +- crates/integrations/datafusion/src/lib.rs | 6 +++--- .../datafusion/src/{schema.rs => schema_provider.rs} | 0 3 files changed, 4 insertions(+), 4 deletions(-) rename crates/integrations/datafusion/src/{catalog.rs => catalog_provider.rs} (98%) rename crates/integrations/datafusion/src/{schema.rs => schema_provider.rs} (100%) diff --git a/crates/integrations/datafusion/src/catalog.rs b/crates/integrations/datafusion/src/catalog_provider.rs similarity index 98% rename from crates/integrations/datafusion/src/catalog.rs rename to crates/integrations/datafusion/src/catalog_provider.rs index 2c6e1ff002..af01779e0d 100644 --- a/crates/integrations/datafusion/src/catalog.rs +++ b/crates/integrations/datafusion/src/catalog_provider.rs @@ -22,7 +22,7 @@ use datafusion::catalog::{CatalogProvider, SchemaProvider}; use futures::future::try_join_all; use iceberg::{Catalog, NamespaceIdent, Result}; -use crate::schema::IcebergSchemaProvider; +use crate::schema_provider::IcebergSchemaProvider; /// Provides an interface to manage and access multiple schemas /// within an Iceberg [`Catalog`]. diff --git a/crates/integrations/datafusion/src/lib.rs b/crates/integrations/datafusion/src/lib.rs index 4b0ea8606d..5bf798892f 100644 --- a/crates/integrations/datafusion/src/lib.rs +++ b/crates/integrations/datafusion/src/lib.rs @@ -15,14 +15,14 @@ // specific language governing permissions and limitations // under the License. -mod catalog; -pub use catalog::*; +mod catalog_provider; +pub use catalog_provider::*; mod error; pub use error::*; pub mod physical_plan; -mod schema; +mod schema_provider; pub mod table; pub use table::table_provider_factory::IcebergTableProviderFactory; pub use table::*; diff --git a/crates/integrations/datafusion/src/schema.rs b/crates/integrations/datafusion/src/schema_provider.rs similarity index 100% rename from crates/integrations/datafusion/src/schema.rs rename to crates/integrations/datafusion/src/schema_provider.rs From a824a158fe8e192205d2b720c11656058de98f1b Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sat, 15 Aug 2026 14:59:49 +0200 Subject: [PATCH 02/11] Introduce the SessionContextResolver trait --- crates/integrations/datafusion/public-api.txt | 2 ++ crates/integrations/datafusion/src/lib.rs | 25 ++++++++++++++++--- 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/crates/integrations/datafusion/public-api.txt b/crates/integrations/datafusion/public-api.txt index e197c2057d..85b7f568f0 100644 --- a/crates/integrations/datafusion/public-api.txt +++ b/crates/integrations/datafusion/public-api.txt @@ -134,5 +134,7 @@ impl core::fmt::Debug for iceberg_datafusion::table_provider_factory::IcebergTab pub fn iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::table::TableProviderFactory for iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory pub fn iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory::create<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, cmd: &'life2 datafusion_expr::logical_plan::ddl::CreateExternalTable) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub trait iceberg_datafusion::SessionContextResolver: core::fmt::Debug + core::marker::Send + core::marker::Sync +pub fn iceberg_datafusion::SessionContextResolver::resolve(&self, session: &dyn datafusion_session::session::Session) -> datafusion_common::error::Result pub fn iceberg_datafusion::from_datafusion_error(error: datafusion_common::error::DataFusionError) -> iceberg::error::Error pub fn iceberg_datafusion::to_datafusion_error(error: iceberg::error::Error) -> datafusion_common::error::DataFusionError diff --git a/crates/integrations/datafusion/src/lib.rs b/crates/integrations/datafusion/src/lib.rs index 5bf798892f..3e3930c004 100644 --- a/crates/integrations/datafusion/src/lib.rs +++ b/crates/integrations/datafusion/src/lib.rs @@ -15,12 +15,11 @@ // specific language governing permissions and limitations // under the License. -mod catalog_provider; -pub use catalog_provider::*; - mod error; pub use error::*; +mod catalog_provider; +pub use catalog_provider::*; pub mod physical_plan; mod schema_provider; pub mod table; @@ -28,3 +27,23 @@ pub use table::table_provider_factory::IcebergTableProviderFactory; pub use table::*; pub(crate) mod task_writer; + +use std::fmt; + +use datafusion::catalog::Session as DFSession; +use datafusion::error::Result as DFResult; +use iceberg::SessionContext; + +/// Resolves an Iceberg [`SessionContext`] from a DataFusion session. +/// +/// The DataFusion integration calls the resolver once while planning each +/// session-aware scan or insert. The returned context is bound to that +/// operation and, for inserts, is retained through transaction commit. +/// Implementations should therefore return a stable Iceberg session identity +/// for repeated operations from the same DataFusion session. +pub trait SessionContextResolver: fmt::Debug + Send + Sync { + /// Returns the Iceberg context associated with `session`. + /// + /// Returning an error aborts planning before the catalog is accessed. + fn resolve(&self, session: &dyn DFSession) -> DFResult; +} From 63ee82678a2475d21cafdd051d5069eedd5eec29 Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sat, 15 Aug 2026 15:01:16 +0200 Subject: [PATCH 03/11] Introduce the CatalogAccess abstraction --- .../datafusion/src/catalog_access.rs | 38 +++++++++++++++++++ crates/integrations/datafusion/src/lib.rs | 1 + 2 files changed, 39 insertions(+) create mode 100644 crates/integrations/datafusion/src/catalog_access.rs diff --git a/crates/integrations/datafusion/src/catalog_access.rs b/crates/integrations/datafusion/src/catalog_access.rs new file mode 100644 index 0000000000..db9135bb61 --- /dev/null +++ b/crates/integrations/datafusion/src/catalog_access.rs @@ -0,0 +1,38 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use iceberg::{Catalog, SessionCatalog}; + +use crate::SessionContextResolver; + +/// Describes how the DataFusion integration accesses an Iceberg catalog. +/// +/// A catalog can either be accessed directly through [`Catalog`] or through +/// [`SessionCatalog`] with an Iceberg [`SessionContext`] derived from the +/// current DataFusion session. Operations for which DataFusion provides no +/// session use an empty Iceberg context. +#[derive(Clone, Debug)] +pub(crate) enum CatalogAccess { + /// A catalog accessed directly through the [`Catalog`] API. + Direct(Arc), + + /// A session-aware catalog and the resolver used to derive its Iceberg + /// context from the current DataFusion session. + SessionAware(Arc, Arc), +} diff --git a/crates/integrations/datafusion/src/lib.rs b/crates/integrations/datafusion/src/lib.rs index 3e3930c004..c7ad5f8124 100644 --- a/crates/integrations/datafusion/src/lib.rs +++ b/crates/integrations/datafusion/src/lib.rs @@ -17,6 +17,7 @@ mod error; pub use error::*; +mod catalog_access; mod catalog_provider; pub use catalog_provider::*; From e0cf986fa86ed66e83a37f46fe2198d663250718 Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sat, 15 Aug 2026 15:16:13 +0200 Subject: [PATCH 04/11] Use CatalogAccess in providers and forward session were available --- crates/integrations/datafusion/public-api.txt | 6 +- .../datafusion/src/catalog_provider.rs | 91 +++++++++++-------- .../datafusion/src/physical_plan/commit.rs | 4 +- .../datafusion/src/schema_provider.rs | 83 ++++++++++------- .../integrations/datafusion/src/table/mod.rs | 65 +++++++++---- 5 files changed, 155 insertions(+), 94 deletions(-) diff --git a/crates/integrations/datafusion/public-api.txt b/crates/integrations/datafusion/public-api.txt index 85b7f568f0..76d5d4f33e 100644 --- a/crates/integrations/datafusion/public-api.txt +++ b/crates/integrations/datafusion/public-api.txt @@ -76,7 +76,7 @@ impl core::fmt::Debug for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::table::TableProvider for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::insert_into<'life0, 'life1, 'async_trait>(&'life0 self, state: &'life1 dyn datafusion_session::session::Session, input: alloc::sync::Arc, _insert_op: datafusion_expr::logical_plan::dml::InsertOp) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait -pub fn iceberg_datafusion::IcebergTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, projection: core::option::Option<&'life2 alloc::vec::Vec>, filters: &'life3 [datafusion_expr::expr::Expr], limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait +pub fn iceberg_datafusion::IcebergTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, state: &'life1 dyn datafusion_session::session::Session, projection: core::option::Option<&'life2 alloc::vec::Vec>, filters: &'life3 [datafusion_expr::expr::Expr], limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait pub fn iceberg_datafusion::IcebergTableProvider::schema(&self) -> arrow_schema::schema::SchemaRef pub fn iceberg_datafusion::IcebergTableProvider::supports_filters_pushdown(&self, filters: &[&datafusion_expr::expr::Expr]) -> datafusion_common::error::Result> pub fn iceberg_datafusion::IcebergTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType @@ -93,7 +93,7 @@ impl datafusion_catalog::table::TableProviderFactory for iceberg_datafusion::tab pub fn iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory::create<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, cmd: &'life2 datafusion_expr::logical_plan::ddl::CreateExternalTable) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait pub struct iceberg_datafusion::IcebergCatalogProvider impl iceberg_datafusion::IcebergCatalogProvider -pub async fn iceberg_datafusion::IcebergCatalogProvider::try_new(client: alloc::sync::Arc) -> iceberg::error::Result +pub async fn iceberg_datafusion::IcebergCatalogProvider::try_new(catalog: alloc::sync::Arc) -> iceberg::error::Result impl core::fmt::Debug for iceberg_datafusion::IcebergCatalogProvider pub fn iceberg_datafusion::IcebergCatalogProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::catalog::CatalogProvider for iceberg_datafusion::IcebergCatalogProvider @@ -120,7 +120,7 @@ impl core::fmt::Debug for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::table::TableProvider for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::insert_into<'life0, 'life1, 'async_trait>(&'life0 self, state: &'life1 dyn datafusion_session::session::Session, input: alloc::sync::Arc, _insert_op: datafusion_expr::logical_plan::dml::InsertOp) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait -pub fn iceberg_datafusion::IcebergTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, projection: core::option::Option<&'life2 alloc::vec::Vec>, filters: &'life3 [datafusion_expr::expr::Expr], limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait +pub fn iceberg_datafusion::IcebergTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, state: &'life1 dyn datafusion_session::session::Session, projection: core::option::Option<&'life2 alloc::vec::Vec>, filters: &'life3 [datafusion_expr::expr::Expr], limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait pub fn iceberg_datafusion::IcebergTableProvider::schema(&self) -> arrow_schema::schema::SchemaRef pub fn iceberg_datafusion::IcebergTableProvider::supports_filters_pushdown(&self, filters: &[&datafusion_expr::expr::Expr]) -> datafusion_common::error::Result> pub fn iceberg_datafusion::IcebergTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType diff --git a/crates/integrations/datafusion/src/catalog_provider.rs b/crates/integrations/datafusion/src/catalog_provider.rs index af01779e0d..d5367f206e 100644 --- a/crates/integrations/datafusion/src/catalog_provider.rs +++ b/crates/integrations/datafusion/src/catalog_provider.rs @@ -20,8 +20,10 @@ use std::sync::Arc; use datafusion::catalog::{CatalogProvider, SchemaProvider}; use futures::future::try_join_all; -use iceberg::{Catalog, NamespaceIdent, Result}; +use iceberg::{Catalog, NamespaceIdent, Result, SessionCatalog, SessionContext}; +use crate::SessionContextResolver; +use crate::catalog_access::CatalogAccess; use crate::schema_provider::IcebergSchemaProvider; /// Provides an interface to manage and access multiple schemas @@ -38,47 +40,36 @@ pub struct IcebergCatalogProvider { } impl IcebergCatalogProvider { - /// Asynchronously tries to construct a new [`IcebergCatalogProvider`] - /// using the given client to fetch and initialize schema providers for - /// each namespace in the Iceberg [`Catalog`]. + /// Asynchronously constructs an [`IcebergCatalogProvider`] from a + /// [`Catalog`], fetching and initializing a schema provider for each + /// namespace. /// - /// This method retrieves the list of namespace names - /// attempts to create a schema provider for each namespace, and - /// collects these providers into a `HashMap`. - pub async fn try_new(client: Arc) -> Result { + /// This method retrieves the namespace names and collects an initialized + /// schema provider for each namespace into a `HashMap`. + pub async fn try_new(catalog: Arc) -> Result { + let direct = CatalogAccess::Direct(catalog); + Self::try_new_with_access(direct).await + } + + async fn try_new_with_access(catalog: CatalogAccess) -> Result { // TODO: // Schemas and providers should be cached and evicted based on time // As of right now; schemas might become stale. - let schema_names: Vec<_> = client - .list_namespaces(None) - .await? - .iter() - .flat_map(|ns| ns.as_ref().clone()) - .collect(); - - let providers = try_join_all( - schema_names - .iter() - .map(|name| { - IcebergSchemaProvider::try_new( - client.clone(), - NamespaceIdent::new(name.clone()), - ) - }) - .collect::>(), - ) - .await?; - - let schemas: HashMap> = schema_names - .into_iter() - .zip(providers) - .map(|(name, provider)| { - let provider = Arc::new(provider) as Arc; - (name, provider) - }) - .collect(); + let schema_names: Vec<_> = match &catalog { + CatalogAccess::Direct(catalog) => catalog.list_namespaces(None).await?, + CatalogAccess::SessionAware(session_catalog, _) => { + session_catalog + .list_namespaces(&SessionContext::empty(), None) + .await? + } + } + .iter() + .flat_map(|ns| ns.as_ref().clone()) + .collect(); - Ok(IcebergCatalogProvider { schemas }) + Ok(IcebergCatalogProvider { + schemas: load_schema_providers(catalog, schema_names).await?, + }) } } @@ -91,3 +82,29 @@ impl CatalogProvider for IcebergCatalogProvider { self.schemas.get(name).cloned() } } + +async fn load_schema_providers( + catalog: CatalogAccess, + schema_names: Vec, +) -> Result>> { + let iceberg_providers = try_join_all( + schema_names + .iter() + .map(|name| { + IcebergSchemaProvider::try_new(catalog.clone(), NamespaceIdent::new(name.clone())) + }) + .collect::>(), + ) + .await?; + + let provider_map = schema_names + .into_iter() + .zip(iceberg_providers) + .map(|(name, iceberg_provider)| { + let provider = Arc::new(iceberg_provider) as Arc; + (name, provider) + }) + .collect(); + + Ok(provider_map) +} diff --git a/crates/integrations/datafusion/src/physical_plan/commit.rs b/crates/integrations/datafusion/src/physical_plan/commit.rs index 9ae8b845ce..7901e8b7ab 100644 --- a/crates/integrations/datafusion/src/physical_plan/commit.rs +++ b/crates/integrations/datafusion/src/physical_plan/commit.rs @@ -42,6 +42,7 @@ use crate::to_datafusion_error; #[derive(Debug)] pub(crate) struct IcebergCommitExec { table: Table, + /// Catalog already bound to the session resolved during insert planning. catalog: Arc, input: Arc, schema: ArrowSchemaRef, @@ -287,6 +288,7 @@ mod tests { use iceberg::{Catalog, CatalogBuilder, NamespaceIdent, TableCreation, TableIdent}; use super::*; + use crate::catalog_access::CatalogAccess; use crate::physical_plan::DATA_FILES_COL_NAME; use crate::table::IcebergTableProvider; @@ -658,7 +660,7 @@ mod tests { ctx.register_table("source_table", source_table)?; let iceberg_table_provider = IcebergTableProvider::try_new( - catalog.clone(), + CatalogAccess::Direct(catalog), namespace.clone(), table_name.to_string(), ) diff --git a/crates/integrations/datafusion/src/schema_provider.rs b/crates/integrations/datafusion/src/schema_provider.rs index 545863f8c6..106a83837c 100644 --- a/crates/integrations/datafusion/src/schema_provider.rs +++ b/crates/integrations/datafusion/src/schema_provider.rs @@ -23,23 +23,27 @@ use datafusion::catalog::SchemaProvider; use datafusion::datasource::TableProvider; use datafusion::error::{DataFusionError, Result as DFResult}; use datafusion::execution::TaskContext; -use datafusion::prelude::SessionContext; +use datafusion::prelude::SessionContext as DFSessionContext; use futures::StreamExt; use futures::future::try_join_all; use iceberg::arrow::arrow_schema_to_schema_auto_assign_ids; use iceberg::inspect::MetadataTableType; use iceberg::spec::FormatVersion; -use iceberg::{Catalog, Error, ErrorKind, NamespaceIdent, Result, TableCreation, TableIdent}; +use iceberg::{ + Error, ErrorKind, NamespaceIdent, Result, SessionContext as IcebergSessionContext, + TableCreation, TableIdent, +}; +use crate::catalog_access::CatalogAccess; use crate::table::IcebergTableProvider; use crate::to_datafusion_error; -/// Represents a [`SchemaProvider`] for the Iceberg [`Catalog`], managing -/// access to table providers within a specific namespace. +/// Represents a [`SchemaProvider`] for an Iceberg catalog, managing table +/// providers within a specific namespace. #[derive(Debug)] pub(crate) struct IcebergSchemaProvider { - /// Reference to the Iceberg catalog - catalog: Arc, + /// Access policy for the Iceberg catalog. + catalog: CatalogAccess, /// The namespace this schema represents namespace: NamespaceIdent, /// A concurrent map where keys are table names @@ -54,28 +58,28 @@ impl IcebergSchemaProvider { /// using the given client to fetch and initialize table providers for /// the provided namespace in the Iceberg [`Catalog`]. /// - /// This method retrieves a list of table names - /// attempts to create a table provider for each table name, and - /// collects these providers into a `HashMap`. - pub(crate) async fn try_new( - client: Arc, - namespace: NamespaceIdent, - ) -> Result { + /// This method retrieves a list of table names, attempts to create a table + /// provider for each name, and collects the providers into a [`DashMap`]. + pub(crate) async fn try_new(catalog: CatalogAccess, namespace: NamespaceIdent) -> Result { // TODO: // Tables and providers should be cached based on table_name // if we have a cache miss; we update our internal cache & check again // As of right now; tables might become stale. - let table_names: Vec<_> = client - .list_tables(&namespace) - .await? - .iter() - .map(|tbl| tbl.name().to_string()) - .collect(); + let table_names: Vec<_> = match &catalog { + CatalogAccess::Direct(catalog) => catalog.list_tables(&namespace).await?, + CatalogAccess::SessionAware(session_catalog, _) => { + let context = IcebergSessionContext::empty(); + session_catalog.list_tables(&context, &namespace).await? + } + } + .iter() + .map(|tbl| tbl.name().to_string()) + .collect(); let providers = try_join_all( table_names .iter() - .map(|name| IcebergTableProvider::try_new(client.clone(), namespace.clone(), name)) + .map(|name| IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), name)) .collect::>(), ) .await?; @@ -86,7 +90,7 @@ impl IcebergSchemaProvider { } Ok(IcebergSchemaProvider { - catalog: client, + catalog, namespace, tables, }) @@ -186,10 +190,18 @@ impl SchemaProvider for IcebergSchemaProvider { .await .map_err(to_datafusion_error)?; - catalog - .create_table(&namespace, table_creation) - .await - .map_err(to_datafusion_error)?; + match &catalog { + CatalogAccess::Direct(catalog) => { + catalog.create_table(&namespace, table_creation).await + } + CatalogAccess::SessionAware(session_catalog, _) => { + let context = IcebergSessionContext::empty(); + session_catalog + .create_table(&context, &namespace, table_creation) + .await + } + } + .map_err(to_datafusion_error)?; // Create a new table provider using the catalog reference let table_provider = IcebergTableProvider::try_new( @@ -232,10 +244,14 @@ impl SchemaProvider for IcebergSchemaProvider { let table_ident = TableIdent::new(namespace, table_name.clone()); // Drop the table from the Iceberg catalog - catalog - .drop_table(&table_ident) - .await - .map_err(to_datafusion_error)?; + match &catalog { + CatalogAccess::Direct(catalog) => catalog.drop_table(&table_ident).await, + CatalogAccess::SessionAware(session_catalog, _) => { + let context = IcebergSessionContext::empty(); + session_catalog.drop_table(&context, &table_ident).await + } + } + .map_err(to_datafusion_error)?; // Remove from local cache and return the removed provider let removed = tables @@ -254,7 +270,7 @@ impl SchemaProvider for IcebergSchemaProvider { /// Verifies that a table provider contains no data by scanning with LIMIT 1. /// Returns an error if the table has any rows. async fn ensure_table_is_empty(table: &Arc) -> Result<()> { - let session_ctx = SessionContext::new(); + let session_ctx = DFSessionContext::new(); let exec_plan = table .scan(&session_ctx.state(), None, &[], Some(1)) .await @@ -317,9 +333,10 @@ mod tests { .await .unwrap(); - let provider = IcebergSchemaProvider::try_new(Arc::new(catalog), namespace) - .await - .unwrap(); + let provider = + IcebergSchemaProvider::try_new(CatalogAccess::Direct(Arc::new(catalog)), namespace) + .await + .unwrap(); (provider, temp_dir) } diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 9de7bcb9c2..6812756822 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -45,9 +45,10 @@ use iceberg::arrow::schema_to_arrow_schema; use iceberg::inspect::MetadataTableType; use iceberg::spec::TableProperties; use iceberg::table::Table; -use iceberg::{Catalog, Error, ErrorKind, NamespaceIdent, Result, TableIdent}; +use iceberg::{Error, ErrorKind, NamespaceIdent, Result, SessionContext, TableIdent}; use metadata_table::IcebergMetadataTableProvider; +use crate::catalog_access::CatalogAccess; use crate::error::to_datafusion_error; use crate::physical_plan::commit::IcebergCommitExec; use crate::physical_plan::project::project_with_partition; @@ -67,7 +68,7 @@ use crate::physical_plan::write::IcebergWriteExec; #[derive(Debug, Clone)] pub struct IcebergTableProvider { /// The catalog that manages this table - catalog: Arc, + catalog: CatalogAccess, /// The table identifier (namespace + name) table_ident: TableIdent, /// A reference-counted arrow `Schema` (cached at construction) @@ -80,14 +81,20 @@ impl IcebergTableProvider { /// Loads the table once to get the initial schema, then stores the catalog /// reference for future metadata refreshes on each operation. pub(crate) async fn try_new( - catalog: Arc, + catalog: CatalogAccess, namespace: NamespaceIdent, name: impl Into, ) -> Result { let table_ident = TableIdent::new(namespace, name.into()); // Load table once to get initial schema - let table = catalog.load_table(&table_ident).await?; + let table = match &catalog { + CatalogAccess::Direct(catalog) => catalog.load_table(&table_ident).await?, + CatalogAccess::SessionAware(session_catalog, _) => { + let context = SessionContext::empty(); + session_catalog.load_table(&context, &table_ident).await? + } + }; let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?); Ok(IcebergTableProvider { @@ -102,7 +109,15 @@ impl IcebergTableProvider { r#type: MetadataTableType, ) -> Result { // Load fresh table metadata for metadata table access - let table = self.catalog.load_table(&self.table_ident).await?; + let table = match &self.catalog { + CatalogAccess::Direct(catalog) => catalog.load_table(&self.table_ident).await?, + CatalogAccess::SessionAware(session_catalog, _) => { + let context = SessionContext::empty(); + session_catalog + .load_table(&context, &self.table_ident) + .await? + } + }; Ok(IcebergMetadataTableProvider { table, r#type }) } } @@ -119,17 +134,22 @@ impl TableProvider for IcebergTableProvider { async fn scan( &self, - _state: &dyn Session, + state: &dyn Session, projection: Option<&Vec>, filters: &[Expr], limit: Option, ) -> DFResult> { // Load fresh table metadata from catalog - let table = self - .catalog - .load_table(&self.table_ident) - .await - .map_err(to_datafusion_error)?; + let table = match &self.catalog { + CatalogAccess::Direct(catalog) => catalog.load_table(&self.table_ident).await, + CatalogAccess::SessionAware(session_catalog, session_context_resolver) => { + let context = session_context_resolver.resolve(state)?; + session_catalog + .load_table(&context, &self.table_ident) + .await + } + } + .map_err(to_datafusion_error)?; // Create scan with fresh metadata (always use current snapshot) Ok(Arc::new(IcebergTableScan::new( @@ -163,11 +183,16 @@ impl TableProvider for IcebergTableProvider { } // Load fresh table metadata from catalog - let table = self - .catalog - .load_table(&self.table_ident) - .await - .map_err(to_datafusion_error)?; + let table = match &self.catalog { + CatalogAccess::Direct(catalog) => catalog.load_table(&self.table_ident).await, + CatalogAccess::SessionAware(session_catalog, session_context_resolver) => { + let context = session_context_resolver.resolve(state)?; + session_catalog + .load_table(&context, &self.table_ident) + .await + } + } + .map_err(to_datafusion_error)?; let partition_spec = table.metadata().default_partition_spec(); @@ -374,7 +399,7 @@ mod tests { static_table.into_table() } - async fn get_test_catalog_and_table() -> (Arc, NamespaceIdent, String, TempDir) { + async fn get_test_catalog_and_table() -> (CatalogAccess, NamespaceIdent, String, TempDir) { let temp_dir = TempDir::new().unwrap(); let warehouse_path = temp_dir.path().to_str().unwrap().to_string(); @@ -414,7 +439,7 @@ mod tests { .unwrap(); ( - Arc::new(catalog), + CatalogAccess::Direct(Arc::new(catalog)), namespace, "test_table".to_string(), temp_dir, @@ -625,7 +650,7 @@ mod tests { async fn get_partitioned_test_catalog_and_table( fanout_enabled: Option, - ) -> (Arc, NamespaceIdent, String, TempDir) { + ) -> (CatalogAccess, NamespaceIdent, String, TempDir) { use iceberg::spec::{Transform, UnboundPartitionSpec}; let temp_dir = TempDir::new().unwrap(); @@ -682,7 +707,7 @@ mod tests { .unwrap(); ( - Arc::new(catalog), + CatalogAccess::Direct(Arc::new(catalog)), namespace, "partitioned_table".to_string(), temp_dir, From 93ba9ecb713ad9568ea64355415cdc3f858d98d6 Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sat, 15 Aug 2026 15:17:20 +0200 Subject: [PATCH 05/11] Add new constructor for session-aware providers --- crates/integrations/datafusion/public-api.txt | 1 + .../datafusion/src/catalog_provider.rs | 14 ++++++++++++++ 2 files changed, 15 insertions(+) diff --git a/crates/integrations/datafusion/public-api.txt b/crates/integrations/datafusion/public-api.txt index 76d5d4f33e..780ad26341 100644 --- a/crates/integrations/datafusion/public-api.txt +++ b/crates/integrations/datafusion/public-api.txt @@ -94,6 +94,7 @@ pub fn iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory:: pub struct iceberg_datafusion::IcebergCatalogProvider impl iceberg_datafusion::IcebergCatalogProvider pub async fn iceberg_datafusion::IcebergCatalogProvider::try_new(catalog: alloc::sync::Arc) -> iceberg::error::Result +pub async fn iceberg_datafusion::IcebergCatalogProvider::try_new_with_session_catalog(catalog: alloc::sync::Arc, resolver: alloc::sync::Arc) -> iceberg::error::Result impl core::fmt::Debug for iceberg_datafusion::IcebergCatalogProvider pub fn iceberg_datafusion::IcebergCatalogProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::catalog::CatalogProvider for iceberg_datafusion::IcebergCatalogProvider diff --git a/crates/integrations/datafusion/src/catalog_provider.rs b/crates/integrations/datafusion/src/catalog_provider.rs index d5367f206e..0b5882b090 100644 --- a/crates/integrations/datafusion/src/catalog_provider.rs +++ b/crates/integrations/datafusion/src/catalog_provider.rs @@ -51,6 +51,20 @@ impl IcebergCatalogProvider { Self::try_new_with_access(direct).await } + /// Creates an [`IcebergCatalogProvider`] backed by a [`SessionCatalog`]. + /// + /// The [`SessionContextResolver`] derives an Iceberg [`SessionContext`] + /// from the current DataFusion session for operations that provide one. + /// Initialization and operations without a DataFusion session use + /// [`SessionContext::empty`]. + pub async fn try_new_with_session_catalog( + catalog: Arc, + resolver: Arc, + ) -> Result { + let session_aware = CatalogAccess::SessionAware(session_catalog, resolver); + Self::try_new_with_access(session_aware).await + } + async fn try_new_with_access(catalog: CatalogAccess) -> Result { // TODO: // Schemas and providers should be cached and evicted based on time From 8d1e9e7e9b7203f046d7696dba93b8c2a3e25528 Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sat, 15 Aug 2026 15:17:41 +0200 Subject: [PATCH 06/11] Match IcebergCommitExec::new's visibility --- crates/integrations/datafusion/src/physical_plan/commit.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/integrations/datafusion/src/physical_plan/commit.rs b/crates/integrations/datafusion/src/physical_plan/commit.rs index 7901e8b7ab..482414672f 100644 --- a/crates/integrations/datafusion/src/physical_plan/commit.rs +++ b/crates/integrations/datafusion/src/physical_plan/commit.rs @@ -51,7 +51,7 @@ pub(crate) struct IcebergCommitExec { } impl IcebergCommitExec { - pub fn new( + pub(crate) fn new( table: Table, catalog: Arc, input: Arc, From 34679090fdb5670dec2f1b83bd072d4a3c94e20a Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sat, 15 Aug 2026 17:18:55 +0200 Subject: [PATCH 07/11] Store fallback_context and reuse it --- .../datafusion/src/catalog_access.rs | 21 +++++--- .../datafusion/src/catalog_provider.rs | 33 +++++++----- .../datafusion/src/schema_provider.rs | 34 +++++++------ .../integrations/datafusion/src/table/mod.rs | 50 +++++++++++-------- 4 files changed, 82 insertions(+), 56 deletions(-) diff --git a/crates/integrations/datafusion/src/catalog_access.rs b/crates/integrations/datafusion/src/catalog_access.rs index db9135bb61..2d1c2d4961 100644 --- a/crates/integrations/datafusion/src/catalog_access.rs +++ b/crates/integrations/datafusion/src/catalog_access.rs @@ -17,22 +17,27 @@ use std::sync::Arc; -use iceberg::{Catalog, SessionCatalog}; +use iceberg::{Catalog, SessionCatalog, SessionContext}; use crate::SessionContextResolver; /// Describes how the DataFusion integration accesses an Iceberg catalog. /// -/// A catalog can either be accessed directly through [`Catalog`] or through -/// [`SessionCatalog`] with an Iceberg [`SessionContext`] derived from the -/// current DataFusion session. Operations for which DataFusion provides no -/// session use an empty Iceberg context. +/// A catalog can either be accessed directly through [`Catalog`] or through a +/// [`SessionCatalog`] bound to an Iceberg [`SessionContext`]. Operations that +/// receive a DataFusion session resolve and bind its context once. Operations +/// for which DataFusion provides no session reuse one anonymous fallback +/// context so session-scoped catalog state remains stable across those calls. #[derive(Clone, Debug)] pub(crate) enum CatalogAccess { /// A catalog accessed directly through the [`Catalog`] API. Direct(Arc), - /// A session-aware catalog and the resolver used to derive its Iceberg - /// context from the current DataFusion session. - SessionAware(Arc, Arc), + /// A session-aware catalog, its resolver, and the stable context used by + /// DataFusion APIs that do not expose a session. + SessionAware { + catalog: Arc, + resolver: Arc, + fallback_context: SessionContext, + }, } diff --git a/crates/integrations/datafusion/src/catalog_provider.rs b/crates/integrations/datafusion/src/catalog_provider.rs index 0b5882b090..92a1ec2fbd 100644 --- a/crates/integrations/datafusion/src/catalog_provider.rs +++ b/crates/integrations/datafusion/src/catalog_provider.rs @@ -26,8 +26,8 @@ use crate::SessionContextResolver; use crate::catalog_access::CatalogAccess; use crate::schema_provider::IcebergSchemaProvider; -/// Provides an interface to manage and access multiple schemas -/// within an Iceberg [`Catalog`]. +/// Provides a DataFusion interface to schemas in an Iceberg [`Catalog`] or +/// [`SessionCatalog`]. /// /// Acts as a centralized catalog provider that aggregates /// multiple [`SchemaProvider`], each associated with distinct namespaces. @@ -53,15 +53,24 @@ impl IcebergCatalogProvider { /// Creates an [`IcebergCatalogProvider`] backed by a [`SessionCatalog`]. /// - /// The [`SessionContextResolver`] derives an Iceberg [`SessionContext`] - /// from the current DataFusion session for operations that provide one. - /// Initialization and operations without a DataFusion session use - /// [`SessionContext::empty`]. + /// The [`SessionContextResolver`] derives an Iceberg session context for + /// scans and inserts. Provider initialization, metadata-table lookup, + /// table registration, and table deregistration do not receive a + /// DataFusion session; they share one anonymous fallback context instead. + /// + /// Namespace and table discovery is performed once during construction and + /// shared by all DataFusion sessions. Catalogs with session-dependent + /// visibility must make the intended discovery set available to the + /// anonymous fallback; discovery is not repeated per DataFusion session. pub async fn try_new_with_session_catalog( catalog: Arc, resolver: Arc, ) -> Result { - let session_aware = CatalogAccess::SessionAware(session_catalog, resolver); + let session_aware = CatalogAccess::SessionAware { + catalog, + resolver, + fallback_context: SessionContext::empty(), + }; Self::try_new_with_access(session_aware).await } @@ -71,11 +80,11 @@ impl IcebergCatalogProvider { // As of right now; schemas might become stale. let schema_names: Vec<_> = match &catalog { CatalogAccess::Direct(catalog) => catalog.list_namespaces(None).await?, - CatalogAccess::SessionAware(session_catalog, _) => { - session_catalog - .list_namespaces(&SessionContext::empty(), None) - .await? - } + CatalogAccess::SessionAware { + catalog, + resolver: _, + fallback_context, + } => catalog.list_namespaces(&fallback_context, None).await?, } .iter() .flat_map(|ns| ns.as_ref().clone()) diff --git a/crates/integrations/datafusion/src/schema_provider.rs b/crates/integrations/datafusion/src/schema_provider.rs index 106a83837c..bb6c7b110f 100644 --- a/crates/integrations/datafusion/src/schema_provider.rs +++ b/crates/integrations/datafusion/src/schema_provider.rs @@ -29,10 +29,7 @@ use futures::future::try_join_all; use iceberg::arrow::arrow_schema_to_schema_auto_assign_ids; use iceberg::inspect::MetadataTableType; use iceberg::spec::FormatVersion; -use iceberg::{ - Error, ErrorKind, NamespaceIdent, Result, SessionContext as IcebergSessionContext, - TableCreation, TableIdent, -}; +use iceberg::{Error, ErrorKind, NamespaceIdent, Result, TableCreation, TableIdent}; use crate::catalog_access::CatalogAccess; use crate::table::IcebergTableProvider; @@ -67,10 +64,11 @@ impl IcebergSchemaProvider { // As of right now; tables might become stale. let table_names: Vec<_> = match &catalog { CatalogAccess::Direct(catalog) => catalog.list_tables(&namespace).await?, - CatalogAccess::SessionAware(session_catalog, _) => { - let context = IcebergSessionContext::empty(); - session_catalog.list_tables(&context, &namespace).await? - } + CatalogAccess::SessionAware { + catalog, + resolver: _, + fallback_context, + } => catalog.list_tables(&fallback_context, &namespace).await?, } .iter() .map(|tbl| tbl.name().to_string()) @@ -194,10 +192,13 @@ impl SchemaProvider for IcebergSchemaProvider { CatalogAccess::Direct(catalog) => { catalog.create_table(&namespace, table_creation).await } - CatalogAccess::SessionAware(session_catalog, _) => { - let context = IcebergSessionContext::empty(); - session_catalog - .create_table(&context, &namespace, table_creation) + CatalogAccess::SessionAware { + catalog, + resolver: _, + fallback_context, + } => { + catalog + .create_table(&fallback_context, &namespace, table_creation) .await } } @@ -246,10 +247,11 @@ impl SchemaProvider for IcebergSchemaProvider { // Drop the table from the Iceberg catalog match &catalog { CatalogAccess::Direct(catalog) => catalog.drop_table(&table_ident).await, - CatalogAccess::SessionAware(session_catalog, _) => { - let context = IcebergSessionContext::empty(); - session_catalog.drop_table(&context, &table_ident).await - } + CatalogAccess::SessionAware { + catalog, + resolver: _, + fallback_context, + } => catalog.drop_table(&fallback_context, &table_ident).await, } .map_err(to_datafusion_error)?; diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 6812756822..f989078169 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -45,7 +45,7 @@ use iceberg::arrow::schema_to_arrow_schema; use iceberg::inspect::MetadataTableType; use iceberg::spec::TableProperties; use iceberg::table::Table; -use iceberg::{Error, ErrorKind, NamespaceIdent, Result, SessionContext, TableIdent}; +use iceberg::{Error, ErrorKind, NamespaceIdent, Result, TableIdent}; use metadata_table::IcebergMetadataTableProvider; use crate::catalog_access::CatalogAccess; @@ -59,9 +59,11 @@ use crate::physical_plan::write::IcebergWriteExec; /// Catalog-backed table provider with automatic metadata refresh. /// -/// This provider loads fresh table metadata from the catalog on every scan and write -/// operation, ensuring you always see the latest table state. Use this when you need -/// write operations or want to see the most up-to-date data. +/// This provider loads fresh table metadata from the catalog on every scan and +/// write operation, ensuring you always see the latest table state. A +/// session-aware provider binds the current DataFusion session for those +/// operations. Initial schema loading and metadata-table lookup do not receive +/// a DataFusion session and use the provider's shared anonymous context. /// /// For read-only access to a specific snapshot without catalog overhead, use /// [`IcebergStaticTableProvider`] instead. @@ -90,10 +92,11 @@ impl IcebergTableProvider { // Load table once to get initial schema let table = match &catalog { CatalogAccess::Direct(catalog) => catalog.load_table(&table_ident).await?, - CatalogAccess::SessionAware(session_catalog, _) => { - let context = SessionContext::empty(); - session_catalog.load_table(&context, &table_ident).await? - } + CatalogAccess::SessionAware { + catalog, + resolver: _, + fallback_context, + } => catalog.load_table(&fallback_context, &table_ident).await?, }; let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?); @@ -111,10 +114,13 @@ impl IcebergTableProvider { // Load fresh table metadata for metadata table access let table = match &self.catalog { CatalogAccess::Direct(catalog) => catalog.load_table(&self.table_ident).await?, - CatalogAccess::SessionAware(session_catalog, _) => { - let context = SessionContext::empty(); - session_catalog - .load_table(&context, &self.table_ident) + CatalogAccess::SessionAware { + catalog, + resolver: _, + fallback_context, + } => { + catalog + .load_table(&fallback_context, &self.table_ident) .await? } }; @@ -142,11 +148,13 @@ impl TableProvider for IcebergTableProvider { // Load fresh table metadata from catalog let table = match &self.catalog { CatalogAccess::Direct(catalog) => catalog.load_table(&self.table_ident).await, - CatalogAccess::SessionAware(session_catalog, session_context_resolver) => { + CatalogAccess::SessionAware { + catalog, + resolver: session_context_resolver, + fallback_context: _, + } => { let context = session_context_resolver.resolve(state)?; - session_catalog - .load_table(&context, &self.table_ident) - .await + catalog.load_table(&context, &self.table_ident).await } } .map_err(to_datafusion_error)?; @@ -185,11 +193,13 @@ impl TableProvider for IcebergTableProvider { // Load fresh table metadata from catalog let table = match &self.catalog { CatalogAccess::Direct(catalog) => catalog.load_table(&self.table_ident).await, - CatalogAccess::SessionAware(session_catalog, session_context_resolver) => { + CatalogAccess::SessionAware { + catalog, + resolver: session_context_resolver, + fallback_context: _, + } => { let context = session_context_resolver.resolve(state)?; - session_catalog - .load_table(&context, &self.table_ident) - .await + catalog.load_table(&context, &self.table_ident).await } } .map_err(to_datafusion_error)?; From 2fc099b73a29199ec09c594748617ec2eb1de527 Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sat, 15 Aug 2026 17:55:34 +0200 Subject: [PATCH 08/11] SessionCatalog to Catalog adapter --- .../datafusion/src/catalog_access.rs | 119 +++++++++++++++++- .../integrations/datafusion/src/table/mod.rs | 20 +-- 2 files changed, 130 insertions(+), 9 deletions(-) diff --git a/crates/integrations/datafusion/src/catalog_access.rs b/crates/integrations/datafusion/src/catalog_access.rs index 2d1c2d4961..95f9bc6573 100644 --- a/crates/integrations/datafusion/src/catalog_access.rs +++ b/crates/integrations/datafusion/src/catalog_access.rs @@ -15,9 +15,15 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; use std::sync::Arc; -use iceberg::{Catalog, SessionCatalog, SessionContext}; +use async_trait::async_trait; +use iceberg::table::Table; +use iceberg::{ + Catalog, Namespace, NamespaceIdent, Result, SessionCatalog, SessionContext, TableCommit, + TableCreation, TableIdent, +}; use crate::SessionContextResolver; @@ -41,3 +47,114 @@ pub(crate) enum CatalogAccess { fallback_context: SessionContext, }, } + +/// Adapts a [`SessionCatalog`] to [`Catalog`] by binding one [`SessionContext`]. +/// +/// Every catalog operation is forwarded to the inner session-aware catalog +/// with the same context. The binding is fixed for the lifetime of this +/// adapter; create another adapter to use a different session context. +#[derive(Debug)] +pub(crate) struct SessionBoundCatalog { + context: SessionContext, + inner: Arc, +} + +impl SessionBoundCatalog { + /// Creates a catalog view of `inner` bound to `context`. + /// + /// The inner catalog receives this context for every operation performed + /// through the returned adapter. + pub fn new(context: SessionContext, inner: Arc) -> Self { + Self { context, inner } + } +} + +#[async_trait] +impl Catalog for SessionBoundCatalog { + async fn list_namespaces( + &self, + parent: Option<&NamespaceIdent>, + ) -> Result> { + self.inner.list_namespaces(&self.context, parent).await + } + + async fn create_namespace( + &self, + namespace: &NamespaceIdent, + properties: HashMap, + ) -> Result { + self.inner + .create_namespace(&self.context, namespace, properties) + .await + } + + async fn get_namespace(&self, namespace: &NamespaceIdent) -> Result { + self.inner.get_namespace(&self.context, namespace).await + } + + async fn namespace_exists(&self, ns: &NamespaceIdent) -> Result { + self.inner.namespace_exists(&self.context, ns).await + } + + async fn update_namespace( + &self, + namespace: &NamespaceIdent, + properties: HashMap, + ) -> Result<()> { + self.inner + .update_namespace(&self.context, namespace, properties) + .await + } + + async fn drop_namespace(&self, namespace: &NamespaceIdent) -> Result<()> { + self.inner.drop_namespace(&self.context, namespace).await + } + + async fn list_tables(&self, namespace: &NamespaceIdent) -> Result> { + self.inner.list_tables(&self.context, namespace).await + } + + async fn create_table( + &self, + namespace: &NamespaceIdent, + creation: TableCreation, + ) -> Result { + self.inner + .create_table(&self.context, namespace, creation) + .await + } + + async fn load_table(&self, table_ident: &TableIdent) -> Result
{ + self.inner.load_table(&self.context, table_ident).await + } + + async fn drop_table(&self, table: &TableIdent) -> Result<()> { + self.inner.drop_table(&self.context, table).await + } + + async fn purge_table(&self, table: &TableIdent) -> Result<()> { + self.inner.purge_table(&self.context, table).await + } + + async fn table_exists(&self, table: &TableIdent) -> Result { + self.inner.table_exists(&self.context, table).await + } + + async fn rename_table(&self, src: &TableIdent, dest: &TableIdent) -> Result<()> { + self.inner.rename_table(&self.context, src, dest).await + } + + async fn register_table( + &self, + table_ident: &TableIdent, + metadata_location: String, + ) -> Result
{ + self.inner + .register_table(&self.context, table_ident, metadata_location) + .await + } + + async fn update_table(&self, commit: TableCommit) -> Result
{ + self.inner.update_table(&self.context, commit).await + } +} diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index f989078169..11f8514fd2 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -48,7 +48,7 @@ use iceberg::table::Table; use iceberg::{Error, ErrorKind, NamespaceIdent, Result, TableIdent}; use metadata_table::IcebergMetadataTableProvider; -use crate::catalog_access::CatalogAccess; +use crate::catalog_access::{CatalogAccess, SessionBoundCatalog}; use crate::error::to_datafusion_error; use crate::physical_plan::commit::IcebergCommitExec; use crate::physical_plan::project::project_with_partition; @@ -190,19 +190,23 @@ impl TableProvider for IcebergTableProvider { ))); } - // Load fresh table metadata from catalog - let table = match &self.catalog { - CatalogAccess::Direct(catalog) => catalog.load_table(&self.table_ident).await, + let catalog = match &self.catalog { + CatalogAccess::Direct(catalog) => Arc::clone(catalog), CatalogAccess::SessionAware { catalog, resolver: session_context_resolver, fallback_context: _, } => { let context = session_context_resolver.resolve(state)?; - catalog.load_table(&context, &self.table_ident).await + Arc::new(SessionBoundCatalog::new(context, Arc::clone(catalog))) } - } - .map_err(to_datafusion_error)?; + }; + + // Load fresh table metadata from catalog + let table = catalog + .load_table(&self.table_ident) + .await + .map_err(to_datafusion_error)?; let partition_spec = table.metadata().default_partition_spec(); @@ -260,7 +264,7 @@ impl TableProvider for IcebergTableProvider { Ok(Arc::new(IcebergCommitExec::new( table, - self.catalog.clone(), + catalog, coalesce_partitions, self.schema.clone(), ))) From 4a0ec7136b5b03cf7d3f05fea3ec2ae20179e8ec Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sat, 15 Aug 2026 18:12:21 +0200 Subject: [PATCH 09/11] Keep SessionBoundCatalog private --- .../datafusion/src/catalog_access.rs | 22 +++++++++++++++++-- .../integrations/datafusion/src/table/mod.rs | 14 ++---------- 2 files changed, 22 insertions(+), 14 deletions(-) diff --git a/crates/integrations/datafusion/src/catalog_access.rs b/crates/integrations/datafusion/src/catalog_access.rs index 95f9bc6573..4695ddbc71 100644 --- a/crates/integrations/datafusion/src/catalog_access.rs +++ b/crates/integrations/datafusion/src/catalog_access.rs @@ -19,6 +19,8 @@ use std::collections::HashMap; use std::sync::Arc; use async_trait::async_trait; +use datafusion::catalog::Session; +use datafusion::error::Result as DFResult; use iceberg::table::Table; use iceberg::{ Catalog, Namespace, NamespaceIdent, Result, SessionCatalog, SessionContext, TableCommit, @@ -48,13 +50,29 @@ pub(crate) enum CatalogAccess { }, } +impl CatalogAccess { + pub(crate) fn with_session(&self, session: &dyn Session) -> DFResult> { + match self { + CatalogAccess::Direct(catalog) => Ok(Arc::clone(catalog)), + CatalogAccess::SessionAware { + catalog, + resolver: session_context_resolver, + .. + } => Ok(Arc::new(SessionBoundCatalog::new( + session_context_resolver.resolve(session)?, + Arc::clone(catalog), + ))), + } + } +} + /// Adapts a [`SessionCatalog`] to [`Catalog`] by binding one [`SessionContext`]. /// /// Every catalog operation is forwarded to the inner session-aware catalog /// with the same context. The binding is fixed for the lifetime of this /// adapter; create another adapter to use a different session context. #[derive(Debug)] -pub(crate) struct SessionBoundCatalog { +struct SessionBoundCatalog { context: SessionContext, inner: Arc, } @@ -64,7 +82,7 @@ impl SessionBoundCatalog { /// /// The inner catalog receives this context for every operation performed /// through the returned adapter. - pub fn new(context: SessionContext, inner: Arc) -> Self { + fn new(context: SessionContext, inner: Arc) -> Self { Self { context, inner } } } diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 11f8514fd2..347758e911 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -48,7 +48,7 @@ use iceberg::table::Table; use iceberg::{Error, ErrorKind, NamespaceIdent, Result, TableIdent}; use metadata_table::IcebergMetadataTableProvider; -use crate::catalog_access::{CatalogAccess, SessionBoundCatalog}; +use crate::catalog_access::CatalogAccess; use crate::error::to_datafusion_error; use crate::physical_plan::commit::IcebergCommitExec; use crate::physical_plan::project::project_with_partition; @@ -190,17 +190,7 @@ impl TableProvider for IcebergTableProvider { ))); } - let catalog = match &self.catalog { - CatalogAccess::Direct(catalog) => Arc::clone(catalog), - CatalogAccess::SessionAware { - catalog, - resolver: session_context_resolver, - fallback_context: _, - } => { - let context = session_context_resolver.resolve(state)?; - Arc::new(SessionBoundCatalog::new(context, Arc::clone(catalog))) - } - }; + let catalog = self.catalog.with_session(state)?; // Load fresh table metadata from catalog let table = catalog From 0208eef3933db3c13f5b47b5ed5fce4d3e51cb05 Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sat, 15 Aug 2026 22:13:45 +0200 Subject: [PATCH 10/11] Factor out catalog access with shared session --- .../datafusion/src/catalog_access.rs | 14 ++++ .../datafusion/src/catalog_provider.rs | 26 +++---- .../datafusion/src/schema_provider.rs | 72 ++++++++----------- .../integrations/datafusion/src/table/mod.rs | 57 ++++++--------- 4 files changed, 76 insertions(+), 93 deletions(-) diff --git a/crates/integrations/datafusion/src/catalog_access.rs b/crates/integrations/datafusion/src/catalog_access.rs index 4695ddbc71..1a995de112 100644 --- a/crates/integrations/datafusion/src/catalog_access.rs +++ b/crates/integrations/datafusion/src/catalog_access.rs @@ -64,6 +64,20 @@ impl CatalogAccess { ))), } } + + pub(crate) fn without_session(&self) -> Arc { + match self { + CatalogAccess::Direct(catalog) => Arc::clone(catalog), + CatalogAccess::SessionAware { + catalog, + fallback_context, + .. + } => Arc::new(SessionBoundCatalog::new( + fallback_context.clone(), + Arc::clone(catalog), + )), + } + } } /// Adapts a [`SessionCatalog`] to [`Catalog`] by binding one [`SessionContext`]. diff --git a/crates/integrations/datafusion/src/catalog_provider.rs b/crates/integrations/datafusion/src/catalog_provider.rs index 92a1ec2fbd..5aedfad134 100644 --- a/crates/integrations/datafusion/src/catalog_provider.rs +++ b/crates/integrations/datafusion/src/catalog_provider.rs @@ -69,6 +69,7 @@ impl IcebergCatalogProvider { let session_aware = CatalogAccess::SessionAware { catalog, resolver, + // One session context that's shared for all query-unrelated catalog operations. fallback_context: SessionContext::empty(), }; Self::try_new_with_access(session_aware).await @@ -78,17 +79,13 @@ impl IcebergCatalogProvider { // TODO: // Schemas and providers should be cached and evicted based on time // As of right now; schemas might become stale. - let schema_names: Vec<_> = match &catalog { - CatalogAccess::Direct(catalog) => catalog.list_namespaces(None).await?, - CatalogAccess::SessionAware { - catalog, - resolver: _, - fallback_context, - } => catalog.list_namespaces(&fallback_context, None).await?, - } - .iter() - .flat_map(|ns| ns.as_ref().clone()) - .collect(); + let schema_names: Vec<_> = catalog + .without_session() + .list_namespaces(None) + .await? + .iter() + .flat_map(|ns| ns.as_ref().clone()) + .collect(); Ok(IcebergCatalogProvider { schemas: load_schema_providers(catalog, schema_names).await?, @@ -107,14 +104,17 @@ impl CatalogProvider for IcebergCatalogProvider { } async fn load_schema_providers( - catalog: CatalogAccess, + catalog_access: CatalogAccess, schema_names: Vec, ) -> Result>> { let iceberg_providers = try_join_all( schema_names .iter() .map(|name| { - IcebergSchemaProvider::try_new(catalog.clone(), NamespaceIdent::new(name.clone())) + IcebergSchemaProvider::try_new( + catalog_access.clone(), + NamespaceIdent::new(name.clone()), + ) }) .collect::>(), ) diff --git a/crates/integrations/datafusion/src/schema_provider.rs b/crates/integrations/datafusion/src/schema_provider.rs index bb6c7b110f..0871f3b9f9 100644 --- a/crates/integrations/datafusion/src/schema_provider.rs +++ b/crates/integrations/datafusion/src/schema_provider.rs @@ -40,7 +40,7 @@ use crate::to_datafusion_error; #[derive(Debug)] pub(crate) struct IcebergSchemaProvider { /// Access policy for the Iceberg catalog. - catalog: CatalogAccess, + catalog_access: CatalogAccess, /// The namespace this schema represents namespace: NamespaceIdent, /// A concurrent map where keys are table names @@ -57,27 +57,28 @@ impl IcebergSchemaProvider { /// /// This method retrieves a list of table names, attempts to create a table /// provider for each name, and collects the providers into a [`DashMap`]. - pub(crate) async fn try_new(catalog: CatalogAccess, namespace: NamespaceIdent) -> Result { + pub(crate) async fn try_new( + catalog_access: CatalogAccess, + namespace: NamespaceIdent, + ) -> Result { // TODO: // Tables and providers should be cached based on table_name // if we have a cache miss; we update our internal cache & check again // As of right now; tables might become stale. - let table_names: Vec<_> = match &catalog { - CatalogAccess::Direct(catalog) => catalog.list_tables(&namespace).await?, - CatalogAccess::SessionAware { - catalog, - resolver: _, - fallback_context, - } => catalog.list_tables(&fallback_context, &namespace).await?, - } - .iter() - .map(|tbl| tbl.name().to_string()) - .collect(); + let table_names: Vec<_> = catalog_access + .without_session() + .list_tables(&namespace) + .await? + .iter() + .map(|tbl| tbl.name().to_string()) + .collect(); let providers = try_join_all( table_names .iter() - .map(|name| IcebergTableProvider::try_new(catalog.clone(), namespace.clone(), name)) + .map(|name| { + IcebergTableProvider::try_new(catalog_access.clone(), namespace.clone(), name) + }) .collect::>(), ) .await?; @@ -88,7 +89,7 @@ impl IcebergSchemaProvider { } Ok(IcebergSchemaProvider { - catalog, + catalog_access, namespace, tables, }) @@ -173,7 +174,7 @@ impl SchemaProvider for IcebergSchemaProvider { .format_version(format_version) .build(); - let catalog = self.catalog.clone(); + let catalog_access = self.catalog_access.clone(); let namespace = self.namespace.clone(); let tables = self.tables.clone(); let name_clone = name.clone(); @@ -188,25 +189,15 @@ impl SchemaProvider for IcebergSchemaProvider { .await .map_err(to_datafusion_error)?; - match &catalog { - CatalogAccess::Direct(catalog) => { - catalog.create_table(&namespace, table_creation).await - } - CatalogAccess::SessionAware { - catalog, - resolver: _, - fallback_context, - } => { - catalog - .create_table(&fallback_context, &namespace, table_creation) - .await - } - } - .map_err(to_datafusion_error)?; + catalog_access + .without_session() + .create_table(&namespace, table_creation) + .await + .map_err(to_datafusion_error)?; - // Create a new table provider using the catalog reference + // Create a new table provider using the catalog access let table_provider = IcebergTableProvider::try_new( - catalog.clone(), + catalog_access, namespace.clone(), name_clone.clone(), ) @@ -233,7 +224,7 @@ impl SchemaProvider for IcebergSchemaProvider { return Ok(None); } - let catalog = self.catalog.clone(); + let catalog = self.catalog_access.without_session(); let namespace = self.namespace.clone(); let tables = self.tables.clone(); let table_name = name.to_string(); @@ -245,15 +236,10 @@ impl SchemaProvider for IcebergSchemaProvider { let table_ident = TableIdent::new(namespace, table_name.clone()); // Drop the table from the Iceberg catalog - match &catalog { - CatalogAccess::Direct(catalog) => catalog.drop_table(&table_ident).await, - CatalogAccess::SessionAware { - catalog, - resolver: _, - fallback_context, - } => catalog.drop_table(&fallback_context, &table_ident).await, - } - .map_err(to_datafusion_error)?; + catalog + .drop_table(&table_ident) + .await + .map_err(to_datafusion_error)?; // Remove from local cache and return the removed provider let removed = tables diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 347758e911..d53933cb00 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -69,8 +69,8 @@ use crate::physical_plan::write::IcebergWriteExec; /// [`IcebergStaticTableProvider`] instead. #[derive(Debug, Clone)] pub struct IcebergTableProvider { - /// The catalog that manages this table - catalog: CatalogAccess, + /// Access to the catalog that manages this table + catalog_access: CatalogAccess, /// The table identifier (namespace + name) table_ident: TableIdent, /// A reference-counted arrow `Schema` (cached at construction) @@ -83,25 +83,21 @@ impl IcebergTableProvider { /// Loads the table once to get the initial schema, then stores the catalog /// reference for future metadata refreshes on each operation. pub(crate) async fn try_new( - catalog: CatalogAccess, + catalog_access: CatalogAccess, namespace: NamespaceIdent, name: impl Into, ) -> Result { let table_ident = TableIdent::new(namespace, name.into()); // Load table once to get initial schema - let table = match &catalog { - CatalogAccess::Direct(catalog) => catalog.load_table(&table_ident).await?, - CatalogAccess::SessionAware { - catalog, - resolver: _, - fallback_context, - } => catalog.load_table(&fallback_context, &table_ident).await?, - }; + let table = catalog_access + .without_session() + .load_table(&table_ident) + .await?; let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?); Ok(IcebergTableProvider { - catalog, + catalog_access, table_ident, schema, }) @@ -112,18 +108,11 @@ impl IcebergTableProvider { r#type: MetadataTableType, ) -> Result { // Load fresh table metadata for metadata table access - let table = match &self.catalog { - CatalogAccess::Direct(catalog) => catalog.load_table(&self.table_ident).await?, - CatalogAccess::SessionAware { - catalog, - resolver: _, - fallback_context, - } => { - catalog - .load_table(&fallback_context, &self.table_ident) - .await? - } - }; + let table = self + .catalog_access + .without_session() + .load_table(&self.table_ident) + .await?; Ok(IcebergMetadataTableProvider { table, r#type }) } } @@ -146,18 +135,12 @@ impl TableProvider for IcebergTableProvider { limit: Option, ) -> DFResult> { // Load fresh table metadata from catalog - let table = match &self.catalog { - CatalogAccess::Direct(catalog) => catalog.load_table(&self.table_ident).await, - CatalogAccess::SessionAware { - catalog, - resolver: session_context_resolver, - fallback_context: _, - } => { - let context = session_context_resolver.resolve(state)?; - catalog.load_table(&context, &self.table_ident).await - } - } - .map_err(to_datafusion_error)?; + let table = self + .catalog_access + .with_session(state)? + .load_table(&self.table_ident) + .await + .map_err(to_datafusion_error)?; // Create scan with fresh metadata (always use current snapshot) Ok(Arc::new(IcebergTableScan::new( @@ -190,7 +173,7 @@ impl TableProvider for IcebergTableProvider { ))); } - let catalog = self.catalog.with_session(state)?; + let catalog = self.catalog_access.with_session(state)?; // Load fresh table metadata from catalog let table = catalog From 391d7fd58b40dfc85f790097e0bb64d4b6b185ca Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sun, 16 Aug 2026 22:59:53 +0200 Subject: [PATCH 11/11] Add example --- Cargo.lock | 3 + crates/examples/Cargo.toml | 7 + crates/examples/README.md | 14 +- .../src/datafusion_session_catalog.rs | 303 ++++++++++++++++++ 4 files changed, 325 insertions(+), 2 deletions(-) create mode 100644 crates/examples/src/datafusion_session_catalog.rs diff --git a/Cargo.lock b/Cargo.lock index ae367e7d8c..15e3a41bbb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3964,9 +3964,12 @@ dependencies = [ name = "iceberg-examples" version = "0.10.0" dependencies = [ + "async-trait", + "datafusion", "futures", "iceberg", "iceberg-catalog-rest", + "iceberg-datafusion", "iceberg-storage-opendal", "tokio", ] diff --git a/crates/examples/Cargo.toml b/crates/examples/Cargo.toml index 0492fed343..5b05b68e3c 100644 --- a/crates/examples/Cargo.toml +++ b/crates/examples/Cargo.toml @@ -26,12 +26,19 @@ rust-version = { workspace = true } version = { workspace = true } [dependencies] +async-trait = { workspace = true } +datafusion = { workspace = true } futures = { workspace = true } iceberg = { workspace = true } iceberg-catalog-rest = { workspace = true } +iceberg-datafusion = { workspace = true } iceberg-storage-opendal = { workspace = true, optional = true } tokio = { workspace = true, features = ["full"] } +[[example]] +name = "datafusion-session-catalog" +path = "src/datafusion_session_catalog.rs" + [[example]] name = "rest-catalog-namespace" path = "src/rest_catalog_namespace.rs" diff --git a/crates/examples/README.md b/crates/examples/README.md index 335d2ea287..40801c69e8 100644 --- a/crates/examples/README.md +++ b/crates/examples/README.md @@ -17,5 +17,15 @@ ~ under the License. --> -Example usage codes for `iceberg-rust`. Currently, these examples can't run directly since it requires setting up of -environments for catalogs, for example, rest catalog server. \ No newline at end of file +Example usage code for `iceberg-rust`. + +The [`datafusion-session-catalog` example](src/datafusion_session_catalog.rs) is +self-contained. It demonstrates how to attach application-specific request metadata to a +DataFusion session, resolve it into an Iceberg `SessionContext`, and query a +session-catalog-backed `IcebergCatalogProvider`: + +```shell +cargo run -p iceberg-examples --example datafusion-session-catalog +``` + +The REST catalog examples require a catalog server and its supporting environment. diff --git a/crates/examples/src/datafusion_session_catalog.rs b/crates/examples/src/datafusion_session_catalog.rs new file mode 100644 index 0000000000..a52c665275 --- /dev/null +++ b/crates/examples/src/datafusion_session_catalog.rs @@ -0,0 +1,303 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Connects a session-aware Iceberg catalog to DataFusion. +//! +//! Run with: +//! +//! ```text +//! cargo run -p iceberg-examples --example datafusion-session-catalog +//! ``` +//! +//! The adapter at the bottom only makes the example self-contained. Applications +//! should pass their own `SessionCatalog` implementation to the provider. + +use std::collections::HashMap; +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::catalog::Session as DataFusionSession; +use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datafusion::prelude::{SessionConfig, SessionContext as DataFusionSessionContext}; +use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalog, MemoryCatalogBuilder}; +use iceberg::spec::{NestedField, PrimitiveType, Schema, Type}; +use iceberg::table::Table; +use iceberg::{ + Catalog, CatalogBuilder, Namespace, NamespaceIdent, Result, SessionCatalog, + SessionContext as IcebergSessionContext, TableCommit, TableCreation, TableIdent, +}; +use iceberg_datafusion::{IcebergCatalogProvider, SessionContextResolver}; + +/// User metadata stored as an application-specific DataFusion extension. +#[derive(Debug)] +struct UserContext { + name: String, +} + +/// Maps the application's DataFusion user context to an Iceberg session. +#[derive(Debug)] +struct UserSessionContextResolver; + +impl SessionContextResolver for UserSessionContextResolver { + fn resolve(&self, session: &dyn DataFusionSession) -> DataFusionResult { + let user = session + .config() + .get_extension::() + .ok_or_else(|| { + DataFusionError::Configuration( + "the DataFusion session has no RequestContext extension".to_string(), + ) + })?; + + Ok(IcebergSessionContext::builder() + // Reusing the DataFusion session ID gives the catalog a stable key + // for session-scoped caches. + .session_id(session.session_id().to_string()) + .identity(user.name.to_string()) + .build()) + } +} + +#[tokio::main] +async fn main() -> std::result::Result<(), Box> { + let catalog = + init_catalog_with_table(TableIdent::from_strs(&["datafusion", "example"])?).await?; + + let session_catalog = Arc::new(ExampleSessionCatalog::new(catalog)); + // Provider construction discovers namespaces and tables with one stable, + // anonymous fallback context. Session-dependent catalogs must make that + // discovery set available to the fallback context. + let provider = IcebergCatalogProvider::try_new_with_session_catalog( + session_catalog, + Arc::new(UserSessionContextResolver), + ) + .await?; + + let config = SessionConfig::new().with_extension(Arc::new(UserContext { + name: "user123".to_string(), + })); + let datafusion = DataFusionSessionContext::new_with_config(config); + datafusion.register_catalog("iceberg", Arc::new(provider)); + + // Planning the scan invokes the resolver and forwards its Iceberg context + // to the session catalog's `load_table` operation. + datafusion + .sql("SELECT COUNT(*) AS event_count FROM iceberg.datafusion.example") + .await? + .show() + .await?; + + Ok(()) +} + +/// A small session-aware wrapper around the in-memory catalog used by this +/// standalone example. +/// +/// Real session catalogs can use the context for authorization, credentials, +/// configuration, and caching. This wrapper logs it, then delegates catalog +/// operations so the example needs no external service. +#[derive(Debug)] +struct ExampleSessionCatalog { + inner: MemoryCatalog, +} + +impl ExampleSessionCatalog { + fn new(inner: MemoryCatalog) -> Self { + Self { inner } + } + + fn log_context(context: &IcebergSessionContext, operation: &str) { + let identity = context.identity().unwrap_or(""); + println!( + "{operation}: session_id={}, identity={identity}", + context.session_id() + ); + } +} + +#[async_trait] +impl SessionCatalog for ExampleSessionCatalog { + async fn list_namespaces( + &self, + context: &IcebergSessionContext, + parent: Option<&NamespaceIdent>, + ) -> Result> { + Self::log_context(context, "list_namespaces"); + self.inner.list_namespaces(parent).await + } + + async fn create_namespace( + &self, + context: &IcebergSessionContext, + namespace: &NamespaceIdent, + properties: HashMap, + ) -> Result { + Self::log_context(context, "create_namespace"); + self.inner.create_namespace(namespace, properties).await + } + + async fn get_namespace( + &self, + context: &IcebergSessionContext, + namespace: &NamespaceIdent, + ) -> Result { + Self::log_context(context, "get_namespace"); + self.inner.get_namespace(namespace).await + } + + async fn namespace_exists( + &self, + context: &IcebergSessionContext, + namespace: &NamespaceIdent, + ) -> Result { + Self::log_context(context, "namespace_exists"); + self.inner.namespace_exists(namespace).await + } + + async fn update_namespace( + &self, + context: &IcebergSessionContext, + namespace: &NamespaceIdent, + properties: HashMap, + ) -> Result<()> { + Self::log_context(context, "update_namespace"); + self.inner.update_namespace(namespace, properties).await + } + + async fn drop_namespace( + &self, + context: &IcebergSessionContext, + namespace: &NamespaceIdent, + ) -> Result<()> { + Self::log_context(context, "drop_namespace"); + self.inner.drop_namespace(namespace).await + } + + async fn list_tables( + &self, + context: &IcebergSessionContext, + namespace: &NamespaceIdent, + ) -> Result> { + Self::log_context(context, "list_tables"); + self.inner.list_tables(namespace).await + } + + async fn create_table( + &self, + context: &IcebergSessionContext, + namespace: &NamespaceIdent, + creation: TableCreation, + ) -> Result
{ + Self::log_context(context, "create_table"); + self.inner.create_table(namespace, creation).await + } + + async fn load_table( + &self, + context: &IcebergSessionContext, + table: &TableIdent, + ) -> Result
{ + Self::log_context(context, "load_table"); + self.inner.load_table(table).await + } + + async fn drop_table(&self, context: &IcebergSessionContext, table: &TableIdent) -> Result<()> { + Self::log_context(context, "drop_table"); + self.inner.drop_table(table).await + } + + async fn purge_table(&self, context: &IcebergSessionContext, table: &TableIdent) -> Result<()> { + Self::log_context(context, "purge_table"); + self.inner.purge_table(table).await + } + + async fn table_exists( + &self, + context: &IcebergSessionContext, + table: &TableIdent, + ) -> Result { + Self::log_context(context, "table_exists"); + self.inner.table_exists(table).await + } + + async fn rename_table( + &self, + context: &IcebergSessionContext, + src: &TableIdent, + dest: &TableIdent, + ) -> Result<()> { + Self::log_context(context, "rename_table"); + self.inner.rename_table(src, dest).await + } + + async fn register_table( + &self, + context: &IcebergSessionContext, + table: &TableIdent, + metadata_location: String, + ) -> Result
{ + Self::log_context(context, "register_table"); + self.inner.register_table(table, metadata_location).await + } + + async fn update_table( + &self, + context: &IcebergSessionContext, + commit: TableCommit, + ) -> Result
{ + Self::log_context(context, "update_table"); + self.inner.update_table(commit).await + } +} + +async fn init_catalog_with_table(table_ident: TableIdent) -> Result { + let catalog = MemoryCatalogBuilder::default() + .load( + "memory", + HashMap::from([( + MEMORY_CATALOG_WAREHOUSE.to_string(), + "memory://session-catalog-example".to_string(), + )]), + ) + .await?; + + catalog + .create_namespace(&table_ident.namespace(), HashMap::new()) + .await?; + catalog + .create_table( + &table_ident.namespace(), + TableCreation::builder() + .name(table_ident.name().to_string()) + .schema( + Schema::builder() + .with_fields(vec![ + NestedField::required( + 1, + "event_id", + Type::Primitive(PrimitiveType::Long), + ) + .into(), + ]) + .build()?, + ) + .build(), + ) + .await?; + + Ok(catalog) +}