yuqi1129 commented on code in PR #13928:
URL: https://github.com/apache/gravitino/pull/13928#discussion_r4310989446
##########
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:
nit: the env fallback is only tested in basic mode. The scenario in #13927
is path-based (Iceberg REST load table), which goes through
`createCredentialsWithSessionPolicy`. A path-context test asserting
`RoleArn=<env role>` and a `Policy=` in the STS request would cover it.
##########
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:
nit: equivalent spellings of the same location are counted as two roots, so
requests that one SAS could serve are rejected. I checked this with a local
probe; each of these fails with "2 roots":
- `.../wh/t` together with `.../wh/t/` (even with `.../wh/t/data` also
present)
- the same path once as `abfs://` and once as `abfss://`
- account host case: `@Account.dfs...` vs `@account.dfs...`
None of these is a regression, since they all failed before too. The
likeliest case in practice is `write.data.path` with a trailing slash that
equals the table location. Deduplicating by `(container, trimSlashes(path))`
before computing roots would cover all three.
Also, the message says "N roots: ..." but then prints every path, not the
roots.
##########
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:
nit: two checks in `findCredentialLocation` are not covered. I verified each
by mutating the code locally:
- Removing `!storageAccountName.equalsIgnoreCase(location.getAccountName())`
still passes, because in this test the first path is in the configured account
and the "different from first" check rejects the request instead. A case with a
single path in another account (for example,
`abfss://[email protected]/wh/t`) would cover it.
- Changing `isDescendantOf` to `path.startsWith(parentPath)` (without the
`/`) also passes. A case asserting that `.../wh/t` plus `.../wh/t2` is rejected
would cover the sibling-prefix boundary.
##########
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:
nit: a blank `s3-role-arn` still fails here, while `getValidatedRoleArn`
already treats a blank value as unset and falls back.
This can happen with the Docker image:
`dev/docker/iceberg-rest-server/rewrite_config.py` writes the key whenever the
env var exists, even if it is empty. For example,
`GRAVITINO_ICEBERG_REST_S3_ROLE_ARN=${S3_ROLE_ARN}` with `S3_ROLE_ARN` unset
produces `s3-role-arn = `, and `s3RoleArn()` then throws on the not-blank check.
Suggestion:
```java
this.roleArn =
StringUtils.isNotBlank(properties.get(S3CredentialConfig.S3_ROLE_ARN.getKey()))
? s3CredentialConfig.s3RoleArn()
: null;
```
Separately, a missing ARN is now reported on the first `generate()` instead
of at `initialize()`. That follows from the fallback design and is fine, but it
may be worth a sentence in the PR description.
--
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]
