roryqi commented on code in PR #13928:
URL: https://github.com/apache/gravitino/pull/13928#discussion_r4311255908
##########
bundles/azure/src/main/java/org/apache/gravitino/abs/credential/ADLSTokenGenerator.java:
##########
@@ -125,6 +106,76 @@ public ADLSTokenCredential generate(CredentialContext
context) {
locationParts.getAccountName(), sasToken,
expiry.toInstant().toEpochMilli());
}
+ UserDelegationKey getUserDelegationKey(OffsetDateTime start, OffsetDateTime
expiry) {
+ ClientSecretCredential clientSecretCredential =
+ new ClientSecretCredentialBuilder()
+ .tenantId(tenantId)
+ .clientId(clientId)
+ .clientSecret(clientSecret)
+ .build();
+
+ DataLakeServiceClient dataLakeServiceClient =
+ new DataLakeServiceClientBuilder()
+ .endpoint(endpoint)
+ .credential(clientSecretCredential)
+ .buildClient();
+ return dataLakeServiceClient.getUserDelegationKey(start, expiry);
+ }
+
+ private ADLSLocationUtils.ADLSLocationParts
findCredentialLocation(Set<String> paths) {
+ if (paths.isEmpty()) {
+ throw new IllegalArgumentException(
+ "ADLS credential context should contain at least one path");
+ }
+
+ List<ADLSLocationUtils.ADLSLocationParts> locations =
+
paths.stream().map(ADLSLocationUtils::parseLocation).collect(Collectors.toList());
+ ADLSLocationUtils.ADLSLocationParts first = locations.get(0);
+ boolean hasDifferentScope =
+ locations.stream()
+ .anyMatch(
+ location ->
+
!storageAccountName.equalsIgnoreCase(location.getAccountName())
+ ||
!first.getAccountName().equalsIgnoreCase(location.getAccountName())
+ ||
!first.getContainer().equalsIgnoreCase(location.getContainer()));
+ if (hasDifferentScope) {
+ throw new IllegalArgumentException(
+ "ADLS paths must belong to the configured storage account and a
single container, but found: "
+ + paths);
+ }
+
+ List<ADLSLocationUtils.ADLSLocationParts> rootLocations =
+ locations.stream()
+ .filter(
+ location ->
+ locations.stream()
+ .noneMatch(
+ possibleParent ->
+ possibleParent != location
+ && isDescendantOf(location,
possibleParent)))
+ .collect(Collectors.toList());
+ if (rootLocations.size() != 1) {
Review Comment:
Fixed in 672e3ae6ac. Equivalent locations are now deduplicated by container
and slash-trimmed path after account/container validation, so abfs/abfss,
account-name case, and trailing-slash variants collapse to one root. I also
changed the error wording to say that the listed values are input paths and
added a regression test covering all three variants.
##########
bundles/aws/src/main/java/org/apache/gravitino/s3/credential/AwsIrsaCredentialGenerator.java:
##########
@@ -64,12 +66,24 @@ public class AwsIrsaCredentialGenerator implements
CredentialGenerator<AwsIrsaCr
private StsClient basicModeStsClient;
private Map<String, StsAssumeRoleWithWebIdentityCredentialsProvider>
basicModeProviders;
+ /** Creates an AWS IRSA credential generator. */
+ public AwsIrsaCredentialGenerator() {
+ this(() -> System.getenv("AWS_ROLE_ARN"));
+ }
+
+ AwsIrsaCredentialGenerator(Supplier<String> environmentRoleArnSupplier) {
+ this.environmentRoleArnSupplier =
Objects.requireNonNull(environmentRoleArnSupplier);
+ }
+
@Override
public void initialize(Map<String, String> properties) {
this.tokenSource = WebIdentityTokenSources.create(properties);
S3CredentialConfig s3CredentialConfig = new S3CredentialConfig(properties);
- this.roleArn = s3CredentialConfig.s3RoleArn();
+ this.roleArn =
+ properties.containsKey(S3CredentialConfig.S3_ROLE_ARN.getKey())
Review Comment:
Fixed in 672e3ae6ac. A blank configured role ARN is now treated as unset
with StringUtils.isNotBlank and falls back to AWS_ROLE_ARN. The test covers the
blank Docker-style value, and the PR description now notes that a missing ARN
is reported when credentials are first requested.
##########
bundles/azure/src/test/java/org/apache/gravitino/abs/credential/TestADLSTokenGenerator.java:
##########
@@ -0,0 +1,139 @@
+/*
+ * 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.abs.credential;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import com.azure.storage.file.datalake.models.UserDelegationKey;
+import java.time.OffsetDateTime;
+import java.util.Arrays;
+import java.util.Base64;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import org.apache.gravitino.credential.ADLSTokenCredential;
+import org.apache.gravitino.credential.PathBasedCredentialContext;
+import org.apache.gravitino.storage.AzureProperties;
+import org.junit.jupiter.api.Test;
+
+class TestADLSTokenGenerator {
+
+ @Test
+ void generatesTokenForParentDirectoryOfMultiplePaths() throws Exception {
+ try (ADLSTokenGenerator generator = new TestableADLSTokenGenerator()) {
+ generator.initialize(createProperties());
+ PathBasedCredentialContext context =
+ new PathBasedCredentialContext(
+ "test-user",
+ Collections.emptySet(),
+ new HashSet<>(
+ Arrays.asList(
+
"abfss://[email protected]/warehouse/table",
+
"abfss://[email protected]/warehouse/table/data",
+
"abfss://[email protected]/warehouse/table/metadata")));
+
+ ADLSTokenCredential credential = generator.generate(context);
+
+ assertEquals("account", credential.accountName());
+ assertTrue(credential.sasToken().contains("sr=d"));
+ assertTrue(credential.sasToken().contains("sdd=2"));
+ }
+ }
+
+ @Test
+ void rejectsPathsWithoutProvidedCommonRoot() throws Exception {
+ String first =
"abfss://[email protected]/warehouse/table/data";
+ String second =
"abfss://[email protected]/warehouse/table/metadata";
+ try (ADLSTokenGenerator generator = new TestableADLSTokenGenerator()) {
+ generator.initialize(createProperties());
+ PathBasedCredentialContext context =
+ new PathBasedCredentialContext(
+ "test-user", Collections.emptySet(), new
HashSet<>(Arrays.asList(first, second)));
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () ->
generator.generate(context));
+
+ assertTrue(exception.getMessage().contains("exactly one root location"));
+ assertTrue(exception.getMessage().contains(first));
+ assertTrue(exception.getMessage().contains(second));
+ }
+ }
+
+ @Test
+ void rejectsPathsFromDifferentContainers() throws Exception {
+ String first =
"abfss://[email protected]/warehouse/table";
+ String second =
"abfss://[email protected]/warehouse/table/data";
+ try (ADLSTokenGenerator generator = new TestableADLSTokenGenerator()) {
+ generator.initialize(createProperties());
+ PathBasedCredentialContext context =
+ new PathBasedCredentialContext(
+ "test-user", Collections.emptySet(), new
HashSet<>(Arrays.asList(first, second)));
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () ->
generator.generate(context));
+
+ assertTrue(exception.getMessage().contains(first));
+ assertTrue(exception.getMessage().contains(second));
+ }
+ }
+
+ @Test
+ void rejectsPathsFromDifferentStorageAccounts() throws Exception {
Review Comment:
Fixed in 672e3ae6ac. Added tests for a single path from an account different
from the configured account and for sibling paths whose names share a string
prefix (table and table2).
##########
bundles/aws/src/test/java/org/apache/gravitino/s3/credential/TestAwsIrsaCredentialGenerator.java:
##########
@@ -84,6 +84,39 @@ void
basicModeUsesConfiguredTokenSourceWithoutSessionPolicy(@TempDir Path dir)
assertFalse(body.contains("Policy="));
}
+ @Test
+ void basicModeFallsBackToEnvironmentRoleArn(@TempDir Path dir) throws
IOException {
Review Comment:
Fixed in 672e3ae6ac. Added a path-based fallback test that removes the
configured ARN and asserts both the environment RoleArn and a Policy parameter
in the STS request.
--
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]
