From 1a8714d159f6f7d2ba3c8f33a1935770e1844750 Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sat, 15 Aug 2026 14:41:53 +0200 Subject: [PATCH 01/12] 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/12] 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/12] 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 612a2bd44f85b9487fc8eec2a8891cd8cfaaacaf Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sat, 15 Aug 2026 15:16:13 +0200 Subject: [PATCH 04/12] 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 | 85 ++++++++++------- .../integrations/datafusion/src/table/mod.rs | 65 +++++++++---- 5 files changed, 156 insertions(+), 95 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..41da3ac8ee 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 @@ -52,30 +56,30 @@ pub(crate) struct IcebergSchemaProvider { impl IcebergSchemaProvider { /// Asynchronously tries to construct a new [`IcebergSchemaProvider`] /// using the given client to fetch and initialize table providers for - /// the provided namespace in the Iceberg [`Catalog`]. + /// the provided namespace in the Iceberg [`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 f5432700eec0a3f0c48175a163ce818ba8fdc39a Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sat, 15 Aug 2026 15:17:20 +0200 Subject: [PATCH 05/12] 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 373e365aa89a1db7c0d280a04307a9b123d2df5e Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sat, 15 Aug 2026 15:17:41 +0200 Subject: [PATCH 06/12] 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 9abe204b46ef54e0461bdece12cc41d0bdc085da Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sat, 15 Aug 2026 17:18:55 +0200 Subject: [PATCH 07/12] 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 41da3ac8ee..ab10c8ae2c 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 394baf036981b6ea1db0e35ded94e48043027bdd Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sat, 15 Aug 2026 17:55:34 +0200 Subject: [PATCH 08/12] 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 c361fe5e3ced00f7a6e1fb943ec139d90e87679b Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sat, 15 Aug 2026 18:12:21 +0200 Subject: [PATCH 09/12] 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 7ed13076409914c9ae96c1dd4fac6ae5de1b0340 Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sat, 15 Aug 2026 22:13:45 +0200 Subject: [PATCH 10/12] 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 ab10c8ae2c..059f467c05 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 1bfcded8477d1c1ea88f60653a8a581928ecdc59 Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sun, 16 Aug 2026 22:59:53 +0200 Subject: [PATCH 11/12] 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..3381de5e01 --- /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 UserContext 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) +} From 39b780895f494e61164c59d92bbe1ced19b6f83d Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Tue, 18 Aug 2026 23:35:26 +0200 Subject: [PATCH 12/12] Add tests --- .../datafusion/src/catalog_provider.rs | 194 +++++++++++ crates/integrations/datafusion/src/lib.rs | 2 + .../integrations/datafusion/src/test_utils.rs | 311 ++++++++++++++++++ 3 files changed, 507 insertions(+) create mode 100644 crates/integrations/datafusion/src/test_utils.rs diff --git a/crates/integrations/datafusion/src/catalog_provider.rs b/crates/integrations/datafusion/src/catalog_provider.rs index 5aedfad134..0d04a8e262 100644 --- a/crates/integrations/datafusion/src/catalog_provider.rs +++ b/crates/integrations/datafusion/src/catalog_provider.rs @@ -131,3 +131,197 @@ async fn load_schema_providers( Ok(provider_map) } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::catalog::CatalogProvider; + use datafusion::datasource::MemTable; + use datafusion::logical_expr::dml::InsertOp; + use datafusion::physical_plan::empty::EmptyExec; + use datafusion::prelude::SessionContext as DFSessionContext; + + use super::*; + use crate::test_utils::{ + FailingSessionContextResolver, FixedSessionContextResolver, create_recording_catalog, + resolved_session_context, + }; + + #[tokio::test] + async fn test_session_aware_scan_uses_resolved_context() { + let (session_catalog, namespace, table_name, _temp_dir) = create_recording_catalog().await; + let resolver = Arc::new(FixedSessionContextResolver::new(resolved_session_context())); + let provider = IcebergCatalogProvider::try_new_with_session_catalog( + session_catalog.clone(), + resolver.clone(), + ) + .await + .unwrap(); + + let bootstrap_calls = session_catalog.calls(); + assert_eq!( + bootstrap_calls + .iter() + .map(|call| call.operation) + .collect::>(), + vec!["list_namespaces", "list_tables", "load_table"] + ); + session_catalog.clear_calls(); + + let fallback_session_id = bootstrap_calls[0].session_id.as_str(); + assert!(bootstrap_calls.iter().all(|call| { + call.session_id == fallback_session_id + && call.identity.is_none() + && call.properties.is_empty() + && call.credential_keys.is_empty() + })); + assert_eq!(resolver.calls(), 0); + + let schema = provider.schema(namespace[0].as_str()).unwrap(); + let table = schema.table(&table_name).await.unwrap().unwrap(); + + let df_context = DFSessionContext::new(); + table + .scan(&df_context.state(), None, &[], None) + .await + .unwrap(); + + assert_eq!(resolver.calls(), 1); + let calls = session_catalog.calls(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].operation, "load_table"); + assert_eq!(calls[0].session_id, "resolved-session"); + assert_eq!(calls[0].identity.as_deref(), Some("test-user")); + assert_eq!( + calls[0].properties.get("test-property").map(String::as_str), + Some("test-value") + ); + assert_eq!(calls[0].credential_keys, vec!["test-token"]); + } + + #[tokio::test] + async fn test_session_resolution_errors_prevent_catalog_access() { + let (session_catalog, namespace, table_name, _temp_dir) = create_recording_catalog().await; + let provider = IcebergCatalogProvider::try_new_with_session_catalog( + session_catalog.clone(), + Arc::new(FailingSessionContextResolver), + ) + .await + .unwrap(); + let schema = provider.schema(namespace[0].as_str()).unwrap(); + let table = schema.table(&table_name).await.unwrap().unwrap(); + session_catalog.clear_calls(); + + let df_context = DFSessionContext::new(); + let error = table + .scan(&df_context.state(), None, &[], None) + .await + .err() + .unwrap(); + + assert!(error.to_string().contains("session resolution failed")); + + let input = Arc::new(EmptyExec::new(table.schema())); + let error = table + .insert_into(&df_context.state(), input, InsertOp::Append) + .await + .err() + .unwrap(); + assert!(error.to_string().contains("session resolution failed")); + assert!(session_catalog.calls().is_empty()); + } + + #[tokio::test] + async fn test_session_aware_insert_reuses_resolved_context_for_commit() { + let (session_catalog, namespace, table_name, _temp_dir) = create_recording_catalog().await; + let resolver = Arc::new(FixedSessionContextResolver::new(resolved_session_context())); + let provider = IcebergCatalogProvider::try_new_with_session_catalog( + session_catalog.clone(), + resolver.clone(), + ) + .await + .unwrap(); + let schema = provider.schema(namespace[0].as_str()).unwrap(); + let table = schema.table(&table_name).await.unwrap().unwrap(); + session_catalog.clear_calls(); + + let df_context = DFSessionContext::new(); + df_context.register_table("test_table", table).unwrap(); + df_context + .sql("INSERT INTO test_table VALUES (1, 'test')") + .await + .unwrap() + .collect() + .await + .unwrap(); + + assert_eq!(resolver.calls(), 1); + let calls = session_catalog.calls(); + let load = calls + .iter() + .find(|call| call.operation == "load_table") + .unwrap(); + let update = calls + .iter() + .find(|call| call.operation == "update_table") + .unwrap(); + assert_eq!(load.session_id, "resolved-session"); + assert_eq!(update.session_id, load.session_id); + assert_eq!(update.identity, load.identity); + assert_eq!(update.properties, load.properties); + assert_eq!(update.credential_keys, load.credential_keys); + } + + #[tokio::test] + async fn test_sessionless_schema_operations_share_fallback_context() { + let (session_catalog, namespace, table_name, _temp_dir) = create_recording_catalog().await; + let resolver = Arc::new(FixedSessionContextResolver::new(resolved_session_context())); + let provider = IcebergCatalogProvider::try_new_with_session_catalog( + session_catalog.clone(), + resolver.clone(), + ) + .await + .unwrap(); + let schema = provider.schema(namespace[0].as_str()).unwrap(); + session_catalog.clear_calls(); + + schema + .table(&format!("{table_name}$snapshots")) + .await + .unwrap() + .unwrap(); + + let metadata_calls = session_catalog.calls(); + assert_eq!(metadata_calls.len(), 1); + assert_eq!(metadata_calls[0].operation, "load_table"); + let fallback_session_id = metadata_calls[0].session_id.clone(); + assert!(metadata_calls.iter().all(|call| { + call.session_id == fallback_session_id + && call.identity.is_none() + && call.properties.is_empty() + && call.credential_keys.is_empty() + })); + session_catalog.clear_calls(); + + let arrow_schema = schema.table(&table_name).await.unwrap().unwrap().schema(); + let empty_batch = RecordBatch::new_empty(arrow_schema.clone()); + let empty_table = MemTable::try_new(arrow_schema, vec![vec![empty_batch]]).unwrap(); + schema + .register_table("registered_table".to_string(), Arc::new(empty_table)) + .unwrap(); + schema.deregister_table("registered_table").unwrap(); + + let calls = session_catalog.calls(); + assert!(calls.iter().any(|call| call.operation == "create_table")); + assert!(calls.iter().any(|call| call.operation == "drop_table")); + assert!(calls.iter().all(|call| { + call.session_id == fallback_session_id + && call.identity.is_none() + && call.properties.is_empty() + && call.credential_keys.is_empty() + })); + assert_eq!(resolver.calls(), 0); + } +} diff --git a/crates/integrations/datafusion/src/lib.rs b/crates/integrations/datafusion/src/lib.rs index c7ad5f8124..d35d4eb1a9 100644 --- a/crates/integrations/datafusion/src/lib.rs +++ b/crates/integrations/datafusion/src/lib.rs @@ -28,6 +28,8 @@ pub use table::table_provider_factory::IcebergTableProviderFactory; pub use table::*; pub(crate) mod task_writer; +#[cfg(test)] +mod test_utils; use std::fmt; diff --git a/crates/integrations/datafusion/src/test_utils.rs b/crates/integrations/datafusion/src/test_utils.rs new file mode 100644 index 0000000000..aba2773e15 --- /dev/null +++ b/crates/integrations/datafusion/src/test_utils.rs @@ -0,0 +1,311 @@ +// 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::collections::HashMap; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use datafusion::catalog::Session as DFSession; +use datafusion::error::{DataFusionError, Result as DFResult}; +use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; +use iceberg::sensitive::SensitiveString; +use iceberg::spec::{NestedField, PrimitiveType, Schema, Type}; +use iceberg::table::Table; +use iceberg::{ + Catalog, CatalogBuilder, Namespace, NamespaceIdent, Result, SessionCatalog, SessionContext, + TableCommit, TableCreation, TableIdent, +}; +use tempfile::TempDir; + +use crate::SessionContextResolver; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct CatalogCall { + pub(crate) operation: &'static str, + pub(crate) session_id: String, + pub(crate) identity: Option, + pub(crate) properties: HashMap, + pub(crate) credential_keys: Vec, +} + +impl CatalogCall { + fn new(operation: &'static str, context: &SessionContext) -> Self { + let mut credential_keys = context.credentials().keys().cloned().collect::>(); + credential_keys.sort(); + Self { + operation, + session_id: context.session_id().to_string(), + identity: context.identity().map(ToString::to_string), + properties: context.properties().clone(), + credential_keys, + } + } +} + +#[derive(Debug)] +pub(crate) struct RecordingSessionCatalog { + inner: Arc, + calls: Mutex>, +} + +impl RecordingSessionCatalog { + fn new(inner: Arc) -> Self { + Self { + inner, + calls: Mutex::new(Vec::new()), + } + } + + fn record(&self, operation: &'static str, context: &SessionContext) { + self.calls + .lock() + .unwrap() + .push(CatalogCall::new(operation, context)); + } + + pub(crate) fn calls(&self) -> Vec { + self.calls.lock().unwrap().clone() + } + + pub(crate) fn clear_calls(&self) { + self.calls.lock().unwrap().clear(); + } +} + +#[async_trait] +impl SessionCatalog for RecordingSessionCatalog { + async fn list_namespaces( + &self, + context: &SessionContext, + parent: Option<&NamespaceIdent>, + ) -> Result> { + self.record("list_namespaces", context); + self.inner.list_namespaces(parent).await + } + + async fn create_namespace( + &self, + context: &SessionContext, + namespace: &NamespaceIdent, + properties: HashMap, + ) -> Result { + self.record("create_namespace", context); + self.inner.create_namespace(namespace, properties).await + } + + async fn get_namespace( + &self, + context: &SessionContext, + namespace: &NamespaceIdent, + ) -> Result { + self.record("get_namespace", context); + self.inner.get_namespace(namespace).await + } + + async fn namespace_exists( + &self, + context: &SessionContext, + namespace: &NamespaceIdent, + ) -> Result { + self.record("namespace_exists", context); + self.inner.namespace_exists(namespace).await + } + + async fn update_namespace( + &self, + context: &SessionContext, + namespace: &NamespaceIdent, + properties: HashMap, + ) -> Result<()> { + self.record("update_namespace", context); + self.inner.update_namespace(namespace, properties).await + } + + async fn drop_namespace( + &self, + context: &SessionContext, + namespace: &NamespaceIdent, + ) -> Result<()> { + self.record("drop_namespace", context); + self.inner.drop_namespace(namespace).await + } + + async fn list_tables( + &self, + context: &SessionContext, + namespace: &NamespaceIdent, + ) -> Result> { + self.record("list_tables", context); + self.inner.list_tables(namespace).await + } + + async fn create_table( + &self, + context: &SessionContext, + namespace: &NamespaceIdent, + creation: TableCreation, + ) -> Result
{ + self.record("create_table", context); + self.inner.create_table(namespace, creation).await + } + + async fn load_table(&self, context: &SessionContext, table: &TableIdent) -> Result
{ + self.record("load_table", context); + self.inner.load_table(table).await + } + + async fn drop_table(&self, context: &SessionContext, table: &TableIdent) -> Result<()> { + self.record("drop_table", context); + self.inner.drop_table(table).await + } + + async fn purge_table(&self, context: &SessionContext, table: &TableIdent) -> Result<()> { + self.record("purge_table", context); + self.inner.purge_table(table).await + } + + async fn table_exists(&self, context: &SessionContext, table: &TableIdent) -> Result { + self.record("table_exists", context); + self.inner.table_exists(table).await + } + + async fn rename_table( + &self, + context: &SessionContext, + src: &TableIdent, + dest: &TableIdent, + ) -> Result<()> { + self.record("rename_table", context); + self.inner.rename_table(src, dest).await + } + + async fn register_table( + &self, + context: &SessionContext, + table: &TableIdent, + metadata_location: String, + ) -> Result
{ + self.record("register_table", context); + self.inner.register_table(table, metadata_location).await + } + + async fn update_table(&self, context: &SessionContext, commit: TableCommit) -> Result
{ + self.record("update_table", context); + self.inner.update_table(commit).await + } +} + +#[derive(Debug)] +pub(crate) struct FixedSessionContextResolver { + context: SessionContext, + calls: AtomicUsize, +} + +impl FixedSessionContextResolver { + pub(crate) fn new(context: SessionContext) -> Self { + Self { + context, + calls: AtomicUsize::new(0), + } + } + + pub(crate) fn calls(&self) -> usize { + self.calls.load(Ordering::SeqCst) + } +} + +impl SessionContextResolver for FixedSessionContextResolver { + fn resolve(&self, _session: &dyn DFSession) -> DFResult { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(self.context.clone()) + } +} + +#[derive(Debug)] +pub(crate) struct FailingSessionContextResolver; + +impl SessionContextResolver for FailingSessionContextResolver { + fn resolve(&self, _session: &dyn DFSession) -> DFResult { + Err(DataFusionError::External( + "session resolution failed".into(), + )) + } +} + +pub(crate) fn resolved_session_context() -> SessionContext { + SessionContext::builder() + .session_id("resolved-session".to_string()) + .identity("test-user".to_string()) + .properties(HashMap::from([( + "test-property".to_string(), + "test-value".to_string(), + )])) + .credentials(HashMap::from([( + "test-token".to_string(), + SensitiveString::from("test-secret".to_string()), + )])) + .build() +} + +pub(crate) async fn create_recording_catalog() -> ( + Arc, + NamespaceIdent, + String, + TempDir, +) { + let temp_dir = TempDir::new().unwrap(); + let warehouse_path = temp_dir.path().to_str().unwrap().to_string(); + let catalog = Arc::new( + MemoryCatalogBuilder::default() + .load( + "memory", + HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse_path.clone())]), + ) + .await + .unwrap(), + ); + + let namespace = NamespaceIdent::new("test_ns".to_string()); + catalog + .create_namespace(&namespace, HashMap::new()) + .await + .unwrap(); + + let schema = Schema::builder() + .with_schema_id(0) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(), + ]) + .build() + .unwrap(); + let table_name = "test_table".to_string(); + let creation = TableCreation::builder() + .name(table_name.clone()) + .location(format!("{warehouse_path}/{table_name}")) + .schema(schema) + .build(); + catalog.create_table(&namespace, creation).await.unwrap(); + + ( + Arc::new(RecordingSessionCatalog::new(catalog)), + namespace, + table_name, + temp_dir, + ) +}