[#10559] feat(trino-connector): Forward Trino session credentials to Gravitino server

clients/client-java/src/main/java/org/apache/gravitino/client/ExtraHeadersProvider.java

@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.gravitino.client;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.Map;
+
+/**
+ * Provider of extra HTTP headers appended to every outgoing request. Implementations must be
+ * thread-safe because {@link #getHeaders()} is called once per HTTP request and may be invoked
+ * concurrently from multiple threads.
+ *
+ * <p>A typical use-case is forwarding per-request context (e.g. the originating user identity) that
+ * is stored in a {@link ThreadLocal} and populated before each request is dispatched.
+ *
+ * @since 1.3.0
+ */
+public interface ExtraHeadersProvider extends Closeable {
+
+  /**
+   * Returns the extra headers to append to the current HTTP request. The returned map must not be
+   * {@code null}; return an empty map when no headers should be added.
+   *
+   * @return a non-null, possibly empty map of header name → header value pairs
+   */
+  Map<String, String> getHeaders();
+
+  /** Default no-op close implementation. Override if resources need to be released. */
+  @Override
+  default void close() throws IOException {}
+}

clients/client-java/src/main/java/org/apache/gravitino/client/GravitinoAdminClient.java

@@ -54,6 +54,7 @@ public class GravitinoAdminClient extends GravitinoClientBase implements Support
    *
    * @param uri The base URI for the Gravitino API.
    * @param authDataProvider The provider of the data which is used for authentication.
+   * @param extraHeadersProvider The provider of per-request extra headers, may be {@code null}.
    * @param checkVersion Whether to check the version of the Gravitino server. Gravitino does not
    *     support the case that the client-side version is higher than the server-side version.
    * @param headers The base header for Gravitino API.
@@ -62,10 +63,11 @@ public class GravitinoAdminClient extends GravitinoClientBase implements Support
   private GravitinoAdminClient(
       String uri,
       AuthDataProvider authDataProvider,
+      ExtraHeadersProvider extraHeadersProvider,
       boolean checkVersion,
       Map<String, String> headers,
       Map<String, String> properties) {
-    super(uri, authDataProvider, checkVersion, headers, properties);
+    super(uri, authDataProvider, extraHeadersProvider, checkVersion, headers, properties);
   }
 
   /**
@@ -259,7 +261,12 @@ public GravitinoAdminClient build() {
       Preconditions.checkArgument(
           uri != null && !uri.isEmpty(), "The argument 'uri' must be a valid URI");
       return new GravitinoAdminClient(
-          uri, authDataProvider, isVersionCheckEnabled(), headers, properties);
+          uri,
+          authDataProvider,
+          extraHeadersProvider,
+          isVersionCheckEnabled(),
+          headers,
+          properties);
     }
   }
 }

clients/client-java/src/main/java/org/apache/gravitino/client/GravitinoClient.java

@@ -99,10 +99,11 @@ private GravitinoClient(
       String uri,
       String metalakeName,
       AuthDataProvider authDataProvider,
+      ExtraHeadersProvider extraHeadersProvider,
       boolean checkVersion,
       Map<String, String> headers,
       Map<String, String> properties) {
-    super(uri, authDataProvider, checkVersion, headers, properties);
+    super(uri, authDataProvider, extraHeadersProvider, checkVersion, headers, properties);
     this.metalake = loadMetalake(metalakeName);
   }
 
@@ -712,7 +713,13 @@ public GravitinoClient build() {
           "The argument 'metalakeName' must be a valid name");
 
       return new GravitinoClient(
-          uri, metalakeName, authDataProvider, isVersionCheckEnabled(), headers, properties);
+          uri,
+          metalakeName,
+          authDataProvider,
+          extraHeadersProvider,
+          isVersionCheckEnabled(),
+          headers,
+          properties);
     }
   }
 }

clients/client-java/src/main/java/org/apache/gravitino/client/GravitinoClientBase.java

@@ -61,13 +61,15 @@ public abstract class GravitinoClientBase implements Closeable {
    *
    * @param uri The base URI for the Gravitino API.
    * @param authDataProvider The provider of the data which is used for authentication.
+   * @param extraHeadersProvider The provider of per-request extra headers, may be {@code null}.
    * @param checkVersion Whether to check the version of the Gravitino server.
    * @param headers The base header of the Gravitino API.
    * @param properties A map of properties (key-value pairs) used to configure the Gravitino client.
    */
   protected GravitinoClientBase(
       String uri,
       AuthDataProvider authDataProvider,
+      ExtraHeadersProvider extraHeadersProvider,
       boolean checkVersion,
       Map<String, String> headers,
       Map<String, String> properties) {
@@ -78,6 +80,7 @@ protected GravitinoClientBase(
           HTTPClient.builder(properties)
               .uri(uri)
               .withAuthDataProvider(authDataProvider)
+              .withExtraHeadersProvider(extraHeadersProvider)
               .withObjectMapper(mapper)
               .withPreConnectHandler(this::checkVersion)
               .withHeaders(headers)
@@ -88,6 +91,7 @@ protected GravitinoClientBase(
           HTTPClient.builder(properties)
               .uri(uri)
               .withAuthDataProvider(authDataProvider)
+              .withExtraHeadersProvider(extraHeadersProvider)
               .withObjectMapper(mapper)
               .withHeaders(headers)
               .build();
@@ -208,6 +212,8 @@ public abstract static class Builder<T> {
     protected String uri;
     /** The authentication provider. */
     protected AuthDataProvider authDataProvider;
+    /** The provider of per-request extra headers (e.g. forwarded user identity). */
+    protected ExtraHeadersProvider extraHeadersProvider;
     /** The check version flag. */
     protected boolean checkVersion = true;
     /** The request base header for the Gravitino API. */
@@ -329,6 +335,19 @@ public Builder<T> withCustomTokenAuth(CustomTokenProvider dataProvider) {
       return this;
     }
 
+    /**
+     * Sets an {@link ExtraHeadersProvider} that appends per-request headers to every outgoing HTTP
+     * call. A common use-case is forwarding the originating user identity (e.g. the Trino session
+     * user) alongside the service-level authentication credentials.
+     *
+     * @param provider the provider, or {@code null} to disable extra-header forwarding
+     * @return This Builder instance for method chaining.
+     */
+    public Builder<T> withExtraHeaders(ExtraHeadersProvider provider) {
+      this.extraHeadersProvider = provider;
+      return this;
+    }
+
     /**
      * Set base header for Gravitino Client.
      *

clients/client-java/src/main/java/org/apache/gravitino/client/HTTPClient.java

@@ -80,6 +80,7 @@ public class HTTPClient implements RESTClient {
   private final CloseableHttpClient httpClient;
   private final ObjectMapper mapper;
   private final AuthDataProvider authDataProvider;
+  private final ExtraHeadersProvider extraHeadersProvider;
 
   // Handler to be executed before connecting to the server.
   private final Runnable beforeConnectHandler;
@@ -103,6 +104,7 @@ enum HandlerStatus {
    * @param baseHeaders A map of base headers to be included in all HTTP requests.
    * @param objectMapper The ObjectMapper used for JSON serialization and deserialization.
    * @param authDataProvider The provider of authentication data.
+   * @param extraHeadersProvider The provider of per-request extra headers, may be {@code null}.
    * @param beforeConnectHandler The function to be executed before connecting to the server.
    * @param properties A map of properties (key-value pairs) used to configure the HTTP client.
    */
@@ -111,6 +113,7 @@ private HTTPClient(
       Map<String, String> baseHeaders,
       ObjectMapper objectMapper,
       AuthDataProvider authDataProvider,
+      ExtraHeadersProvider extraHeadersProvider,
       Runnable beforeConnectHandler,
       Map<String, String> properties) {
     this.uri = uri;
@@ -130,6 +133,7 @@ private HTTPClient(
 
     this.httpClient = clientBuilder.build();
     this.authDataProvider = authDataProvider;
+    this.extraHeadersProvider = extraHeadersProvider;
 
     if (beforeConnectHandler == null) {
       handlerStatus = HandlerStatus.Finished;
@@ -360,10 +364,15 @@ private <T> T execute(
     } else {
       addRequestHeaders(request, headers, ContentType.APPLICATION_JSON.getMimeType());
     }
-    if (authDataProvider != null) {
-      request.setHeader(
-          AuthConstants.HTTP_HEADER_AUTHORIZATION,
-          new String(authDataProvider.getTokenData(), StandardCharsets.UTF_8));
+    if (authDataProvider != null && authDataProvider.hasTokenData()) {
+      byte[] tokenData = authDataProvider.getTokenData();
+      if (tokenData != null) {
+        request.setHeader(
+            AuthConstants.HTTP_HEADER_AUTHORIZATION, new String(tokenData, StandardCharsets.UTF_8));
+      }
+    }
+    if (extraHeadersProvider != null) {
+      extraHeadersProvider.getHeaders().forEach(request::setHeader);
     }
 
     try (CloseableHttpResponse response = httpClient.execute(request)) {
@@ -705,6 +714,9 @@ public void close() throws IOException {
     if (authDataProvider != null) {
       authDataProvider.close();
     }
+    if (extraHeadersProvider != null) {
+      extraHeadersProvider.close();
+    }
     httpClient.close(CloseMode.GRACEFUL);
   }
 
@@ -758,6 +770,7 @@ public static class Builder {
     private String uri;
     private ObjectMapper mapper = ObjectMapperProvider.objectMapper();
     private AuthDataProvider authDataProvider;
+    private ExtraHeadersProvider extraHeadersProvider;
     private Runnable beforeConnectHandler;
 
     private Builder(Map<String, String> properties) {
@@ -833,6 +846,19 @@ public Builder withAuthDataProvider(AuthDataProvider authDataProvider) {
       return this;
     }
 
+    /**
+     * Sets the ExtraHeadersProvider for the HTTP client. The provider's {@link
+     * ExtraHeadersProvider#getHeaders()} is called on every outgoing request to append additional
+     * headers (e.g. a forwarded user-identity header).
+     *
+     * @param extraHeadersProvider the provider, or {@code null} to disable extra headers
+     * @return This Builder instance for method chaining.
+     */
+    public Builder withExtraHeadersProvider(ExtraHeadersProvider extraHeadersProvider) {
+      this.extraHeadersProvider = extraHeadersProvider;
+      return this;
+    }
+
     /**
      * Builds and returns an instance of the HTTPClient with the configured options.
      *
@@ -841,7 +867,13 @@ public Builder withAuthDataProvider(AuthDataProvider authDataProvider) {
     public HTTPClient build() {
 
       return new HTTPClient(
-          uri, baseHeaders, mapper, authDataProvider, beforeConnectHandler, properties);
+          uri,
+          baseHeaders,
+          mapper,
+          authDataProvider,
+          extraHeadersProvider,
+          beforeConnectHandler,
+          properties);
     }
   }
 

clients/client-java/src/test/java/org/apache/gravitino/client/TestGravitinoClientBuilder.java

@@ -101,7 +101,7 @@ private MockGravitinoClient(
       AuthDataProvider authDataProvider,
       Map<String, String> headers,
       Map<String, String> properties) {
-    super(uri, authDataProvider, false, headers, properties);
+    super(uri, authDataProvider, null, false, headers, properties);
     this.headers = headers;
     this.authDataProvider = authDataProvider;
     this.properties = properties;

clients/client-java/src/test/java/org/apache/gravitino/client/TestVersionCheck.java

@@ -43,13 +43,25 @@ public class TestVersionCheck extends TestBase {
 
   private static class TestClient extends GravitinoClientBase {
     protected TestClient(String uri, AuthDataProvider authDataProvider, boolean checkVersion) {
-      super(uri, authDataProvider, checkVersion, Collections.emptyMap(), Collections.emptyMap());
+      super(
+          uri,
+          authDataProvider,
+          null,
+          checkVersion,
+          Collections.emptyMap(),
+          Collections.emptyMap());
     }
   }
 
   private static class TestAdminClient extends GravitinoClientBase {
     protected TestAdminClient(String uri, AuthDataProvider authDataProvider, boolean checkVersion) {
-      super(uri, authDataProvider, checkVersion, Collections.emptyMap(), Collections.emptyMap());
+      super(
+          uri,
+          authDataProvider,
+          null,
+          checkVersion,
+          Collections.emptyMap(),
+          Collections.emptyMap());
     }
 
     private int listMetalakesCount() {

docs/trino-connector/authentication.md

@@ -31,7 +31,7 @@ gravitino.user=admin
 
 | Property                          | Description                                                          | Default value    | Required                               | Since version |
 |-----------------------------------|----------------------------------------------------------------------|------------------|----------------------------------------|---------------|
-| `gravitino.client.authType`       | Authentication type: `simple`, `oauth2`, or `kerberos`              | (none)           | No                                     | 1.3.0         |
+| `gravitino.client.authType`       | Authentication type: `simple`, `oauth2`, or `kerberos`                 | (none)        | No                                     | 1.3.0         |
 | `gravitino.user`                  | Username for simple authentication                                   | (none)           | No (uses system user if not specified) | 1.3.0         |
 
 ### OAuth2 Authentication
@@ -57,7 +57,7 @@ gravitino.client.oauth2.scope=gravitino
 
 | Property                             | Description                                            | Default value | Required                    | Since version |
 |--------------------------------------|--------------------------------------------------------|---------------|-----------------------------|---------------|
-| `gravitino.client.authType`          | Authentication type: `simple`, `oauth2`, or `kerberos` | (none)        | Yes (to enable OAuth2)      | 1.3.0         |
+| `gravitino.client.authType`          | Authentication type: `simple`, `oauth2`, or `kerberos` | (none)        | Yes (to enable OAuth2)           | 1.3.0         |
 | `gravitino.client.oauth2.serverUri`  | OAuth2 server URI                                      | (none)        | Yes if authType is `oauth2` | 1.3.0         |
 | `gravitino.client.oauth2.credential` | OAuth2 credentials in format `client_id:client_secret` | (none)        | Yes if authType is `oauth2` | 1.3.0         |
 | `gravitino.client.oauth2.path`       | OAuth2 token endpoint path                             | (none)        | Yes if authType is `oauth2` | 1.3.0         |
@@ -88,6 +88,7 @@ gravitino.client.kerberos.keytabFilePath=/path/to/user.keytab
 | `gravitino.client.kerberos.principal`      | Kerberos principal  | (none)        | Yes if authType is `kerberos`           | 1.3.0         |
 | `gravitino.client.kerberos.keytabFilePath` | Path to keytab file | (none)        | No (uses ticket cache if not specified) | 1.3.0         |
 
+
 ### Example: Connecting to OAuth-protected Gravitino Server
 
 This example shows how to configure the Trino connector to connect to a Gravitino server protected by OAuth authentication.
@@ -123,6 +124,51 @@ gravitino.client.oauth2.scope=test
 SHOW CATALOGS;
 ```
 
+### Session Credential Forwarding
+
+Setting `gravitino.client.session.forwardUser=true` enables per-query credential forwarding from Trino to Gravitino. The behavior depends on the configured `authType`:
+
+- **`authType=simple`**: The Trino session username is encoded as a `Basic` credential and forwarded per-request, so each Trino user is visible in the Gravitino audit log instead of the shared `gravitino.user`.
+- **`authType=oauth2`**: The Bearer token from Trino session extra credentials is forwarded per-request to Gravitino, enabling per-user OAuth2 authorization.
+
+**Configuration for simple auth (forward session user):**
+
+```properties
+connector.name=gravitino
+gravitino.metalake=metalake
+gravitino.uri=http://localhost:8090
+
+gravitino.client.authType=simple
+gravitino.client.session.forwardUser=true
+```
+
+**Configuration for OAuth2 (forward session Bearer token):**
+
+```properties
+connector.name=gravitino
+gravitino.metalake=metalake
+gravitino.uri=http://localhost:8090
+
+gravitino.client.authType=oauth2
+gravitino.client.session.forwardUser=true
+gravitino.client.oauth2.token.credentialKey=gravitino.token
+```
+
+**Query execution with OAuth2 token:**
+
+```sql
+-- Pass an OAuth2 token through Trino extra credentials
+SET SESSION gravitino.extra_credentials = 'gravitino.token=<your-oauth2-token>';
+SHOW SCHEMAS IN my_catalog;
+```
+
+**Configuration properties:**
+
+| Property                                        | Description                                                                                                                              | Default value | Required                                                 | Since version |
+|-------------------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------|---------------|----------------------------------------------------------|---------------|
+| `gravitino.client.session.forwardUser`          | When `true`, forwards the Trino session user (`simple`) or Bearer token (`oauth2`) to Gravitino per-request                              | `false`       | No                                                       | 1.9.0         |
+| `gravitino.client.oauth2.token.credentialKey`   | Key name in Trino extra credentials that holds the Bearer token. Only used when `authType=oauth2` and `forwardUser=true`                  | (none)        | Yes if `authType=oauth2` and `forwardUser=true`          | 1.9.0         |
+
 ### Notes
 
 - The Gravitino server must be configured with the corresponding authentication mechanism enabled.

integration-test-common/docker-script/init/trino/config/catalog/gravitino.properties

@@ -21,3 +21,5 @@ connector.name = gravitino
 gravitino.uri = http://GRAVITINO_HOST_IP:GRAVITINO_HOST_PORT
 gravitino.metalake = GRAVITINO_METALAKE_NAME
 gravitino.trino.skip-version-validation=true
+gravitino.client.authType = simple
+gravitino.client.session.forwardUser = true