Browse Source
* Refactoring `Sagemaker` task plugin with connections managed in connection center. --------- Co-authored-by: Eric Gao <ericgao.apache@gmail.com>3.2.1-prepare
chenrj
1 year ago
committed by
GitHub
29 changed files with 729 additions and 22 deletions
@ -0,0 +1,50 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?> |
||||
<!-- |
||||
~ 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. |
||||
--> |
||||
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" |
||||
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd"> |
||||
<modelVersion>4.0.0</modelVersion> |
||||
<parent> |
||||
<groupId>org.apache.dolphinscheduler</groupId> |
||||
<artifactId>dolphinscheduler-datasource-plugin</artifactId> |
||||
<version>dev-SNAPSHOT</version> |
||||
</parent> |
||||
|
||||
<artifactId>dolphinscheduler-datasource-sagemaker</artifactId> |
||||
<packaging>jar</packaging> |
||||
<name>${project.artifactId}</name> |
||||
|
||||
<dependencies> |
||||
<dependency> |
||||
<groupId>org.apache.dolphinscheduler</groupId> |
||||
<artifactId>dolphinscheduler-spi</artifactId> |
||||
<scope>provided</scope> |
||||
</dependency> |
||||
|
||||
<dependency> |
||||
<groupId>org.apache.dolphinscheduler</groupId> |
||||
<artifactId>dolphinscheduler-datasource-api</artifactId> |
||||
<version>${project.version}</version> |
||||
</dependency> |
||||
|
||||
<dependency> |
||||
<groupId>com.amazonaws</groupId> |
||||
<artifactId>aws-java-sdk-sagemaker</artifactId> |
||||
</dependency> |
||||
</dependencies> |
||||
|
||||
</project> |
@ -0,0 +1,65 @@
|
||||
/* |
||||
* 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.dolphinscheduler.plugin.datasource.sagemaker; |
||||
|
||||
import static com.google.common.base.Preconditions.checkNotNull; |
||||
|
||||
import lombok.extern.slf4j.Slf4j; |
||||
|
||||
import com.amazonaws.auth.AWSCredentialsProvider; |
||||
import com.amazonaws.auth.AWSStaticCredentialsProvider; |
||||
import com.amazonaws.auth.BasicAWSCredentials; |
||||
import com.amazonaws.services.sagemaker.AmazonSageMaker; |
||||
import com.amazonaws.services.sagemaker.AmazonSageMakerClientBuilder; |
||||
import com.amazonaws.services.sagemaker.model.ListNotebookInstancesRequest; |
||||
|
||||
@Slf4j |
||||
public class SagemakerClientWrapper implements AutoCloseable { |
||||
|
||||
private AmazonSageMaker amazonSageMaker; |
||||
|
||||
public SagemakerClientWrapper(String accessKey, String secretAccessKey, String region) { |
||||
checkNotNull(accessKey, "sagemaker accessKey cannot be null"); |
||||
checkNotNull(secretAccessKey, "sagemaker secretAccessKey cannot be null"); |
||||
checkNotNull(region, "sagemaker region cannot be null"); |
||||
|
||||
final BasicAWSCredentials basicAWSCredentials = new BasicAWSCredentials(accessKey, secretAccessKey); |
||||
final AWSCredentialsProvider awsCredentialsProvider = new AWSStaticCredentialsProvider(basicAWSCredentials); |
||||
// create a SageMaker client
|
||||
amazonSageMaker = AmazonSageMakerClientBuilder.standard().withCredentials(awsCredentialsProvider) |
||||
.withRegion(region).build(); |
||||
} |
||||
|
||||
public boolean checkConnect() { |
||||
try { |
||||
// If listing notebook instances fails, an exception will be thrown directly
|
||||
ListNotebookInstancesRequest request = new ListNotebookInstancesRequest(); |
||||
amazonSageMaker.listNotebookInstances(request); |
||||
log.info("sagemaker client connects to server successfully"); |
||||
return true; |
||||
} catch (Exception e) { |
||||
log.info("sagemaker client failed to connect to the server"); |
||||
} |
||||
return false; |
||||
} |
||||
|
||||
@Override |
||||
public void close() { |
||||
|
||||
} |
||||
} |
@ -0,0 +1,37 @@
|
||||
/* |
||||
* 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.dolphinscheduler.plugin.datasource.sagemaker; |
||||
|
||||
import org.apache.dolphinscheduler.spi.datasource.AdHocDataSourceClient; |
||||
import org.apache.dolphinscheduler.spi.datasource.BaseConnectionParam; |
||||
import org.apache.dolphinscheduler.spi.datasource.DataSourceChannel; |
||||
import org.apache.dolphinscheduler.spi.datasource.PooledDataSourceClient; |
||||
import org.apache.dolphinscheduler.spi.enums.DbType; |
||||
|
||||
public class SagemakerDataSourceChannel implements DataSourceChannel { |
||||
|
||||
@Override |
||||
public AdHocDataSourceClient createAdHocDataSourceClient(BaseConnectionParam baseConnectionParam, DbType dbType) { |
||||
throw new UnsupportedOperationException("Sagemaker AdHocDataSourceClient is not supported"); |
||||
} |
||||
|
||||
@Override |
||||
public PooledDataSourceClient createPooledDataSourceClient(BaseConnectionParam baseConnectionParam, DbType dbType) { |
||||
throw new UnsupportedOperationException("Sagemaker AdHocDataSourceClient is not supported"); |
||||
} |
||||
} |
@ -0,0 +1,38 @@
|
||||
/* |
||||
* 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.dolphinscheduler.plugin.datasource.sagemaker; |
||||
|
||||
import org.apache.dolphinscheduler.spi.datasource.DataSourceChannel; |
||||
import org.apache.dolphinscheduler.spi.datasource.DataSourceChannelFactory; |
||||
|
||||
import com.google.auto.service.AutoService; |
||||
|
||||
@AutoService(DataSourceChannelFactory.class) |
||||
public class SagemakerDataSourceChannelFactory implements DataSourceChannelFactory { |
||||
|
||||
@Override |
||||
public DataSourceChannel create() { |
||||
return new SagemakerDataSourceChannel(); |
||||
} |
||||
|
||||
@Override |
||||
public String getName() { |
||||
return "sagemaker"; |
||||
} |
||||
|
||||
} |
@ -0,0 +1,35 @@
|
||||
/* |
||||
* 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.dolphinscheduler.plugin.datasource.sagemaker.param; |
||||
|
||||
import org.apache.dolphinscheduler.spi.datasource.ConnectionParam; |
||||
|
||||
import lombok.Data; |
||||
|
||||
import com.fasterxml.jackson.annotation.JsonInclude; |
||||
|
||||
@Data |
||||
@JsonInclude(JsonInclude.Include.NON_NULL) |
||||
public class SagemakerConnectionParam implements ConnectionParam { |
||||
|
||||
protected String userName; |
||||
|
||||
protected String password; |
||||
|
||||
protected String awsRegion; |
||||
} |
@ -0,0 +1,34 @@
|
||||
/* |
||||
* 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.dolphinscheduler.plugin.datasource.sagemaker.param; |
||||
|
||||
import org.apache.dolphinscheduler.plugin.datasource.api.datasource.BaseDataSourceParamDTO; |
||||
import org.apache.dolphinscheduler.spi.enums.DbType; |
||||
|
||||
import lombok.Data; |
||||
|
||||
@Data |
||||
public class SagemakerDataSourceParamDTO extends BaseDataSourceParamDTO { |
||||
|
||||
protected String awsRegion; |
||||
|
||||
@Override |
||||
public DbType getType() { |
||||
return DbType.SAGEMAKER; |
||||
} |
||||
} |
@ -0,0 +1,136 @@
|
||||
/* |
||||
* 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.dolphinscheduler.plugin.datasource.sagemaker.param; |
||||
|
||||
import org.apache.dolphinscheduler.common.utils.JSONUtils; |
||||
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.sagemaker.SagemakerClientWrapper; |
||||
import org.apache.dolphinscheduler.spi.datasource.ConnectionParam; |
||||
import org.apache.dolphinscheduler.spi.enums.DbType; |
||||
|
||||
import org.apache.commons.lang3.StringUtils; |
||||
|
||||
import java.sql.Connection; |
||||
import java.text.MessageFormat; |
||||
|
||||
import lombok.extern.slf4j.Slf4j; |
||||
|
||||
import com.google.auto.service.AutoService; |
||||
|
||||
@AutoService(DataSourceProcessor.class) |
||||
@Slf4j |
||||
public class SagemakerDataSourceProcessor implements DataSourceProcessor { |
||||
|
||||
@Override |
||||
public BaseDataSourceParamDTO castDatasourceParamDTO(String paramJson) { |
||||
return JSONUtils.parseObject(paramJson, SagemakerDataSourceParamDTO.class); |
||||
} |
||||
|
||||
@Override |
||||
public void checkDatasourceParam(BaseDataSourceParamDTO datasourceParamDTO) { |
||||
SagemakerDataSourceParamDTO sageMakerDataSourceParamDTO = (SagemakerDataSourceParamDTO) datasourceParamDTO; |
||||
if (StringUtils.isEmpty(sageMakerDataSourceParamDTO.getUserName()) |
||||
|| StringUtils.isEmpty(sageMakerDataSourceParamDTO.getPassword()) |
||||
|| StringUtils.isEmpty(sageMakerDataSourceParamDTO.getAwsRegion())) { |
||||
throw new IllegalArgumentException("sagemaker datasource param is not valid"); |
||||
} |
||||
} |
||||
|
||||
@Override |
||||
public String getDatasourceUniqueId(ConnectionParam connectionParam, DbType dbType) { |
||||
SagemakerConnectionParam baseConnectionParam = (SagemakerConnectionParam) connectionParam; |
||||
return MessageFormat.format("{0}@{1}@{2}@{3}", dbType.getDescp(), |
||||
PasswordUtils.encodePassword(baseConnectionParam.getUserName()), |
||||
PasswordUtils.encodePassword(baseConnectionParam.getPassword()), |
||||
PasswordUtils.encodePassword(baseConnectionParam.getAwsRegion())); |
||||
} |
||||
|
||||
// SageMaker
|
||||
@Override |
||||
public BaseDataSourceParamDTO createDatasourceParamDTO(String connectionJson) { |
||||
SagemakerConnectionParam connectionParams = (SagemakerConnectionParam) createConnectionParams(connectionJson); |
||||
SagemakerDataSourceParamDTO sagemakerDataSourceParamDTO = new SagemakerDataSourceParamDTO(); |
||||
|
||||
sagemakerDataSourceParamDTO.setUserName(connectionParams.getUserName()); |
||||
sagemakerDataSourceParamDTO.setPassword(connectionParams.getPassword()); |
||||
sagemakerDataSourceParamDTO.setAwsRegion(connectionParams.getAwsRegion()); |
||||
return sagemakerDataSourceParamDTO; |
||||
} |
||||
|
||||
@Override |
||||
public SagemakerConnectionParam createConnectionParams(BaseDataSourceParamDTO datasourceParam) { |
||||
SagemakerDataSourceParamDTO sageMakerDataSourceParam = (SagemakerDataSourceParamDTO) datasourceParam; |
||||
SagemakerConnectionParam sageMakerConnectionParam = new SagemakerConnectionParam(); |
||||
sageMakerConnectionParam.setUserName(sageMakerDataSourceParam.getUserName()); |
||||
sageMakerConnectionParam.setPassword(sageMakerDataSourceParam.getPassword()); |
||||
sageMakerConnectionParam.setAwsRegion(sageMakerDataSourceParam.getAwsRegion()); |
||||
|
||||
return sageMakerConnectionParam; |
||||
} |
||||
|
||||
@Override |
||||
public ConnectionParam createConnectionParams(String connectionJson) { |
||||
return JSONUtils.parseObject(connectionJson, SagemakerConnectionParam.class); |
||||
} |
||||
|
||||
@Override |
||||
public String getDatasourceDriver() { |
||||
return ""; |
||||
} |
||||
|
||||
@Override |
||||
public String getValidationQuery() { |
||||
return ""; |
||||
} |
||||
|
||||
@Override |
||||
public String getJdbcUrl(ConnectionParam connectionParam) { |
||||
return ""; |
||||
} |
||||
|
||||
@Override |
||||
public Connection getConnection(ConnectionParam connectionParam) { |
||||
return null; |
||||
} |
||||
|
||||
@Override |
||||
public boolean checkDataSourceConnectivity(ConnectionParam connectionParam) { |
||||
SagemakerConnectionParam baseConnectionParam = (SagemakerConnectionParam) connectionParam; |
||||
try ( |
||||
SagemakerClientWrapper sagemakerClientWrapper = |
||||
new SagemakerClientWrapper(baseConnectionParam.userName, |
||||
baseConnectionParam.password, baseConnectionParam.awsRegion)) { |
||||
return sagemakerClientWrapper.checkConnect(); |
||||
} catch (Exception e) { |
||||
log.error("sagemaker client failed to connect to the server", e); |
||||
return false; |
||||
} |
||||
} |
||||
|
||||
@Override |
||||
public DbType getDbType() { |
||||
return DbType.SAGEMAKER; |
||||
} |
||||
|
||||
@Override |
||||
public DataSourceProcessor create() { |
||||
return new SagemakerDataSourceProcessor(); |
||||
} |
||||
} |
@ -0,0 +1,115 @@
|
||||
/* |
||||
* 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.dolphinscheduler.plugin.datasource.sagemaker; |
||||
|
||||
import org.apache.dolphinscheduler.plugin.datasource.sagemaker.param.SagemakerConnectionParam; |
||||
import org.apache.dolphinscheduler.plugin.datasource.sagemaker.param.SagemakerDataSourceParamDTO; |
||||
import org.apache.dolphinscheduler.plugin.datasource.sagemaker.param.SagemakerDataSourceProcessor; |
||||
import org.apache.dolphinscheduler.spi.enums.DbType; |
||||
|
||||
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.MockedConstruction; |
||||
import org.mockito.Mockito; |
||||
import org.mockito.junit.jupiter.MockitoExtension; |
||||
|
||||
@ExtendWith(MockitoExtension.class) |
||||
|
||||
public class SagemakerDataSourceProcessorTest { |
||||
|
||||
private SagemakerDataSourceProcessor sagemakerDataSourceProcessor; |
||||
|
||||
private String connectJson = |
||||
"{\"userName\":\"access key\",\"password\":\"secret access key\",\"awsRegion\":\"region\"}"; |
||||
|
||||
@BeforeEach |
||||
public void init() { |
||||
sagemakerDataSourceProcessor = new SagemakerDataSourceProcessor(); |
||||
} |
||||
|
||||
@Test |
||||
void testCheckDatasourceParam() { |
||||
SagemakerDataSourceParamDTO sagemakerDataSourceParamDTO = new SagemakerDataSourceParamDTO(); |
||||
Assertions.assertThrows(IllegalArgumentException.class, |
||||
() -> sagemakerDataSourceProcessor.checkDatasourceParam(sagemakerDataSourceParamDTO)); |
||||
sagemakerDataSourceParamDTO.setUserName("access key"); |
||||
Assertions.assertThrows(IllegalArgumentException.class, |
||||
() -> sagemakerDataSourceProcessor.checkDatasourceParam(sagemakerDataSourceParamDTO)); |
||||
sagemakerDataSourceParamDTO.setPassword("secret access key"); |
||||
Assertions.assertThrows(IllegalArgumentException.class, |
||||
() -> sagemakerDataSourceProcessor.checkDatasourceParam(sagemakerDataSourceParamDTO)); |
||||
sagemakerDataSourceParamDTO.setAwsRegion("region"); |
||||
|
||||
Assertions |
||||
.assertDoesNotThrow( |
||||
() -> sagemakerDataSourceProcessor.checkDatasourceParam(sagemakerDataSourceParamDTO)); |
||||
} |
||||
|
||||
@Test |
||||
void testGetDatasourceUniqueId() { |
||||
SagemakerConnectionParam sagemakerConnectionParam = new SagemakerConnectionParam(); |
||||
sagemakerConnectionParam.setUserName("access key"); |
||||
sagemakerConnectionParam.setPassword("secret access key"); |
||||
sagemakerConnectionParam.setAwsRegion("region"); |
||||
Assertions.assertEquals("sagemaker@access key@secret access key@region", |
||||
sagemakerDataSourceProcessor.getDatasourceUniqueId(sagemakerConnectionParam, DbType.SAGEMAKER)); |
||||
|
||||
} |
||||
|
||||
@Test |
||||
void testCreateDatasourceParamDTO() { |
||||
SagemakerDataSourceParamDTO sagemakerDataSourceParamDTO = |
||||
(SagemakerDataSourceParamDTO) sagemakerDataSourceProcessor.createDatasourceParamDTO(connectJson); |
||||
Assertions.assertEquals("access key", sagemakerDataSourceParamDTO.getUserName()); |
||||
Assertions.assertEquals("secret access key", sagemakerDataSourceParamDTO.getPassword()); |
||||
Assertions.assertEquals("region", sagemakerDataSourceParamDTO.getAwsRegion()); |
||||
} |
||||
|
||||
@Test |
||||
void testCreateConnectionParams() { |
||||
SagemakerDataSourceParamDTO sagemakerDataSourceParamDTO = |
||||
(SagemakerDataSourceParamDTO) sagemakerDataSourceProcessor.createDatasourceParamDTO(connectJson); |
||||
SagemakerConnectionParam sagemakerConnectionParam = |
||||
sagemakerDataSourceProcessor.createConnectionParams(sagemakerDataSourceParamDTO); |
||||
Assertions.assertEquals("access key", sagemakerConnectionParam.getUserName()); |
||||
Assertions.assertEquals("secret access key", sagemakerConnectionParam.getPassword()); |
||||
Assertions.assertEquals("region", sagemakerConnectionParam.getAwsRegion()); |
||||
} |
||||
|
||||
@Test |
||||
void testTestConnection() { |
||||
SagemakerDataSourceParamDTO sagemakerDataSourceParamDTO = |
||||
(SagemakerDataSourceParamDTO) sagemakerDataSourceProcessor.createDatasourceParamDTO(connectJson); |
||||
SagemakerConnectionParam connectionParam = |
||||
sagemakerDataSourceProcessor.createConnectionParams(sagemakerDataSourceParamDTO); |
||||
Assertions.assertFalse(sagemakerDataSourceProcessor.checkDataSourceConnectivity(connectionParam)); |
||||
|
||||
try ( |
||||
MockedConstruction<SagemakerClientWrapper> sshClientWrapperMockedConstruction = |
||||
Mockito.mockConstruction(SagemakerClientWrapper.class, (mock, context) -> { |
||||
Mockito.when( |
||||
mock.checkConnect()) |
||||
.thenReturn(true); |
||||
})) { |
||||
Assertions.assertTrue(sagemakerDataSourceProcessor.checkDataSourceConnectivity(connectionParam)); |
||||
} |
||||
|
||||
} |
||||
} |
@ -0,0 +1,45 @@
|
||||
/* |
||||
* 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.dolphinscheduler.plugin.task.sagemaker; |
||||
|
||||
import java.io.Serializable; |
||||
|
||||
/** |
||||
* master/worker task transport |
||||
*/ |
||||
|
||||
public class SagemakerTaskExecutionContext implements Serializable { |
||||
|
||||
/** |
||||
* connectionParams |
||||
*/ |
||||
private String connectionParams; |
||||
|
||||
public String getConnectionParams() { |
||||
return connectionParams; |
||||
} |
||||
|
||||
public void setConnectionParams(String connectionParams) { |
||||
this.connectionParams = connectionParams; |
||||
} |
||||
|
||||
@Override |
||||
public String toString() { |
||||
return "SagemakerTaskExecutionContext{" + "connectionParams='" + connectionParams + '\'' + '}'; |
||||
} |
||||
} |
Loading…
Reference in new issue