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

[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!