Skip to content

Commit 375dcc4

Browse files
CosminLazarCosminLazar
authored
[Enhancement] Support Azure Workload Identity authentication for Azure Data Lake Storage Gen2 (backport #62754) (#62882)
Signed-off-by: Cosmin Constantin Lazar <[email protected]>
1 parent ef06261 commit 375dcc4

File tree

4 files changed

+46
-4
lines changed

4 files changed

+46
-4
lines changed

‎fe/fe-core/src/main/java/com/starrocks/credential/azure/AzureCloudConfigurationProvider.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
import static com.starrocks.connector.share.credential.CloudConfigurationConstants.AZURE_ADLS2_OAUTH2_CLIENT_ID;
3131
import static com.starrocks.connector.share.credential.CloudConfigurationConstants.AZURE_ADLS2_OAUTH2_CLIENT_SECRET;
3232
import static com.starrocks.connector.share.credential.CloudConfigurationConstants.AZURE_ADLS2_OAUTH2_TENANT_ID;
33+
import static com.starrocks.connector.share.credential.CloudConfigurationConstants.AZURE_ADLS2_OAUTH2_TOKEN_FILE;
3334
import static com.starrocks.connector.share.credential.CloudConfigurationConstants.AZURE_ADLS2_OAUTH2_USE_MANAGED_IDENTITY;
3435
import static com.starrocks.connector.share.credential.CloudConfigurationConstants.AZURE_ADLS2_SAS_TOKEN;
3536
import static com.starrocks.connector.share.credential.CloudConfigurationConstants.AZURE_ADLS2_SHARED_KEY;
@@ -95,7 +96,8 @@ public CloudConfiguration build(Map<String, String> properties) {
9596
properties.getOrDefault(AZURE_ADLS2_SHARED_KEY, ""),
9697
properties.getOrDefault(AZURE_ADLS2_SAS_TOKEN, ""),
9798
properties.getOrDefault(AZURE_ADLS2_OAUTH2_CLIENT_SECRET, ""),
98-
properties.getOrDefault(AZURE_ADLS2_OAUTH2_CLIENT_ENDPOINT, "")
99+
properties.getOrDefault(AZURE_ADLS2_OAUTH2_CLIENT_ENDPOINT, ""),
100+
properties.getOrDefault(AZURE_ADLS2_OAUTH2_TOKEN_FILE, "")
99101
);
100102
if (adls2.validate()) {
101103
return new AzureCloudConfiguration(adls2);

‎fe/fe-core/src/main/java/com/starrocks/credential/azure/AzureStorageCloudCredential.java‎

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import org.apache.hadoop.fs.azurebfs.constants.ConfigurationKeys;
3030
import org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider;
3131
import org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider;
32+
import org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider;
3233
import org.apache.logging.log4j.LogManager;
3334
import org.apache.logging.log4j.Logger;
3435

@@ -40,7 +41,7 @@
4041
// For Azure Data Lake Gen1 (adl://)
4142
// We support Managed Service Identity & Service Principal
4243
// For Azure Data Lake Gen2 (abfs:// & abfss://)
43-
// We support Managed Identity & Shared Key & Service Principal
44+
// We support Managed Identity & Shared Key & Service Principal & Workload Identity
4445
abstract class AzureStorageCloudCredential implements CloudCredential {
4546

4647
public static final Logger LOG = LogManager.getLogger(AzureStorageCloudCredential.class);
@@ -231,10 +232,11 @@ class AzureADLS2CloudCredential extends AzureStorageCloudCredential {
231232
private final String sasToken;
232233
private final String oauth2ClientSecret;
233234
private final String oauth2ClientEndpoint;
235+
private final String oauth2TokenFile;
234236

235237
public AzureADLS2CloudCredential(String endpoint, boolean oauth2ManagedIdentity, String oauth2TenantId, String oauth2ClientId,
236238
String storageAccount, String sharedKey, String sasToken, String oauth2ClientSecret,
237-
String oauth2ClientEndpoint) {
239+
String oauth2ClientEndpoint, String oauth2TokenFile) {
238240
Preconditions.checkNotNull(endpoint);
239241
Preconditions.checkNotNull(oauth2TenantId);
240242
Preconditions.checkNotNull(oauth2ClientId);
@@ -243,6 +245,7 @@ public AzureADLS2CloudCredential(String endpoint, boolean oauth2ManagedIdentity,
243245
Preconditions.checkNotNull(sasToken);
244246
Preconditions.checkNotNull(oauth2ClientSecret);
245247
Preconditions.checkNotNull(oauth2ClientEndpoint);
248+
Preconditions.checkNotNull(oauth2TokenFile);
246249

247250
this.endpoint = endpoint;
248251
this.oauth2ManagedIdentity = oauth2ManagedIdentity;
@@ -253,6 +256,7 @@ public AzureADLS2CloudCredential(String endpoint, boolean oauth2ManagedIdentity,
253256
this.sasToken = sasToken;
254257
this.oauth2ClientSecret = oauth2ClientSecret;
255258
this.oauth2ClientEndpoint = oauth2ClientEndpoint;
259+
this.oauth2TokenFile = oauth2TokenFile;
256260

257261
tryGenerateConfigurationMap();
258262
}
@@ -315,6 +319,18 @@ void tryGenerateConfigurationMap() {
315319
oauth2ClientSecret);
316320
generatedConfigurationMap.put(createConfigKey(ConfigurationKeys.FS_AZURE_ACCOUNT_OAUTH_CLIENT_ENDPOINT),
317321
oauth2ClientEndpoint);
322+
} else if (!oauth2TokenFile.isEmpty() && !oauth2TenantId.isEmpty() && !oauth2ClientId.isEmpty()) {
323+
generatedConfigurationMap.put(createConfigKey(ConfigurationKeys.FS_AZURE_ACCOUNT_AUTH_TYPE_PROPERTY_NAME),
324+
"OAuth");
325+
generatedConfigurationMap.put(
326+
createConfigKey(ConfigurationKeys.FS_AZURE_ACCOUNT_TOKEN_PROVIDER_TYPE_PROPERTY_NAME),
327+
WorkloadIdentityTokenProvider.class.getName());
328+
generatedConfigurationMap.put(createConfigKey(ConfigurationKeys.FS_AZURE_ACCOUNT_OAUTH_TOKEN_FILE),
329+
oauth2TokenFile);
330+
generatedConfigurationMap.put(createConfigKey(ConfigurationKeys.FS_AZURE_ACCOUNT_OAUTH_CLIENT_ID),
331+
oauth2ClientId);
332+
generatedConfigurationMap.put(createConfigKey(ConfigurationKeys.FS_AZURE_ACCOUNT_OAUTH_MSI_TENANT),
333+
oauth2TenantId);
318334
}
319335
}
320336

@@ -328,6 +344,7 @@ public String toCredString() {
328344
", sharedKey='" + sharedKey + '\'' +
329345
", oauth2ClientSecret='" + oauth2ClientSecret + '\'' +
330346
", oauth2ClientEndpoint='" + oauth2ClientEndpoint + '\'' +
347+
", oauth2TokenFile='" + oauth2TokenFile + '\'' +
331348
'}';
332349
}
333350

‎fe/fe-core/src/test/java/com/starrocks/credential/CloudConfigurationFactoryTest.java‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -246,6 +246,7 @@ public void testAzureADLS2CloudConfiguration() {
246246
put(CloudConfigurationConstants.AZURE_ADLS2_OAUTH2_CLIENT_ENDPOINT, "XX");
247247
put(CloudConfigurationConstants.AZURE_ADLS2_OAUTH2_CLIENT_SECRET, "XX");
248248
put(CloudConfigurationConstants.AZURE_ADLS2_OAUTH2_USE_MANAGED_IDENTITY, "XX");
249+
put(CloudConfigurationConstants.AZURE_ADLS2_OAUTH2_TOKEN_FILE, "XX");
249250
}
250251
};
251252
CloudConfiguration cc = CloudConfigurationFactory.buildCloudConfigurationForStorage(map);
@@ -259,7 +260,7 @@ public void testAzureADLS2CloudConfiguration() {
259260
"AzureCloudConfiguration{resources='', jars='', hdpuser='', " +
260261
"cred=AzureADLS2CloudCredential{oauth2ManagedIdentity=false, " +
261262
"oauth2TenantId='XX', oauth2ClientId='XX', storageAccount='XX', sharedKey='XX', " +
262-
"oauth2ClientSecret='XX', oauth2ClientEndpoint='XX'}}");
263+
"oauth2ClientSecret='XX', oauth2ClientEndpoint='XX', oauth2TokenFile='XX'}}");
263264
}
264265

265266
@Test
@@ -305,6 +306,27 @@ public void testAzureADLS2Oauth2() {
305306
Assertions.assertEquals("client-id", conf.get("fs.azure.account.oauth2.client.id"));
306307
}
307308

309+
@Test
310+
public void testAzureADLS2WorkloadIdentity() {
311+
Map<String, String> map = new HashMap<>() {
312+
{
313+
put(CloudConfigurationConstants.AZURE_ADLS2_OAUTH2_CLIENT_ID, "client-id");
314+
put(CloudConfigurationConstants.AZURE_ADLS2_OAUTH2_TENANT_ID, "tenant-id");
315+
put(CloudConfigurationConstants.AZURE_ADLS2_OAUTH2_TOKEN_FILE, "/path/to/token");
316+
}
317+
};
318+
319+
CloudConfiguration cc = CloudConfigurationFactory.buildCloudConfigurationForStorage(map);
320+
Assertions.assertEquals(cc.getCloudType(), CloudType.AZURE);
321+
Configuration conf = new Configuration();
322+
cc.applyToConfiguration(conf);
323+
Assertions.assertEquals("OAuth", conf.get("fs.azure.account.auth.type"));
324+
Assertions.assertEquals("org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider",
325+
conf.get("fs.azure.account.oauth.provider.type"));
326+
Assertions.assertEquals("tenant-id", conf.get("fs.azure.account.oauth2.msi.tenant"));
327+
Assertions.assertEquals("/path/to/token", conf.get("fs.azure.account.oauth2.token.file"));
328+
}
329+
308330
@Test
309331
public void testGCPCloudConfiguration() {
310332
Map<String, String> map = new HashMap<String, String>() {

‎java-extensions/hadoop-ext/src/main/java/com/starrocks/connector/share/credential/CloudConfigurationConstants.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,7 @@ public class CloudConfigurationConstants {
103103
public static final String AZURE_ADLS2_SAS_TOKEN = "azure.adls2.sas_token";
104104
public static final String AZURE_ADLS2_OAUTH2_CLIENT_SECRET = "azure.adls2.oauth2_client_secret";
105105
public static final String AZURE_ADLS2_OAUTH2_CLIENT_ENDPOINT = "azure.adls2.oauth2_client_endpoint";
106+
public static final String AZURE_ADLS2_OAUTH2_TOKEN_FILE = "azure.adls2.oauth2_token_file";
106107

107108
// Credential for Google Cloud Platform (GCP)
108109
// For Google Cloud Storage (GCS)

0 commit comments

Comments
 (0)