Skip to content

Commit 5e72665

Browse files
author
Sreesh Maheshwar
committed
Iceberg Azure credential vending
1 parent 4a74c68 commit 5e72665

File tree

5 files changed

+134
-19
lines changed

5 files changed

+134
-19
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: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
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+
public static final String EXTRA_SAS_TOKEN_PROPERTY_PREFIX = "internal$account_sas$";
19+
20+
private AzureFileSystemConstants() {}
21+
}

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

Lines changed: 6 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,9 @@ public static HttpClient createAzureHttpClient(ConnectionProvider connectionProv
143143
.eventLoopGroup(eventLoopGroup)
144144
.build();
145145
}
146+
147+
private static AzureAuth withVendedAuth(ConnectorIdentity identity, AzureAuth defaultAuth)
148+
{
149+
return identity.getExtraCredentials().isEmpty() ? defaultAuth : new AzureVendedAuth(identity.getExtraCredentials(), defaultAuth);
150+
}
146151
}

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

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
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+
21+
public final class AzureVendedAuth
22+
implements AzureAuth
23+
{
24+
private final Map<String, String> accountSasTokens;
25+
private final AzureAuth fallbackAuth;
26+
27+
public AzureVendedAuth(Map<String, String> accountSasTokens, AzureAuth fallbackAuth)
28+
{
29+
this.accountSasTokens = accountSasTokens;
30+
this.fallbackAuth = fallbackAuth;
31+
}
32+
33+
@Override
34+
public void setAuth(String storageAccount, BlobContainerClientBuilder builder)
35+
{
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+
}
43+
}
44+
45+
@Override
46+
public void setAuth(String storageAccount, DataLakeServiceClientBuilder builder)
47+
{
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+
}
55+
}
56+
}

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

Lines changed: 50 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,11 @@
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;
2425

26+
import static io.trino.filesystem.azure.AzureFileSystemConstants.EXTRA_SAS_TOKEN_PROPERTY_PREFIX;
2527
import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_ACCESS_KEY_PROPERTY;
2628
import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_SECRET_KEY_PROPERTY;
2729
import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_SESSION_TOKEN_PROPERTY;
@@ -34,6 +36,8 @@ public class IcebergRestCatalogFileSystemFactory
3436
private static final String VENDED_S3_SECRET_KEY = "s3.secret-access-key";
3537
private static final String VENDED_S3_SESSION_TOKEN = "s3.session-token";
3638

39+
private static final String VENDED_ADLS_SAS_TOKEN_PREFIX = "adls.sas-token.";
40+
3741
private final TrinoFileSystemFactory fileSystemFactory;
3842
private final boolean vendedCredentialsEnabled;
3943

@@ -47,25 +51,54 @@ public IcebergRestCatalogFileSystemFactory(TrinoFileSystemFactory fileSystemFact
4751
@Override
4852
public TrinoFileSystem create(ConnectorIdentity identity, Map<String, String> fileIoProperties)
4953
{
50-
if (vendedCredentialsEnabled &&
51-
fileIoProperties.containsKey(VENDED_S3_ACCESS_KEY) &&
52-
fileIoProperties.containsKey(VENDED_S3_SECRET_KEY) &&
53-
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);
54+
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+
}
6784
}
6885

6986
return fileSystemFactory.create(identity);
7087
}
88+
89+
private static Map<String, String> getAzureCredentials(Map<String, String> fileIoProperties)
90+
{
91+
ImmutableMap.Builder<String, String> azureCredentialBuilder = ImmutableMap.builder();
92+
93+
PropertyUtil.propertiesWithPrefix(fileIoProperties, VENDED_ADLS_SAS_TOKEN_PREFIX)
94+
.forEach((host, token) -> {
95+
String storageAccount = host.contains(".") ? host.substring(0, host.indexOf('.')) : host;
96+
97+
if (!storageAccount.isEmpty() && !token.isEmpty()) {
98+
azureCredentialBuilder.put(EXTRA_SAS_TOKEN_PROPERTY_PREFIX + storageAccount, token);
99+
}
100+
});
101+
102+
return azureCredentialBuilder.build();
103+
}
71104
}

0 commit comments

Comments
 (0)