Skip to content

Commit d35e5f5

Browse files
orencclorenccl
authored
[#5894] feat(iceberg): support Azure account key credential (#5938)
### What changes were proposed in this pull request? Support Azure account key credential ### Why are the changes needed? Fix: #5894 ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Unit Test IcebergRESTADLSAccountKeyIT at iceberg 1.6.0
1 parent 2a5838f commit d35e5f5

File tree

12 files changed

+374
-46
lines changed

12 files changed

+374
-46
lines changed

‎api/src/main/java/org/apache/gravitino/credential/AzureAccountKeyCredential.java

+109
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.gravitino.credential;
21+
22+
import com.google.common.base.Preconditions;
23+
import com.google.common.collect.ImmutableMap;
24+
import java.util.Map;
25+
import org.apache.commons.lang3.StringUtils;
26+
27+
/** Azure account key credential. */
28+
public class AzureAccountKeyCredential implements Credential {
29+
30+
/** Azure account key credential type. */
31+
public static final String AZURE_ACCOUNT_KEY_CREDENTIAL_TYPE = "azure-account-key";
32+
/** Azure storage account name */
33+
public static final String GRAVITINO_AZURE_STORAGE_ACCOUNT_NAME = "azure-storage-account-name";
34+
/** Azure storage account key */
35+
public static final String GRAVITINO_AZURE_STORAGE_ACCOUNT_KEY = "azure-storage-account-key";
36+
37+
private String accountName;
38+
private String accountKey;
39+
40+
/**
41+
* Constructs an instance of {@link AzureAccountKeyCredential}.
42+
*
43+
* @param accountName The Azure account name.
44+
* @param accountKey The Azure account key.
45+
*/
46+
public AzureAccountKeyCredential(String accountName, String accountKey) {
47+
validate(accountName, accountKey);
48+
this.accountName = accountName;
49+
this.accountKey = accountKey;
50+
}
51+
52+
/**
53+
* This is the constructor that is used by credential factory to create an instance of credential
54+
* according to the credential information.
55+
*/
56+
public AzureAccountKeyCredential() {}
57+
58+
@Override
59+
public String credentialType() {
60+
return AZURE_ACCOUNT_KEY_CREDENTIAL_TYPE;
61+
}
62+
63+
@Override
64+
public long expireTimeInMs() {
65+
return 0;
66+
}
67+
68+
@Override
69+
public Map<String, String> credentialInfo() {
70+
return (new ImmutableMap.Builder<String, String>())
71+
.put(GRAVITINO_AZURE_STORAGE_ACCOUNT_NAME, accountName)
72+
.put(GRAVITINO_AZURE_STORAGE_ACCOUNT_KEY, accountKey)
73+
.build();
74+
}
75+
76+
@Override
77+
public void initialize(Map<String, String> credentialInfo, long expireTimeInMS) {
78+
String accountName = credentialInfo.get(GRAVITINO_AZURE_STORAGE_ACCOUNT_NAME);
79+
String accountKey = credentialInfo.get(GRAVITINO_AZURE_STORAGE_ACCOUNT_KEY);
80+
validate(accountName, accountKey);
81+
this.accountName = accountName;
82+
this.accountKey = accountKey;
83+
}
84+
85+
/**
86+
* Get Azure account name
87+
*
88+
* @return The Azure account name
89+
*/
90+
public String accountName() {
91+
return accountName;
92+
}
93+
94+
/**
95+
* Get Azure account key
96+
*
97+
* @return The Azure account key
98+
*/
99+
public String accountKey() {
100+
return accountKey;
101+
}
102+
103+
private void validate(String accountName, String accountKey) {
104+
Preconditions.checkArgument(
105+
StringUtils.isNotBlank(accountName), "Azure account name should not be empty.");
106+
Preconditions.checkArgument(
107+
StringUtils.isNotBlank(accountKey), "Azure account key should not be empty.");
108+
}
109+
}

‎api/src/main/resources/META-INF/services/org.apache.gravitino.credential.Credential

+1
Original file line numberDiff line numberDiff line change
@@ -23,3 +23,4 @@ org.apache.gravitino.credential.GCSTokenCredential
2323
org.apache.gravitino.credential.OSSTokenCredential
2424
org.apache.gravitino.credential.OSSSecretKeyCredential
2525
org.apache.gravitino.credential.ADLSTokenCredential
26+
org.apache.gravitino.credential.AzureAccountKeyCredential

‎bundles/azure-bundle/src/main/java/org/apache/gravitino/abs/credential/ADLSTokenProvider.java

+7-7
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@
3838
import org.apache.gravitino.credential.CredentialContext;
3939
import org.apache.gravitino.credential.CredentialProvider;
4040
import org.apache.gravitino.credential.PathBasedCredentialContext;
41-
import org.apache.gravitino.credential.config.ADLSCredentialConfig;
41+
import org.apache.gravitino.credential.config.AzureCredentialConfig;
4242

4343
/** Generates ADLS token to access ADLS data. */
4444
public class ADLSTokenProvider implements CredentialProvider {
@@ -51,14 +51,14 @@ public class ADLSTokenProvider implements CredentialProvider {
5151

5252
@Override
5353
public void initialize(Map<String, String> properties) {
54-
ADLSCredentialConfig adlsCredentialConfig = new ADLSCredentialConfig(properties);
55-
this.storageAccountName = adlsCredentialConfig.storageAccountName();
56-
this.tenantId = adlsCredentialConfig.tenantId();
57-
this.clientId = adlsCredentialConfig.clientId();
58-
this.clientSecret = adlsCredentialConfig.clientSecret();
54+
AzureCredentialConfig azureCredentialConfig = new AzureCredentialConfig(properties);
55+
this.storageAccountName = azureCredentialConfig.storageAccountName();
56+
this.tenantId = azureCredentialConfig.tenantId();
57+
this.clientId = azureCredentialConfig.clientId();
58+
this.clientSecret = azureCredentialConfig.clientSecret();
5959
this.endpoint =
6060
String.format("https://%s.%s", storageAccountName, ADLSTokenCredential.ADLS_DOMAIN);
61-
this.tokenExpireSecs = adlsCredentialConfig.tokenExpireInSecs();
61+
this.tokenExpireSecs = azureCredentialConfig.adlsTokenExpireInSecs();
6262
}
6363

6464
@Override

‎bundles/azure-bundle/src/main/java/org/apache/gravitino/abs/credential/AzureAccountKeyProvider.java

+54
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.gravitino.abs.credential;
21+
22+
import java.util.Map;
23+
import org.apache.gravitino.credential.AzureAccountKeyCredential;
24+
import org.apache.gravitino.credential.Credential;
25+
import org.apache.gravitino.credential.CredentialConstants;
26+
import org.apache.gravitino.credential.CredentialContext;
27+
import org.apache.gravitino.credential.CredentialProvider;
28+
import org.apache.gravitino.credential.config.AzureCredentialConfig;
29+
30+
/** Generates Azure account key to access data. */
31+
public class AzureAccountKeyProvider implements CredentialProvider {
32+
private String accountName;
33+
private String accountKey;
34+
35+
@Override
36+
public void initialize(Map<String, String> properties) {
37+
AzureCredentialConfig azureCredentialConfig = new AzureCredentialConfig(properties);
38+
this.accountName = azureCredentialConfig.storageAccountName();
39+
this.accountKey = azureCredentialConfig.storageAccountKey();
40+
}
41+
42+
@Override
43+
public void close() {}
44+
45+
@Override
46+
public String credentialType() {
47+
return CredentialConstants.AZURE_ACCOUNT_KEY_CREDENTIAL_PROVIDER_TYPE;
48+
}
49+
50+
@Override
51+
public Credential getCredential(CredentialContext context) {
52+
return new AzureAccountKeyCredential(accountName, accountKey);
53+
}
54+
}

‎bundles/azure-bundle/src/main/resources/META-INF/services/org.apache.gravitino.credential.CredentialProvider

+2-1
Original file line numberDiff line numberDiff line change
@@ -16,4 +16,5 @@
1616
# specific language governing permissions and limitations
1717
# under the License.
1818
#
19-
org.apache.gravitino.abs.credential.ADLSTokenProvider
19+
org.apache.gravitino.abs.credential.ADLSTokenProvider
20+
org.apache.gravitino.abs.credential.AzureAccountKeyProvider

‎catalogs/catalog-common/src/main/java/org/apache/gravitino/credential/CredentialConstants.java

+2
Original file line numberDiff line numberDiff line change
@@ -32,5 +32,7 @@ public class CredentialConstants {
3232
public static final String ADLS_TOKEN_CREDENTIAL_PROVIDER_TYPE = "adls-token";
3333
public static final String ADLS_TOKEN_EXPIRE_IN_SECS = "adls-token-expire-in-secs";
3434

35+
public static final String AZURE_ACCOUNT_KEY_CREDENTIAL_PROVIDER_TYPE = "azure-account-key";
36+
3537
private CredentialConstants() {}
3638
}

‎common/src/main/java/org/apache/gravitino/credential/CredentialPropertyUtils.java

+23-8
Original file line numberDiff line numberDiff line change
@@ -33,12 +33,19 @@ public class CredentialPropertyUtils {
3333
@VisibleForTesting static final String ICEBERG_S3_SECRET_ACCESS_KEY = "s3.secret-access-key";
3434
@VisibleForTesting static final String ICEBERG_S3_TOKEN = "s3.session-token";
3535
@VisibleForTesting static final String ICEBERG_GCS_TOKEN = "gcs.oauth2.token";
36-
@VisibleForTesting static final String ICEBERG_ADLS_TOKEN = "adls.sas-token";
3736

3837
@VisibleForTesting static final String ICEBERG_OSS_ACCESS_KEY_ID = "client.access-key-id";
3938
@VisibleForTesting static final String ICEBERG_OSS_ACCESS_KEY_SECRET = "client.access-key-secret";
4039
@VisibleForTesting static final String ICEBERG_OSS_SECURITY_TOKEN = "client.security-token";
4140

41+
@VisibleForTesting static final String ICEBERG_ADLS_TOKEN = "adls.sas-token";
42+
43+
@VisibleForTesting
44+
static final String ICEBERG_ADLS_ACCOUNT_NAME = "adls.auth.shared-key.account.name";
45+
46+
@VisibleForTesting
47+
static final String ICEBERG_ADLS_ACCOUNT_KEY = "adls.auth.shared-key.account.key";
48+
4249
private static Map<String, String> icebergCredentialPropertyMap =
4350
ImmutableMap.of(
4451
GCSTokenCredential.GCS_TOKEN_NAME,
@@ -54,7 +61,11 @@ public class CredentialPropertyUtils {
5461
OSSTokenCredential.GRAVITINO_OSS_SESSION_ACCESS_KEY_ID,
5562
ICEBERG_OSS_ACCESS_KEY_ID,
5663
OSSTokenCredential.GRAVITINO_OSS_SESSION_SECRET_ACCESS_KEY,
57-
ICEBERG_OSS_ACCESS_KEY_SECRET);
64+
ICEBERG_OSS_ACCESS_KEY_SECRET,
65+
AzureAccountKeyCredential.GRAVITINO_AZURE_STORAGE_ACCOUNT_NAME,
66+
ICEBERG_ADLS_ACCOUNT_NAME,
67+
AzureAccountKeyCredential.GRAVITINO_AZURE_STORAGE_ACCOUNT_KEY,
68+
ICEBERG_ADLS_ACCOUNT_KEY);
5869

5970
/**
6071
* Transforms a specific credential into a map of Iceberg properties.
@@ -63,19 +74,22 @@ public class CredentialPropertyUtils {
6374
* @return a map of Iceberg properties derived from the credential
6475
*/
6576
public static Map<String, String> toIcebergProperties(Credential credential) {
77+
if (credential instanceof S3TokenCredential
78+
|| credential instanceof S3SecretKeyCredential
79+
|| credential instanceof OSSTokenCredential
80+
|| credential instanceof OSSSecretKeyCredential
81+
|| credential instanceof AzureAccountKeyCredential) {
82+
return transformProperties(credential.credentialInfo(), icebergCredentialPropertyMap);
83+
}
84+
6685
if (credential instanceof GCSTokenCredential) {
6786
Map<String, String> icebergGCSCredentialProperties =
6887
transformProperties(credential.credentialInfo(), icebergCredentialPropertyMap);
6988
icebergGCSCredentialProperties.put(
7089
"gcs.oauth2.token-expires-at", String.valueOf(credential.expireTimeInMs()));
7190
return icebergGCSCredentialProperties;
7291
}
73-
if (credential instanceof S3TokenCredential || credential instanceof S3SecretKeyCredential) {
74-
return transformProperties(credential.credentialInfo(), icebergCredentialPropertyMap);
75-
}
76-
if (credential instanceof OSSTokenCredential || credential instanceof OSSSecretKeyCredential) {
77-
return transformProperties(credential.credentialInfo(), icebergCredentialPropertyMap);
78-
}
92+
7993
if (credential instanceof ADLSTokenCredential) {
8094
ADLSTokenCredential adlsCredential = (ADLSTokenCredential) credential;
8195
String sasTokenKey =
@@ -87,6 +101,7 @@ public static Map<String, String> toIcebergProperties(Credential credential) {
87101
icebergADLSCredentialProperties.put(sasTokenKey, adlsCredential.sasToken());
88102
return icebergADLSCredentialProperties;
89103
}
104+
90105
return credential.toProperties();
91106
}
92107

‎common/src/test/java/org/apache/gravitino/credential/TestCredentialFactory.java

+27
Original file line numberDiff line numberDiff line change
@@ -165,4 +165,31 @@ void testADLSTokenCredential() {
165165
Assertions.assertEquals(sasToken, adlsTokenCredential.sasToken());
166166
Assertions.assertEquals(expireTime, adlsTokenCredential.expireTimeInMs());
167167
}
168+
169+
@Test
170+
void testAzureAccountKeyCredential() {
171+
String storageAccountName = "storage-account-name";
172+
String storageAccountKey = "storage-account-key";
173+
174+
Map<String, String> azureAccountKeyCredentialInfo =
175+
ImmutableMap.of(
176+
AzureAccountKeyCredential.GRAVITINO_AZURE_STORAGE_ACCOUNT_NAME,
177+
storageAccountName,
178+
AzureAccountKeyCredential.GRAVITINO_AZURE_STORAGE_ACCOUNT_KEY,
179+
storageAccountKey);
180+
long expireTime = 0;
181+
Credential credential =
182+
CredentialFactory.create(
183+
AzureAccountKeyCredential.AZURE_ACCOUNT_KEY_CREDENTIAL_TYPE,
184+
azureAccountKeyCredentialInfo,
185+
expireTime);
186+
Assertions.assertEquals(
187+
AzureAccountKeyCredential.AZURE_ACCOUNT_KEY_CREDENTIAL_TYPE, credential.credentialType());
188+
Assertions.assertInstanceOf(AzureAccountKeyCredential.class, credential);
189+
190+
AzureAccountKeyCredential azureAccountKeyCredential = (AzureAccountKeyCredential) credential;
191+
Assertions.assertEquals(storageAccountName, azureAccountKeyCredential.accountName());
192+
Assertions.assertEquals(storageAccountKey, azureAccountKeyCredential.accountKey());
193+
Assertions.assertEquals(expireTime, azureAccountKeyCredential.expireTimeInMs());
194+
}
168195
}

‎core/src/main/java/org/apache/gravitino/credential/config/ADLSCredentialConfig.java ‎core/src/main/java/org/apache/gravitino/credential/config/AzureCredentialConfig.java

+3-3
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
import org.apache.gravitino.credential.CredentialConstants;
3030
import org.apache.gravitino.storage.AzureProperties;
3131

32-
public class ADLSCredentialConfig extends Config {
32+
public class AzureCredentialConfig extends Config {
3333

3434
public static final ConfigEntry<String> AZURE_STORAGE_ACCOUNT_NAME =
3535
new ConfigBuilder(AzureProperties.GRAVITINO_AZURE_STORAGE_ACCOUNT_NAME)
@@ -79,7 +79,7 @@ public class ADLSCredentialConfig extends Config {
7979
.intConf()
8080
.createWithDefault(3600);
8181

82-
public ADLSCredentialConfig(Map<String, String> properties) {
82+
public AzureCredentialConfig(Map<String, String> properties) {
8383
super(false);
8484
loadFromMap(properties, k -> true);
8585
}
@@ -110,7 +110,7 @@ public String clientSecret() {
110110
}
111111

112112
@NotNull
113-
public Integer tokenExpireInSecs() {
113+
public Integer adlsTokenExpireInSecs() {
114114
return this.get(ADLS_TOKEN_EXPIRE_IN_SECS);
115115
}
116116
}

0 commit comments

Comments
 (0)