Skip to content

Commit 8c08a0d

Browse files
smaheshwar-pltrsmaheshwar-pltr
committed
Iceberg Azure credential vending
1 parent 4a74c68 commit 8c08a0d

File tree

6 files changed

+150
-18
lines changed

6 files changed

+150
-18
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
import com.azure.storage.file.datalake.DataLakeServiceClientBuilder;
1818

1919
public sealed interface AzureAuth
20-
permits AzureAuthAccessKey, AzureAuthDefault, AzureAuthOauth
20+
permits AzureAuthAccessKey, AzureAuthDefault, AzureAuthOauth, AzureVendedAuth
2121
{
2222
void setAuth(String storageAccount, BlobContainerClientBuilder builder);
2323

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package io.trino.filesystem.azure;
15+
16+
public final class AzureFileSystemConstants
17+
{
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+
*/
26+
public static final String EXTRA_SAS_TOKEN_PROPERTY_PREFIX = "internal$account_sas$";
27+
28+
private AzureFileSystemConstants() {}
29+
}

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,7 @@ public void destroy()
127127
@Override
128128
public TrinoFileSystem create(ConnectorIdentity identity)
129129
{
130-
return new AzureFileSystem(httpClient, concurrencyPolicy, uploadExecutor, tracingOptions, auth, endpoint, readBlockSize, writeBlockSize, maxWriteConcurrency, maxSingleUploadSize, multipart);
130+
return new AzureFileSystem(httpClient, concurrencyPolicy, uploadExecutor, tracingOptions, withVendedAuth(identity, auth), endpoint, readBlockSize, writeBlockSize, maxWriteConcurrency, maxSingleUploadSize, multipart);
131131
}
132132

133133
public static HttpClient createAzureHttpClient(ConnectionProvider connectionProvider, EventLoopGroup eventLoopGroup, HttpClientOptions clientOptions)
@@ -143,4 +143,13 @@ public static HttpClient createAzureHttpClient(ConnectionProvider connectionProv
143143
.eventLoopGroup(eventLoopGroup)
144144
.build();
145145
}
146+
147+
private static AzureAuth withVendedAuth(ConnectorIdentity identity, AzureAuth defaultAuth)
148+
{
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;
154+
}
146155
}

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

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package io.trino.filesystem.azure;
15+
16+
import com.azure.storage.blob.BlobContainerClientBuilder;
17+
import com.azure.storage.file.datalake.DataLakeServiceClientBuilder;
18+
19+
import java.util.Map;
20+
import java.util.Optional;
21+
22+
public final class AzureVendedAuth
23+
implements AzureAuth
24+
{
25+
private final Map<String, String> sasTokens;
26+
private final AzureAuth fallbackAuth;
27+
28+
public AzureVendedAuth(Map<String, String> sasTokens, AzureAuth fallbackAuth)
29+
{
30+
this.sasTokens = sasTokens;
31+
this.fallbackAuth = fallbackAuth;
32+
}
33+
34+
@Override
35+
public void setAuth(String storageAccount, BlobContainerClientBuilder builder)
36+
{
37+
getSasToken(storageAccount)
38+
.ifPresentOrElse(builder::sasToken, () -> fallbackAuth.setAuth(storageAccount, builder));
39+
}
40+
41+
@Override
42+
public void setAuth(String storageAccount, DataLakeServiceClientBuilder builder)
43+
{
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));
51+
}
52+
}

‎plugin/trino-iceberg/pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,11 @@
122122
<artifactId>trino-filesystem-manager</artifactId>
123123
</dependency>
124124

125+
<dependency>
126+
<groupId>io.trino</groupId>
127+
<artifactId>trino-filesystem-azure</artifactId>
128+
</dependency>
129+
125130
<dependency>
126131
<groupId>io.trino</groupId>
127132
<artifactId>trino-filesystem-s3</artifactId>

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

Lines changed: 53 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,13 @@
1919
import io.trino.filesystem.TrinoFileSystemFactory;
2020
import io.trino.plugin.iceberg.IcebergFileSystemFactory;
2121
import io.trino.spi.security.ConnectorIdentity;
22+
import org.apache.iceberg.util.PropertyUtil;
2223

2324
import java.util.Map;
25+
import java.util.Optional;
2426

