diff --git a/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureAuth.java b/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureAuth.java index 3d6da6e4ed99..6c9ab2e1c098 100644 --- a/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureAuth.java +++ b/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureAuth.java @@ -17,7 +17,7 @@ import com.azure.storage.file.datalake.DataLakeServiceClientBuilder; public sealed interface AzureAuth - permits AzureAuthAccessKey, AzureAuthDefault, AzureAuthOauth + permits AzureAuthAccessKey, AzureAuthDefault, AzureAuthOauth, AzureVendedAuth { void setAuth(String storageAccount, BlobContainerClientBuilder builder); diff --git a/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureFileSystemConstants.java b/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureFileSystemConstants.java new file mode 100644 index 000000000000..09f44435d527 --- /dev/null +++ b/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureFileSystemConstants.java @@ -0,0 +1,29 @@ +/* + * Licensed 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 io.trino.filesystem.azure; + +public final class AzureFileSystemConstants +{ + /** + * Internal property enabling {@link AzureVendedAuth} on the filesystem when set to true. + */ + public static final String EXTRA_USE_VENDED_TOKEN = "internal$use_vended_token"; + + /** + * Internal prefix for SAS token property keys, mapping storage accounts to their SAS tokens. + */ + public static final String EXTRA_SAS_TOKEN_PROPERTY_PREFIX = "internal$account_sas$"; + + private AzureFileSystemConstants() {} +} diff --git a/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureFileSystemFactory.java b/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureFileSystemFactory.java index 0c1152235b66..4c83a129c2eb 100644 --- a/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureFileSystemFactory.java +++ b/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureFileSystemFactory.java @@ -127,7 +127,7 @@ public void destroy() @Override public TrinoFileSystem create(ConnectorIdentity identity) { - return new AzureFileSystem(httpClient, concurrencyPolicy, uploadExecutor, tracingOptions, auth, endpoint, readBlockSize, writeBlockSize, maxWriteConcurrency, maxSingleUploadSize, multipart); + return new AzureFileSystem(httpClient, concurrencyPolicy, uploadExecutor, tracingOptions, withVendedAuth(identity, auth), endpoint, readBlockSize, writeBlockSize, maxWriteConcurrency, maxSingleUploadSize, multipart); } public static HttpClient createAzureHttpClient(ConnectionProvider connectionProvider, EventLoopGroup eventLoopGroup, HttpClientOptions clientOptions) @@ -143,4 +143,13 @@ public static HttpClient createAzureHttpClient(ConnectionProvider connectionProv .eventLoopGroup(eventLoopGroup) .build(); } + + private static AzureAuth withVendedAuth(ConnectorIdentity identity, AzureAuth defaultAuth) + { + if (identity.getExtraCredentials().containsKey(AzureFileSystemConstants.EXTRA_USE_VENDED_TOKEN) && + identity.getExtraCredentials().get(AzureFileSystemConstants.EXTRA_USE_VENDED_TOKEN).equalsIgnoreCase("true")) { + return new AzureVendedAuth(identity.getExtraCredentials(), defaultAuth); + } + return defaultAuth; + } } diff --git a/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureVendedAuth.java b/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureVendedAuth.java new file mode 100644 index 000000000000..a41270472111 --- /dev/null +++ b/lib/trino-filesystem-azure/src/main/java/io/trino/filesystem/azure/AzureVendedAuth.java @@ -0,0 +1,52 @@ +/* + * Licensed 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 io.trino.filesystem.azure; + +import com.azure.storage.blob.BlobContainerClientBuilder; +import com.azure.storage.file.datalake.DataLakeServiceClientBuilder; + +import java.util.Map; +import java.util.Optional; + +public final class AzureVendedAuth + implements AzureAuth +{ + private final Map sasTokens; + private final AzureAuth fallbackAuth; + + public AzureVendedAuth(Map sasTokens, AzureAuth fallbackAuth) + { + this.sasTokens = sasTokens; + this.fallbackAuth = fallbackAuth; + } + + @Override + public void setAuth(String storageAccount, BlobContainerClientBuilder builder) + { + getSasToken(storageAccount) + .ifPresentOrElse(builder::sasToken, () -> fallbackAuth.setAuth(storageAccount, builder)); + } + + @Override + public void setAuth(String storageAccount, DataLakeServiceClientBuilder builder) + { + getSasToken(storageAccount) + .ifPresentOrElse(builder::sasToken, () -> fallbackAuth.setAuth(storageAccount, builder)); + } + + public Optional getSasToken(String storageAccount) + { + return Optional.ofNullable(sasTokens.get(AzureFileSystemConstants.EXTRA_SAS_TOKEN_PROPERTY_PREFIX + storageAccount)); + } +} diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogFileSystemFactory.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogFileSystemFactory.java index b5e777c7ee15..e772d6677a98 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogFileSystemFactory.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/rest/IcebergRestCatalogFileSystemFactory.java @@ -19,9 +19,13 @@ import io.trino.filesystem.TrinoFileSystemFactory; import io.trino.plugin.iceberg.IcebergFileSystemFactory; import io.trino.spi.security.ConnectorIdentity; +import org.apache.iceberg.util.PropertyUtil; import java.util.Map; +import java.util.Optional; +import static io.trino.filesystem.azure.AzureFileSystemConstants.EXTRA_SAS_TOKEN_PROPERTY_PREFIX; +import static io.trino.filesystem.azure.AzureFileSystemConstants.EXTRA_USE_VENDED_TOKEN; import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_ACCESS_KEY_PROPERTY; import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_SECRET_KEY_PROPERTY; import static io.trino.filesystem.s3.S3FileSystemConstants.EXTRA_CREDENTIALS_SESSION_TOKEN_PROPERTY; @@ -34,6 +38,8 @@ public class IcebergRestCatalogFileSystemFactory private static final String VENDED_S3_SECRET_KEY = "s3.secret-access-key"; private static final String VENDED_S3_SESSION_TOKEN = "s3.session-token"; + private static final String VENDED_ADLS_SAS_TOKEN_PREFIX = "adls.sas-token."; + private final TrinoFileSystemFactory fileSystemFactory; private final boolean vendedCredentialsEnabled; @@ -47,25 +53,56 @@ public IcebergRestCatalogFileSystemFactory(TrinoFileSystemFactory fileSystemFact @Override public TrinoFileSystem create(ConnectorIdentity identity, Map fileIoProperties) { - if (vendedCredentialsEnabled && - fileIoProperties.containsKey(VENDED_S3_ACCESS_KEY) && + if (vendedCredentialsEnabled) { + return fileSystemFactory.create( + getVendedS3Identity(identity, fileIoProperties) + .or(() -> getVendedAzureIdentity(identity, fileIoProperties)) + .orElse(identity)); + } + + return fileSystemFactory.create(identity); + } + + private static Optional getVendedS3Identity(ConnectorIdentity identity, Map fileIoProperties) + { + if (fileIoProperties.containsKey(VENDED_S3_ACCESS_KEY) && fileIoProperties.containsKey(VENDED_S3_SECRET_KEY) && fileIoProperties.containsKey(VENDED_S3_SESSION_TOKEN)) { - // Do not include original credentials as they should not be used in vended mode - ConnectorIdentity identityWithExtraCredentials = ConnectorIdentity.forUser(identity.getUser()) - .withGroups(identity.getGroups()) - .withPrincipal(identity.getPrincipal()) - .withEnabledSystemRoles(identity.getEnabledSystemRoles()) - .withConnectorRole(identity.getConnectorRole()) - .withExtraCredentials(ImmutableMap.builder() - .put(EXTRA_CREDENTIALS_ACCESS_KEY_PROPERTY, fileIoProperties.get(VENDED_S3_ACCESS_KEY)) - .put(EXTRA_CREDENTIALS_SECRET_KEY_PROPERTY, fileIoProperties.get(VENDED_S3_SECRET_KEY)) - .put(EXTRA_CREDENTIALS_SESSION_TOKEN_PROPERTY, fileIoProperties.get(VENDED_S3_SESSION_TOKEN)) - .buildOrThrow()) - .build(); - return fileSystemFactory.create(identityWithExtraCredentials); + return Optional.of(getVendedIdentity(identity, ImmutableMap.builder() + .put(EXTRA_CREDENTIALS_ACCESS_KEY_PROPERTY, fileIoProperties.get(VENDED_S3_ACCESS_KEY)) + .put(EXTRA_CREDENTIALS_SECRET_KEY_PROPERTY, fileIoProperties.get(VENDED_S3_SECRET_KEY)) + .put(EXTRA_CREDENTIALS_SESSION_TOKEN_PROPERTY, fileIoProperties.get(VENDED_S3_SESSION_TOKEN)) + .buildOrThrow())); } + return Optional.empty(); + } - return fileSystemFactory.create(identity); + private static Optional getVendedAzureIdentity(ConnectorIdentity identity, Map fileIoProperties) + { + ImmutableMap.Builder azureCredentialBuilder = ImmutableMap.builder(); + PropertyUtil.propertiesWithPrefix(fileIoProperties, VENDED_ADLS_SAS_TOKEN_PREFIX) + .forEach((host, token) -> { + String storageAccount = host.contains(".") ? host.substring(0, host.indexOf('.')) : host; + + if (!storageAccount.isEmpty() && !token.isEmpty()) { + azureCredentialBuilder.put(EXTRA_SAS_TOKEN_PROPERTY_PREFIX + storageAccount, token); + azureCredentialBuilder.put(EXTRA_USE_VENDED_TOKEN, "true"); + } + }); + + Map azureCredentials = azureCredentialBuilder.buildKeepingLast(); + return azureCredentials.isEmpty() ? Optional.empty() : Optional.of(getVendedIdentity(identity, azureCredentials)); + } + + private static ConnectorIdentity getVendedIdentity(ConnectorIdentity identity, Map extraCredentials) + { + // Do not include original credentials as they should not be used in vended mode + return ConnectorIdentity.forUser(identity.getUser()) + .withGroups(identity.getGroups()) + .withPrincipal(identity.getPrincipal()) + .withEnabledSystemRoles(identity.getEnabledSystemRoles()) + .withConnectorRole(identity.getConnectorRole()) + .withExtraCredentials(extraCredentials) + .build(); } }