feat(rest_catalog::auth) Contextual Sessions by DerGut · Pull Request #3081 · apache/iceberg-rust

DerGut · 2026-08-26T22:21:28Z

Which issue does this PR close?

Another piece that's working towards #2774.

This PR bridges the RestSessionCatalog's SessionContexts (introduced with #2920) with the AuthManager (introduced with #2838 thanks again @plusplusjiajia).
The RestSessionCatalog will be able to forward session context to an auth manager to generate "contextual" auth sessions. These auth sessions may contain credentials specific to a query in a multi-tenant system, i.e. a specific tenant's credentials.

What changes are included in this PR?

This PR adds a new pub trait method to the iceberg_catalog_rest::AuthManager with a default implementation. The new AuthManager::contextual_session allows implementors to convert per catalog operation SessionContext into HTTP-level authentication metadata.

async fn contextual_session(
        &self,
        _context: &SessionContext,
        catalog_session: Arc<dyn AuthSession>,
    ) -> Result<Arc<dyn AuthSession>> {
        Ok(catalog_session)
    }

The RestSessionCatalog receives updates to forward its SessionContext to its underlying RestClient which now holds an dyn AuthManager to authenticate any incoming session. Previously the AuthManager was only used once to retrieve a catalog session.

Are these changes tested?

Added a test test_contextual_session_authenticates_each_catalog_request to assert that the catalog session is passed as the parent, session context reaches the AuthManager, requests are modified with auth headers.

The OAuth2Manager::contextual_session implementation in #3170 verifies the trait with an actual implementation.

AI Disclosure

I've used Codex to help me with setting up this PR.

requests

DerGut · 2026-09-08T18:14:27Z

This PR should now be ready for review, it also has it's API validated by #3170 but doesn't depend on it to be merged in any way 👍
cc @CTTY

CTTY

Thanks for having this! Left some questions

CTTY · 2026-09-08T22:20:36Z

+    /// derive authentication for its [`SessionContext`].
+    auth_manager: Arc<dyn AuthManager>,
+    /// The catalog-wide session passed to [`AuthManager::contextual_session`].
+    catalog_session: Arc<dyn AuthSession>,


Storing catalog session here is fine but feels a bit excessive. Why not build contextual sessions in the catalog, and then pass them to RestClient::query_catalog?

Really only because the catalog is already doing quite a lot and the RestClient isn't.
I agree that conceptually it belongs closer to the catalog though.
I can make the change tomorrow 👌

There's actually something more to it. The catalog session depends on the /v1/config handshake and on the init session. Because the RestClient is initialized lazily, it's not trivial to simply move the field to the RestSessionCatalog.

I see three options:

keep them in the OnceCell<RestClient> and re-use the lazy intialization

introduce new OnceCell<Arc<dyn AuthManager>> and OnceCell<Arc<dyn AuthSession>> and have their initialization depend on the rest client's initialization

use a wrapper to encapsulate the rest client, auth manager and catalog session into a single lazy initialization

My thoughts on them:
2. is spreading lazy initialization logic around the code where it's really only one lazy chain that could be simplified
3. I could see this one work out but I haven't found a good name yet. I also am concerned that we're building a deep tree of encapsulating types: RestCatalog -> RestSessionCatalog -> <SomeNewType> -> RestClient -> HttpClient. The distinction between the RestClient and HttpClient isn't too well-defined yet (it's essentially an http client with rest-catalog specific config). Maybe this could be an opportunity to give it a more specific responsibility to save us from another layer -> that would get us to a hybrid of 1. and 3.

This makes sense, thanks for the detailed explanation!

I think keeping it as is for now is good. Maybe we need to rename RestClient to something else in the future since it represents more than just a client

I'd rename it to RestCatalogClient for now to emphasize it's specificity to the catalog, will push a commit in a few minutes.
I have some refactors in mind that can go on top of this PR

the auth manager has no dependency on sessions and shouldn't have to be loaded lazily -> it should be constructed by the builder

going a step further, auth manager resolution only requires the RestConfig's properties, I'm exploring whether we shouldn't move the auth manager to the config and change it from a TypedBuilder pattern to a custom builder that returns a Result<RestConfig> with a resolved auth manager. This would avoid an Arc<dyn AuthManager> on the catalog just to pass it to RestCatalogClient initialization (and then not use it again)

Rename commit -> 8785ac8

Let me know what you think 🙏

CTTY · 2026-09-08T22:21:13Z