27+
import static io.trino.filesystem.azure.AzureFileSystemConstants.EXTRA_SAS_TOKEN_PROPERTY_PREFIX;
28+
import static io.trino.filesystem.azure.AzureFileSystemConstants.EXTRA_USE_VENDED_TOKEN;
2529
import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_ACCESS_KEY_PROPERTY;
2630
import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_SECRET_KEY_PROPERTY;
2731
import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_SESSION_TOKEN_PROPERTY;
@@ -34,6 +38,8 @@ public class IcebergRestCatalogFileSystemFactory
3438
private static final String VENDED_S3_SECRET_KEY = "s3.secret-access-key";
3539
private static final String VENDED_S3_SESSION_TOKEN = "s3.session-token";
3640

41+
private static final String VENDED_ADLS_SAS_TOKEN_PREFIX = "adls.sas-token.";
42+
3743
private final TrinoFileSystemFactory fileSystemFactory;
3844
private final boolean vendedCredentialsEnabled;
3945

@@ -47,25 +53,56 @@ public IcebergRestCatalogFileSystemFactory(TrinoFileSystemFactory fileSystemFact
4753
@Override
4854
public TrinoFileSystem create(ConnectorIdentity identity, Map<String, String> fileIoProperties)
4955
{
50-
if (vendedCredentialsEnabled &&
51-
fileIoProperties.containsKey(VENDED_S3_ACCESS_KEY) &&
56+
if (vendedCredentialsEnabled) {
57+
return fileSystemFactory.create(
58+
getVendedS3Identity(identity, fileIoProperties)
59+
.or(() -> getVendedAzureIdentity(identity, fileIoProperties))
60+
.orElse(identity));
61+
}
62+
63+
return fileSystemFactory.create(identity);
64+
}
65+
66+
private static Optional<ConnectorIdentity> getVendedS3Identity(ConnectorIdentity identity, Map<String, String> fileIoProperties)
67+
{
68+
if (fileIoProperties.containsKey(VENDED_S3_ACCESS_KEY) &&
5269
fileIoProperties.containsKey(VENDED_S3_SECRET_KEY) &&
5370
fileIoProperties.containsKey(VENDED_S3_SESSION_TOKEN)) {
54-
// Do not include original credentials as they should not be used in vended mode
55-
ConnectorIdentity identityWithExtraCredentials = ConnectorIdentity.forUser(identity.getUser())
56-
.withGroups(identity.getGroups())
57-
.withPrincipal(identity.getPrincipal())
58-
.withEnabledSystemRoles(identity.getEnabledSystemRoles())
59-
.withConnectorRole(identity.getConnectorRole())
60-
.withExtraCredentials(ImmutableMap.<String, String>builder()
61-
.put(EXTRA_CREDENTIALS_ACCESS_KEY_PROPERTY, fileIoProperties.get(VENDED_S3_ACCESS_KEY))
62-
.put(EXTRA_CREDENTIALS_SECRET_KEY_PROPERTY, fileIoProperties.get(VENDED_S3_SECRET_KEY))
63-
.put(EXTRA_CREDENTIALS_SESSION_TOKEN_PROPERTY, fileIoProperties.get(VENDED_S3_SESSION_TOKEN))
64-
.buildOrThrow())
65-
.build();
66-
return fileSystemFactory.create(identityWithExtraCredentials);
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()));
6776
}
77+
return Optional.empty();
78+
}
6879

69-
return fileSystemFactory.create(identity);
80+
private static Optional<ConnectorIdentity> getVendedAzureIdentity(ConnectorIdentity identity, Map<String, String> fileIoProperties)
81+
{
82+
ImmutableMap.Builder<String, String> azureCredentialBuilder = ImmutableMap.builder();
83+
PropertyUtil.propertiesWithPrefix(fileIoProperties, VENDED_ADLS_SAS_TOKEN_PREFIX)
84+
.forEach((host, token) -> {
85+
String storageAccount = host.contains(".") ? host.substring(0, host.indexOf('.')) : host;
86+
87+
if (!storageAccount.isEmpty() && !token.isEmpty()) {
88+
azureCredentialBuilder.put(EXTRA_SAS_TOKEN_PROPERTY_PREFIX + storageAccount, token);
89+
azureCredentialBuilder.put(EXTRA_USE_VENDED_TOKEN, "true");
90+
}
91+
});
92+
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();
70107
}
71108
}

0 commit comments

Comments
 (0)