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(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!("%s: %s", 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?

⏳ Tests are coming.

AI Disclosure