Skip to content

Commit 44d5a9b

Browse files
author
Sreesh Maheshwar
committed
Improvements
1 parent 5e72665 commit 44d5a9b

File tree

4 files changed

+62
-50
lines changed

4 files changed

+62
-50
lines changed

‎lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureFileSystemConstants.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,14 @@
1515

1616
public final class AzureFileSystemConstants
1717
{
18+
/**
19+
* Internal property enabling {@link AzureVendedAuth} on the filesystem when set to true.
20+
*/
21+
public static final String EXTRA_USE_VENDED_TOKEN = "internal$use_vended_token";
22+
23+
/**
24+
* Internal prefix for SAS token property keys, mapping storage accounts to their SAS tokens.
25+
*/
1826
public static final String EXTRA_SAS_TOKEN_PROPERTY_PREFIX = "internal$account_sas$";
1927

2028
private AzureFileSystemConstants() {}

‎lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureFileSystemFactory.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,10 @@ public static HttpClient createAzureHttpClient(ConnectionProvider connectionProv
146146

147147
private static AzureAuth withVendedAuth(ConnectorIdentity identity, AzureAuth defaultAuth)
148148
{
149-
return identity.getExtraCredentials().isEmpty() ? defaultAuth : new AzureVendedAuth(identity.getExtraCredentials(), defaultAuth);
149+
if (identity.getExtraCredentials().containsKey(AzureFileSystemConstants.EXTRA_USE_VENDED_TOKEN) &&
150+
identity.getExtraCredentials().get(AzureFileSystemConstants.EXTRA_USE_VENDED_TOKEN).equalsIgnoreCase("true")) {
151+
return new AzureVendedAuth(identity.getExtraCredentials(), defaultAuth);
152+
}
153+
return defaultAuth;
150154
}
151155
}

‎lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureVendedAuth.java‎

Lines changed: 13 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -17,40 +17,36 @@
1717
import com.azure.storage.file.datalake.DataLakeServiceClientBuilder;
1818

1919
import java.util.Map;
20+
import java.util.Optional;
2021

2122
public final class AzureVendedAuth
2223
implements AzureAuth
2324
{
24-
private final Map<String, String> accountSasTokens;
25+
private final Map<String, String> sasTokens;
2526
private final AzureAuth fallbackAuth;
2627

27-
public AzureVendedAuth(Map<String, String> accountSasTokens, AzureAuth fallbackAuth)
28+
public AzureVendedAuth(Map<String, String> sasTokens, AzureAuth fallbackAuth)
2829
{
29-
this.accountSasTokens = accountSasTokens;
30+
this.sasTokens = sasTokens;
3031
this.fallbackAuth = fallbackAuth;
3132
}
3233

3334
@Override
3435
public void setAuth(String storageAccount, BlobContainerClientBuilder builder)
3536
{
36-
String sasToken = accountSasTokens.get(AzureFileSystemConstants.EXTRA_SAS_TOKEN_PROPERTY_PREFIX + storageAccount);
37-
if (sasToken == null) {
38-
fallbackAuth.setAuth(storageAccount, builder);
39-
}
40-
else {
41-
builder.sasToken(sasToken);
42-
}
37+
getSasToken(storageAccount)
38+
.ifPresentOrElse(builder::sasToken, () -> fallbackAuth.setAuth(storageAccount, builder));
4339
}
4440

4541
@Override
4642
public void setAuth(String storageAccount, DataLakeServiceClientBuilder builder)
4743
{
48-
String sasToken = accountSasTokens.get(AzureFileSystemConstants.EXTRA_SAS_TOKEN_PROPERTY_PREFIX + storageAccount);
49-
if (sasToken == null) {
50-
fallbackAuth.setAuth(storageAccount, builder);
51-
}
52-
else {
53-
builder.sasToken(sasToken);
54-
}
44+
getSasToken(storageAccount)
45+
.ifPresentOrElse(builder::sasToken, () -> fallbackAuth.setAuth(storageAccount, builder));
46+
}
47+
48+
public Optional<String> getSasToken(String storageAccount)
49+
{
50+
return Optional.ofNullable(sasTokens.get(AzureFileSystemConstants.EXTRA_SAS_TOKEN_PROPERTY_PREFIX + storageAccount));
5551
}
5652
}

‎plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogFileSystemFactory.java‎

Lines changed: 36 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,10 @@
2222
import org.apache.iceberg.util.PropertyUtil;
2323

2424
import java.util.Map;
25+
import java.util.Optional;
2526

2627
import static io.trino.filesystem.azure.AzureFileSystemConstants.EXTRA_SAS_TOKEN_PROPERTY_PREFIX;
28+
import static io.trino.filesystem.azure.AzureFileSystemConstants.EXTRA_USE_VENDED_TOKEN;
2729
import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_ACCESS_KEY_PROPERTY;
2830
import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_SECRET_KEY_PROPERTY;
2931
import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_SESSION_TOKEN_PROPERTY;
@@ -52,53 +54,55 @@ public IcebergRestCatalogFileSystemFactory(TrinoFileSystemFactory fileSystemFact
5254
public TrinoFileSystem create(ConnectorIdentity identity, Map<String, String> fileIoProperties)
5355
{
5456
if (vendedCredentialsEnabled) {
55-
ImmutableMap.Builder<String, String> overriddenCredentialsBuilder = ImmutableMap.builder();
56-
57-
if (fileIoProperties.containsKey(VENDED_S3_ACCESS_KEY) &&
58-
fileIoProperties.containsKey(VENDED_S3_SECRET_KEY) &&
59-
fileIoProperties.containsKey(VENDED_S3_SESSION_TOKEN)) {
60-
// S3 vended credentials
61-
overriddenCredentialsBuilder
62-
.put(EXTRA_CREDENTIALS_ACCESS_KEY_PROPERTY, fileIoProperties.get(VENDED_S3_ACCESS_KEY))
63-
.put(EXTRA_CREDENTIALS_SECRET_KEY_PROPERTY, fileIoProperties.get(VENDED_S3_SECRET_KEY))
64-
.put(EXTRA_CREDENTIALS_SESSION_TOKEN_PROPERTY, fileIoProperties.get(VENDED_S3_SESSION_TOKEN));
65-
}
66-
else {
67-
// Azure vended credentials
68-
overriddenCredentialsBuilder.putAll(getAzureCredentials(fileIoProperties));
69-
}
70-
71-
Map<String, String> overriddenCredentials = overriddenCredentialsBuilder.buildOrThrow();
72-
if (!overriddenCredentials.isEmpty()) {
73-
// Do not include original credentials as they should not be used in vended mode
74-
ConnectorIdentity identityWithExtraCredentials = ConnectorIdentity
75-
.forUser(identity.getUser())
76-
.withGroups(identity.getGroups())
77-
.withPrincipal(identity.getPrincipal())
78-
.withEnabledSystemRoles(identity.getEnabledSystemRoles())
79-
.withConnectorRole(identity.getConnectorRole())
80-
.withExtraCredentials(overriddenCredentials).build();
81-
82-
return fileSystemFactory.create(identityWithExtraCredentials);
83-
}
57+
return fileSystemFactory.create(
58+
getVendedS3Identity(identity, fileIoProperties)
59+
.or(() -> getVendedAzureIdentity(identity, fileIoProperties))
60+
.orElse(identity));
8461
}
8562

8663
return fileSystemFactory.create(identity);
8764
}
8865

89-
private static Map<String, String> getAzureCredentials(Map<String, String> fileIoProperties)
66+
private static Optional<ConnectorIdentity> getVendedS3Identity(ConnectorIdentity identity, Map<String, String> fileIoProperties)
9067
{
91-
ImmutableMap.Builder<String, String> azureCredentialBuilder = ImmutableMap.builder();
68+
if (fileIoProperties.containsKey(VENDED_S3_ACCESS_KEY) &&
69+
fileIoProperties.containsKey(VENDED_S3_SECRET_KEY) &&
70+
fileIoProperties.containsKey(VENDED_S3_SESSION_TOKEN)) {
71+
return Optional.of(getVendedIdentity(identity, ImmutableMap.<String, String>builder()
72+
.put(EXTRA_CREDENTIALS_ACCESS_KEY_PROPERTY, fileIoProperties.get(VENDED_S3_ACCESS_KEY))
73+
.put(EXTRA_CREDENTIALS_SECRET_KEY_PROPERTY, fileIoProperties.get(VENDED_S3_SECRET_KEY))
74+
.put(EXTRA_CREDENTIALS_SESSION_TOKEN_PROPERTY, fileIoProperties.get(VENDED_S3_SESSION_TOKEN))
75+
.buildOrThrow()));
76+
}
77+
return Optional.empty();
78+
}
9279

80+
private static Optional<ConnectorIdentity> getVendedAzureIdentity(ConnectorIdentity identity, Map<String, String> fileIoProperties)
81+
{
82+
ImmutableMap.Builder<String, String> azureCredentialBuilder = ImmutableMap.builder();
9383
PropertyUtil.propertiesWithPrefix(fileIoProperties, VENDED_ADLS_SAS_TOKEN_PREFIX)
9484
.forEach((host, token) -> {
9585
String storageAccount = host.contains(".") ? host.substring(0, host.indexOf('.')) : host;
9686

9787
if (!storageAccount.isEmpty() && !token.isEmpty()) {
9888
azureCredentialBuilder.put(EXTRA_SAS_TOKEN_PROPERTY_PREFIX + storageAccount, token);
89+
azureCredentialBuilder.put(EXTRA_USE_VENDED_TOKEN, "true");
9990
}
10091
});
10192

102-
return azureCredentialBuilder.build();
93+
Map<String, String> azureCredentials = azureCredentialBuilder.buildKeepingLast();
94+
return azureCredentials.isEmpty() ? Optional.empty() : Optional.of(getVendedIdentity(identity, azureCredentials));
95+
}
96+
97+
private static ConnectorIdentity getVendedIdentity(ConnectorIdentity identity, Map<String, String> extraCredentials)
98+
{
99+
// Do not include original credentials as they should not be used in vended mode
100+
return ConnectorIdentity.forUser(identity.getUser())
101+
.withGroups(identity.getGroups())
102+
.withPrincipal(identity.getPrincipal())
103+
.withEnabledSystemRoles(identity.getEnabledSystemRoles())
104+
.withConnectorRole(identity.getConnectorRole())
105+
.withExtraCredentials(extraCredentials)
106+
.build();
103107
}
104108
}

0 commit comments

Comments
 (0)