Skip to content

Commit

Permalink
[#5894] feat(iceberg): support Azure account key credential (#5938)
Browse files Browse the repository at this point in the history
### 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
  • Loading branch information
orenccl authored Dec 22, 2024
1 parent 2a5838f commit d35e5f5
Show file tree
Hide file tree
Showing 12 changed files with 374 additions and 46 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,109 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you 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 org.apache.gravitino.credential;

import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableMap;
import java.util.Map;
import org.apache.commons.lang3.StringUtils;

/** Azure account key credential. */
public class AzureAccountKeyCredential implements Credential {

/** Azure account key credential type. */
public static final String AZURE_ACCOUNT_KEY_CREDENTIAL_TYPE = "azure-account-key";
/** Azure storage account name */
public static final String GRAVITINO_AZURE_STORAGE_ACCOUNT_NAME = "azure-storage-account-name";
/** Azure storage account key */
public static final String GRAVITINO_AZURE_STORAGE_ACCOUNT_KEY = "azure-storage-account-key";

private String accountName;
private String accountKey;

/**
* Constructs an instance of {@link AzureAccountKeyCredential}.
*
* @param accountName The Azure account name.
* @param accountKey The Azure account key.
*/
public AzureAccountKeyCredential(String accountName, String accountKey) {
validate(accountName, accountKey);
this.accountName = accountName;
this.accountKey = accountKey;
}

/**
* This is the constructor that is used by credential factory to create an instance of credential
* according to the credential information.
*/
public AzureAccountKeyCredential() {}

@Override
public String credentialType() {
return AZURE_ACCOUNT_KEY_CREDENTIAL_TYPE;
}

@Override
public long expireTimeInMs() {
return 0;
}

@Override
public Map<String, String> credentialInfo() {
return (new ImmutableMap.Builder<String, String>())
.put(GRAVITINO_AZURE_STORAGE_ACCOUNT_NAME, accountName)
.put(GRAVITINO_AZURE_STORAGE_ACCOUNT_KEY, accountKey)
.build();
}

@Override
public void initialize(Map<String, String> credentialInfo, long expireTimeInMS) {
String accountName = credentialInfo.get(GRAVITINO_AZURE_STORAGE_ACCOUNT_NAME);
String accountKey = credentialInfo.get(GRAVITINO_AZURE_STORAGE_ACCOUNT_KEY);
validate(accountName, accountKey);
this.accountName = accountName;
this.accountKey = accountKey;
}

/**
* Get Azure account name
*
* @return The Azure account name
*/
public String accountName() {
return accountName;
}

/**
* Get Azure account key
*
* @return The Azure account key
*/
public String accountKey() {
return accountKey;
}

private void validate(String accountName, String accountKey) {
Preconditions.checkArgument(
StringUtils.isNotBlank(accountName), "Azure account name should not be empty.");
Preconditions.checkArgument(
StringUtils.isNotBlank(accountKey), "Azure account key should not be empty.");
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -23,3 +23,4 @@ org.apache.gravitino.credential.GCSTokenCredential
org.apache.gravitino.credential.OSSTokenCredential
org.apache.gravitino.credential.OSSSecretKeyCredential
org.apache.gravitino.credential.ADLSTokenCredential
org.apache.gravitino.credential.AzureAccountKeyCredential
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@
import org.apache.gravitino.credential.CredentialContext;
import org.apache.gravitino.credential.CredentialProvider;
import org.apache.gravitino.credential.PathBasedCredentialContext;
import org.apache.gravitino.credential.config.ADLSCredentialConfig;
import org.apache.gravitino.credential.config.AzureCredentialConfig;

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

@Override
public void initialize(Map<String, String> properties) {
ADLSCredentialConfig adlsCredentialConfig = new ADLSCredentialConfig(properties);
this.storageAccountName = adlsCredentialConfig.storageAccountName();
this.tenantId = adlsCredentialConfig.tenantId();
this.clientId = adlsCredentialConfig.clientId();
this.clientSecret = adlsCredentialConfig.clientSecret();
AzureCredentialConfig azureCredentialConfig = new AzureCredentialConfig(properties);
this.storageAccountName = azureCredentialConfig.storageAccountName();
this.tenantId = azureCredentialConfig.tenantId();
this.clientId = azureCredentialConfig.clientId();
this.clientSecret = azureCredentialConfig.clientSecret();
this.endpoint =
String.format("https://%s.%s", storageAccountName, ADLSTokenCredential.ADLS_DOMAIN);
this.tokenExpireSecs = adlsCredentialConfig.tokenExpireInSecs();
this.tokenExpireSecs = azureCredentialConfig.adlsTokenExpireInSecs();
}

@Override
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you 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 org.apache.gravitino.abs.credential;

import java.util.Map;
import org.apache.gravitino.credential.AzureAccountKeyCredential;
import org.apache.gravitino.credential.Credential;
import org.apache.gravitino.credential.CredentialConstants;
import org.apache.gravitino.credential.CredentialContext;
import org.apache.gravitino.credential.CredentialProvider;
import org.apache.gravitino.credential.config.AzureCredentialConfig;

/** Generates Azure account key to access data. */
public class AzureAccountKeyProvider implements CredentialProvider {
private String accountName;
private String accountKey;

@Override
public void initialize(Map<String, String> properties) {
AzureCredentialConfig azureCredentialConfig = new AzureCredentialConfig(properties);
this.accountName = azureCredentialConfig.storageAccountName();
this.accountKey = azureCredentialConfig.storageAccountKey();
}

@Override
public void close() {}

@Override
public String credentialType() {
return CredentialConstants.AZURE_ACCOUNT_KEY_CREDENTIAL_PROVIDER_TYPE;
}

@Override
public Credential getCredential(CredentialContext context) {
return new AzureAccountKeyCredential(accountName, accountKey);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -16,4 +16,5 @@
# specific language governing permissions and limitations
# under the License.
#
org.apache.gravitino.abs.credential.ADLSTokenProvider
org.apache.gravitino.abs.credential.ADLSTokenProvider
org.apache.gravitino.abs.credential.AzureAccountKeyProvider
Original file line number Diff line number Diff line change
Expand Up @@ -32,5 +32,7 @@ public class CredentialConstants {
public static final String ADLS_TOKEN_CREDENTIAL_PROVIDER_TYPE = "adls-token";
public static final String ADLS_TOKEN_EXPIRE_IN_SECS = "adls-token-expire-in-secs";

public static final String AZURE_ACCOUNT_KEY_CREDENTIAL_PROVIDER_TYPE = "azure-account-key";

private CredentialConstants() {}
}
Original file line number Diff line number Diff line change
Expand Up @@ -33,12 +33,19 @@ public class CredentialPropertyUtils {
@VisibleForTesting static final String ICEBERG_S3_SECRET_ACCESS_KEY = "s3.secret-access-key";
@VisibleForTesting static final String ICEBERG_S3_TOKEN = "s3.session-token";
@VisibleForTesting static final String ICEBERG_GCS_TOKEN = "gcs.oauth2.token";
@VisibleForTesting static final String ICEBERG_ADLS_TOKEN = "adls.sas-token";

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

@VisibleForTesting static final String ICEBERG_ADLS_TOKEN = "adls.sas-token";

@VisibleForTesting
static final String ICEBERG_ADLS_ACCOUNT_NAME = "adls.auth.shared-key.account.name";

@VisibleForTesting
static final String ICEBERG_ADLS_ACCOUNT_KEY = "adls.auth.shared-key.account.key";

private static Map<String, String> icebergCredentialPropertyMap =
ImmutableMap.of(
GCSTokenCredential.GCS_TOKEN_NAME,
Expand All @@ -54,7 +61,11 @@ public class CredentialPropertyUtils {
OSSTokenCredential.GRAVITINO_OSS_SESSION_ACCESS_KEY_ID,
ICEBERG_OSS_ACCESS_KEY_ID,
OSSTokenCredential.GRAVITINO_OSS_SESSION_SECRET_ACCESS_KEY,
ICEBERG_OSS_ACCESS_KEY_SECRET);
ICEBERG_OSS_ACCESS_KEY_SECRET,
AzureAccountKeyCredential.GRAVITINO_AZURE_STORAGE_ACCOUNT_NAME,
ICEBERG_ADLS_ACCOUNT_NAME,
AzureAccountKeyCredential.GRAVITINO_AZURE_STORAGE_ACCOUNT_KEY,
ICEBERG_ADLS_ACCOUNT_KEY);

/**
* Transforms a specific credential into a map of Iceberg properties.
Expand All @@ -63,19 +74,22 @@ public class CredentialPropertyUtils {
* @return a map of Iceberg properties derived from the credential
*/
public static Map<String, String> toIcebergProperties(Credential credential) {
if (credential instanceof S3TokenCredential
|| credential instanceof S3SecretKeyCredential
|| credential instanceof OSSTokenCredential
|| credential instanceof OSSSecretKeyCredential
|| credential instanceof AzureAccountKeyCredential) {
return transformProperties(credential.credentialInfo(), icebergCredentialPropertyMap);
}

if (credential instanceof GCSTokenCredential) {
Map<String, String> icebergGCSCredentialProperties =
transformProperties(credential.credentialInfo(), icebergCredentialPropertyMap);
icebergGCSCredentialProperties.put(
"gcs.oauth2.token-expires-at", String.valueOf(credential.expireTimeInMs()));
return icebergGCSCredentialProperties;
}
if (credential instanceof S3TokenCredential || credential instanceof S3SecretKeyCredential) {
return transformProperties(credential.credentialInfo(), icebergCredentialPropertyMap);
}
if (credential instanceof OSSTokenCredential || credential instanceof OSSSecretKeyCredential) {
return transformProperties(credential.credentialInfo(), icebergCredentialPropertyMap);
}

if (credential instanceof ADLSTokenCredential) {
ADLSTokenCredential adlsCredential = (ADLSTokenCredential) credential;
String sasTokenKey =
Expand All @@ -87,6 +101,7 @@ public static Map<String, String> toIcebergProperties(Credential credential) {
icebergADLSCredentialProperties.put(sasTokenKey, adlsCredential.sasToken());
return icebergADLSCredentialProperties;
}

return credential.toProperties();
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -165,4 +165,31 @@ void testADLSTokenCredential() {
Assertions.assertEquals(sasToken, adlsTokenCredential.sasToken());
Assertions.assertEquals(expireTime, adlsTokenCredential.expireTimeInMs());
}

@Test
void testAzureAccountKeyCredential() {
String storageAccountName = "storage-account-name";
String storageAccountKey = "storage-account-key";

Map<String, String> azureAccountKeyCredentialInfo =
ImmutableMap.of(
AzureAccountKeyCredential.GRAVITINO_AZURE_STORAGE_ACCOUNT_NAME,
storageAccountName,
AzureAccountKeyCredential.GRAVITINO_AZURE_STORAGE_ACCOUNT_KEY,
storageAccountKey);
long expireTime = 0;
Credential credential =
CredentialFactory.create(
AzureAccountKeyCredential.AZURE_ACCOUNT_KEY_CREDENTIAL_TYPE,
azureAccountKeyCredentialInfo,
expireTime);
Assertions.assertEquals(
AzureAccountKeyCredential.AZURE_ACCOUNT_KEY_CREDENTIAL_TYPE, credential.credentialType());
Assertions.assertInstanceOf(AzureAccountKeyCredential.class, credential);

AzureAccountKeyCredential azureAccountKeyCredential = (AzureAccountKeyCredential) credential;
Assertions.assertEquals(storageAccountName, azureAccountKeyCredential.accountName());
Assertions.assertEquals(storageAccountKey, azureAccountKeyCredential.accountKey());
Assertions.assertEquals(expireTime, azureAccountKeyCredential.expireTimeInMs());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
import org.apache.gravitino.credential.CredentialConstants;
import org.apache.gravitino.storage.AzureProperties;

public class ADLSCredentialConfig extends Config {
public class AzureCredentialConfig extends Config {

public static final ConfigEntry<String> AZURE_STORAGE_ACCOUNT_NAME =
new ConfigBuilder(AzureProperties.GRAVITINO_AZURE_STORAGE_ACCOUNT_NAME)
Expand Down Expand Up @@ -79,7 +79,7 @@ public class ADLSCredentialConfig extends Config {
.intConf()
.createWithDefault(3600);

public ADLSCredentialConfig(Map<String, String> properties) {
public AzureCredentialConfig(Map<String, String> properties) {
super(false);
loadFromMap(properties, k -> true);
}
Expand Down Expand Up @@ -110,7 +110,7 @@ public String clientSecret() {
}

@NotNull
public Integer tokenExpireInSecs() {
public Integer adlsTokenExpireInSecs() {
return this.get(ADLS_TOKEN_EXPIRE_IN_SECS);
}
}
Loading

0 comments on commit d35e5f5

Please sign in to comment.