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

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`**: A dedicated Gravitino client is created per Trino session user, so each user is visible in the Gravitino audit log instead of the shared `gravitino.user`.
+- **`authType=oauth2`**: A dedicated Gravitino client is created per Bearer token from the Trino session extra credentials, 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.3.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.3.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

trino-connector/integration-test/src/test/java/org/apache/gravitino/trino/connector/integration/test/TrinoQueryTestTool.java

@@ -20,6 +20,7 @@
 
 import java.io.File;
 import java.util.Arrays;
+import java.util.concurrent.CountDownLatch;
 import org.apache.commons.cli.CommandLine;
 import org.apache.commons.cli.CommandLineParser;
 import org.apache.commons.cli.HelpFormatter;
@@ -102,6 +103,11 @@ public static void main(String[] args) throws Exception {
               + "The JAR file under ${trino_connector_dir} will be copied into the test image, "
               + "the default value is ${project_root}/trino-connector/trino-connector/build/libs.");
 
+      options.addOption(
+          "env_only",
+          false,
+          "Start the environment (Gravitino + Trino) and keep it running for manual testing. Press Ctrl+C to shutdown.");
+
       options.addOption("help", false, "Print this help message");
 
       CommandLineParser parser = new PosixParser();
@@ -254,6 +260,26 @@ public static void main(String[] args) throws Exception {
         return;
       }
 
+      if (commandLine.hasOption("env_only")) {
+        CountDownLatch shutdownLatch = new CountDownLatch(1);
+        Runtime.getRuntime().addShutdownHook(new Thread(shutdownLatch::countDown));
+
+        System.out.println("=======================================================");
+        System.out.println("Environment is ready for manual testing.");
+        System.out.println("  Gravitino URI : " + TrinoQueryITBase.gravitinoUri);
+        System.out.println("  Trino URI     : " + TrinoQueryITBase.trinoUri);
+        System.out.println("Connect to Trino CLI:");
+        System.out.println("  docker exec -it trino-ci-trino trino");
+        System.out.println("Press Ctrl+C to shutdown the environment.");
+        System.out.println("=======================================================");
+        try {
+          shutdownLatch.await();
+        } catch (InterruptedException e) {
+          Thread.currentThread().interrupt();
+        }
+        return;
+      }
+
       if (testSet == null) {
         testerRunner.testSql();
       } else {

trino-connector/integration-test/trino-test-tools/trino_integration_test.sh

@@ -26,6 +26,35 @@ export HADOOP_USER_NAME=anonymous
 echo $GRAVITINO_ROOT_DIR
 cd $GRAVITINO_ROOT_DIR
 
+PID_FILE="$GRAVITINO_ROOT_DIR/integration-test-common/build/trino-test-env.pid"
+
+if [ "$1" = "--stop" ]; then
+  if [ -f "$PID_FILE" ]; then
+    GRADLE_PID=$(cat "$PID_FILE")
+    if kill -0 "$GRADLE_PID" 2>/dev/null; then
+      PGID=$(ps -o pgid= -p "$GRADLE_PID" 2>/dev/null | tr -d ' ')
+      if [[ -n "$PGID" && "$PGID" =~ ^[0-9]+$ ]]; then
+        echo "Stopping environment (PGID: $PGID)..."
+        kill -TERM -- "-$PGID" 2>/dev/null
+        for i in $(seq 1 60); do
+          kill -0 "$GRADLE_PID" 2>/dev/null || break
+          sleep 1
+        done
+      else
+        echo "Could not determine process group for PID $GRADLE_PID, stopping containers directly..."
+      fi
+    else
+      echo "Environment process ($GRADLE_PID) is no longer running, stopping containers directly..."
+    fi
+    "$GRAVITINO_ROOT_DIR/integration-test-common/docker-script/shutdown.sh"
+    rm -f "$PID_FILE"
+  else
+    echo "No running environment found, stopping containers directly..."
+    "$GRAVITINO_ROOT_DIR/integration-test-common/docker-script/shutdown.sh"
+  fi
+  exit 0
+fi
+
 # Parse --auto_patch and --trino_version from arguments.
 # --auto_patch is consumed here and not forwarded to Gradle.
 # --trino_version is forwarded to Gradle and also used for patch selection.
@@ -106,4 +135,34 @@ if [ "$auto_patch" = true ]; then
     fi
 fi
 
-./gradlew :trino-connector:integration-test:TrinoTest -PappArgs="\"$args\""
+if echo "$args" | grep -q -- '--env_only'; then
+  LOG_FILE="$GRAVITINO_ROOT_DIR/integration-test-common/build/trino-test-env.log"
+  mkdir -p "$(dirname "$PID_FILE")" "$(dirname "$LOG_FILE")"
+
+  ./gradlew :trino-connector:integration-test:TrinoTest -PappArgs="\"$args\"" > "$LOG_FILE" 2>&1 &
+  GRADLE_PID=$!
+  echo $GRADLE_PID > "$PID_FILE"
+
+  echo "Starting environment, please wait..."
+  while kill -0 "$GRADLE_PID" 2>/dev/null; do
+    if grep -q "Press Ctrl+C to shutdown" "$LOG_FILE" 2>/dev/null; then
+      break
+    fi
+    sleep 1
+  done
+
+  if ! kill -0 "$GRADLE_PID" 2>/dev/null; then
+    echo "Environment failed to start. Check logs at: $LOG_FILE"
+    rm -f "$PID_FILE"
+    exit 1
+  fi
+
+  grep -A 6 "=======================================================" "$LOG_FILE" | head -8
+  echo ""
+  echo "Environment is running in background (PID: $GRADLE_PID)"
+  echo "Logs  : $LOG_FILE"
+  echo "Stop  : ./trino-connector/integration-test/trino-test-tools/trino_integration_test.sh --stop"
+  exit 0
+else
+  ./gradlew :trino-connector:integration-test:TrinoTest -PappArgs="\"$args\""
+fi

trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/GravitinoConnector.java

@@ -19,8 +19,12 @@
 package org.apache.gravitino.trino.connector;
 
 import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED;
+import static io.trino.spi.StandardErrorCode.PERMISSION_DENIED;
 
 import com.google.common.base.Preconditions;
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.cache.RemovalNotification;
 import io.trino.spi.TrinoException;
 import io.trino.spi.connector.Connector;
 import io.trino.spi.connector.ConnectorAccessControl;
@@ -36,11 +40,20 @@
 import io.trino.spi.session.PropertyMetadata;
 import io.trino.spi.transaction.IsolationLevel;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.gravitino.NameIdentifier;
+import org.apache.gravitino.client.GravitinoAdminClient;
+import org.apache.gravitino.client.GravitinoMetalake;
 import org.apache.gravitino.trino.connector.catalog.CatalogConnectorContext;
 import org.apache.gravitino.trino.connector.catalog.CatalogConnectorMetadata;
 import org.apache.gravitino.trino.connector.catalog.CatalogConnectorMetadataAdapter;
+import org.apache.gravitino.trino.connector.security.GravitinoAuthProvider;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * GravitinoConnector serves as the entry point for operations on the connector managed by Trino and
@@ -49,9 +62,15 @@
  */
 public class GravitinoConnector implements Connector {
 
+  private static final Logger LOG = LoggerFactory.getLogger(GravitinoConnector.class);
+
   private final NameIdentifier catalogIdentifier;
   protected final CatalogConnectorContext catalogConnectorContext;
   private final CatalogConnectorMetadata connectorMetadata;
+  private final boolean forwardUser;
+  private final GravitinoAuthProvider.AuthType authType;
+  private final String oauth2CredentialKey;
+  private final Cache<String, UserSession> perUserSessionCache;
 
   /**
    * Constructs a new GravitinoConnector with the specified catalog identifier and catalog connector
@@ -64,6 +83,39 @@ public GravitinoConnector(CatalogConnectorContext catalogConnectorContext) {
     this.catalogConnectorContext = catalogConnectorContext;
     this.connectorMetadata =
         new CatalogConnectorMetadata(catalogConnectorContext.getMetalake(), this.catalogIdentifier);
+
+    GravitinoConfig config = catalogConnectorContext.getConfig();
+    Map<String, String> clientConfig = config.getClientConfig();
+    this.forwardUser =
+        Boolean.parseBoolean(
+            clientConfig.getOrDefault(GravitinoAuthProvider.FORWARD_SESSION_USER_KEY, "false"));
+    String authTypeStr = clientConfig.getOrDefault(GravitinoAuthProvider.AUTH_TYPE_KEY, "none");
+    this.authType = GravitinoAuthProvider.parseAuthType(authTypeStr);
+    this.oauth2CredentialKey = clientConfig.get(GravitinoAuthProvider.OAUTH2_TOKEN_CREDENTIAL_KEY);
+
+    if (forwardUser && authType == GravitinoAuthProvider.AuthType.OAUTH2) {
+      Preconditions.checkArgument(
+          StringUtils.isNotBlank(oauth2CredentialKey),
+          "oauth2 with forwardUser=true requires '%s' to be set",
+          GravitinoAuthProvider.OAUTH2_TOKEN_CREDENTIAL_KEY);
+    }
+
+    if (forwardUser) {
+      this.perUserSessionCache =
+          CacheBuilder.newBuilder()
+              .maximumSize(500)
+              .expireAfterAccess(1, TimeUnit.HOURS)
+              .removalListener(
+                  (RemovalNotification<String, UserSession> notification) -> {
+                    UserSession session = notification.getValue();
+                    if (session != null) {
+                      session.close();
+                    }
+                  })
+              .build();
+    } else {
+      this.perUserSessionCache = null;
+    }
   }
 
   @Override
@@ -90,6 +142,45 @@ public ConnectorMetadata getMetadata(
     ConnectorMetadata internalMetadata =
         internalConnector.getMetadata(session, gravitinoTransactionHandle.getInternalHandle());
     Preconditions.checkArgument(internalMetadata != null, "Internal metadata must not be null");
+
+    if (forwardUser) {
+      String credKey =
+          authType == GravitinoAuthProvider.AuthType.OAUTH2
+              ? "oauth2:" + session.getIdentity().getExtraCredentials().get(oauth2CredentialKey)
+              : "simple:" + session.getUser();
+      UserSession userSession;
+      try {
+        userSession =
+            perUserSessionCache.get(
+                credKey,
+                () -> {
+                  GravitinoAdminClient userClient =
+                      GravitinoAuthProvider.buildForSession(
+                          catalogConnectorContext.getConfig(), session);
+                  GravitinoMetalake userMetalake =
+                      userClient.loadMetalake(catalogConnectorContext.getMetalake().name());
+                  CatalogConnectorMetadata userMetadata =
+                      new CatalogConnectorMetadata(userMetalake, catalogIdentifier);
+                  return new UserSession(userClient, userMetadata);
+                });
+      } catch (ExecutionException e) {
+        Throwable cause = e.getCause();
+        LOG.warn(
+            "Failed to create per-user Gravitino client for user '{}': {}",
+            session.getUser(),
+            cause.getMessage());
+        throw new TrinoException(
+            PERMISSION_DENIED,
+            "Failed to authenticate user '"
+                + session.getUser()
+                + "' with Gravitino: "
+                + cause.getMessage(),
+            cause);
+      }
+      return createGravitinoMetadata(
+          userSession.metadata, catalogConnectorContext.getMetadataAdapter(), internalMetadata);
+    }
+
     return createGravitinoMetadata(
         connectorMetadata, catalogConnectorContext.getMetadataAdapter(), internalMetadata);
   }
@@ -179,8 +270,30 @@ public ConnectorNodePartitioningProvider getNodePartitioningProvider() {
   }
 
   public void shutdown() {
+    if (forwardUser) {
+      perUserSessionCache.invalidateAll();
+    }
     Connector internalConnector = catalogConnectorContext.getInternalConnector();
     internalConnector.shutdown();
     catalogConnectorContext.close();
   }
+
+  /** Holds a per-user {@link GravitinoAdminClient} together with its derived metadata. */
+  private static final class UserSession {
+    final GravitinoAdminClient client;
+    final CatalogConnectorMetadata metadata;
+
+    UserSession(GravitinoAdminClient client, CatalogConnectorMetadata metadata) {
+      this.client = client;
+      this.metadata = metadata;
+    }
+
+    void close() {
+      try {
+        client.close();
+      } catch (Exception e) {
+        LOG.warn("Failed to close GravitinoAdminClient", e);
+      }
+    }
+  }
 }

trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/GravitinoMetadata.java

@@ -313,11 +313,7 @@ public void setColumnComment(
       ColumnHandle column,
       Optional<String> comment) {
     String columnName = getColumnName(column);
-
-    String commentString = "";
-    if (comment.isPresent() && !StringUtils.isBlank(comment.get())) {
-      commentString = comment.get();
-    }
+    String commentString = comment.filter(c -> !StringUtils.isBlank(c)).orElse("");
     catalogConnectorMetadata.setColumnComment(getTableName(tableHandle), columnName, commentString);
   }