diff --git a/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/datasource/aliyunserverlessspark/AliyunServerlessSparkClientWrapper.java b/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/datasource/aliyunserverlessspark/AliyunServerlessSparkClientWrapper.java index 078597d33e..45883aa8c4 100644 --- a/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/datasource/aliyunserverlessspark/AliyunServerlessSparkClientWrapper.java +++ b/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/datasource/aliyunserverlessspark/AliyunServerlessSparkClientWrapper.java @@ -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); } diff --git a/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/datasource/aliyunserverlessspark/AliyunServerlessSparkUtils.java b/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/datasource/aliyunserverlessspark/AliyunServerlessSparkUtils.java index 45799d172e..91928020a8 100644 --- a/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/datasource/aliyunserverlessspark/AliyunServerlessSparkUtils.java +++ b/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/datasource/aliyunserverlessspark/AliyunServerlessSparkUtils.java @@ -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); } diff --git a/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/datasource/aliyunserverlessspark/param/AliyunServerlessSparkDataSourceProcessor.java b/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/datasource/aliyunserverlessspark/param/AliyunServerlessSparkDataSourceProcessor.java index c2579c6eae..5a3c76d482 100644 --- a/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/datasource/aliyunserverlessspark/param/AliyunServerlessSparkDataSourceProcessor.java +++ b/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/datasource/aliyunserverlessspark/param/AliyunServerlessSparkDataSourceProcessor.java @@ -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; diff --git a/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-aliyunserverlessspark/src/test/java/org/apache/dolphinscheduler/plugin/datasource/aliyunserverlessspark/AliyunServerlessSparkDataSourceProcessorTest.java b/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-aliyunserverlessspark/src/test/java/org/apache/dolphinscheduler/plugin/datasource/aliyunserverlessspark/AliyunServerlessSparkDataSourceProcessorTest.java index 50aae38487..9973ef10f5 100644 --- a/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-aliyunserverlessspark/src/test/java/org/apache/dolphinscheduler/plugin/datasource/aliyunserverlessspark/AliyunServerlessSparkDataSourceProcessorTest.java +++ b/dolphinscheduler-datasource-plugin/dolphinscheduler-datasource-aliyunserverlessspark/src/test/java/org/apache/dolphinscheduler/plugin/datasource/aliyunserverlessspark/AliyunServerlessSparkDataSourceProcessorTest.java @@ -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 = 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)); } } } diff --git a/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkParameters.java b/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkParameters.java index 8e69dc735f..cf684a4629 100644 --- a/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkParameters.java +++ b/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkParameters.java @@ -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; diff --git a/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkTask.java b/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkTask.java index d6d60db030..891e1df725 100644 --- a/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkTask.java +++ b/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkTask.java @@ -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 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 entryPointArguments = StringUtils.isEmpty(aliyunServerlessSparkParameters.getEntryPointArguments()) ? - Collections.emptyList() : Arrays.asList(aliyunServerlessSparkParameters.getEntryPointArguments().split("#")); + List 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; diff --git a/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkTaskChannel.java b/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkTaskChannel.java index 0f67c94457..de91cc3515 100644 --- a/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkTaskChannel.java +++ b/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkTaskChannel.java @@ -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) { diff --git a/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkTaskChannelFactory.java b/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkTaskChannelFactory.java index 908c562e2e..a10e65a026 100644 --- a/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkTaskChannelFactory.java +++ b/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/main/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkTaskChannelFactory.java @@ -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"; diff --git a/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/test/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkTaskTest.java b/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/test/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkTaskTest.java index 1a360ea8a8..9a120ba344 100644 --- a/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/test/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkTaskTest.java +++ b/dolphinscheduler-task-plugin/dolphinscheduler-task-aliyunserverlessspark/src/test/java/org/apache/dolphinscheduler/plugin/task/aliyunserverlessspark/AliyunServerlessSparkTaskTest.java @@ -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(); + } + }