+    /// Builder-supplied override retained so lazy client initialization can
+    /// clone it into the runtime state. When absent, a manager is resolved from
+    /// `rest.auth.type` during initialization.
+    auth_manager_override: Option<Arc<dyn AuthManager>>,


naming is hard and I like the new name more!

CTTY · 2026-09-08T22:22:19Z

-    async fn check_exists_via_head(&self, client: &RestClient, url: String) -> Result<bool> {
+    async fn check_exists_via_head(
+        &self,
+        context: &SessionContext,


Same here, seems like we can pass the pre-built contextual session to this function directly

CTTY · 2026-09-08T22:25:28Z


        let request = HttpRequest::build(request_builder)?;
-        let http_response = client.query_catalog(request).await?;
+        let session = client.contextual_session(context).await?;


We can build the contextual session using the auth manager here

CTTY

LGTM!

CTTY · 2026-09-11T23:13:12Z

+    /// derive authentication for its [`SessionContext`].
+    auth_manager: Arc<dyn AuthManager>,
+    /// The catalog-wide session passed to [`AuthManager::contextual_session`].
+    catalog_session: Arc<dyn AuthSession>,


This makes sense, thanks for the detailed explanation!

I think keeping it as is for now is good. Maybe we need to rename RestClient to something else in the future since it represents more than just a client

alexanderbianchi · 2026-09-12T23:21:08Z

@CTTY Looks like auto-merge was disabled

laskoviymishka

Looks good to me too — coming in late after CTTY's approval, so mostly a fresh set of eyes on the public API. The core move is right: keeping HttpClient session-free and deriving auth per operation through contextual_session, and it nicely fixes the two existence-check paths that were dropping the context before.

Nothing blocks the merge for me. Left a few inline notes — a couple worth tidying while it's open (the double HeaderMap clone, the _context underscore) and a nudge to exercise more than list_namespaces in the test, since check_exists_via_head is the path that was silently broken and the one where the wrong tenant's session leaks existence. The rest are just forward-looking notes for the auth series.

laskoviymishka · 2026-09-14T18:19:20Z

    }

+    #[tokio::test]
+    async fn test_contextual_session_authenticates_catalog_operation() {


This is the one thing I'd want before merge: the new wiring touches all twelve operations, but only list_namespaces is exercised here.

The gap I care about most is check_exists_via_head — both table_exists and the HEAD branch of namespace_exists route through it, and that path silently ignored the context before this PR. An existence check sent with the wrong tenant's session is how "does table X exist" leaks across tenants, so I'd really like a HEAD test with match_header on the contextual session for both.

While we're in there, a test where contextual_session returns Err (asserting the op fails before any HTTP request goes out) plus one write-op test would cover the parts most likely to break silently on a mechanical slip. The existing ContextManager harness makes each ~15 lines. wdyt?

laskoviymishka · 2026-09-14T18:19:21Z

+        request: HttpRequest,
+    ) -> Result<HttpResponse> {
+        self.http_client
+            .with_auth_session(session)


I think this quietly doubles the per-request HeaderMap clone. with_auth_session does Self { auth_session, ..self.clone() }, which clones extra_headers, and then HttpClient::query_catalog clones extra_headers again when it extends the request. Before this PR the session was baked in at init and we paid one clone per request.

An internal method that authenticates in place keeps it to a single clone:

pub(crate) async fn query_catalog_with_session( &self, session: &dyn AuthSession, mut request: HttpRequest, ) -> Result<HttpResponse> { session.authenticate(&mut request).await?; let mut inner = request.into_inner(); inner.headers_mut().extend(self.extra_headers.clone()); HttpResponse::read(self.client.execute(inner).await?).await }

Not a blocker, but it's a regression on the hot path for single-context callers, so I'd fix it here.

laskoviymishka · 2026-09-14T18:19:21Z

+    /// context may therefore return the previously cached session.
+    async fn contextual_session(
+        &self,
+        _context: &SessionContext,


Small one, but the underscore here fights the doc. The docstring tells implementors to key their cache off SessionContext::session_id, yet _context renders in the generated public API and in rust-analyzer as "this argument is unused" — the opposite message.

I'd drop the underscore from the trait signature and silence the default body instead:

async fn contextual_session( &self, context: &SessionContext, catalog_session: Arc<dyn AuthSession>, ) -> Result<Arc<dyn AuthSession>> { let _ = context; Ok(catalog_session) }

laskoviymishka · 2026-09-14T18:19:21Z

+    /// [`SessionContext::session_id`] and are responsible for eviction and
+    /// releasing any associated resources. Reusing a session ID with different
+    /// context may therefore return the previously cached session.
+    async fn contextual_session(


Not blocking — more a design question for the series. The doc asks implementors to cache sessions and "release any associated resources," but the trait gives them no signal for when to do that. Java's AuthManager extends AutoCloseable and OAuth2Manager.close() calls sessionCache.invalidateAll() for exactly this.

In a long-lived multi-tenant catalog, a per-context cache grows unbounded with no teardown hook. Worth a default async fn close(&self) -> Result<()> { Ok(()) } now, or at least a doc note that Drop is the intended cleanup point? Happy to leave it for a follow-up if that's the plan.

laskoviymishka · 2026-09-14T18:19:21Z

+    /// The catalog calls this method only after [`Self::catalog_session`] has
+    /// succeeded. `catalog_session` is the catalog session returned by this
+    /// manager. If the context does not require different authentication,
+    /// implementations should return `catalog_session` unchanged.


Two things the implementor guidance leaves unsaid that I think are worth a sentence each, since the OAuth2 impl in #3170 will hit both.

First, contextual_session gets no HttpClient — an implementation that needs one for a token exchange has to stash a clone of the client passed to catalog_session. That's the intended pattern (Java holds the client internally too), but nothing here says so.

Second, the catalog doesn't serialize concurrent calls with the same session_id, so a caching impl without its own synchronization can race and build several sessions for one context. A note that impls must guard concurrent creation would save someone that bug.

laskoviymishka · 2026-09-14T18:19:21Z

        let endpoint = client.config.namespaces_endpoint();
        let mut namespaces = Vec::new();
        let mut next_token = None;
+        let session = client.contextual_session(context).await?;


One subtle thing — we grab the contextual session once here and Arc::clone it across every page. If an implementation hands back a short-lived token and a listing spans more pages than the token's TTL, the later pages go out with stale auth and surface as a surprise 401 deep in the loop rather than at the call.

Probably fine for now since no built-in manager does short-lived contextual tokens yet, but worth either re-deriving per page or a doc line telling impls a session may be reused across an operation's requests. wdyt? (Same applies to list_tables.)

DerGut mentioned this pull request Aug 27, 2026

Session/ Request/ Auth Context Propagation #2774

Open

DerGut marked this pull request as ready for review August 27, 2026 13:24

DerGut mentioned this pull request Aug 27, 2026

refactor(auth) Emphasize AuthManagers aren't shared between catalogs #3086

Merged

DerGut added 3 commits September 4, 2026 00:11

Add trait method AuthManager::contextual_session

e4d4f91

Implement contextual session handling

45d29c6

Produce an auth_session per operation and re-use it for multiple Http

d539f65

requests

DerGut force-pushed the contextual-session branch from be83e3d to d539f65 Compare September 4, 2026 08:37

DerGut mentioned this pull request Sep 7, 2026

feat: Oauth2 contextual session #3170

Open

CTTY reviewed Sep 8, 2026

View reviewed changes

This was referenced Sep 9, 2026

Draft: resolve Iceberg session catalogs per DataFusion query alexanderbianchi/iceberg-rust#1

Draft

feat[iceberg]: add catalog and schema providers datafusion-contrib/datafusion-distributed#722

Draft

CTTY approved these changes Sep 11, 2026

View reviewed changes

CTTY enabled auto-merge September 11, 2026 23:15

auto-merge was automatically disabled September 11, 2026 23:55
Head branch was pushed to by a user without write access

Rename client to RestCatalogClient

8785ac8

DerGut force-pushed the contextual-session branch from dd6aaaa to 8785ac8 Compare September 12, 2026 00:16

laskoviymishka approved these changes Sep 14, 2026

View reviewed changes

Uh oh!

Conversation

DerGut commented Aug 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

What changes are included in this PR?

Are these changes tested?

AI Disclosure

Uh oh!

DerGut commented Sep 8, 2026

Uh oh!

CTTY left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

DerGut Sep 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

DerGut Sep 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

DerGut Sep 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

CTTY left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alexanderbianchi commented Sep 12, 2026

Uh oh!

laskoviymishka left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

DerGut commented Aug 26, 2026 •

edited

Loading

DerGut Sep 9, 2026 •

edited

Loading

DerGut Sep 11, 2026 •

edited

Loading

DerGut Sep 11, 2026 •

edited

Loading