[#10559] feat(trino-connector): Forward Trino session credentials to Gravitino server#10730
Open
diqiu50 wants to merge 19 commits intoapache:mainfrom
Open
[#10559] feat(trino-connector): Forward Trino session credentials to Gravitino server#10730diqiu50 wants to merge 19 commits intoapache:mainfrom
diqiu50 wants to merge 19 commits intoapache:mainfrom
Conversation
…OAuth2 token to Gravitino server - Add ExtraHeadersProvider interface to client-java for dynamic per-request HTTP headers - Add oauth2_token auth type: forwards Bearer token from Trino session extra credentials - Add forwardUser option: forwards Trino session username via X-Gravitino-User header for all auth types - Use ThreadLocal pattern (TrinoSessionAuthProvider, TrinoUserHeaderProvider) to propagate session context through beginQuery/cleanupQuery lifecycle - Propagate TrinoSessionContext through CatalogConnectorContext -> GravitinoConnector -> GravitinoMetadata - Add tests for new auth scenarios and TrinoUserHeaderProvider - Update authentication.md with new configuration options
…User option The X-Gravitino-User header was never read by the Gravitino server (AuthenticationFilter/SimpleAuthenticator only reads Authorization). In a unified OAuth2 setup the token already carries user identity, making the extra header both redundant and forgeable. Remove TrinoUserHeaderProvider, the gravitino.client.session.forwardUser config key, and the related Session User Forwarding documentation section. The oauth2_token Bearer token forwarding (TrinoSessionAuthProvider) is unaffected.
…on test script Port --env_only and --stop from gravitino-internal. Allows starting the full Gravitino + Trino environment in the background for manual testing, then stopping it cleanly without running any test cases.
…e session management Root cause: ThreadLocal token set in beginQuery (coordinator thread) but HTTP requests made in SplitRunner threads, causing token to be null. Changes: 1. Move session apply/clear from beginQuery to each catalogConnectorMetadata call using callWithSession/runWithSession wrappers 2. Add explicit error logging when OAuth2 token or username is missing 3. Optimize performance: - Restore beginQuery/cleanupQuery session management as primary mechanism - Cache Base64-encoded tokens per username (ConcurrentHashMap) - Skip redundant applySession if token already set Performance: Reduced from 100+ ThreadLocal ops per query to 1 set + 1 clear. Verified: alice/bob users correctly forwarded to Gravitino server.
…areCatalogMetadata wrapper Replace the callWithSession/runWithSession pattern with a cleaner wrapper-based approach: - Introduce SessionAwareCatalogMetadata wrapper class that automatically handles session apply/clear around each operation - Use method overloading (Supplier<T> vs Runnable) for elegant handling of both returning and void methods - Remove callWithSession/runWithSession helper methods from GravitinoMetadata - Update all 6 version-specific modules (435-439, 440-445, 446-451, 452-468, 469-472, 473-478) Benefits: - Cleaner code: net reduction of 69 lines - Better separation of concerns: session management encapsulated in wrapper - Easier to maintain: no need to wrap each method call manually - Fixes potential bugs: version-specific addColumn methods now properly wrapped Testing: - All unit tests pass - Manual testing verified user forwarding works correctly - Concurrent queries show no session cross-contamination
Contributor
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
This PR adds per-query credential forwarding from Trino to Gravitino via a new gravitino.client.session.forwardUser connector option, enabling Gravitino to authenticate/authorize using each Trino user’s identity rather than shared static credentials.
Changes:
- Add session-aware auth forwarding (
TrinoSessionAuthProvider+TrinoSessionContext) and wrap Gravitino metadata calls withSessionAwareCatalogMetadata. - Extend Java client HTTP layer to support per-request extra headers via
ExtraHeadersProvider. - Update integration tooling and documentation, and add/adjust unit tests.
Reviewed changes
Copilot reviewed 33 out of 33 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
| trino-connector/trino-connector/src/test/java/org/apache/gravitino/trino/connector/catalog/TestGravitinoAuthProvider.java | Adds build-result tests for session forwarding and changes package. |
| trino-connector/trino-connector/src/test/java/org/apache/gravitino/trino/connector/TestGravitinoMetadataGetSystemTable.java | Updates tests to use SessionAwareCatalogMetadata. |
| trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/security/TrinoSessionContext.java | Introduces a facade for applying/clearing session context. |
| trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/security/TrinoSessionAuthProvider.java | Implements per-thread, per-session auth forwarding via CustomTokenProvider. |
| trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/catalog/SessionAwareCatalogMetadata.java | Wraps metadata calls to apply/clear session forwarding around each call. |
| trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/catalog/GravitinoAuthProvider.java | Builds client + optional session context; adds new config keys; changes package/public API. |
| trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/catalog/CatalogConnectorManager.java | Uses new GravitinoAuthProvider.build() and propagates TrinoSessionContext. |
| trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/catalog/CatalogConnectorContext.java | Stores and exposes optional TrinoSessionContext; plumbs through builder. |
| trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/GravitinoMetadata.java | Switches to SessionAwareCatalogMetadata and passes ConnectorSession through. |
| trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/GravitinoConnector.java | Wraps metadata delegate in SessionAwareCatalogMetadata. |
| trino-connector/trino-connector-473-478/src/main/java/org/apache/gravitino/trino/connector/GravitinoMetadata478.java | Passes session into catalog metadata calls (session-aware wrapper). |
| trino-connector/trino-connector-473-478/src/main/java/org/apache/gravitino/trino/connector/GravitinoConnector478.java | Updates factory to accept SessionAwareCatalogMetadata. |
| trino-connector/trino-connector-469-472/src/main/java/org/apache/gravitino/trino/connector/GravitinoMetadata469.java | Passes session into catalog metadata calls. |
| trino-connector/trino-connector-469-472/src/main/java/org/apache/gravitino/trino/connector/GravitinoConnector469.java | Updates factory signature for SessionAwareCatalogMetadata. |
| trino-connector/trino-connector-452-468/src/main/java/org/apache/gravitino/trino/connector/GravitinoMetadata452.java | Passes session into catalog metadata calls. |
| trino-connector/trino-connector-452-468/src/main/java/org/apache/gravitino/trino/connector/GravitinoConnector452.java | Updates factory signature for SessionAwareCatalogMetadata. |
| trino-connector/trino-connector-446-451/src/main/java/org/apache/gravitino/trino/connector/GravitinoMetadata446.java | Passes session into catalog metadata calls. |
| trino-connector/trino-connector-446-451/src/main/java/org/apache/gravitino/trino/connector/GravitinoConnector446.java | Updates factory signature for SessionAwareCatalogMetadata. |
| trino-connector/trino-connector-440-445/src/main/java/org/apache/gravitino/trino/connector/GravitinoMetadata440.java | Passes session into catalog metadata calls. |
| trino-connector/trino-connector-440-445/src/main/java/org/apache/gravitino/trino/connector/GravitinoConnector440.java | Updates factory signature for SessionAwareCatalogMetadata. |
| trino-connector/trino-connector-435-439/src/main/java/org/apache/gravitino/trino/connector/GravitinoMetadata435.java | Passes session into catalog metadata calls. |
| trino-connector/trino-connector-435-439/src/main/java/org/apache/gravitino/trino/connector/GravitinoConnector435.java | Updates factory signature for SessionAwareCatalogMetadata. |
| trino-connector/integration-test/trino-test-tools/trino_integration_test.sh | Adds --stop and --env_only runner behavior for manual testing. |
| trino-connector/integration-test/src/test/java/org/apache/gravitino/trino/connector/integration/test/TrinoQueryTestTool.java | Adds env_only option to keep env alive for manual testing. |
| integration-test-common/docker-script/init/trino/config/catalog/gravitino.properties | Enables simple auth + forwardUser in integration test catalog config. |
| docs/trino-connector/authentication.md | Documents session credential forwarding config and adds new properties. |
| clients/client-java/src/test/java/org/apache/gravitino/client/TestVersionCheck.java | Updates constructor call for new ExtraHeadersProvider param. |
| clients/client-java/src/test/java/org/apache/gravitino/client/TestGravitinoClientBuilder.java | Updates constructor call for new ExtraHeadersProvider param. |
| clients/client-java/src/main/java/org/apache/gravitino/client/HTTPClient.java | Adds ExtraHeadersProvider; guards against missing auth token data. |
| clients/client-java/src/main/java/org/apache/gravitino/client/GravitinoClientBase.java | Plumbs ExtraHeadersProvider through client builders. |
| clients/client-java/src/main/java/org/apache/gravitino/client/GravitinoClient.java | Plumbs ExtraHeadersProvider into concrete client. |
| clients/client-java/src/main/java/org/apache/gravitino/client/GravitinoAdminClient.java | Plumbs ExtraHeadersProvider into admin client. |
| clients/client-java/src/main/java/org/apache/gravitino/client/ExtraHeadersProvider.java | Introduces interface for per-request extra headers. |
Comments suppressed due to low confidence (1)
trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/security/TrinoSessionContext.java:1
- This Javadoc describes a per-query lifecycle (
beginQuery/cleanupQuery), but the introduced wrapper (SessionAwareCatalogMetadata) applies/clears the session around each metadata call. These two approaches conflict (especially given nested-call concerns). Update the Javadoc to match the actual lifecycle used in this PR, or refactor to consistently use per-query apply/clear and ensure thread/context propagation across all call sites.
roryqi
reviewed
Apr 9, 2026
…h per-user GravitinoClient cache Replace the ThreadLocal-based credential forwarding (SessionAwareCatalogMetadata + TrinoSessionAuthProvider) with a cleaner per-user GravitinoAdminClient approach: - Each unique session credential (username for simple, token for oauth2) gets its own GravitinoAdminClient built at client creation time via GravitinoAuthProvider.buildForSession() - Per-user clients and their CatalogConnectorMetadata are cached in GravitinoConnector (max 500 entries, 1h TTL) with a RemovalListener that closes evicted clients - Auth config (forwardUser, authType, credentialKey) cached as constructor fields to avoid HashMap allocation on every query - Remove SessionAwareCatalogMetadata, TrinoSessionAuthProvider, TrinoSessionContext - Consolidate GravitinoAuthProvider: extract removeAuthSpecificKeys() and parseAuthType() helpers, simplify build() to return GravitinoAdminClient directly (drop BuildResult wrapper)
Code Coverage Report
Files
|
- Fix Builder.clone() to propagate GravitinoConfig (was causing NPE when catalogs are cloned via DefaultCatalogConnectorFactory) - Add Preconditions.checkArgument(config != null) in Builder.build() - Validate oauth2CredentialKey at GravitinoConnector construction time when authType=OAUTH2 and forwardUser=true, preventing NPE in hot path - Replace non-atomic getIfPresent/put with Cache.get(key, Callable) to eliminate race condition in per-user session cache - Wrap buildForSession/loadMetalake failures as TrinoException(PERMISSION_DENIED) with user-facing message and logging instead of leaking internal errors - Fix UserSession.close() to log warnings instead of throwing RuntimeException, which was aborting invalidateAll() evictions during shutdown - Fix RemovalListener to use lambda instead of anonymous class - Remove stale ThreadLocal reference from ExtraHeadersProvider Javadoc - Update buildClient() deprecation comment to be accurate - Fix authentication.md: "per-request" → "per-user" client description
ExtraHeadersProvider was introduced for the old ThreadLocal-based session forwarding approach. With the redesign to per-user GravitinoAdminClient caching, credentials are baked in at client construction time and per-request header injection is no longer needed. - Remove ExtraHeadersProvider interface and all plumbing through HTTPClient, GravitinoClientBase, GravitinoClient, GravitinoAdminClient - Revert test changes that adapted to the new constructor signatures - Fix authentication.md: change "Since version" from 1.9.0 to 1.3.0
- Always call shutdown.sh after killing the process group, so Docker containers are stopped even if the Gradle process died unexpectedly - Add PGID numeric guard before issuing kill to process group, falling back to shutdown.sh directly when PGID cannot be resolved
- GravitinoConnector: change authType field from raw String to
AuthType enum, parsed once at construction via parseAuthType();
use enum reference comparison instead of string equals("OAUTH2")
- GravitinoAuthProvider: make parseAuthType() public so callers can
parse auth type without duplicating the logic; remove redundant
getTokenData() override in StaticTokenProvider — base class already
returns schemeName + " " + getCustomTokenInfo() which is identical
- CatalogConnectorContext.Builder: fix checkArgument messages from
"X is not null" to "X must not be null"
- CatalogConnectorManager: replace hardcoded "gravitino.client.authType"
string with GravitinoAuthProvider.AUTH_TYPE_KEY constant; fix
checkArgument message
roryqi
reviewed
Apr 13, 2026
| case SIMPLE: | ||
| builder.withSimpleAuth(session.getUser()); | ||
| break; | ||
| case OAUTH2: |
Contributor
There was a problem hiding this comment.
I don't figure out to support OAUTH for the session user. Maybe we can throw an UnSupportedException now.
OAuth2 forwardUser support is removed until a proper token exchange approach is defined. Only simple auth session forwarding is supported. - Remove OAUTH2_TOKEN_CREDENTIAL_KEY constant and StaticTokenProvider - buildForSession() now only supports authType=simple; throws for others - GravitinoConnector: remove oauth2CredentialKey field, cache key is always "simple:<user>" - Remove OAuth2 forwarding tests and doc section
- buildForSession() now throws UnsupportedOperationException for non-simple auth types - Add startup validation: forwardUser=true with non-simple authType throws TrinoException at construction time - Fix CatalogConnectorManager.createCatalogConnectorContext() to use per-request config instead of global this.config - Add tests for startup validation and UnsupportedOperationException
…le cache params - Extract auth validation and cache construction into GravitinoConnector.buildSessionCache() - Add SESSION_CACHE_MAX_SIZE_KEY and SESSION_CACHE_EXPIRE_AFTER_ACCESS_SECONDS_KEY config options (gravitino.client.session.cache.maxSize, gravitino.client.session.cache.expireAfterAccessSeconds) - Remove hardcoded 500/1h values; defaults remain 500 / 3600s - Document new config properties in authentication.md
- Move SESSION_CACHE_MAX_SIZE_KEY / SESSION_CACHE_EXPIRE_AFTER_ACCESS_SECONDS_KEY from GravitinoAuthProvider to GravitinoConnector (they are connector-level, not auth concerns) - Replace explicit removal of 3 session.* keys in removeAuthSpecificKeys() with a single prefix-based removeIf() covering all gravitino.client.session.* keys - Add parseLongConfig() helper with TrinoException on NumberFormatException - Remove unnecessary Javadoc from buildSessionCache() - Make UserSession.close() idempotent with volatile closed flag
…moval - Add SESSION_CACHE_MAX_SIZE_KEY, SESSION_CACHE_EXPIRE_AFTER_ACCESS_SECONDS_KEY public constants and ConfigEntry definitions in GravitinoConfig, consistent with project pattern - Add getSessionCacheMaxSize() and getSessionCacheExpireAfterAccessSeconds() typed getters with NumberFormatException handling in parseLongConfigEntry() - GravitinoConnector.buildSessionCache() takes GravitinoConfig and uses typed getters; remove private key constants and parseLongConfig() helper - Revert removeAuthSpecificKeys() to explicit removal referencing GravitinoConfig constants
… method - Change session cache config keys to gravitino.trino.session-cache.* prefix (not Gravitino client config, no gravitino.client. prefix needed) - Add GravitinoConfig.isForwardUser() getter; simplify GravitinoConnector constructor - Extract per-user session cache lookup into resolveSessionMetadata(); getMetadata() is now a single ternary plus one createGravitinoMetadata() call - Remove SESSION_CACHE_*_KEY removes from removeAuthSpecificKeys() since these keys no longer appear in getClientConfig() output - Revert unrelated GravitinoMetadata.setColumnComment() style change
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
Add per-query credential forwarding from Trino to Gravitino via a new
gravitino.client.session.forwardUseroption. When enabled, the connector forwards the Trino session identity (username or Bearer token) to Gravitino on every request instead of using static credentials.Why are the changes needed?
Fix: #10559
In a multi-user deployment, each Trino user should authenticate with their own identity so Gravitino can apply per-user authorization policies.
Does this PR introduce any user-facing change?
New connector property:
gravitino.client.session.forwardUser=true— forwards the Trino session identity to Gravitino per request. Behavior depends onauthType: forwards the session username forsimple, or a Bearer token from extra credentials foroauth2.How was this patch tested?
TestGravitinoAuthProvider