Datafusion session integration

[DerGut]

Which issue does this PR close?

This is another PR in my series to close #2774.

Today, the Datafusion Session (which is already available for each query) terminates at the Iceberg catalog boundary. Scans for example, use a shared dyn Catalog and then call the catalog.load_table(&self.table_ident) function on it, which doesn't support any way of context propagation.

Our downstream REST catalog requires this context to make authorization, rate limiting and shard routing decisions.

What changes are included in this PR?

This PR starts to use the freshly introduced SessionCatalog trait in our Datafusion {Catalog, Schema, Table}Providers.
A user will now be able to provide forward Datafusion query context to their Iceberg catalog by providing 1) a dyn SessionCatalog (like the RestSessionCatalog introduced by #2920) and 2) a custom dyn SessionContextResolver implementation.

Public API

The public API is extended with two new symbols:

1. a new constructor IcebergCatalogProvider::try_new_with_session_catalog
2. a new trait to allow users to extract relevant metadata from their Datafusion query session, and translate it into an Iceberg SessionContext (which the session catalog accepts)

pub trait SessionContextResolver {
    fn resolve(&self, session: &dyn DFSession) -> DFResult<SessionContext>;
}

A possible implementation of this trait may look like

struct CustomUserContext{
    name: String,
    id: String,
    auth_token: String,
}

struct CustomUserContextResolver {}

impl SessionContextResolver for CustomUserContextResolver {
    fn resolve(&self, session: &dyn DFSession) -> DFResult<SessionContext> {
        let user = session
            .config()
            .get_extension::<CustomUserContext>()
            .ok_or_else(|| {
                DataFusionError::Configuration(
                    "the DataFusion session has no CustomUserContext extension".to_string(),
                )
            })?;

        Ok(SessionContext::builder()
            // Reusing the DataFusion session ID gives the catalog a stable key
            // for session-scoped caches.
            .session_id(session.session_id().to_string())
            .identity(format!("{}: {}", user.name.to_string(), user.id.to_string()))
            .credentials(HashMap::from([(
                "token".to_string(),
                SensitiveString::from(user.auth_token.to_string()),
            )]))
            .build())
    }
}

Why a new Trait?

This is necessary because Datafusion doesn't have a canonical way of encoding query context (in contrast to Trino's ConnectorSession). Instead, it propagates arbitrary types via its SessionConfig's extension mechanism.

This leaves us with two ways to shape a Datafusion SessionConfig's extension into a SessionContext:

1. make users provide an Iceberg-defined extension
2. make users provide an implementation to parse their own types

Option 1. has a meaningful drawback: a Datafusion instance that connects to multiple data sources/ catalog providers (and supports joins between those) shouldn't use a dedicated query context for each, but one user-defined one that can be interpreted by each data source.

Note on RestSessionCatalog

Since the REST catalog implementations abstract the HTTP protocol away, there's another layer missing to specify how an Iceberg SessionContext can be used to enrich HTTP requests with the provided metadata. The newly introduced AuthManager trait (via #2838) can be used for that.

The Implementation

CatalogAccess Enum

I'd like to keep a way for users to create CatalogProviders from plain Catalogs in case they don't deal with sessions. Removing that constructor would be breaking anyway.

Again, I saw two options to do this:

1. provide two sets of implementations: next to {Catalog, Schema, Table}Provider we'd have something like {Catalog, Schema, Table}SessionProviders
2. support two constructors and back a single implementation set by a common abstraction

For this draft, I figured that the overhead of two sets of public APIs, in addition to the duplicate code (or a similar common abstraction to 1. to reduce some duplication) makes 2. seem simpler. So that's what I went for.

Are these changes tested?

Included an example datafusion_session_catalog.rs and a test suite under catalog_provider.rs to exercise that metadata set on the Datafusion session make it into the SessionCatalog..

AI Disclosure

Used help from Codex and Claude to prototype different designs and generate tests.

[CTTY]

I'm not sure I understand why do we need to have a separate trait to customize this. I thought the ability to resolve Datafusion context should be an extension of SessionCatalog.

[DerGut]

Thanks again for the quick review!

Essentially, the SessionContextResolver is an adapter between the query engine and Iceberg. The SessionCatalog is pure Iceberg and required to be provided. The SessionContextResolver is specific to how Datafusion works (it accepts a datafusion::catalog::Session) and cannot easily be abstracted further and tied to the SessionCatalog.
In the Java implementation, we don't have an equivalent because its only query engine integration (in the apache/iceberg repo at least) is Spark which is single-tenant and doesn't need/ use a SessionCatalog.
A better example in the Java world would be the Trino-Iceberg connector, which defines its own equivalent of the SessionContextResolver -> included in the TrinoRestCatalog.

Note that Trino is different in the sense that a Trino ConnectorSession is already structured, and so a default implementation can be provided. This is very different in the Datafusion world, where users provide their custom types (more on this in the PR description).

[CTTY]

The current design looks correct function wise, but I have some concerns and maybe we can improve this:

• We don't know when to call with_session and when to call without_session, my expectation is iceberg always tries to resolve the session while we have the query context, and the result depends on whether the underlying catalog type is a session catalog. This way we don't need to call something like "without_session" when we don't have the query context (e.g. IcebergTableProvider)
• There are many abstractions hanging around and it's hard for users to figure out what to implement. SessionContextResolver feels like something that should come with SessionCatalog implementation or extension. Having two of them in parallel in APIs like below will shift the responsibility of resolving df session to users.

