fix format error

This commit is contained in:
EricGao888 2024-06-05 20:04:36 +08:00
parent 66618f009c
commit 7370c1e46b
9 changed files with 268 additions and 85 deletions

View File

@ -18,7 +18,9 @@
package org.apache.dolphinscheduler.plugin.datasource.aliyunserverlessspark;
import static com.google.common.base.Preconditions.checkNotNull;
import lombok.extern.slf4j.Slf4j;
import com.aliyun.emr_serverless_spark20230808.Client;
import com.aliyun.teaopenapi.models.Config;
@ -28,17 +30,17 @@ public class AliyunServerlessSparkClientWrapper implements AutoCloseable {
private Client aliyunServerlessSparkClient;
public AliyunServerlessSparkClientWrapper(
String accessKeyId,
String accessKeySecret,
String regionId)
throws Exception {
String accessKeyId,
String accessKeySecret,
String regionId)
throws Exception {
checkNotNull(accessKeyId, accessKeySecret, regionId);
String endpoint = String.format("emr-serverless-spark.%s.aliyuncs.com", regionId);
Config config = new Config()
.setEndpoint(endpoint)
.setAccessKeyId(accessKeyId)
.setAccessKeySecret(accessKeySecret);
.setEndpoint(endpoint)
.setAccessKeyId(accessKeyId)
.setAccessKeySecret(accessKeySecret);
aliyunServerlessSparkClient = new Client(config);
}

View File

@ -31,9 +31,9 @@ public class AliyunServerlessSparkUtils {
public static Client getAliyunServerlessSparkClient(AliyunServerlessSparkConnectionParam connectionParam) throws Exception {
String endpoint = String.format("emr-serverless-spark.%s.aliyuncs.com", connectionParam.getRegionId());
Config config = new Config()
.setEndpoint(endpoint)
.setAccessKeyId(connectionParam.getAccessKeyId())
.setAccessKeySecret(connectionParam.getAccessKeySecret());
.setEndpoint(endpoint)
.setAccessKeyId(connectionParam.getAccessKeyId())
.setAccessKeySecret(connectionParam.getAccessKeySecret());
return new Client(config);
}

View File

@ -18,11 +18,11 @@
package org.apache.dolphinscheduler.plugin.datasource.aliyunserverlessspark.param;
import org.apache.dolphinscheduler.common.utils.JSONUtils;
import org.apache.dolphinscheduler.plugin.datasource.aliyunserverlessspark.AliyunServerlessSparkClientWrapper;
import org.apache.dolphinscheduler.plugin.datasource.api.datasource.AbstractDataSourceProcessor;
import org.apache.dolphinscheduler.plugin.datasource.api.datasource.BaseDataSourceParamDTO;
import org.apache.dolphinscheduler.plugin.datasource.api.datasource.DataSourceProcessor;
import org.apache.dolphinscheduler.plugin.datasource.api.utils.PasswordUtils;
import org.apache.dolphinscheduler.plugin.datasource.aliyunserverlessspark.AliyunServerlessSparkClientWrapper;
import org.apache.dolphinscheduler.spi.datasource.ConnectionParam;
import org.apache.dolphinscheduler.spi.enums.DbType;
@ -46,29 +46,33 @@ public class AliyunServerlessSparkDataSourceProcessor extends AbstractDataSource
@Override
public void checkDatasourceParam(BaseDataSourceParamDTO datasourceParamDTO) {
AliyunServerlessSparkDataSourceParamDTO aliyunServerlessSparkDataSourceParamDTO = (AliyunServerlessSparkDataSourceParamDTO) datasourceParamDTO;
AliyunServerlessSparkDataSourceParamDTO aliyunServerlessSparkDataSourceParamDTO =
(AliyunServerlessSparkDataSourceParamDTO) datasourceParamDTO;
if (StringUtils.isEmpty(aliyunServerlessSparkDataSourceParamDTO.getRegionId()) ||
StringUtils.isEmpty(aliyunServerlessSparkDataSourceParamDTO.getAccessKeyId()) ||
StringUtils.isEmpty(aliyunServerlessSparkDataSourceParamDTO.getRegionId())) {
StringUtils.isEmpty(aliyunServerlessSparkDataSourceParamDTO.getAccessKeyId()) ||
StringUtils.isEmpty(aliyunServerlessSparkDataSourceParamDTO.getRegionId())) {
throw new IllegalArgumentException("spark datasource param is not valid");
}
}
@Override
public String getDatasourceUniqueId(ConnectionParam connectionParam, DbType dbType) {
AliyunServerlessSparkConnectionParam baseConnectionParam = (AliyunServerlessSparkConnectionParam) connectionParam;
AliyunServerlessSparkConnectionParam baseConnectionParam =
(AliyunServerlessSparkConnectionParam) connectionParam;
return MessageFormat.format(
"{0}@{1}@{2}@{3}",
dbType.getName(),
baseConnectionParam.getRegionId(),
PasswordUtils.encodePassword(baseConnectionParam.getAccessKeyId()),
PasswordUtils.encodePassword(baseConnectionParam.getAccessKeySecret()));
"{0}@{1}@{2}@{3}",
dbType.getName(),
baseConnectionParam.getRegionId(),
PasswordUtils.encodePassword(baseConnectionParam.getAccessKeyId()),
PasswordUtils.encodePassword(baseConnectionParam.getAccessKeySecret()));
}
@Override
public BaseDataSourceParamDTO createDatasourceParamDTO(String connectionJson) {
AliyunServerlessSparkConnectionParam connectionParams = (AliyunServerlessSparkConnectionParam) createConnectionParams(connectionJson);
AliyunServerlessSparkDataSourceParamDTO aliyunServerlessSparkDataSourceParamDTO = new AliyunServerlessSparkDataSourceParamDTO();
AliyunServerlessSparkConnectionParam connectionParams =
(AliyunServerlessSparkConnectionParam) createConnectionParams(connectionJson);
AliyunServerlessSparkDataSourceParamDTO aliyunServerlessSparkDataSourceParamDTO =
new AliyunServerlessSparkDataSourceParamDTO();
aliyunServerlessSparkDataSourceParamDTO.setAccessKeyId(connectionParams.getAccessKeyId());
aliyunServerlessSparkDataSourceParamDTO.setAccessKeySecret(connectionParams.getAccessKeySecret());
@ -78,10 +82,13 @@ public class AliyunServerlessSparkDataSourceProcessor extends AbstractDataSource
@Override
public AliyunServerlessSparkConnectionParam createConnectionParams(BaseDataSourceParamDTO datasourceParam) {
AliyunServerlessSparkDataSourceParamDTO aliyunServerlessSparkDataSourceParamDTO = (AliyunServerlessSparkDataSourceParamDTO) datasourceParam;
AliyunServerlessSparkConnectionParam aliyunServerlessSparkConnectionParam = new AliyunServerlessSparkConnectionParam();
AliyunServerlessSparkDataSourceParamDTO aliyunServerlessSparkDataSourceParamDTO =
(AliyunServerlessSparkDataSourceParamDTO) datasourceParam;
AliyunServerlessSparkConnectionParam aliyunServerlessSparkConnectionParam =
new AliyunServerlessSparkConnectionParam();
aliyunServerlessSparkConnectionParam.setAccessKeyId(aliyunServerlessSparkDataSourceParamDTO.getAccessKeyId());
aliyunServerlessSparkConnectionParam.setAccessKeySecret(aliyunServerlessSparkDataSourceParamDTO.getAccessKeySecret());
aliyunServerlessSparkConnectionParam
.setAccessKeySecret(aliyunServerlessSparkDataSourceParamDTO.getAccessKeySecret());
aliyunServerlessSparkConnectionParam.setRegionId(aliyunServerlessSparkDataSourceParamDTO.getRegionId());
return aliyunServerlessSparkConnectionParam;
@ -114,18 +121,18 @@ public class AliyunServerlessSparkDataSourceProcessor extends AbstractDataSource
@Override
public boolean checkDataSourceConnectivity(ConnectionParam connectionParam) {
AliyunServerlessSparkConnectionParam baseConnectionParam = (AliyunServerlessSparkConnectionParam) connectionParam;
AliyunServerlessSparkConnectionParam baseConnectionParam =
(AliyunServerlessSparkConnectionParam) connectionParam;
try (
AliyunServerlessSparkClientWrapper aliyunServerlessSparkClientWrapper =
AliyunServerlessSparkClientWrapper aliyunServerlessSparkClientWrapper =
new AliyunServerlessSparkClientWrapper(
baseConnectionParam.getAccessKeyId(),
baseConnectionParam.getAccessKeySecret(),
baseConnectionParam.getRegionId())
) {
baseConnectionParam.getAccessKeyId(),
baseConnectionParam.getAccessKeySecret(),
baseConnectionParam.getRegionId())) {
return aliyunServerlessSparkClientWrapper.checkConnect(
baseConnectionParam.getAccessKeyId(),
baseConnectionParam.getAccessKeySecret(),
baseConnectionParam.getRegionId());
baseConnectionParam.getAccessKeyId(),
baseConnectionParam.getAccessKeySecret(),
baseConnectionParam.getRegionId());
} catch (Exception e) {
log.error("spark client failed to connect to the server", e);
return false;

View File

@ -24,7 +24,6 @@ import org.apache.dolphinscheduler.spi.enums.DbType;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.MockedConstruction;
@ -46,30 +45,36 @@ public class AliyunServerlessSparkDataSourceProcessorTest {
@Test
void testCheckDatasourceParam() {
AliyunServerlessSparkDataSourceParamDTO aliyunServerlessSparkDataSourceParamDTO = new AliyunServerlessSparkDataSourceParamDTO();
AliyunServerlessSparkDataSourceParamDTO aliyunServerlessSparkDataSourceParamDTO =
new AliyunServerlessSparkDataSourceParamDTO();
aliyunServerlessSparkDataSourceParamDTO.setRegionId("cn-hangzhou");
Assertions.assertThrows(IllegalArgumentException.class,
() -> aliyunServerlessSparkDataSourceProcessor.checkDatasourceParam(aliyunServerlessSparkDataSourceParamDTO));
() -> aliyunServerlessSparkDataSourceProcessor
.checkDatasourceParam(aliyunServerlessSparkDataSourceParamDTO));
aliyunServerlessSparkDataSourceParamDTO.setAccessKeyId("mockAccessKeyId");
aliyunServerlessSparkDataSourceParamDTO.setAccessKeySecret("mockAccessKeySecret");
Assertions
.assertDoesNotThrow(() -> aliyunServerlessSparkDataSourceProcessor.checkDatasourceParam(aliyunServerlessSparkDataSourceParamDTO));
.assertDoesNotThrow(() -> aliyunServerlessSparkDataSourceProcessor
.checkDatasourceParam(aliyunServerlessSparkDataSourceParamDTO));
}
@Test
void testGetDatasourceUniqueId() {
AliyunServerlessSparkConnectionParam aliyunServerlessSparkConnectionParam = new AliyunServerlessSparkConnectionParam();
AliyunServerlessSparkConnectionParam aliyunServerlessSparkConnectionParam =
new AliyunServerlessSparkConnectionParam();
aliyunServerlessSparkConnectionParam.setRegionId("cn-hangzhou");
aliyunServerlessSparkConnectionParam.setAccessKeyId("mockAccessKeyId");
aliyunServerlessSparkConnectionParam.setAccessKeySecret("mockAccessKeySecret");
Assertions.assertEquals("aliyun_serverless_spark@cn-hangzhou@mockAccessKeyId@mockAccessKeySecret",
aliyunServerlessSparkDataSourceProcessor.getDatasourceUniqueId(aliyunServerlessSparkConnectionParam, DbType.ALIYUN_SERVERLESS_SPARK));
aliyunServerlessSparkDataSourceProcessor.getDatasourceUniqueId(aliyunServerlessSparkConnectionParam,
DbType.ALIYUN_SERVERLESS_SPARK));
}
@Test
void testCreateDatasourceParamDTO() {
AliyunServerlessSparkDataSourceParamDTO aliyunServerlessSparkDataSourceParamDTO =
(AliyunServerlessSparkDataSourceParamDTO) aliyunServerlessSparkDataSourceProcessor.createDatasourceParamDTO(connectJson);
(AliyunServerlessSparkDataSourceParamDTO) aliyunServerlessSparkDataSourceProcessor
.createDatasourceParamDTO(connectJson);
Assertions.assertEquals("cn-hangzhou", aliyunServerlessSparkDataSourceParamDTO.getRegionId());
Assertions.assertEquals("mockAccessKeyId", aliyunServerlessSparkDataSourceParamDTO.getAccessKeyId());
Assertions.assertEquals("mockAccessKeySecret", aliyunServerlessSparkDataSourceParamDTO.getAccessKeySecret());
@ -78,9 +83,11 @@ public class AliyunServerlessSparkDataSourceProcessorTest {
@Test
void testCreateConnectionParams() {
AliyunServerlessSparkDataSourceParamDTO aliyunServerlessSparkDataSourceParamDTO =
(AliyunServerlessSparkDataSourceParamDTO) aliyunServerlessSparkDataSourceProcessor.createDatasourceParamDTO(connectJson);
(AliyunServerlessSparkDataSourceParamDTO) aliyunServerlessSparkDataSourceProcessor
.createDatasourceParamDTO(connectJson);
AliyunServerlessSparkConnectionParam aliyunServerlessSparkConnectionParam =
aliyunServerlessSparkDataSourceProcessor.createConnectionParams(aliyunServerlessSparkDataSourceParamDTO);
aliyunServerlessSparkDataSourceProcessor
.createConnectionParams(aliyunServerlessSparkDataSourceParamDTO);
Assertions.assertEquals("cn-hangzhou", aliyunServerlessSparkConnectionParam.getRegionId());
Assertions.assertEquals("mockAccessKeyId", aliyunServerlessSparkConnectionParam.getAccessKeyId());
Assertions.assertEquals("mockAccessKeySecret", aliyunServerlessSparkConnectionParam.getAccessKeySecret());
@ -89,18 +96,22 @@ public class AliyunServerlessSparkDataSourceProcessorTest {
@Test
void testTestConnection() {
AliyunServerlessSparkDataSourceParamDTO aliyunServerlessSparkDataSourceParamDTO =
(AliyunServerlessSparkDataSourceParamDTO) aliyunServerlessSparkDataSourceProcessor.createDatasourceParamDTO(connectJson);
(AliyunServerlessSparkDataSourceParamDTO) aliyunServerlessSparkDataSourceProcessor
.createDatasourceParamDTO(connectJson);
AliyunServerlessSparkConnectionParam connectionParam =
aliyunServerlessSparkDataSourceProcessor.createConnectionParams(aliyunServerlessSparkDataSourceParamDTO);
aliyunServerlessSparkDataSourceProcessor
.createConnectionParams(aliyunServerlessSparkDataSourceParamDTO);
Assertions.assertTrue(aliyunServerlessSparkDataSourceProcessor.checkDataSourceConnectivity(connectionParam));
try (
MockedConstruction<AliyunServerlessSparkClientWrapper> AliyunServerlessSparkClientWrapper =
Mockito.mockConstruction(AliyunServerlessSparkClientWrapper.class, (mock, context) -> {
Mockito.when(
mock.checkConnect(connectionParam.getAccessKeyId(), connectionParam.getAccessKeySecret(), connectionParam.getRegionId()))
mock.checkConnect(connectionParam.getAccessKeyId(),
connectionParam.getAccessKeySecret(), connectionParam.getRegionId()))
.thenReturn(true);
})) {
Assertions.assertTrue(aliyunServerlessSparkDataSourceProcessor.checkDataSourceConnectivity(connectionParam));
Assertions
.assertTrue(aliyunServerlessSparkDataSourceProcessor.checkDataSourceConnectivity(connectionParam));
}
}
}

View File

@ -16,18 +16,17 @@
*/
package org.apache.dolphinscheduler.plugin.task.aliyunserverlessspark;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.dolphinscheduler.plugin.task.api.enums.ResourceType;
import org.apache.dolphinscheduler.plugin.task.api.parameters.AbstractParameters;
import org.apache.dolphinscheduler.plugin.task.api.parameters.resource.ResourceParametersHelper;
import java.util.List;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
@Data
@Slf4j
public class AliyunServerlessSparkParameters extends AbstractParameters {
// spark job configurations
private String workspaceId;
private String resourceQueueId;

View File

@ -31,7 +31,7 @@ import org.apache.dolphinscheduler.plugin.task.api.parameters.resource.DataSourc
import org.apache.dolphinscheduler.plugin.task.api.parameters.resource.ResourceParametersHelper;
import org.apache.dolphinscheduler.spi.enums.DbType;
import com.aliyun.emr_serverless_spark20230808.Client;
import org.apache.commons.lang3.StringUtils;
import java.util.Arrays;
import java.util.Collections;
@ -41,8 +41,7 @@ import java.util.Map;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import com.aliyun.emr_serverless_spark20230808.Client;
import com.aliyun.emr_serverless_spark20230808.models.CancelJobRunRequest;
import com.aliyun.emr_serverless_spark20230808.models.GetJobRunRequest;
import com.aliyun.emr_serverless_spark20230808.models.GetJobRunResponse;
@ -88,11 +87,12 @@ public class AliyunServerlessSparkTask extends AbstractRemoteTask {
}
ResourceParametersHelper resourceParametersHelper = taskExecutionContext.getResourceParametersHelper();
DataSourceParameters dataSourceParameters = (DataSourceParameters) resourceParametersHelper.getResourceParameters(ResourceType.DATASOURCE, aliyunServerlessSparkParameters.getDatasource());
DataSourceParameters dataSourceParameters = (DataSourceParameters) resourceParametersHelper
.getResourceParameters(ResourceType.DATASOURCE, aliyunServerlessSparkParameters.getDatasource());
aliyunServerlessSparkConnectionParam = (AliyunServerlessSparkConnectionParam) DataSourceUtils
.buildConnectionParams(
DbType.valueOf(aliyunServerlessSparkParameters.getType()),
dataSourceParameters.getConnectionParams());
.buildConnectionParams(
DbType.valueOf(aliyunServerlessSparkParameters.getType()),
dataSourceParameters.getConnectionParams());
accessKeyId = aliyunServerlessSparkConnectionParam.getAccessKeyId();
accessKeySecret = aliyunServerlessSparkConnectionParam.getAccessKeySecret();
@ -114,15 +114,17 @@ public class AliyunServerlessSparkTask extends AbstractRemoteTask {
StartJobRunRequest startJobRunRequest = buildStartJobRunRequest(aliyunServerlessSparkParameters);
RuntimeOptions runtime = new RuntimeOptions();
Map<String, String> headers = new HashMap<>();
StartJobRunResponse startJobRunResponse = aliyunServerlessSparkClient.startJobRunWithOptions(aliyunServerlessSparkParameters.getWorkspaceId(), startJobRunRequest, headers, runtime);
log.info("[debug111] aliyunServerlessSparkParameters - {}", aliyunServerlessSparkParameters);
StartJobRunResponse startJobRunResponse = aliyunServerlessSparkClient.startJobRunWithOptions(
aliyunServerlessSparkParameters.getWorkspaceId(), startJobRunRequest, headers, runtime);
jobRunId = startJobRunResponse.getBody().getJobRunId();
setAppIds(jobRunId);
log.info("Successfully submitted serverless spark job, jobRunId - {}", jobRunId);
// todo: deal with null
while(!RunState.isFinal(currentState)) {
while (!RunState.isFinal(currentState)) {
GetJobRunRequest getJobRunRequest = buildGetJobRunRequest(aliyunServerlessSparkParameters);
GetJobRunResponse getJobRunResponse = aliyunServerlessSparkClient.getJobRun(aliyunServerlessSparkParameters.getWorkspaceId(), jobRunId, getJobRunRequest);
GetJobRunResponse getJobRunResponse = aliyunServerlessSparkClient
.getJobRun(aliyunServerlessSparkParameters.getWorkspaceId(), jobRunId, getJobRunRequest);
currentState = RunState.valueOf(getJobRunResponse.getBody().getJobRun().getState());
log.info("job - {} state - {}", jobRunId, currentState);
Thread.sleep(10 * 1000L);
@ -146,7 +148,7 @@ public class AliyunServerlessSparkTask extends AbstractRemoteTask {
}
private int mapFinalStateToExitCode(RunState state) {
protected int mapFinalStateToExitCode(RunState state) {
switch (state) {
case Success:
return TaskConstants.EXIT_CODE_SUCCESS;
@ -166,7 +168,8 @@ public class AliyunServerlessSparkTask extends AbstractRemoteTask {
public void cancelApplication() throws TaskException {
CancelJobRunRequest cancelJobRunRequest = buildCancelJobRunRequest(aliyunServerlessSparkParameters);
try {
aliyunServerlessSparkClient.cancelJobRun(aliyunServerlessSparkParameters.getWorkspaceId(), jobRunId, cancelJobRunRequest);
aliyunServerlessSparkClient.cancelJobRun(aliyunServerlessSparkParameters.getWorkspaceId(), jobRunId,
cancelJobRunRequest);
} catch (Exception e) {
log.error("Failed to cancel serverless spark job run", e);
}
@ -177,23 +180,26 @@ public class AliyunServerlessSparkTask extends AbstractRemoteTask {
return Collections.emptyList();
}
private Client buildAliyunServerlessSparkClient(String accessKeyId, String accessKeySecret, String regionId) throws Exception {
protected Client buildAliyunServerlessSparkClient(String accessKeyId, String accessKeySecret,
String regionId) throws Exception {
String endpoint = String.format("emr-serverless-spark.%s.aliyuncs.com", regionId);
Config config = new Config()
.setEndpoint(endpoint)
.setAccessKeyId(accessKeyId)
.setAccessKeySecret(accessKeySecret);
.setEndpoint(endpoint)
.setAccessKeyId(accessKeyId)
.setAccessKeySecret(accessKeySecret);
return new Client(config);
}
private StartJobRunRequest buildStartJobRunRequest(AliyunServerlessSparkParameters aliyunServerlessSparkParameters) {
protected StartJobRunRequest buildStartJobRunRequest(AliyunServerlessSparkParameters aliyunServerlessSparkParameters) {
StartJobRunRequest startJobRunRequest = new StartJobRunRequest();
startJobRunRequest.setRegionId(regionId);
startJobRunRequest.setResourceQueueId(aliyunServerlessSparkParameters.getResourceQueueId());
startJobRunRequest.setCodeType(aliyunServerlessSparkParameters.getCodeType());
startJobRunRequest.setName(aliyunServerlessSparkParameters.getJobName());
String engineReleaseVersion = aliyunServerlessSparkParameters.getEngineReleaseVersion();
engineReleaseVersion = StringUtils.isEmpty(engineReleaseVersion) ? "esr-2.1-native (Spark 3.3.1, Scala 2.12, Native Runtime)" : engineReleaseVersion;
engineReleaseVersion =
StringUtils.isEmpty(engineReleaseVersion) ? "esr-2.1-native (Spark 3.3.1, Scala 2.12, Native Runtime)"
: engineReleaseVersion;
startJobRunRequest.setReleaseVersion(engineReleaseVersion);
Tag envTag = new Tag();
envTag.setKey("environment");
@ -203,25 +209,26 @@ public class AliyunServerlessSparkTask extends AbstractRemoteTask {
workflowTag.setKey("workflow");
workflowTag.setValue("true");
startJobRunRequest.setTags(Arrays.asList(envTag, workflowTag));
List<String> entryPointArguments = StringUtils.isEmpty(aliyunServerlessSparkParameters.getEntryPointArguments()) ?
Collections.emptyList() : Arrays.asList(aliyunServerlessSparkParameters.getEntryPointArguments().split("#"));
List<String> entryPointArguments =
StringUtils.isEmpty(aliyunServerlessSparkParameters.getEntryPointArguments()) ? Collections.emptyList()
: Arrays.asList(aliyunServerlessSparkParameters.getEntryPointArguments().split("#"));
JobDriver.JobDriverSparkSubmit jobDriverSparkSubmit = new JobDriver.JobDriverSparkSubmit()
.setEntryPoint(aliyunServerlessSparkParameters.getEntryPoint())
.setEntryPointArguments(entryPointArguments)
.setSparkSubmitParameters(aliyunServerlessSparkParameters.getSparkSubmitParameters());
.setEntryPoint(aliyunServerlessSparkParameters.getEntryPoint())
.setEntryPointArguments(entryPointArguments)
.setSparkSubmitParameters(aliyunServerlessSparkParameters.getSparkSubmitParameters());
JobDriver jobDriver = new com.aliyun.emr_serverless_spark20230808.models.JobDriver()
.setSparkSubmit(jobDriverSparkSubmit);
.setSparkSubmit(jobDriverSparkSubmit);
startJobRunRequest.setJobDriver(jobDriver);
return startJobRunRequest;
}
private GetJobRunRequest buildGetJobRunRequest(AliyunServerlessSparkParameters aliyunServerlessSparkParameters) {
protected GetJobRunRequest buildGetJobRunRequest(AliyunServerlessSparkParameters aliyunServerlessSparkParameters) {
GetJobRunRequest getJobRunRequest = new GetJobRunRequest();
getJobRunRequest.setRegionId(regionId);
return getJobRunRequest;
}
private CancelJobRunRequest buildCancelJobRunRequest(AliyunServerlessSparkParameters aliyunServerlessSparkParameters) {
protected CancelJobRunRequest buildCancelJobRunRequest(AliyunServerlessSparkParameters aliyunServerlessSparkParameters) {
CancelJobRunRequest cancelJobRunRequest = new CancelJobRunRequest();
cancelJobRunRequest.setRegionId(regionId);
return cancelJobRunRequest;

View File

@ -25,6 +25,7 @@ import org.apache.dolphinscheduler.plugin.task.api.parameters.ParametersNode;
import org.apache.dolphinscheduler.plugin.task.api.parameters.resource.ResourceParametersHelper;
public class AliyunServerlessSparkTaskChannel implements TaskChannel {
@Override
public void cancelApplication(boolean status) {

View File

@ -26,6 +26,7 @@ import com.google.auto.service.AutoService;
@AutoService(TaskChannelFactory.class)
public class AliyunServerlessSparkTaskChannelFactory implements TaskChannelFactory {
@Override
public String getName() {
return "ALIYUN_SERVERLESS_SPARK";

View File

@ -16,36 +16,191 @@
*/
package org.apache.dolphinscheduler.plugin.task.aliyunserverlessspark;
import org.apache.dolphinscheduler.plugin.task.api.TaskException;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import org.apache.dolphinscheduler.plugin.task.api.TaskCallBack;
import org.apache.dolphinscheduler.plugin.task.api.TaskConstants;
import org.apache.dolphinscheduler.plugin.task.api.TaskExecutionContext;
import org.apache.dolphinscheduler.plugin.task.api.enums.ResourceType;
import org.apache.dolphinscheduler.plugin.task.api.parameters.resource.DataSourceParameters;
import org.apache.dolphinscheduler.plugin.task.api.parameters.resource.ResourceParametersHelper;
import org.apache.dolphinscheduler.spi.enums.DbType;
import java.util.Collections;
import lombok.extern.slf4j.Slf4j;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.InjectMocks;
import org.mockito.Mock;
import org.mockito.Spy;
import org.mockito.junit.jupiter.MockitoExtension;
import org.mockito.junit.jupiter.MockitoSettings;
import org.mockito.quality.Strictness;
import com.aliyun.emr_serverless_spark20230808.Client;
import com.aliyun.emr_serverless_spark20230808.models.CancelJobRunRequest;
import com.aliyun.emr_serverless_spark20230808.models.CancelJobRunResponse;
import com.aliyun.emr_serverless_spark20230808.models.GetJobRunRequest;
import com.aliyun.emr_serverless_spark20230808.models.GetJobRunResponse;
import com.aliyun.emr_serverless_spark20230808.models.GetJobRunResponseBody;
import com.aliyun.emr_serverless_spark20230808.models.StartJobRunRequest;
import com.aliyun.emr_serverless_spark20230808.models.StartJobRunResponse;
import com.aliyun.emr_serverless_spark20230808.models.Tag;
import com.aliyun.tea.TeaException;
import com.aliyun.emr_serverless_spark20230808.models.StartJobRunResponseBody;
@Slf4j
@ExtendWith(MockitoExtension.class)
@MockitoSettings(strictness = Strictness.LENIENT)
public class AliyunServerlessSparkTaskTest {
@Mock
private TaskExecutionContext mockTaskExecutionContext;
@Mock
private Client mockAliyunServerlessSparkClient;
@Mock
private ResourceParametersHelper mockResourceParametersHelper;
@Mock
private TaskCallBack mockTaskCallBack;
@Mock
private StartJobRunRequest mockStartJobRunRequest;
@Mock
private StartJobRunResponse mockStartJobRunResponse;
@Mock
private GetJobRunRequest mockGetJobRunRequest;
@Mock
private GetJobRunResponse mockGetJobRunResponse;
@Mock
private CancelJobRunRequest mockCancelJobRunRequest;
@Mock
private CancelJobRunResponse mockCancelJobRunResponse;
@InjectMocks
@Spy
private AliyunServerlessSparkTask aliyunServerlessSparkTask;
private static final String mockAccessKeyId = "mockAccessKeyId";
private static final String mockAccessKeySecret = "mockAccessKeySecret";
private static final String mockRegionId = "cn-hangzhou";
private static final int mockDatasourceId = 1;
private static final String taskParamsString =
"{\"localParams\":[],\"resourceList\":[],\"workspaceId\":\"w-ae42e9c929275cc5\",\"resourceQueueId\":\"root_queue\",\"codeType\":\"JAR\",\"jobName\":\"spark\",\"entryPoint\":\"oss://datadev-oss-hdfs-test/spark-resource/examples/jars/spark-examples_2.12-3.3.1.jar\",\"entryPointArguments\":\"10\",\"sparkSubmitParameters\":\"--class org.apache.spark.examples.SparkPi --conf spark.executor.cores=4 --conf spark.executor.memory=20g --conf spark.driver.cores=4 --conf spark.driver.memory=8g --conf spark.executor.instances=1\",\"isProduction\":true,\"type\":\"ALIYUN_SERVERLESS_SPARK\",\"datasource\":1}";
private static final String connectionParamsString =
"{\"accessKeyId\":\"mockAccessKeyId\",\"accessKeySecret\":\"mockAccessKeySecret\",\"regionId\":\"cn-hangzhou\",\"password\":\"\"}";
private static final String mockJobRunId = "jr-f6a1d0dd17d6b8a3";
private static final String mockWorkspaceId = "w-ae42e9c929275cc5";
private static final String mockResourceQueueId = "root_queue";
private static final String mockSparkSubmitParameters =
"--class org.apache.spark.examples.SparkPi --conf spark.executor.cores=4 --conf spark.executor.memory=20g --conf spark.driver.cores=4 --conf spark.driver.memory=8g --conf spark.executor.instances=1";
private static final String mockEntryPoint =
"oss://datadev-oss-hdfs-test/spark-resource/examples/jars/spark-examples_2.12-3.3.1.jar";
private static final String mockEntryPointArguments = "10";
@BeforeEach
public void before() {
when(mockTaskExecutionContext.getTaskParams()).thenReturn(taskParamsString);
DataSourceParameters dataSourceParameters = new DataSourceParameters();
dataSourceParameters.setConnectionParams(connectionParamsString);
dataSourceParameters.setType(DbType.ALIYUN_SERVERLESS_SPARK);
when(mockResourceParametersHelper.getResourceParameters(any(), any())).thenReturn(dataSourceParameters);
when(mockTaskExecutionContext.getResourceParametersHelper()).thenReturn(mockResourceParametersHelper);
Assertions.assertDoesNotThrow(
() -> when(aliyunServerlessSparkTask.buildAliyunServerlessSparkClient(any(), any(), any()))
.thenReturn(mockAliyunServerlessSparkClient));
}
public void testInit() {
@Test
public void testInit() throws Exception {
aliyunServerlessSparkTask.init();
verify(mockTaskExecutionContext).getTaskParams();
verify(mockResourceParametersHelper).getResourceParameters(ResourceType.DATASOURCE, mockDatasourceId);
verify(aliyunServerlessSparkTask).buildAliyunServerlessSparkClient(mockAccessKeyId, mockAccessKeySecret,
mockRegionId);
}
@Test
public void testHandle() {
doReturn(mockStartJobRunRequest).when(aliyunServerlessSparkTask).buildStartJobRunRequest(any());
StartJobRunResponseBody startJobRunResponseBody = new StartJobRunResponseBody();
startJobRunResponseBody.setJobRunId(mockJobRunId);
doReturn(startJobRunResponseBody).when(mockStartJobRunResponse).getBody();
Assertions.assertDoesNotThrow(
() -> doReturn(mockStartJobRunResponse).when(mockAliyunServerlessSparkClient)
.startJobRunWithOptions(any(), any(), any(), any()));
doReturn(mockGetJobRunRequest).when(aliyunServerlessSparkTask).buildGetJobRunRequest(any());
GetJobRunResponseBody getJobRunResponseBody = new GetJobRunResponseBody();
GetJobRunResponseBody.GetJobRunResponseBodyJobRun jobRun =
new GetJobRunResponseBody.GetJobRunResponseBodyJobRun();
jobRun.setState(RunState.Success.name());
getJobRunResponseBody.setJobRun(jobRun);
doReturn(getJobRunResponseBody).when(mockGetJobRunResponse).getBody();
Assertions.assertDoesNotThrow(
() -> doReturn(mockGetJobRunResponse).when(mockAliyunServerlessSparkClient).getJobRun(any(), any(),
any()));
aliyunServerlessSparkTask.init();
aliyunServerlessSparkTask.handle(mockTaskCallBack);
verify(aliyunServerlessSparkTask).setAppIds(mockJobRunId);
verify(aliyunServerlessSparkTask).setExitStatusCode(TaskConstants.EXIT_CODE_SUCCESS);
}
@Test
public void testCancelApplication() throws Exception {
doReturn(mockCancelJobRunRequest).when(aliyunServerlessSparkTask).buildCancelJobRunRequest(any());
Assertions.assertDoesNotThrow(
() -> doReturn(mockCancelJobRunResponse).when(mockAliyunServerlessSparkClient).cancelJobRun(any(),
any(), any()));
aliyunServerlessSparkTask.init();
aliyunServerlessSparkTask.cancelApplication();
verify(aliyunServerlessSparkTask).buildCancelJobRunRequest(any());
verify(mockAliyunServerlessSparkClient).cancelJobRun(eq(mockWorkspaceId), any(), eq(mockCancelJobRunRequest));
}
@Test
public void testBuildStartJobRunRequest() {
AliyunServerlessSparkParameters mockAliyunServerlessSparkParameters =
mock(AliyunServerlessSparkParameters.class);
doReturn(mockResourceQueueId).when(mockAliyunServerlessSparkParameters).getResourceQueueId();
doReturn("JAR").when(mockAliyunServerlessSparkParameters).getCodeType();
doReturn("ds-test").when(mockAliyunServerlessSparkParameters).getJobName();
doReturn(mockSparkSubmitParameters).when(mockAliyunServerlessSparkParameters).getSparkSubmitParameters();
doReturn(mockEntryPoint).when(mockAliyunServerlessSparkParameters).getEntryPoint();
doReturn(mockEntryPointArguments).when(mockAliyunServerlessSparkParameters).getEntryPointArguments();
aliyunServerlessSparkTask.buildStartJobRunRequest(mockAliyunServerlessSparkParameters);
verify(mockAliyunServerlessSparkParameters).getResourceQueueId();
verify(mockAliyunServerlessSparkParameters).getCodeType();
verify(mockAliyunServerlessSparkParameters).getJobName();
verify(mockAliyunServerlessSparkParameters).getEngineReleaseVersion();
verify(mockAliyunServerlessSparkParameters).isProduction();
}
}