Copilot commented on code in PR #12698:
URL: https://github.com/apache/gravitino/pull/12698#discussion_r3877379496
##########
trino-connector/trino-connector/src/main/java/org/apache/gravitino/trino/connector/GravitinoConnector.java:
##########
@@ -225,37 +228,82 @@ public void shutdown() {
catalogConnectorContext.close();
}
- private CatalogConnectorMetadata resolveSessionMetadata(ConnectorSession
session) {
- String credKey = "simple:" + session.getUser();
+ @VisibleForTesting
+ CatalogConnectorMetadata resolveSessionMetadata(ConnectorSession session) {
+ String authType =
+ catalogConnectorContext
+ .getConfig()
+ .getClientConfig()
+ .getOrDefault(GravitinoAuthProvider.AUTH_TYPE_KEY, "simple");
+ String credentialKey =
+ catalogConnectorContext
+ .getConfig()
+ .getClientConfig()
+ .getOrDefault(
+ GravitinoAuthProvider.USER_TOKEN_CREDENTIAL_KEY,
+ GravitinoAuthProvider.DEFAULT_USER_TOKEN_CREDENTIAL_KEY);
+ String token =
session.getIdentity().getExtraCredentials().get(credentialKey);
+ String credKey = sessionCacheKey(authType, session.getUser(), token);
try {
return perUserSessionCache.get(
credKey,
() -> {
- GravitinoAdminClient userClient =
- GravitinoAuthProvider.buildForSession(
- catalogConnectorContext.getConfig(), session);
+ GravitinoAdminClient userClient = buildAuthClient(session);
GravitinoMetalake userMetalake =
userClient.loadMetalake(catalogConnectorContext.getMetalake().name());
return new UserSession(
userClient, new CatalogConnectorMetadata(userMetalake,
catalogIdentifier));
})
.metadata;
- } catch (ExecutionException e) {
- Throwable cause = e.getCause();
+ } catch (ExecutionException | UncheckedExecutionException e) {
+ Throwable cause = e.getCause() == null ? e : e.getCause();
LOG.warn(
- "Failed to create per-user Gravitino client for user '{}': {}",
- session.getUser(),
- cause.getMessage());
+ "Failed to create per-user Gravitino client for user '{}'",
session.getUser(), cause);
+ if (cause instanceof TrinoException) {
+ // Already carries a specific Trino error code (e.g. from
buildForSession); re-wrapping
+ // would swallow it.
+ throw (TrinoException) cause;
+ }
+ if (cause instanceof IllegalArgumentException
+ || cause instanceof UnsupportedOperationException) {
+ throw new TrinoException(
+ PERMISSION_DENIED,
+ "Failed to authenticate user '"
+ + session.getUser()
+ + "' with Gravitino: "
+ + cause.getMessage(),
+ cause);
+ }
throw new TrinoException(
- PERMISSION_DENIED,
- "Failed to authenticate user '"
+ GENERIC_INTERNAL_ERROR,
+ "Unexpected error while creating per-user Gravitino client for user
'"
+ session.getUser()
- + "' with Gravitino: "
+ + "': "
+ cause.getMessage(),
cause);
}
}
+ /**
+ * Builds the per-user Gravitino admin client for a forwarded session.
Extracted as an overridable
+ * seam so tests can substitute the client-building behavior without needing
to mock the static
+ * {@link GravitinoAuthProvider#buildForSession}.
+ *
+ * @param session the current Trino connector session
+ * @return the per-user Gravitino admin client
+ */
+ @VisibleForTesting
+ GravitinoAdminClient buildAuthClient(ConnectorSession session) {
+ return
GravitinoAuthProvider.buildForSession(catalogConnectorContext.getConfig(),
session);
+ }
+
+ @VisibleForTesting
+ static String sessionCacheKey(String authType, String user, String token) {
+ String tokenPart =
+ StringUtils.isBlank(token) ? "" : ":" +
Integer.toHexString(token.hashCode());
+ return authType + ":" + user + tokenPart;
+ }
Review Comment:
`sessionCacheKey` currently incorporates the session extra-credential token
hash for all auth types, even though `authType=simple` does not use the token
to authenticate. This can cause unnecessary cache churn / cache growth whenever
an unrelated extra-credential changes. Also, using only `token.hashCode()` can
(rarely) collide, which could incorrectly reuse a cached client across rotated
tokens.
##########
trino-connector/trino-connector/src/test/java/org/apache/gravitino/trino/connector/catalog/iceberg/TestIcebergCatalogPropertyConverter.java:
##########
@@ -57,6 +57,33 @@ public void testHiveBackendProperty() {
"Missing required property for Hive backend: [uri]");
}
+ @Test
+ public void testRestBackendProperty() {
+ PropertyConverter propertyConverter = new
IcebergCatalogPropertyConverter();
+ Map<String, String> gravitinoIcebergConfig =
+ ImmutableMap.<String, String>builder()
+ .put("uri", "http://localhost:9001/iceberg";)
+ .put("catalog-backend", "rest")
+ .put("warehouse", "gt_iceberg_rest")
+ .build();
+ Map<String, String> restBackendConfig =
+ propertyConverter.gravitinoToEngineProperties(gravitinoIcebergConfig);
+
+ Assertions.assertEquals(restBackendConfig.get("iceberg.catalog.type"),
"rest");
+ Assertions.assertEquals(
+ restBackendConfig.get("iceberg.rest-catalog.uri"),
"http://localhost:9001/iceberg";);
+ Assertions.assertEquals(
+ restBackendConfig.get("iceberg.rest-catalog.warehouse"),
"gt_iceberg_rest");
+
+ Map<String, String> wrongMap = Maps.newHashMap(gravitinoIcebergConfig);
+ wrongMap.remove("uri");
+
+ Assertions.assertThrows(
+ TrinoException.class,
+ () -> propertyConverter.gravitinoToEngineProperties(wrongMap),
+ "Missing required property for Rest backend: [uri]");
+ }
Review Comment:
`Assertions.assertThrows`'s third argument is only a failure message, so
this test currently does not verify the thrown exception message for the
missing-`uri` case. Capturing the exception and asserting on `getMessage()`
makes the test actually validate the intended error output.
##########
trino-connector/trino-connector/src/test/java/org/apache/gravitino/trino/connector/security/TestGravitinoAuthProvider.java:
##########
@@ -259,6 +261,63 @@ public void
testBuildForSessionThrowsForNonSimpleAuthType() {
() -> GravitinoAuthProvider.buildForSession(config, session));
}
+ @Test
+ public void testBuildForSessionOAuth2() {
+ GravitinoConfig config =
+ buildConfig(
+ ImmutableMap.of(
+ GravitinoAuthProvider.AUTH_TYPE_KEY, "oauth2",
+ GravitinoAuthProvider.FORWARD_SESSION_USER_KEY, "true"));
+
+ ConnectorSession session = mock(ConnectorSession.class);
+ when(session.getUser()).thenReturn("alice");
+ when(session.getIdentity())
+ .thenReturn(
+ ConnectorIdentity.forUser("alice")
+ .withExtraCredentials(ImmutableMap.of("token",
"forwarded-user-jwt"))
+ .build());
+
+ GravitinoAdminClient client =
GravitinoAuthProvider.buildForSession(config, session);
+ assertNotNull(client);
+ }
+
+ @Test
+ public void testBuildForSessionOAuth2WithCustomCredentialKey() {
+ GravitinoConfig config =
+ buildConfig(
+ ImmutableMap.of(
+ GravitinoAuthProvider.AUTH_TYPE_KEY, "oauth2",
+ GravitinoAuthProvider.FORWARD_SESSION_USER_KEY, "true",
+ GravitinoAuthProvider.USER_TOKEN_CREDENTIAL_KEY,
"access_token"));
+
+ ConnectorSession session = mock(ConnectorSession.class);
+ when(session.getUser()).thenReturn("alice");
+ when(session.getIdentity())
+ .thenReturn(
+ ConnectorIdentity.forUser("alice")
+ .withExtraCredentials(ImmutableMap.of("access_token",
"alice-jwt"))
+ .build());
+
+ GravitinoAdminClient client =
GravitinoAuthProvider.buildForSession(config, session);
+ assertNotNull(client);
+ }
+
+ @Test
+ public void testBuildForSessionOAuth2ThrowsWhenTokenMissing() {
+ GravitinoConfig config =
+ buildConfig(
+ ImmutableMap.of(
+ GravitinoAuthProvider.AUTH_TYPE_KEY, "oauth2",
+ GravitinoAuthProvider.FORWARD_SESSION_USER_KEY, "true"));
+
+ ConnectorSession session = mock(ConnectorSession.class);
+ when(session.getUser()).thenReturn("alice");
+
when(session.getIdentity()).thenReturn(ConnectorIdentity.forUser("alice").build());
+
+ assertThrows(
+ TrinoException.class, () ->
GravitinoAuthProvider.buildForSession(config, session));
+ }
Review Comment:
This test only asserts that an exception of type `TrinoException` is thrown
when the forwarded token is missing, but it doesn't validate the error code or
message. Adding assertions for the expected error code (e.g.
`GRAVITINO_ILLEGAL_ARGUMENT`) and that the message mentions the missing
credential key would better lock in the intended behavior and prevent
regressions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
