diff --git a/Cargo.lock b/Cargo.lock index d6b897a8cf..80abc6ef52 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3966,9 +3966,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 27e7c9969d..8f832e9c2e 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..350b4657b5 --- /dev/null +++ b/crates/examples/src/datafusion_session_catalog.rs @@ -0,0 +1,269 @@ +// 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::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, IcebergOptions}; + +#[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).await?; + + let mut iceberg_options = IcebergOptions::default(); + iceberg_options.identity = Some("user123".to_string()); + + let config = SessionConfig::new().with_extension(Arc::new(iceberg_options)); + let datafusion = DataFusionSessionContext::new_with_config(config); + datafusion.register_catalog("iceberg", Arc::new(provider)); + + // Planning the scan derives an Iceberg context from the DataFusion session + // and its IcebergOptions, then forwards it 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) +} diff --git a/crates/integrations/datafusion/public-api.txt b/crates/integrations/datafusion/public-api.txt index e197c2057d..59bd6fea60 100644 --- a/crates/integrations/datafusion/public-api.txt +++ b/crates/integrations/datafusion/public-api.txt @@ -75,8 +75,8 @@ pub fn iceberg_datafusion::IcebergTableProvider::clone(&self) -> iceberg_datafus 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::insert_into<'life0, 'life1, 'async_trait>(&'life0 self, session: &'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, session: &'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,12 +93,30 @@ 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 +pub async fn iceberg_datafusion::IcebergCatalogProvider::try_new_with_session_catalog(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 pub fn iceberg_datafusion::IcebergCatalogProvider::schema(&self, name: &str) -> core::option::Option> pub fn iceberg_datafusion::IcebergCatalogProvider::schema_names(&self) -> alloc::vec::Vec +#[non_exhaustive] pub struct iceberg_datafusion::IcebergOptions +pub iceberg_datafusion::IcebergOptions::identity: core::option::Option +impl core::clone::Clone for iceberg_datafusion::IcebergOptions +pub fn iceberg_datafusion::IcebergOptions::clone(&self) -> iceberg_datafusion::IcebergOptions +impl core::default::Default for iceberg_datafusion::IcebergOptions +pub fn iceberg_datafusion::IcebergOptions::default() -> Self +impl core::fmt::Debug for iceberg_datafusion::IcebergOptions +pub fn iceberg_datafusion::IcebergOptions::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl datafusion_common::config::ConfigField for iceberg_datafusion::IcebergOptions +pub fn iceberg_datafusion::IcebergOptions::set(&mut self, key: &str, value: &str) -> datafusion_common::error::Result<()> +pub fn iceberg_datafusion::IcebergOptions::visit(&self, v: &mut V, _key_prefix: &str, _description: &'static str) +impl datafusion_common::config::ExtensionOptions for iceberg_datafusion::IcebergOptions +pub fn iceberg_datafusion::IcebergOptions::as_any(&self) -> &dyn core::any::Any +pub fn iceberg_datafusion::IcebergOptions::as_any_mut(&mut self) -> &mut dyn core::any::Any +pub fn iceberg_datafusion::IcebergOptions::cloned(&self) -> alloc::boxed::Box +pub fn iceberg_datafusion::IcebergOptions::entries(&self) -> alloc::vec::Vec +pub fn iceberg_datafusion::IcebergOptions::set(&mut self, key: &str, value: &str) -> datafusion_common::error::Result<()> pub struct iceberg_datafusion::IcebergStaticTableProvider impl iceberg_datafusion::IcebergStaticTableProvider pub async fn iceberg_datafusion::IcebergStaticTableProvider::try_new_from_table(table: iceberg::table::Table) -> iceberg::error::Result @@ -119,8 +137,8 @@ pub fn iceberg_datafusion::IcebergTableProvider::clone(&self) -> iceberg_datafus 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::insert_into<'life0, 'life1, 'async_trait>(&'life0 self, session: &'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, session: &'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.rs b/crates/integrations/datafusion/src/catalog.rs deleted file mode 100644 index 2c6e1ff002..0000000000 --- a/crates/integrations/datafusion/src/catalog.rs +++ /dev/null @@ -1,93 +0,0 @@ -// 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::Arc; - -use datafusion::catalog::{CatalogProvider, SchemaProvider}; -use futures::future::try_join_all; -use iceberg::{Catalog, NamespaceIdent, Result}; - -use crate::schema::IcebergSchemaProvider; - -/// Provides an interface to manage and access multiple schemas -/// within an Iceberg [`Catalog`]. -/// -/// Acts as a centralized catalog provider that aggregates -/// multiple [`SchemaProvider`], each associated with distinct namespaces. -#[derive(Debug)] -pub struct IcebergCatalogProvider { - /// A `HashMap` where keys are namespace names - /// and values are dynamic references to objects implementing the - /// [`SchemaProvider`] trait. - schemas: HashMap>, -} - -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`]. - /// - /// 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 { - // 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(); - - Ok(IcebergCatalogProvider { schemas }) - } -} - -impl CatalogProvider for IcebergCatalogProvider { - fn schema_names(&self) -> Vec { - self.schemas.keys().cloned().collect() - } - - fn schema(&self, name: &str) -> Option> { - self.schemas.get(name).cloned() - } -} diff --git a/crates/integrations/datafusion/src/catalog_adapter.rs b/crates/integrations/datafusion/src/catalog_adapter.rs new file mode 100644 index 0000000000..4e8b7be169 --- /dev/null +++ b/crates/integrations/datafusion/src/catalog_adapter.rs @@ -0,0 +1,278 @@ +// 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::Arc; + +use async_trait::async_trait; +use datafusion::catalog::Session; +use iceberg::table::Table; +use iceberg::{ + Catalog, Namespace, NamespaceIdent, Result, SessionCatalog, SessionContext, TableCommit, + TableCreation, TableIdent, +}; + +use crate::options::resolve_session_context; + +/// 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(Clone, Debug)] +pub(crate) struct SessionBindingCatalogAdapter { + context: SessionContext, + inner: Arc, +} + +impl SessionBindingCatalogAdapter { + /// Creates a catalog view of `inner` bound to `context`. + /// + /// The inner catalog receives this context for every operation performed + /// through the returned adapter. + pub(crate) fn new(context: SessionContext, inner: Arc) -> Self { + Self { context, inner } + } + + /// Adapts a plain, session-unaware catalog to a [`SessionBindingCatalogAdapter`]. + /// + /// The bound session context is never used, but simply ignored. + pub(crate) fn new_without_context(catalog: Arc) -> Self { + let session_catalog = SessionDroppingCatalogAdapter::new(catalog); + Self::new(SessionContext::empty(), Arc::new(session_catalog)) + } + + /// Overwrites the already bound, usually shared fallback session, with + /// a provided session, usually from a DataFusion query. + pub(crate) fn with_session( + self: &Arc, + session: &dyn Session, + ) -> Arc { + match resolve_session_context(session) { + None => Arc::clone(self), + Some(context) => Arc::new(SessionBindingCatalogAdapter::new( + context, + Arc::clone(&self.inner), + )), + } + } +} + +#[async_trait] +impl Catalog for SessionBindingCatalogAdapter { + 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 + } +} + +/// A wrapper around a [`Catalog`] to provide a [`SessionCatalog`] API by +/// ignoring any passed [`SessionContext`]. +#[derive(Debug)] +struct SessionDroppingCatalogAdapter { + inner: Arc, +} + +impl SessionDroppingCatalogAdapter { + fn new(inner: Arc) -> Self { + Self { inner } + } +} + +#[async_trait] +impl SessionCatalog for SessionDroppingCatalogAdapter { + async fn list_namespaces( + &self, + _: &SessionContext, + parent: Option<&NamespaceIdent>, + ) -> Result> { + self.inner.list_namespaces(parent).await + } + + async fn create_namespace( + &self, + _: &SessionContext, + namespace: &NamespaceIdent, + properties: HashMap, + ) -> Result { + self.inner.create_namespace(namespace, properties).await + } + + async fn get_namespace( + &self, + _: &SessionContext, + namespace: &NamespaceIdent, + ) -> Result { + self.inner.get_namespace(namespace).await + } + + async fn namespace_exists(&self, _: &SessionContext, ns: &NamespaceIdent) -> Result { + self.inner.namespace_exists(ns).await + } + + async fn update_namespace( + &self, + _: &SessionContext, + namespace: &NamespaceIdent, + properties: HashMap, + ) -> Result<()> { + self.inner.update_namespace(namespace, properties).await + } + + async fn drop_namespace(&self, _: &SessionContext, namespace: &NamespaceIdent) -> Result<()> { + self.inner.drop_namespace(namespace).await + } + + async fn list_tables( + &self, + _: &SessionContext, + namespace: &NamespaceIdent, + ) -> Result> { + self.inner.list_tables(namespace).await + } + + async fn create_table( + &self, + _: &SessionContext, + namespace: &NamespaceIdent, + creation: TableCreation, + ) -> Result
{ + self.inner.create_table(namespace, creation).await + } + + async fn load_table(&self, _: &SessionContext, table_ident: &TableIdent) -> Result
{ + self.inner.load_table(table_ident).await + } + + async fn drop_table(&self, _: &SessionContext, table: &TableIdent) -> Result<()> { + self.inner.drop_table(table).await + } + + async fn purge_table(&self, _: &SessionContext, table: &TableIdent) -> Result<()> { + self.inner.purge_table(table).await + } + + async fn table_exists(&self, _: &SessionContext, table: &TableIdent) -> Result { + self.inner.table_exists(table).await + } + + async fn rename_table( + &self, + _: &SessionContext, + src: &TableIdent, + dest: &TableIdent, + ) -> Result<()> { + self.inner.rename_table(src, dest).await + } + + async fn register_table( + &self, + _: &SessionContext, + table_ident: &TableIdent, + metadata_location: String, + ) -> Result
{ + self.inner + .register_table(table_ident, metadata_location) + .await + } + + async fn update_table(&self, _: &SessionContext, commit: TableCommit) -> Result
{ + self.inner.update_table(commit).await + } +} diff --git a/crates/integrations/datafusion/src/catalog_provider.rs b/crates/integrations/datafusion/src/catalog_provider.rs new file mode 100644 index 0000000000..98d9522091 --- /dev/null +++ b/crates/integrations/datafusion/src/catalog_provider.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. + +use std::collections::HashMap; +use std::sync::Arc; + +use datafusion::catalog::{CatalogProvider, SchemaProvider}; +use futures::future::try_join_all; +use iceberg::{Catalog, NamespaceIdent, Result, SessionCatalog, SessionContext}; + +use crate::catalog_adapter::SessionBindingCatalogAdapter; +use crate::schema_provider::IcebergSchemaProvider; + +/// 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. +#[derive(Debug)] +pub struct IcebergCatalogProvider { + /// A `HashMap` where keys are namespace names + /// and values are dynamic references to objects implementing the + /// [`SchemaProvider`] trait. + schemas: HashMap>, +} + +impl IcebergCatalogProvider { + /// Asynchronously constructs an [`IcebergCatalogProvider`] from a + /// [`Catalog`], fetching and initializing a schema provider for each + /// namespace. + /// + /// 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 session_binding_catalog = SessionBindingCatalogAdapter::new_without_context(catalog); + Self::try_new_with_binding_catalog(Arc::new(session_binding_catalog)).await + } + + /// Creates an [`IcebergCatalogProvider`] backed by a [`SessionCatalog`]. + /// + /// Each DataFusion session that has [`IcebergOptions`] configured will + /// propagate an Iceberg [`SessionContext`] 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) -> Result { + let shared_fallback_context = SessionContext::empty(); + let session_bound = SessionBindingCatalogAdapter::new(shared_fallback_context, catalog); + Self::try_new_with_binding_catalog(Arc::new(session_bound)).await + } + + async fn try_new_with_binding_catalog( + catalog: Arc, + ) -> 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<_> = catalog + .list_namespaces(None) + .await? + .iter() + .flat_map(|ns| ns.as_ref().clone()) + .collect(); + + Ok(IcebergCatalogProvider { + schemas: load_schema_providers(catalog, schema_names).await?, + }) + } +} + +impl CatalogProvider for IcebergCatalogProvider { + fn schema_names(&self) -> Vec { + self.schemas.keys().cloned().collect() + } + + fn schema(&self, name: &str) -> Option> { + self.schemas.get(name).cloned() + } +} + +async fn load_schema_providers( + catalog: Arc, + schema_names: Vec, +) -> Result>> { + let iceberg_providers = try_join_all( + schema_names + .iter() + .map(|name| { + IcebergSchemaProvider::try_new( + Arc::clone(&catalog), + 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) +} + +#[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::create_recording_catalog; + + #[tokio::test] + async fn test_session_aware_scan_uses_resolved_context() { + let (session_catalog, namespace, table_name, _temp_dir) = create_recording_catalog().await; + let provider = + IcebergCatalogProvider::try_new_with_session_catalog(session_catalog.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() + })); + + 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(); + + 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()) + .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 provider = + IcebergCatalogProvider::try_new_with_session_catalog(session_catalog.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(); + + 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 provider = + IcebergCatalogProvider::try_new_with_session_catalog(session_catalog.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() + })); + } +} diff --git a/crates/integrations/datafusion/src/lib.rs b/crates/integrations/datafusion/src/lib.rs index 4b0ea8606d..76cea9c816 100644 --- a/crates/integrations/datafusion/src/lib.rs +++ b/crates/integrations/datafusion/src/lib.rs @@ -15,16 +15,20 @@ // specific language governing permissions and limitations // under the License. -mod catalog; -pub use catalog::*; - mod error; pub use error::*; +mod catalog_adapter; +mod catalog_provider; +pub use options::IcebergOptions; +mod options; +pub use catalog_provider::*; pub mod physical_plan; -mod schema; +mod schema_provider; pub mod table; pub use table::table_provider_factory::IcebergTableProviderFactory; pub use table::*; pub(crate) mod task_writer; +#[cfg(test)] +mod test_utils; diff --git a/crates/integrations/datafusion/src/options.rs b/crates/integrations/datafusion/src/options.rs new file mode 100644 index 0000000000..c35c60bbda --- /dev/null +++ b/crates/integrations/datafusion/src/options.rs @@ -0,0 +1,52 @@ +// 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 datafusion::catalog::Session as DFSession; +use datafusion::common::extensions_options; +use iceberg::SessionContext; + +extensions_options! { + /// Iceberg-specific DataFusion options. + /// + /// It does deliberately not implement [`ConfigExtension`](datafusion::config::ConfigExtension) + /// to prevent SQL users from setting unverified authentication properties + /// such as: + /// + /// ```sql + /// SET iceberg.identity = 'alice'; + /// ``` + pub struct IcebergOptions { + /// Optional identity used when deriving the Iceberg session context. + pub identity: Option, default = None + } +} + +/// Derives an Iceberg session context from a DataFusion session and its +/// configured [`IcebergOptions`], if registered. +pub(crate) fn resolve_session_context(session: &dyn DFSession) -> Option { + let options = session.config().get_extension::()?; + + let builder = SessionContext::builder().session_id(session.session_id().to_string()); + + let context = if let Some(identity) = &options.identity { + builder.identity(identity.to_string()).build() + } else { + builder.build() + }; + + Some(context) +} diff --git a/crates/integrations/datafusion/src/physical_plan/commit.rs b/crates/integrations/datafusion/src/physical_plan/commit.rs index f1d93679e0..610489073a 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, @@ -50,7 +51,7 @@ pub(crate) struct IcebergCommitExec { } impl IcebergCommitExec { - pub fn new( + pub(crate) fn new( table: Table, catalog: Arc, input: Arc, @@ -287,6 +288,7 @@ mod tests { use iceberg::{Catalog, CatalogBuilder, NamespaceIdent, TableCreation, TableIdent}; use super::*; + use crate::catalog_adapter::SessionBindingCatalogAdapter; use crate::physical_plan::DATA_FILES_COL_NAME; use crate::table::IcebergTableProvider; @@ -658,8 +660,9 @@ mod tests { let source_table = Arc::new(MemTable::try_new(Arc::clone(&arrow_schema), partitions)?); ctx.register_table("source_table", source_table)?; + let session_binding_catalog = SessionBindingCatalogAdapter::new_without_context(catalog); let iceberg_table_provider = IcebergTableProvider::try_new( - catalog.clone(), + Arc::new(session_binding_catalog), namespace.clone(), table_name.to_string(), ) diff --git a/crates/integrations/datafusion/src/schema.rs b/crates/integrations/datafusion/src/schema_provider.rs similarity index 91% rename from crates/integrations/datafusion/src/schema.rs rename to crates/integrations/datafusion/src/schema_provider.rs index 545863f8c6..bfd1a55294 100644 --- a/crates/integrations/datafusion/src/schema.rs +++ b/crates/integrations/datafusion/src/schema_provider.rs @@ -23,7 +23,7 @@ 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; @@ -31,15 +31,16 @@ use iceberg::inspect::MetadataTableType; use iceberg::spec::FormatVersion; use iceberg::{Catalog, Error, ErrorKind, NamespaceIdent, Result, TableCreation, TableIdent}; +use crate::catalog_adapter::SessionBindingCatalogAdapter; 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, + /// The Iceberg catalog with access to session-aware and session-unaware APIs. + catalog: Arc, /// The namespace this schema represents namespace: NamespaceIdent, /// A concurrent map where keys are table names @@ -52,20 +53,19 @@ 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`. + /// 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( - client: Arc, + catalog: Arc, 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 + let table_names: Vec<_> = catalog .list_tables(&namespace) .await? .iter() @@ -75,7 +75,9 @@ impl IcebergSchemaProvider { let providers = try_join_all( table_names .iter() - .map(|name| IcebergTableProvider::try_new(client.clone(), namespace.clone(), name)) + .map(|name| { + IcebergTableProvider::try_new(Arc::clone(&catalog), namespace.clone(), name) + }) .collect::>(), ) .await?; @@ -86,7 +88,7 @@ impl IcebergSchemaProvider { } Ok(IcebergSchemaProvider { - catalog: client, + catalog, namespace, tables, }) @@ -191,14 +193,11 @@ impl SchemaProvider for IcebergSchemaProvider { .await .map_err(to_datafusion_error)?; - // Create a new table provider using the catalog reference - let table_provider = IcebergTableProvider::try_new( - catalog.clone(), - namespace.clone(), - name_clone.clone(), - ) - .await - .map_err(to_datafusion_error)?; + // Create a new table provider using the catalog access + let table_provider = + IcebergTableProvider::try_new(catalog, namespace.clone(), name_clone.clone()) + .await + .map_err(to_datafusion_error)?; // Store the new table provider tables.insert(name_clone, Arc::new(table_provider)); @@ -220,7 +219,7 @@ impl SchemaProvider for IcebergSchemaProvider { return Ok(None); } - let catalog = self.catalog.clone(); + let catalog = Arc::clone(&self.catalog); let namespace = self.namespace.clone(); let tables = self.tables.clone(); let table_name = name.to_string(); @@ -254,7 +253,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,7 +316,9 @@ mod tests { .await .unwrap(); - let provider = IcebergSchemaProvider::try_new(Arc::new(catalog), namespace) + let session_binding_catalog = + SessionBindingCatalogAdapter::new_without_context(Arc::new(catalog)); + let provider = IcebergSchemaProvider::try_new(Arc::new(session_binding_catalog), namespace) .await .unwrap(); diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 9de7bcb9c2..ef7cd38fe2 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -48,6 +48,7 @@ use iceberg::table::Table; use iceberg::{Catalog, Error, ErrorKind, NamespaceIdent, Result, TableIdent}; use metadata_table::IcebergMetadataTableProvider; +use crate::catalog_adapter::SessionBindingCatalogAdapter; use crate::error::to_datafusion_error; use crate::physical_plan::commit::IcebergCommitExec; use crate::physical_plan::project::project_with_partition; @@ -58,16 +59,18 @@ 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. #[derive(Debug, Clone)] pub struct IcebergTableProvider { - /// The catalog that manages this table - catalog: Arc, + /// Access to the catalog that manages this table + catalog: Arc, /// The table identifier (namespace + name) table_ident: TableIdent, /// A reference-counted arrow `Schema` (cached at construction) @@ -80,7 +83,7 @@ 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: Arc, namespace: NamespaceIdent, name: impl Into, ) -> Result { @@ -119,7 +122,7 @@ impl TableProvider for IcebergTableProvider { async fn scan( &self, - _state: &dyn Session, + session: &dyn Session, projection: Option<&Vec>, filters: &[Expr], limit: Option, @@ -127,6 +130,7 @@ impl TableProvider for IcebergTableProvider { // Load fresh table metadata from catalog let table = self .catalog + .with_session(session) .load_table(&self.table_ident) .await .map_err(to_datafusion_error)?; @@ -152,7 +156,7 @@ impl TableProvider for IcebergTableProvider { async fn insert_into( &self, - state: &dyn Session, + session: &dyn Session, input: Arc, _insert_op: InsertOp, ) -> DFResult> { @@ -162,9 +166,10 @@ impl TableProvider for IcebergTableProvider { ))); } + let catalog = self.catalog.with_session(session); + // Load fresh table metadata from catalog - let table = self - .catalog + let table = catalog .load_table(&self.table_ident) .await .map_err(to_datafusion_error)?; @@ -179,8 +184,8 @@ impl TableProvider for IcebergTableProvider { }; // Step 2: Repartition for parallel processing - let target_partitions = - NonZeroUsize::new(state.config().target_partitions()).ok_or_else(|| { + let target_partitions = NonZeroUsize::new(session.config().target_partitions()) + .ok_or_else(|| { DataFusionError::Configuration( "target_partitions must be greater than 0".to_string(), ) @@ -225,7 +230,7 @@ impl TableProvider for IcebergTableProvider { Ok(Arc::new(IcebergCommitExec::new( table, - self.catalog.clone(), + catalog, coalesce_partitions, self.schema.clone(), ))) @@ -348,7 +353,7 @@ mod tests { use datafusion::common::Column; use datafusion::physical_plan::ExecutionPlan; - use datafusion::prelude::SessionContext; + use datafusion::prelude::SessionContext as DFSessionContext; use iceberg::io::FileIO; use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; use iceberg::spec::{NestedField, PrimitiveType, Schema, Type}; @@ -374,7 +379,12 @@ mod tests { static_table.into_table() } - async fn get_test_catalog_and_table() -> (Arc, NamespaceIdent, String, TempDir) { + async fn get_test_catalog_and_table() -> ( + Arc, + NamespaceIdent, + String, + TempDir, + ) { let temp_dir = TempDir::new().unwrap(); let warehouse_path = temp_dir.path().to_str().unwrap().to_string(); @@ -413,8 +423,11 @@ mod tests { .await .unwrap(); + let session_bound_catalog = + SessionBindingCatalogAdapter::new_without_context(Arc::new(catalog)); + ( - Arc::new(catalog), + Arc::new(session_bound_catalog), namespace, "test_table".to_string(), temp_dir, @@ -429,7 +442,7 @@ mod tests { let table_provider = IcebergStaticTableProvider::try_new_from_table(table.clone()) .await .unwrap(); - let ctx = SessionContext::new(); + let ctx = DFSessionContext::new(); ctx.register_table("mytable", Arc::new(table_provider)) .unwrap(); let df = ctx.sql("SELECT * FROM mytable").await.unwrap(); @@ -455,7 +468,7 @@ mod tests { IcebergStaticTableProvider::try_new_from_table_snapshot(table.clone(), snapshot_id) .await .unwrap(); - let ctx = SessionContext::new(); + let ctx = DFSessionContext::new(); ctx.register_table("mytable", Arc::new(table_provider)) .unwrap(); let df = ctx.sql("SELECT * FROM mytable").await.unwrap(); @@ -479,7 +492,7 @@ mod tests { let table_provider = IcebergStaticTableProvider::try_new_from_table(table.clone()) .await .unwrap(); - let ctx = SessionContext::new(); + let ctx = DFSessionContext::new(); ctx.register_table("mytable", Arc::new(table_provider)) .unwrap(); @@ -502,7 +515,7 @@ mod tests { let table_provider = IcebergStaticTableProvider::try_new_from_table(table.clone()) .await .unwrap(); - let ctx = SessionContext::new(); + let ctx = DFSessionContext::new(); ctx.register_table("mytable", Arc::new(table_provider)) .unwrap(); @@ -540,7 +553,7 @@ mod tests { .await .unwrap(); - let ctx = SessionContext::new(); + let ctx = DFSessionContext::new(); ctx.register_table("test_table", Arc::new(provider)) .unwrap(); @@ -566,7 +579,7 @@ mod tests { .await .unwrap(); - let ctx = SessionContext::new(); + let ctx = DFSessionContext::new(); ctx.register_table("test_table", Arc::new(provider)) .unwrap(); @@ -593,7 +606,7 @@ mod tests { .await .unwrap(); - let ctx = SessionContext::new(); + let ctx = DFSessionContext::new(); ctx.register_table("test_table", Arc::new(provider)) .unwrap(); @@ -625,7 +638,12 @@ mod tests { async fn get_partitioned_test_catalog_and_table( fanout_enabled: Option, - ) -> (Arc, NamespaceIdent, String, TempDir) { + ) -> ( + Arc, + NamespaceIdent, + String, + TempDir, + ) { use iceberg::spec::{Transform, UnboundPartitionSpec}; let temp_dir = TempDir::new().unwrap(); @@ -681,8 +699,11 @@ mod tests { .await .unwrap(); + let session_binding_catalog = + SessionBindingCatalogAdapter::new_without_context(Arc::new(catalog)); + ( - Arc::new(catalog), + Arc::new(session_binding_catalog), namespace, "partitioned_table".to_string(), temp_dir, @@ -710,7 +731,7 @@ mod tests { let provider = IcebergTableProvider::try_new(catalog, namespace, table_name) .await .unwrap(); - let ctx = SessionContext::new(); + let ctx = DFSessionContext::new(); for (insert_op, expected_message) in [ ( @@ -753,7 +774,7 @@ mod tests { .await .unwrap(); - let ctx = SessionContext::new(); + let ctx = DFSessionContext::new(); let input_schema = provider.schema(); let input = Arc::new(EmptyExec::new(input_schema)) as Arc; @@ -785,7 +806,7 @@ mod tests { .await .unwrap(); - let ctx = SessionContext::new(); + let ctx = DFSessionContext::new(); let input_schema = provider.schema(); let input = Arc::new(EmptyExec::new(input_schema)) as Arc; @@ -811,7 +832,7 @@ mod tests { .await .unwrap(); - let ctx = SessionContext::new(); + let ctx = DFSessionContext::new(); let state = ctx.state(); // Test scan with limit @@ -844,7 +865,7 @@ mod tests { .await .unwrap(); - let ctx = SessionContext::new(); + let ctx = DFSessionContext::new(); let state = ctx.state(); // Test scan with limit @@ -872,7 +893,7 @@ mod tests { .await .unwrap(); - let ctx = SessionContext::new(); + let ctx = DFSessionContext::new(); let state = ctx.state(); // Test scan without limit diff --git a/crates/integrations/datafusion/src/test_utils.rs b/crates/integrations/datafusion/src/test_utils.rs new file mode 100644 index 0000000000..e4540cf497 --- /dev/null +++ b/crates/integrations/datafusion/src/test_utils.rs @@ -0,0 +1,253 @@ +// 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::{Arc, Mutex}; + +use async_trait::async_trait; +use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; +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; + +#[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 + } +} + +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, + ) +}