pub async fn try_new_with_session_catalog(
        catalog: Arc<dyn SessionCatalog>,
        resolver: Arc<dyn SessionContextResolver>,
    )

• Do we plan to have default implementations for SessionContextResolver? How do users use RestSessionCatalog out of the box?

[DerGut]

Sorry, I forgot to better explain this in the PR description.

There's really two changes in this PR that I was thinking of contributing as separate PRs but ended up including here in separate commits instead.

1. integrate the SessionCatalog with the Iceberg{Catalog, Schema, Table}Provider implementations (first set of commits)
2. add a SessionBoundCatalog to allow passing a SessionCatalog to the IcebergCommitExec and use it across the providers to save lines

The total diff obscures this. On the one hand, the addition of the SessionBoundCatalog is necessary to commit an insert, but on the other, its been extended as a convenience to reduce repetetive statements like this:

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?,
        };

into something like

let table = catalog_access
            .without_session()
            .load_table(&table_ident)
            .await?;

We don't know when to call with_session and when to call without_session

The API is essentially: whenever we have a datafusion::catalog::Session in scope, use CatalogAccess:with_session, if not, use CatalogAccess::without_session. But the match statement above would be equivalent.

[DerGut]

Do we plan to have default implementations for SessionContextResolver? How do users use RestSessionCatalog out of the box?

Unfortunately, due to Datafusion's nature of handling query sessions, there's little we can provide as default implementations. A datafusion::catalog::Session does provide a Session::session_id which we can use to populate the SessionContext::session_id. But IMO an otherwise empty context has no use, so I don't see a point in providing a default implementation for that.

Also, I'm afraid that since all Datafusion query metadata is user-defined, there's no reason for users to use a SessionCatalog over a Catalog unless they also provide a mechanism to extract+translate that metadata.

[DerGut]

Let me ask some coworkers who know Datafusion very deeply tomorrow, to get another opinion!

[CTTY]

1. Yes, SessionBoundCatalog makes sense to me, let's break it down to smaller PRs and move with that first
2. Instead of CatalogAccess, I was thinking of something like a trait DataFusionCatalog that provides and a default implementation should contain a non-session catalog, so we could use the same trait/concept across catalog/schema/table providers. If users wish to use SessionCatalog, they can implement the context resolving logic there as well. This way we don't have to juggle with two traits (SessionCatalog and ContextResolver) at the same time: these two traits have to be used together anyway, and it feels odd that users have to juggle with them in parallel.

The trait can define a function:

fn catalog_for_session(&self, session: Option<&dyn Session>) -> DFResult<Arc<dyn Catalog>>

I haven't got much chance to think about the caching though, will need to spend a bit more time

[gabotechs]

I've skimmed through the PR, and the proposed solution looks fairly convoluted for something that should be potentially simple.

For example, in DataFusion, there should always be a Session present, so I'm not sure in which situation it makes sense to not pass it.

I'll give it some though, and will try to come up with an actual suggestion.

[DerGut]

On 1.: Sounds good, I'll create a follow-up!

On 2.:
I would like for users to just provide a single parameter but I think the more fitting design depends on what users actually use as a SessionCatalog.

If users are implementing their own SessionCatalog, your approach might be easier. They just add another method to their existing type to make it implement the new trait.

My current version is based on the assumption that the vast majority of users will use a RestSessionCatalog. If we ask users to provide a DataFusionCatalog instead, they will have to provide a dyn SessionCatalog and a dyn SessionContextResolver to some wrapper type we provide, causing them to go through one layer of indirection.

IcebergCatalogProvider::try_new(DataFusionCatalogAdapter::new(rest_session_catalog, my_context_resolver));

Both approaches will work, but I feel like we should be optimizing for the more common one.

My opinion might be biased by our internal usage. Do you have a better assessment of how frequently the community uses custom non-REST catalog implementations?

[DerGut]

in DataFusion, there should always be a Session present, so I'm not sure in which situation it makes sense to not pass it.

Most Iceberg users do not generally deal with sessions and therefore use a simplified Catalog trait that is unaware of sessions.
I think in the DataFusion integration we can discuss whether or not we want to keep the old constructor that's based on that trait. This implementation preserves that constructor.

[gabotechs]

I think there's an opportunity of reducing a lot of the convolution by using DataFusion ConfigOption extensions. For example:

datafusion::common::extensions_options! {
    pub struct IcebergOptions {
        /// The [IcebergSessionContext] `identity` field.
        pub identity: Option<String>, default = None
    }
}

With this, we could automatically map IcebergOptions to the relevant fields of IcebergSessionContext inside this project, without exposing this detail to users.

From a public API standpoint, this crate would just offer this IcebergOptions as a native DataFusion ConfigOptions implementation, and under the hood this can be wired up internally to an IcebergSessionContext.

This is the most DataFusion native way of threading custom config across the callstack.

[DerGut]

That would indeed help the complexity of the integration implementation a lot. But this also means that users now have to set Iceberg-specific extension options in addition to their possibly already existing custom extension options, right?

Definitely good to know that this is the DataFusion-way of doing things, this is exactly the context I was missing, thanks! 🙇

[gabotechs]

I think we can afford to leave this as a normal Arc<dyn SessionCatalog>, and just wrap it in the appropriate places with a SessionBoundCatalog implementation that automatically enriches the IcebergSessionContext under the hood based on whatever is present in the DataFusionSessionConfig.