This is an automated email from the ASF dual-hosted git repository. majin1102 pushed a commit to branch emr/initialize-network-integration in repository https://gitbox.apache.org/repos/asf/amoro.git
commit 6c69b7f94451243d1dfefdb0c42a0eab30ecbf39 Author: majin.nathan <[email protected]> AuthorDate: Mon Aug 3 22:48:37 2026 +0800 feat: add LAS Spark SQL integration scaffold --- amoro-ams/pom.xml | 7 + .../org/apache/amoro/server/las/LasHmsClient.java | 63 ++++++++ .../org/apache/amoro/server/las/LasIamClient.java | 70 ++++++++ .../amoro/server/las/LasIntegrationConfig.java | 42 ++++- .../amoro/server/las/LasIntegrationContext.java | 92 +++++++++-- .../apache/amoro/server/las/LasRestExtension.java | 34 +++- .../apache/amoro/server/las/LasTenantContext.java | 79 +++++++++ .../server/las/ServerlessSparkSqlManager.java | 156 ++++++++++++++++++ .../server/las/TestLasIntegrationContext.java | 14 +- .../amoro/server/las/TestLasSparkSqlFlow.java | 179 +++++++++++++++++++++ .../main/java/org/apache/amoro/hive/HMSClient.java | 3 + .../java/org/apache/amoro/hive/HMSClientImpl.java | 5 + dist/src/main/amoro-bin/conf/config.yaml | 13 +- 13 files changed, 727 insertions(+), 30 deletions(-) diff --git a/amoro-ams/pom.xml b/amoro-ams/pom.xml index 143d2c06f..3d76352d9 100644 --- a/amoro-ams/pom.xml +++ b/amoro-ams/pom.xml @@ -86,6 +86,13 @@ <groupId>com.volcengine.emr.serverless</groupId> <artifactId>serverless-sdk-query</artifactId> <version>${serverless-sdk-query.version}</version> + <exclusions> + <!-- SQLTask does not use the SDK's TOS upload path. TOS access is configured via Proton. --> + <exclusion> + <groupId>com.volcengine</groupId> + <artifactId>ve-tos-java-sdk</artifactId> + </exclusion> + </exclusions> </dependency> <dependency> diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/las/LasHmsClient.java b/amoro-ams/src/main/java/org/apache/amoro/server/las/LasHmsClient.java new file mode 100644 index 000000000..6d30d18b8 --- /dev/null +++ b/amoro-ams/src/main/java/org/apache/amoro/server/las/LasHmsClient.java @@ -0,0 +1,63 @@ +/* + * 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.amoro.server.las; + +import org.apache.amoro.hive.HMSClient; +import org.apache.amoro.hive.HMSClientPool; +import org.apache.hadoop.hive.metastore.api.Table; +import org.apache.thrift.TException; + +import java.util.List; +import java.util.function.Function; + +/** Small HMS3 SDK facade that keeps catalog selection explicit without hiding HMS data types. */ +public final class LasHmsClient { + + private final HMSClientPool discoveryPool; + private final Function<String, HMSClientPool> catalogPoolFactory; + + public LasHmsClient(LasIntegrationContext context) { + this(context.newHmsClientPool(), context::newHmsClientPool); + } + + LasHmsClient(HMSClientPool discoveryPool, Function<String, HMSClientPool> catalogPoolFactory) { + this.discoveryPool = discoveryPool; + this.catalogPoolFactory = catalogPoolFactory; + } + + public List<String> listCatalogs() throws TException, InterruptedException { + return discoveryPool.run(HMSClient::getCatalogs); + } + + public List<String> listDatabases(String catalogName) throws TException, InterruptedException { + return catalogPoolFactory.apply(catalogName).run(HMSClient::getAllDatabases); + } + + public List<String> listTables(String catalogName, String databaseName) + throws TException, InterruptedException { + return catalogPoolFactory.apply(catalogName).run(client -> client.getAllTables(databaseName)); + } + + public Table loadTable(String catalogName, String databaseName, String tableName) + throws TException, InterruptedException { + return catalogPoolFactory + .apply(catalogName) + .run(client -> client.getTable(databaseName, tableName)); + } +} diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/las/LasIamClient.java b/amoro-ams/src/main/java/org/apache/amoro/server/las/LasIamClient.java new file mode 100644 index 000000000..7086783f2 --- /dev/null +++ b/amoro-ams/src/main/java/org/apache/amoro/server/las/LasIamClient.java @@ -0,0 +1,70 @@ +/* + * 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.amoro.server.las; + +import bytedance.olap.iam.Credential; +import bytedance.olap.iam.IAMService; +import bytedance.olap.iam.IamException; +import bytedance.olap.iam.ServiceInfo; +import bytedance.olap.iam.cache.AssumeRoleCredentialCache; +import bytedance.olap.iam.http.ClientConfiguration; +import bytedance.olap.iam.http.model.AssumeRoleResponse.Credentials; + +import java.io.IOException; + +/** IAM SDK owner that caches tenant role credentials and closes the underlying HTTP client. */ +public final class LasIamClient implements AutoCloseable { + + private final IAMService iamService; + private final AssumeRoleCredentialCache credentialCache; + private final String roleSessionName; + + public LasIamClient(LasIntegrationContext context) { + this(context, newIamService(context), context.bootstrapCredential()); + } + + LasIamClient(LasIntegrationContext context, IAMService iamService, Credential credential) { + this.iamService = iamService; + this.roleSessionName = context.iamRoleSessionName(); + this.credentialCache = + new AssumeRoleCredentialCache( + iamService, + credential, + context.iamAssumeRoleTtl().getSeconds(), + context.iamCredentialCacheSize()); + } + + public Credentials assumeRole(String roleTrn) throws IamException { + return credentialCache.get(roleTrn, roleSessionName); + } + + @Override + public void close() throws IOException { + iamService.close(); + } + + private static IAMService newIamService(LasIntegrationContext context) { + ClientConfiguration configuration = new ClientConfiguration(); + configuration.setConnectionTimeout(Math.toIntExact(context.connectTimeout().toMillis())); + configuration.setSocketTimeout(Math.toIntExact(context.readTimeout().toMillis())); + String endpointHost = context.iamEndpoint().getAuthority(); + ServiceInfo serviceInfo = new ServiceInfo(context.iamEndpoint().getScheme(), endpointHost); + return new IAMService(serviceInfo, configuration, context.region()); + } +} diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/las/LasIntegrationConfig.java b/amoro-ams/src/main/java/org/apache/amoro/server/las/LasIntegrationConfig.java index cd119ee09..92852712d 100644 --- a/amoro-ams/src/main/java/org/apache/amoro/server/las/LasIntegrationConfig.java +++ b/amoro-ams/src/main/java/org/apache/amoro/server/las/LasIntegrationConfig.java @@ -52,6 +52,42 @@ public final class LasIntegrationConfig { .noDefaultValue() .withDescription("IAM endpoint used to obtain short-lived workload credentials."); + public static final ConfigOption<String> IAM_BOOTSTRAP_ACCESS_KEY = + ConfigOptions.key(PREFIX + "iam.bootstrap-access-key") + .stringType() + .noDefaultValue() + .withDescription("Access key of the AMS workload identity used to call AssumeRole."); + + public static final ConfigOption<String> IAM_BOOTSTRAP_SECRET_KEY = + ConfigOptions.key(PREFIX + "iam.bootstrap-secret-key") + .stringType() + .noDefaultValue() + .withDescription("Secret key of the AMS workload identity used to call AssumeRole."); + + public static final ConfigOption<String> IAM_BOOTSTRAP_SESSION_TOKEN = + ConfigOptions.key(PREFIX + "iam.bootstrap-session-token") + .stringType() + .defaultValue("") + .withDescription("Optional session token of the AMS workload identity."); + + public static final ConfigOption<String> IAM_ROLE_SESSION_NAME = + ConfigOptions.key(PREFIX + "iam.role-session-name") + .stringType() + .defaultValue("AmoroAssumeRoleSession") + .withDescription("IAM role session name used by AMS."); + + public static final ConfigOption<Duration> IAM_ASSUME_ROLE_TTL = + ConfigOptions.key(PREFIX + "iam.assume-role-ttl") + .durationType() + .defaultValue(Duration.ofHours(1)) + .withDescription("Lifetime of credentials returned by IAM AssumeRole."); + + public static final ConfigOption<Integer> IAM_CREDENTIAL_CACHE_SIZE = + ConfigOptions.key(PREFIX + "iam.credential-cache-size") + .intType() + .defaultValue(1000) + .withDescription("Maximum number of role credentials cached by AMS."); + public static final ConfigOption<String> EMR_SERVERLESS_ENDPOINT = ConfigOptions.key(PREFIX + "emr-serverless-endpoint") .stringType() @@ -82,12 +118,6 @@ public final class LasIntegrationConfig { .defaultValue(Duration.ofSeconds(30)) .withDescription("Socket read timeout for LAS/EMR service clients."); - public static final ConfigOption<String> OPTIMIZER_JAR_URI = - ConfigOptions.key(PREFIX + "optimizer-jar-uri") - .stringType() - .noDefaultValue() - .withDescription("URI of the optimizer jar submitted to EMR Serverless."); - public static final ConfigOption<Boolean> CROSS_VPC_ENABLED = ConfigOptions.key(PREFIX + "cross-vpc.enabled") .booleanType() diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/las/LasIntegrationContext.java b/amoro-ams/src/main/java/org/apache/amoro/server/las/LasIntegrationContext.java index ca4c6f76c..94dcaf911 100644 --- a/amoro-ams/src/main/java/org/apache/amoro/server/las/LasIntegrationContext.java +++ b/amoro-ams/src/main/java/org/apache/amoro/server/las/LasIntegrationContext.java @@ -18,6 +18,8 @@ package org.apache.amoro.server.las; +import bytedance.olap.iam.Credential; +import bytedance.olap.iam.http.model.AssumeRoleResponse.Credentials; import org.apache.amoro.config.ConfigOption; import org.apache.amoro.config.Configurations; import org.apache.amoro.hive.CachedHiveClientPool; @@ -38,21 +40,19 @@ import java.util.Set; public final class LasIntegrationContext { private static final Set<String> HTTP_SCHEMES = new HashSet<>(Arrays.asList("http", "https")); - private static final Set<String> OPTIMIZER_JAR_SCHEMES = - new HashSet<>(Arrays.asList("tos", "http", "https")); - + private final Configurations configurations; private final boolean enabled; private final URI hmsUri; private final URI tosEndpoint; private final URI iamEndpoint; private final URI emrServerlessEndpoint; - private final URI optimizerJarUri; private final String region; private final String emrServerlessService; private final Duration connectTimeout; private final Duration readTimeout; private LasIntegrationContext(Configurations configurations) { + this.configurations = configurations; this.enabled = configurations.getBoolean(LasIntegrationConfig.ENABLED); this.region = configurations.getString(LasIntegrationConfig.REGION); this.emrServerlessService = @@ -65,7 +65,6 @@ public final class LasIntegrationContext { this.tosEndpoint = null; this.iamEndpoint = null; this.emrServerlessEndpoint = null; - this.optimizerJarUri = null; return; } @@ -75,12 +74,18 @@ public final class LasIntegrationContext { this.iamEndpoint = requiredUri(configurations, LasIntegrationConfig.IAM_ENDPOINT, HTTP_SCHEMES); this.emrServerlessEndpoint = requiredUri(configurations, LasIntegrationConfig.EMR_SERVERLESS_ENDPOINT, HTTP_SCHEMES); - this.optimizerJarUri = - requiredUri(configurations, LasIntegrationConfig.OPTIMIZER_JAR_URI, OPTIMIZER_JAR_SCHEMES); requiredString(configurations, LasIntegrationConfig.REGION); requiredString(configurations, LasIntegrationConfig.EMR_SERVERLESS_SERVICE); + requiredString(configurations, LasIntegrationConfig.IAM_BOOTSTRAP_ACCESS_KEY); + requiredString(configurations, LasIntegrationConfig.IAM_BOOTSTRAP_SECRET_KEY); + requiredString(configurations, LasIntegrationConfig.IAM_ROLE_SESSION_NAME); positiveDuration(configurations, LasIntegrationConfig.CONNECT_TIMEOUT); positiveDuration(configurations, LasIntegrationConfig.READ_TIMEOUT); + positiveDuration(configurations, LasIntegrationConfig.IAM_ASSUME_ROLE_TTL); + if (configurations.getInteger(LasIntegrationConfig.IAM_CREDENTIAL_CACHE_SIZE) <= 0) { + throw new IllegalArgumentException( + LasIntegrationConfig.IAM_CREDENTIAL_CACHE_SIZE.key() + " must be greater than zero"); + } if (configurations.getBoolean(LasIntegrationConfig.CROSS_VPC_ENABLED)) { requiredString(configurations, LasIntegrationConfig.CROSS_VPC_ACCOUNT_ID); @@ -99,22 +104,52 @@ public final class LasIntegrationContext { } public CachedHiveClientPool newHmsClientPool() { + return newHmsClientPool(null); + } + + public CachedHiveClientPool newHmsClientPool(String catalogName) { ensureEnabled(); Configuration configuration = new Configuration(); configuration.set("hive.metastore.uris", hmsUri.toString()); + if (StringUtils.isNotBlank(catalogName)) { + configuration.set("metastore.catalog.default", catalogName); + } TableMetaStore metaStore = TableMetaStore.builder().withConfiguration(configuration).build(); return new CachedHiveClientPool(metaStore, Maps.newHashMap()); } - public Configuration newTosConfiguration() { + public Configuration newTosConfiguration(Credentials credentials) { ensureEnabled(); + if (credentials == null) { + throw new IllegalArgumentException("TOS credentials are required"); + } Configuration configuration = new Configuration(false); configuration.set("fs.AbstractFileSystem.tos.impl", "io.proton.fs.ProtonFS"); configuration.set("fs.tos.impl", "io.proton.fs.ProtonFileSystem"); configuration.set("fs.tos.endpoint", tosEndpoint.toString()); + configuration.set("proton.cache.enable", "false"); + configuration.set( + "mapreduce.outputcommitter.factory.class", "io.proton.commit.CommitterFactory"); + configuration.set( + "fs.tos.credentials.provider", "io.proton.common.object.auth.SimpleCredentialsProvider"); + configuration.set("fs.tos.access-key-id", credentials.getAccessKeyId()); + configuration.set("fs.tos.secret-access-key", credentials.getSecretAccessKey()); + configuration.set("fs.tos.session-token", credentials.getSessionToken()); + configuration.set("fs.tos.http.maxConnections", "1024"); return configuration; } + public Credential bootstrapCredential() { + ensureEnabled(); + String accessKey = configurations.getString(LasIntegrationConfig.IAM_BOOTSTRAP_ACCESS_KEY); + String secretKey = configurations.getString(LasIntegrationConfig.IAM_BOOTSTRAP_SECRET_KEY); + String sessionToken = + configurations.getString(LasIntegrationConfig.IAM_BOOTSTRAP_SESSION_TOKEN); + return StringUtils.isBlank(sessionToken) + ? new Credential(accessKey, secretKey) + : new Credential(accessKey, secretKey, sessionToken); + } + public URI hmsUri() { ensureEnabled(); return hmsUri; @@ -130,11 +165,6 @@ public final class LasIntegrationContext { return emrServerlessEndpoint; } - public URI optimizerJarUri() { - ensureEnabled(); - return optimizerJarUri; - } - public String region() { return region; } @@ -151,6 +181,42 @@ public final class LasIntegrationContext { return readTimeout; } + public String iamRoleSessionName() { + ensureEnabled(); + return configurations.getString(LasIntegrationConfig.IAM_ROLE_SESSION_NAME); + } + + public Duration iamAssumeRoleTtl() { + ensureEnabled(); + return configurations.get(LasIntegrationConfig.IAM_ASSUME_ROLE_TTL); + } + + public int iamCredentialCacheSize() { + ensureEnabled(); + return configurations.getInteger(LasIntegrationConfig.IAM_CREDENTIAL_CACHE_SIZE); + } + + public boolean crossVpcEnabled() { + ensureEnabled(); + return configurations.getBoolean(LasIntegrationConfig.CROSS_VPC_ENABLED); + } + + public String crossVpcAccountId() { + return configurations.getString(LasIntegrationConfig.CROSS_VPC_ACCOUNT_ID); + } + + public String crossVpcVpcId() { + return configurations.getString(LasIntegrationConfig.CROSS_VPC_VPC_ID); + } + + public String crossVpcSubnetIds() { + return configurations.getString(LasIntegrationConfig.CROSS_VPC_SUBNET_IDS); + } + + public String crossVpcSecurityGroupId() { + return configurations.getString(LasIntegrationConfig.CROSS_VPC_SECURITY_GROUP_ID); + } + private void ensureEnabled() { if (!enabled) { throw new IllegalStateException("LAS/EMR integration is disabled"); diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/las/LasRestExtension.java b/amoro-ams/src/main/java/org/apache/amoro/server/las/LasRestExtension.java index 29676349f..224d29170 100644 --- a/amoro-ams/src/main/java/org/apache/amoro/server/las/LasRestExtension.java +++ b/amoro-ams/src/main/java/org/apache/amoro/server/las/LasRestExtension.java @@ -37,14 +37,20 @@ public class LasRestExtension implements RestExtension { private static final Logger LOG = LoggerFactory.getLogger(LasRestExtension.class); private final LasIntegrationContext integrationContext; + private final LasHmsClient hmsClient; + private final ServerlessSparkSqlManager sparkSqlManager; private final CatalogManager catalogManager; private final TableManager tableManager; LasRestExtension( LasIntegrationContext integrationContext, + LasHmsClient hmsClient, + ServerlessSparkSqlManager sparkSqlManager, CatalogManager catalogManager, TableManager tableManager) { this.integrationContext = integrationContext; + this.hmsClient = hmsClient; + this.sparkSqlManager = sparkSqlManager; this.catalogManager = catalogManager; this.tableManager = tableManager; LOG.info("LAS/EMR integration initialized, enabled={}", integrationContext.enabled()); @@ -54,10 +60,11 @@ public class LasRestExtension implements RestExtension { public EndpointGroup endpoints() { return () -> { // Intentionally empty. Add management-plane routes here under /api/ams/v1/las when the - // OpenAPI contract is ready. Controllers should receive integrationContext, catalogManager, - // and tableManager from this extension instead of constructing HMS, TOS, IAM, or EMR clients - // themselves. Routes registered here automatically pass through the existing AMS REST - // authentication filter; do not add management-plane routes to the URL whitelist. + // OpenAPI contract is ready. Controllers should receive integrationContext, hmsClient, + // sparkSqlManager, catalogManager, and tableManager from this extension instead of + // constructing HMS, TOS, IAM, or EMR clients themselves. Routes registered here + // automatically pass through the existing AMS REST authentication filter; do not add + // management-plane routes to the URL whitelist. }; } @@ -77,6 +84,7 @@ public class LasRestExtension implements RestExtension { private Configurations serviceConfig; private CatalogManager catalogManager; private TableManager tableManager; + private LasIamClient iamClient; @Override public RestExtensionFactory withServiceConfig(Configurations serviceConfig) { @@ -101,8 +109,16 @@ public class LasRestExtension implements RestExtension { Preconditions.checkNotNull(serviceConfig, "serviceConfig is required"); Preconditions.checkNotNull(catalogManager, "catalogManager is required"); Preconditions.checkNotNull(tableManager, "tableManager is required"); + LasIntegrationContext context = LasIntegrationContext.initialize(serviceConfig); + LasHmsClient hmsClient = null; + ServerlessSparkSqlManager sparkSqlManager = null; + if (context.enabled()) { + iamClient = new LasIamClient(context); + hmsClient = new LasHmsClient(context); + sparkSqlManager = new ServerlessSparkSqlManager(context, iamClient); + } return new LasRestExtension( - LasIntegrationContext.initialize(serviceConfig), catalogManager, tableManager); + context, hmsClient, sparkSqlManager, catalogManager, tableManager); } @Override @@ -112,6 +128,14 @@ public class LasRestExtension implements RestExtension { @Override public void close() { + if (iamClient != null) { + try { + iamClient.close(); + } catch (Exception e) { + LOG.warn("Failed to close LAS IAM client", e); + } + iamClient = null; + } LOG.info("Closing LAS/EMR integration extension"); } diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/las/LasTenantContext.java b/amoro-ams/src/main/java/org/apache/amoro/server/las/LasTenantContext.java new file mode 100644 index 000000000..56e1e0976 --- /dev/null +++ b/amoro-ams/src/main/java/org/apache/amoro/server/las/LasTenantContext.java @@ -0,0 +1,79 @@ +/* + * 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.amoro.server.las; + +import org.apache.commons.lang3.StringUtils; + +/** Tenant-scoped identity and resource information required by one LAS operation. */ +public final class LasTenantContext { + + private final String accountId; + private final String catalogName; + private final String queueName; + private final String submitRoleTrn; + private final String dataRoleTrn; + + public LasTenantContext( + String accountId, + String catalogName, + String queueName, + String submitRoleTrn, + String dataRoleTrn) { + this.accountId = required("accountId", accountId); + this.catalogName = required("catalogName", catalogName); + this.queueName = required("queueName", queueName); + this.submitRoleTrn = requiredRole("submitRoleTrn", submitRoleTrn); + this.dataRoleTrn = requiredRole("dataRoleTrn", dataRoleTrn); + } + + public String accountId() { + return accountId; + } + + public String catalogName() { + return catalogName; + } + + public String queueName() { + return queueName; + } + + public String submitRoleTrn() { + return submitRoleTrn; + } + + public String dataRoleTrn() { + return dataRoleTrn; + } + + private static String required(String name, String value) { + if (StringUtils.isBlank(value)) { + throw new IllegalArgumentException(name + " is required"); + } + return value; + } + + private static String requiredRole(String name, String value) { + String role = required(name, value); + if (!role.startsWith("trn:iam:") || !role.contains(":role/")) { + throw new IllegalArgumentException(name + " must be an IAM role TRN"); + } + return role; + } +} diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/las/ServerlessSparkSqlManager.java b/amoro-ams/src/main/java/org/apache/amoro/server/las/ServerlessSparkSqlManager.java new file mode 100644 index 000000000..b386ed212 --- /dev/null +++ b/amoro-ams/src/main/java/org/apache/amoro/server/las/ServerlessSparkSqlManager.java @@ -0,0 +1,156 @@ +/* + * 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.amoro.server.las; + +import bytedance.olap.iam.IamException; +import bytedance.olap.iam.http.model.AssumeRoleResponse.Credentials; +import com.volcengine.emr.serverless.Job; +import com.volcengine.emr.serverless.SQLTask; +import com.volcengine.emr.serverless.ServerlessClientOption; +import com.volcengine.emr.serverless.ServerlessQueryClient; +import org.apache.commons.lang3.StringUtils; +import org.apache.hadoop.conf.Configuration; + +import java.util.LinkedHashMap; +import java.util.Map; + +/** Submits asynchronous Spark SQL jobs to a tenant queue with tenant-scoped credentials. */ +public final class ServerlessSparkSqlManager { + + private static final String SPARK_HADOOP_PREFIX = "spark.hadoop."; + + private final LasIntegrationContext context; + private final LasIamClient iamClient; + private final ServerlessClientFactory clientFactory; + + public ServerlessSparkSqlManager(LasIntegrationContext context, LasIamClient iamClient) { + this(context, iamClient, credentials -> newServerlessClient(context, credentials)); + } + + ServerlessSparkSqlManager( + LasIntegrationContext context, + LasIamClient iamClient, + ServerlessClientFactory clientFactory) { + this.context = context; + this.iamClient = iamClient; + this.clientFactory = clientFactory; + } + + public String submit( + LasTenantContext tenant, + String taskName, + String sparkSql, + Map<String, String> customSparkConf) + throws IamException { + Credentials submitCredentials = iamClient.assumeRole(tenant.submitRoleTrn()); + Credentials dataCredentials = iamClient.assumeRole(tenant.dataRoleTrn()); + SQLTask task = buildTask(tenant, taskName, sparkSql, customSparkConf, dataCredentials); + return clientFactory.create(submitCredentials).executeSQL(task).getId(); + } + + public Job getJob(LasTenantContext tenant, String jobId) throws IamException { + return submissionClient(tenant).getJob(required("jobId", jobId)); + } + + public void cancelJob(LasTenantContext tenant, String jobId) throws IamException { + submissionClient(tenant).cancelJob(required("jobId", jobId)); + } + + SQLTask buildTask( + LasTenantContext tenant, + String taskName, + String sparkSql, + Map<String, String> customSparkConf, + Credentials dataCredentials) { + required("taskName", taskName); + required("sparkSql", sparkSql); + + Map<String, String> managedConf = managedSparkConf(tenant, dataCredentials); + Map<String, String> taskConf = new LinkedHashMap<>(); + if (customSparkConf != null) { + customSparkConf.forEach( + (key, value) -> { + if (managedConf.containsKey(key)) { + throw new IllegalArgumentException("Custom Spark conf cannot override " + key); + } + taskConf.put(key, value); + }); + } + taskConf.putAll(managedConf); + + return SQLTask.builder(sparkSql) + .name(taskName) + .queue(tenant.queueName()) + .addConf(taskConf) + .sync(false) + .build(); + } + + private ServerlessQueryClient submissionClient(LasTenantContext tenant) throws IamException { + return clientFactory.create(iamClient.assumeRole(tenant.submitRoleTrn())); + } + + private Map<String, String> managedSparkConf( + LasTenantContext tenant, Credentials dataCredentials) { + Map<String, String> conf = new LinkedHashMap<>(); + conf.put("spark.hadoop.hive.metastore.uris", context.hmsUri().toString()); + conf.put("spark.hadoop.metastore.catalog.default", tenant.catalogName()); + conf.put("spark.hive.metastore.catalog.default", tenant.catalogName()); + + Configuration tos = context.newTosConfiguration(dataCredentials); + tos.forEach(entry -> conf.put(SPARK_HADOOP_PREFIX + entry.getKey(), entry.getValue())); + + if (context.crossVpcEnabled()) { + conf.put("serverless.cross.vpc.access.enabled", "true"); + conf.put("serverless.cross.vpc.accountId", context.crossVpcAccountId()); + conf.put("serverless.cross.vpc.vpc.id", context.crossVpcVpcId()); + conf.put("serverless.cross.vpc.subnet.ids", context.crossVpcSubnetIds()); + conf.put("serverless.cross.vpc.security.group.id", context.crossVpcSecurityGroupId()); + } + return conf; + } + + private static ServerlessQueryClient newServerlessClient( + LasIntegrationContext context, Credentials credentials) { + ServerlessClientOption options = + ServerlessClientOption.builder( + credentials.getAccessKeyId(), + credentials.getSecretAccessKey(), + credentials.getSessionToken()) + .endpoint(context.emrServerlessEndpoint().toString()) + .service(context.emrServerlessService()) + .region(context.region()) + .connectionTimeoutMs(Math.toIntExact(context.connectTimeout().toMillis())) + .socketTimeoutMs(Math.toIntExact(context.readTimeout().toMillis())) + .build(); + return new ServerlessQueryClient(options); + } + + private static String required(String name, String value) { + if (StringUtils.isBlank(value)) { + throw new IllegalArgumentException(name + " is required"); + } + return value; + } + + @FunctionalInterface + interface ServerlessClientFactory { + ServerlessQueryClient create(Credentials credentials); + } +} diff --git a/amoro-ams/src/test/java/org/apache/amoro/server/las/TestLasIntegrationContext.java b/amoro-ams/src/test/java/org/apache/amoro/server/las/TestLasIntegrationContext.java index fef97b5e8..fdeebaf5e 100644 --- a/amoro-ams/src/test/java/org/apache/amoro/server/las/TestLasIntegrationContext.java +++ b/amoro-ams/src/test/java/org/apache/amoro/server/las/TestLasIntegrationContext.java @@ -18,6 +18,7 @@ package org.apache.amoro.server.las; +import bytedance.olap.iam.http.model.AssumeRoleResponse.Credentials; import org.apache.amoro.config.Configurations; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -45,8 +46,13 @@ public class TestLasIntegrationContext { Assertions.assertEquals("https://emr-serverless";, context.emrServerlessEndpoint().toString()); Assertions.assertEquals("cn-beijing", context.region()); Assertions.assertNotNull(context.newHmsClientPool()); + Credentials credentials = new Credentials(); + credentials.setAccessKeyId("data-ak"); + credentials.setSecretAccessKey("data-sk"); + credentials.setSessionToken("data-token"); Assertions.assertEquals( - "https://tos-cn-beijing.volces.com";, context.newTosConfiguration().get("fs.tos.endpoint")); + "https://tos-cn-beijing.volces.com";, + context.newTosConfiguration(credentials).get("fs.tos.endpoint")); } @Test @@ -72,15 +78,15 @@ public class TestLasIntegrationContext { exception.getMessage().contains(LasIntegrationConfig.CROSS_VPC_ACCOUNT_ID.key())); } - private static Configurations validConfigurations() { + static Configurations validConfigurations() { Map<String, String> values = new HashMap<>(); values.put(LasIntegrationConfig.ENABLED.key(), "true"); values.put(LasIntegrationConfig.HMS_URI.key(), "thrift://hms-service:9083"); values.put(LasIntegrationConfig.TOS_ENDPOINT.key(), "https://tos-cn-beijing.volces.com";); values.put(LasIntegrationConfig.IAM_ENDPOINT.key(), "https://iam.volcengineapi.com";); values.put(LasIntegrationConfig.EMR_SERVERLESS_ENDPOINT.key(), "https://emr-serverless";); - values.put( - LasIntegrationConfig.OPTIMIZER_JAR_URI.key(), "tos://amoro-artifacts/amoro-optimizer.jar"); + values.put(LasIntegrationConfig.IAM_BOOTSTRAP_ACCESS_KEY.key(), "bootstrap-ak"); + values.put(LasIntegrationConfig.IAM_BOOTSTRAP_SECRET_KEY.key(), "bootstrap-sk"); return Configurations.fromMap(values); } } diff --git a/amoro-ams/src/test/java/org/apache/amoro/server/las/TestLasSparkSqlFlow.java b/amoro-ams/src/test/java/org/apache/amoro/server/las/TestLasSparkSqlFlow.java new file mode 100644 index 000000000..074149ec3 --- /dev/null +++ b/amoro-ams/src/test/java/org/apache/amoro/server/las/TestLasSparkSqlFlow.java @@ -0,0 +1,179 @@ +/* + * 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.amoro.server.las; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import bytedance.olap.iam.Credential; +import bytedance.olap.iam.IAMService; +import bytedance.olap.iam.http.model.AssumeRoleResponse; +import bytedance.olap.iam.http.model.AssumeRoleResponse.Credentials; +import com.volcengine.emr.serverless.Job; +import com.volcengine.emr.serverless.SQLTask; +import com.volcengine.emr.serverless.ServerlessQueryClient; +import org.apache.amoro.client.ClientPool; +import org.apache.amoro.hive.HMSClient; +import org.apache.amoro.hive.HMSClientPool; +import org.apache.hadoop.hive.metastore.api.Table; +import org.apache.thrift.TException; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; + +import java.util.Collections; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; + +/** Executable example of the HMS -> IAM -> EMR Serverless Spark SQL scaffold. */ +public class TestLasSparkSqlFlow { + + private static final String CATALOG = "123456789@las"; + private static final String SUBMIT_ROLE = "trn:iam::123456789:role/EmrJobSubmitRole"; + private static final String DATA_ROLE = "trn:iam::123456789:role/LasDataAccessRole"; + + @Test + public void testCompleteSparkSqlFlow() throws Exception { + LasIntegrationContext integration = + LasIntegrationContext.initialize(TestLasIntegrationContext.validConfigurations()); + + HMSClient hmsSdk = mock(HMSClient.class); + when(hmsSdk.getCatalogs()).thenReturn(Collections.singletonList(CATALOG)); + when(hmsSdk.getAllDatabases()).thenReturn(Collections.singletonList("analytics")); + when(hmsSdk.getAllTables("analytics")).thenReturn(Collections.singletonList("orders")); + Table hmsTable = new Table(); + hmsTable.setDbName("analytics"); + hmsTable.setTableName("orders"); + when(hmsSdk.getTable("analytics", "orders")).thenReturn(hmsTable); + HMSClientPool hmsPool = directPool(hmsSdk); + LasHmsClient hms = new LasHmsClient(hmsPool, ignored -> hmsPool); + + Assertions.assertEquals(Collections.singletonList(CATALOG), hms.listCatalogs()); + Assertions.assertEquals(Collections.singletonList("analytics"), hms.listDatabases(CATALOG)); + Assertions.assertEquals( + Collections.singletonList("orders"), hms.listTables(CATALOG, "analytics")); + Assertions.assertSame(hmsTable, hms.loadTable(CATALOG, "analytics", "orders")); + + IAMService iamSdk = mock(IAMService.class); + when(iamSdk.assumeRole(any(Credential.class), anyString(), anyString(), anyInt(), isNull())) + .thenAnswer( + invocation -> + assumeRoleResponse( + invocation.getArgument(1), + SUBMIT_ROLE.equals(invocation.getArgument(1)) ? "submit" : "data")); + LasIamClient iam = + new LasIamClient(integration, iamSdk, new Credential("bootstrap-ak", "bootstrap-sk")); + + ServerlessQueryClient serverlessSdk = mock(ServerlessQueryClient.class); + Job submitted = mock(Job.class); + when(submitted.getId()).thenReturn("job-20260803-0001"); + when(serverlessSdk.executeSQL(any(SQLTask.class))).thenReturn(submitted); + Job running = mock(Job.class); + when(serverlessSdk.getJob(submitted.getId())).thenReturn(running); + AtomicReference<Credentials> clientCredentials = new AtomicReference<>(); + ServerlessSparkSqlManager sparkSql = + new ServerlessSparkSqlManager( + integration, + iam, + credentials -> { + clientCredentials.set(credentials); + return serverlessSdk; + }); + + LasTenantContext tenant = + new LasTenantContext( + "123456789", CATALOG, "tenant-production-queue", SUBMIT_ROLE, DATA_ROLE); + String sql = "CALL spark_catalog.system.rewrite_data_files(table => 'analytics.orders')"; + String jobId = + sparkSql.submit( + tenant, + "compact-analytics-orders", + sql, + Collections.singletonMap("spark.sql.shuffle.partitions", "200")); + + Assertions.assertEquals(submitted.getId(), jobId); + Assertions.assertEquals("submit-ak", clientCredentials.get().getAccessKeyId()); + ArgumentCaptor<SQLTask> taskCaptor = ArgumentCaptor.forClass(SQLTask.class); + verify(serverlessSdk).executeSQL(taskCaptor.capture()); + SQLTask task = taskCaptor.getValue(); + Assertions.assertEquals(sql, task.getQuery()); + Assertions.assertEquals("tenant-production-queue", task.getQueue().orElse(null)); + Assertions.assertFalse(task.isSync()); + Map<String, String> taskConf = task.getConf(); + Assertions.assertEquals(CATALOG, taskConf.get("spark.hive.metastore.catalog.default")); + Assertions.assertEquals( + "thrift://hms-service:9083", taskConf.get("spark.hadoop.hive.metastore.uris")); + Assertions.assertEquals( + "io.proton.common.object.auth.SimpleCredentialsProvider", + taskConf.get("spark.hadoop.fs.tos.credentials.provider")); + Assertions.assertEquals("data-ak", taskConf.get("spark.hadoop.fs.tos.access-key-id")); + Assertions.assertEquals("data-token", taskConf.get("spark.hadoop.fs.tos.session-token")); + Assertions.assertEquals("200", taskConf.get("spark.sql.shuffle.partitions")); + + Assertions.assertSame(running, sparkSql.getJob(tenant, jobId)); + sparkSql.cancelJob(tenant, jobId); + verify(serverlessSdk).cancelJob(jobId); + + // Submit and data credentials are each loaded once; status/cancel reuse the submit-role cache. + verify(iamSdk, times(2)) + .assumeRole(any(Credential.class), anyString(), anyString(), anyInt(), isNull()); + iam.close(); + verify(iamSdk).close(); + } + + private static AssumeRoleResponse assumeRoleResponse(String roleTrn, String credentialPrefix) { + Credentials credentials = new Credentials(); + credentials.setAccessKeyId(credentialPrefix + "-ak"); + credentials.setSecretAccessKey(credentialPrefix + "-sk"); + credentials.setSessionToken(credentialPrefix + "-token"); + credentials.setCurrentTime("2026-08-03T13:00:00+00:00"); + credentials.setExpiredTime("2099-08-03T14:00:00+00:00"); + + AssumeRoleResponse.ResultBean result = new AssumeRoleResponse.ResultBean(); + result.setCredentials(credentials); + AssumeRoleResponse.AssumedRoleUser user = new AssumeRoleResponse.AssumedRoleUser(); + user.setTrn(roleTrn); + result.setAssumedRoleUser(user); + + AssumeRoleResponse response = new AssumeRoleResponse(); + response.setResult(result); + return response; + } + + private static HMSClientPool directPool(HMSClient client) { + return new HMSClientPool() { + @Override + public <R> R run(ClientPool.Action<R, HMSClient, TException> action) throws TException { + return action.run(client); + } + + @Override + public <R> R run(ClientPool.Action<R, HMSClient, TException> action, boolean retry) + throws TException { + return action.run(client); + } + }; + } +} diff --git a/amoro-common/src/main/java/org/apache/amoro/hive/HMSClient.java b/amoro-common/src/main/java/org/apache/amoro/hive/HMSClient.java index 140324abc..70c721c41 100644 --- a/amoro-common/src/main/java/org/apache/amoro/hive/HMSClient.java +++ b/amoro-common/src/main/java/org/apache/amoro/hive/HMSClient.java @@ -88,4 +88,7 @@ public interface HMSClient { InvocationTargetException, ClassNotFoundException; List<Table> getTableObjectsByName(String dbName, List<String> tableNames) throws TException; + + /** List catalogs exposed by HMS3. */ + List<String> getCatalogs() throws TException; } diff --git a/amoro-common/src/main/java/org/apache/amoro/hive/HMSClientImpl.java b/amoro-common/src/main/java/org/apache/amoro/hive/HMSClientImpl.java index 85ddb6833..01b2e3de7 100644 --- a/amoro-common/src/main/java/org/apache/amoro/hive/HMSClientImpl.java +++ b/amoro-common/src/main/java/org/apache/amoro/hive/HMSClientImpl.java @@ -187,4 +187,9 @@ public class HMSClientImpl implements HMSClient { throws TException { return getClient().getTableObjectsByName(dbName, tableNames); } + + @Override + public List<String> getCatalogs() throws TException { + return getClient().getCatalogs(); + } } diff --git a/dist/src/main/amoro-bin/conf/config.yaml b/dist/src/main/amoro-bin/conf/config.yaml index d505e3823..c14d19089 100644 --- a/dist/src/main/amoro-bin/conf/config.yaml +++ b/dist/src/main/amoro-bin/conf/config.yaml @@ -75,6 +75,9 @@ ams: # Values may also be supplied through environment variables, for example: # AMS_LAS_INTEGRATION_ENABLED=true # AMS_LAS_INTEGRATION_HMS__URI=thrift://hms-service.<namespace>.svc.cluster.local:9083 + # Keep IAM bootstrap credentials in secret-backed environment variables: + # AMS_LAS_INTEGRATION_IAM_BOOTSTRAP__ACCESS__KEY and + # AMS_LAS_INTEGRATION_IAM_BOOTSTRAP__SECRET__KEY. las: integration: enabled: false @@ -82,11 +85,17 @@ ams: # tos-endpoint: https://tos-cn-beijing.volces.com # iam-endpoint: https://iam.volcengineapi.com # emr-serverless-endpoint: https://open.volcengineapi.com - # optimizer-jar-uri: tos://<bucket>/<path>/amoro-optimizer.jar region: cn-beijing emr-serverless-service: emr_serverless connect-timeout: 30s read-timeout: 30s + iam: + # bootstrap-access-key: supplied by a secret-backed environment variable + # bootstrap-secret-key: supplied by a secret-backed environment variable + # bootstrap-session-token: optional, supplied with temporary bootstrap credentials + role-session-name: AmoroAssumeRoleSession + assume-role-ttl: 1h + credential-cache-size: 1000 cross-vpc: enabled: false # account-id: <account-id> @@ -140,7 +149,7 @@ ams: # Support for encrypted sensitive configuration items shade: identifier: default # Built-in support for default/base64. Defaults to "default", indicating no encryption - sensitive-keywords: admin-password;database.password;http-server.authorization.ldap-role-mapping.bind-password + sensitive-keywords: admin-password;database.password;http-server.authorization.ldap-role-mapping.bind-password;las.integration.iam.bootstrap-access-key;las.integration.iam.bootstrap-secret-key;las.integration.iam.bootstrap-session-token overview-cache: refresh-interval: 3min # 3 min
