forked from raystack/dagger
-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add capability to dagger to read python udfs from Ali(oss) and Tencen…
…t(cosn) storage services Given the configuration provided correctly. Set the below environment variables accordingly to access the files stored in the respective bucket. Ali(oss) - OSS_ACCESS_KEY_ID - OSS_ACCESS_KEY_SECRET Tencent(cos) - COS_SECRET_ID - COS_SECRET_KEY - COS_REGION
- Loading branch information
rajuGT
committed
Oct 15, 2024
1 parent
715a65d
commit d0bb427
Showing
11 changed files
with
425 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
66 changes: 66 additions & 0 deletions
66
...src/main/java/com/gotocompany/dagger/functions/udfs/python/file/source/cos/CosClient.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,66 @@ | ||
package com.gotocompany.dagger.functions.udfs.python.file.source.cos; | ||
|
||
import com.qcloud.cos.COSClient; | ||
import com.qcloud.cos.ClientConfig; | ||
import com.qcloud.cos.auth.BasicCOSCredentials; | ||
import com.qcloud.cos.auth.COSCredentials; | ||
import com.qcloud.cos.model.COSObject; | ||
import com.qcloud.cos.model.COSObjectInputStream; | ||
import com.qcloud.cos.region.Region; | ||
import com.qcloud.cos.utils.IOUtils; | ||
|
||
import java.io.IOException; | ||
import java.util.Arrays; | ||
import java.util.List; | ||
import java.util.stream.Collectors; | ||
|
||
public class CosClient { | ||
|
||
// TODO find better way to initialize clients | ||
private final String ENV_COS_SECRET_ID = "COS_SECRET_ID"; | ||
private final String ENV_COS_SECRET_KEY = "COS_SECRET_KEY"; | ||
private final String ENV_COS_REGION = "COS_REGION"; | ||
|
||
private final COSClient libCosClient; | ||
|
||
/** | ||
* Instantiates a new Cos client. | ||
*/ | ||
public CosClient() { | ||
String secretID = System.getenv(ENV_COS_SECRET_ID); | ||
String secretKey = System.getenv(ENV_COS_SECRET_KEY); | ||
String region = System.getenv(ENV_COS_REGION); // ap-singapore | ||
|
||
COSCredentials credentials = new BasicCOSCredentials(secretID, secretKey); | ||
ClientConfig clientConfig = new ClientConfig(new Region(region)); | ||
libCosClient = new COSClient(credentials, clientConfig); | ||
} | ||
|
||
/** | ||
* Instantiates a new Cos client. | ||
* This constructor used for unit test purposes. | ||
* | ||
* @param libCosClient the storage | ||
*/ | ||
public CosClient(COSClient libCosClient) { | ||
this.libCosClient = libCosClient; | ||
} | ||
|
||
/** | ||
* Get file byte [ ]. | ||
* | ||
* @param pythonFile the python file | ||
* @return the byte [ ] | ||
*/ | ||
public byte[] getFile(String pythonFile) throws IOException { | ||
List<String> file = Arrays.asList(pythonFile.replace("cosn://", "").split("/")); | ||
|
||
String bucketName = file.get(0); | ||
String objectName = file.stream().skip(1).collect(Collectors.joining("/")); | ||
|
||
COSObject cosObject = libCosClient.getObject(bucketName, objectName); | ||
try (COSObjectInputStream inputStream = cosObject.getObjectContent()) { | ||
return IOUtils.toByteArray(inputStream); | ||
} | ||
} | ||
} |
47 changes: 47 additions & 0 deletions
47
...main/java/com/gotocompany/dagger/functions/udfs/python/file/source/cos/CosFileSource.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,47 @@ | ||
package com.gotocompany.dagger.functions.udfs.python.file.source.cos; | ||
|
||
import com.gotocompany.dagger.functions.udfs.python.file.source.FileSource; | ||
|
||
import java.io.IOException; | ||
|
||
public class CosFileSource implements FileSource { | ||
|
||
private CosClient cosClient; | ||
private final String pythonFile; | ||
|
||
/** | ||
* Instantiates a new Cos file source. | ||
* | ||
* @param pythonFile the python file | ||
*/ | ||
public CosFileSource(String pythonFile) { | ||
this.pythonFile = pythonFile; | ||
} | ||
|
||
/** | ||
* Instantiates a new Cos file source. | ||
* | ||
* @param pythonFile the python file | ||
*/ | ||
public CosFileSource(String pythonFile, CosClient cosClient) { | ||
this.pythonFile = pythonFile; | ||
this.cosClient = cosClient; | ||
} | ||
|
||
@Override | ||
public byte[] getObjectFile() throws IOException { | ||
return getCosClient().getFile(pythonFile); | ||
} | ||
|
||
/** | ||
* Gets cos client. | ||
* | ||
* @return the cos client | ||
*/ | ||
private CosClient getCosClient() { | ||
if (this.cosClient == null) { | ||
this.cosClient = new CosClient(); | ||
} | ||
return this.cosClient; | ||
} | ||
} |
61 changes: 61 additions & 0 deletions
61
...src/main/java/com/gotocompany/dagger/functions/udfs/python/file/source/oss/OssClient.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,61 @@ | ||
package com.gotocompany.dagger.functions.udfs.python.file.source.oss; | ||
|
||
import com.aliyun.core.utils.IOUtils; | ||
import com.aliyun.oss.OSS; | ||
import com.aliyun.oss.OSSClientBuilder; | ||
import com.aliyun.oss.common.auth.CredentialsProviderFactory; | ||
import com.aliyun.oss.model.OSSObject; | ||
import com.aliyuncs.exceptions.ClientException; | ||
|
||
import java.io.IOException; | ||
import java.io.InputStream; | ||
import java.util.Arrays; | ||
import java.util.List; | ||
import java.util.stream.Collectors; | ||
|
||
public class OssClient { | ||
|
||
// TODO refactor to take this value from the configuration | ||
private static final String endpoint = "oss-cn-hangzhou.aliyuncs.com"; | ||
|
||
private final OSS libOssClient; | ||
|
||
/** | ||
* Instantiates a new Oss client. | ||
*/ | ||
public OssClient() { | ||
try { | ||
libOssClient = new OSSClientBuilder().build(endpoint, CredentialsProviderFactory.newEnvironmentVariableCredentialsProvider()); | ||
} catch (ClientException e) { | ||
throw new RuntimeException("failed to initialise oss client", e); | ||
} | ||
} | ||
|
||
/** | ||
* Instantiates a new OSS client. | ||
* This constructor used for unit test purposes. | ||
* | ||
* @param libOssClient the storage | ||
*/ | ||
public OssClient(OSS libOssClient) { | ||
this.libOssClient = libOssClient; | ||
} | ||
|
||
/** | ||
* Get file byte [ ]. | ||
* | ||
* @param pythonFile the python file | ||
* @return the byte [ ] | ||
*/ | ||
public byte[] getFile(String pythonFile) throws IOException { | ||
List<String> file = Arrays.asList(pythonFile.replace("oss://", "").split("/")); | ||
|
||
String bucketName = file.get(0); | ||
String objectName = file.stream().skip(1).collect(Collectors.joining("/")); | ||
|
||
OSSObject ossObject = libOssClient.getObject(bucketName, objectName); | ||
try (InputStream inputStream = ossObject.getObjectContent()) { | ||
return IOUtils.toByteArray(inputStream); | ||
} | ||
} | ||
} |
47 changes: 47 additions & 0 deletions
47
...main/java/com/gotocompany/dagger/functions/udfs/python/file/source/oss/OssFileSource.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,47 @@ | ||
package com.gotocompany.dagger.functions.udfs.python.file.source.oss; | ||
|
||
import com.gotocompany.dagger.functions.udfs.python.file.source.FileSource; | ||
|
||
import java.io.IOException; | ||
|
||
public class OssFileSource implements FileSource { | ||
|
||
private OssClient ossClient; | ||
private final String pythonFile; | ||
|
||
/** | ||
* Instantiates a new Oss file source. | ||
* | ||
* @param pythonFile the python file | ||
*/ | ||
public OssFileSource(String pythonFile) { | ||
this.pythonFile = pythonFile; | ||
} | ||
|
||
/** | ||
* Instantiates a new Oss file source. | ||
* | ||
* @param pythonFile the python file | ||
*/ | ||
public OssFileSource(String pythonFile, OssClient ossClient) { | ||
this.pythonFile = pythonFile; | ||
this.ossClient = ossClient; | ||
} | ||
|
||
@Override | ||
public byte[] getObjectFile() throws IOException { | ||
return getOssClient().getFile(pythonFile); | ||
} | ||
|
||
/** | ||
* Gets oss client. | ||
* | ||
* @return the oss client | ||
*/ | ||
private OssClient getOssClient() { | ||
if (this.ossClient == null) { | ||
this.ossClient = new OssClient(); | ||
} | ||
return this.ossClient; | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
52 changes: 52 additions & 0 deletions
52
...test/java/com/gotocompany/dagger/functions/udfs/python/file/source/cos/CosClientTest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,52 @@ | ||
package com.gotocompany.dagger.functions.udfs.python.file.source.cos; | ||
|
||
import com.qcloud.cos.COSClient; | ||
import com.qcloud.cos.model.COSObject; | ||
import com.qcloud.cos.model.COSObjectInputStream; | ||
import org.apache.http.client.methods.HttpRequestBase; | ||
import org.junit.Assert; | ||
import org.junit.Before; | ||
import org.junit.Test; | ||
import org.mockito.Mock; | ||
import org.mockito.Mockito; | ||
|
||
import java.io.ByteArrayInputStream; | ||
import java.io.IOException; | ||
import java.util.Arrays; | ||
|
||
import static org.mockito.Mockito.*; | ||
import static org.mockito.MockitoAnnotations.initMocks; | ||
|
||
public class CosClientTest { | ||
|
||
@Mock | ||
private COSClient libCosClient; | ||
|
||
@Mock | ||
private COSObject cosObject; | ||
|
||
@Before | ||
public void setup() { | ||
initMocks(this); | ||
} | ||
|
||
@Test | ||
public void shouldGetObjectFile() throws IOException { | ||
HttpRequestBase mockRequest = Mockito.mock(HttpRequestBase.class); | ||
|
||
String pythonFile = "cosn://bucket_name/path/to/file/python_udf.zip"; | ||
String bucketName = "bucket_name"; | ||
String objectName = "path/to/file/python_udf.zip"; | ||
String expectedValue = Arrays.toString("objectFile".getBytes()); | ||
|
||
when(libCosClient.getObject(bucketName, objectName)).thenReturn(cosObject); | ||
when(cosObject.getObjectContent()).thenReturn(new COSObjectInputStream(new ByteArrayInputStream("objectFile".getBytes()), mockRequest)); | ||
|
||
CosClient cosClient = new CosClient(libCosClient); | ||
byte[] actualValue = cosClient.getFile(pythonFile); | ||
|
||
verify(libCosClient, times(1)).getObject(bucketName, objectName); | ||
verify(cosObject, times(1)).getObjectContent(); | ||
Assert.assertEquals(expectedValue, Arrays.toString(actualValue)); | ||
} | ||
} |
38 changes: 38 additions & 0 deletions
38
.../java/com/gotocompany/dagger/functions/udfs/python/file/source/cos/CosFileSourceTest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,38 @@ | ||
package com.gotocompany.dagger.functions.udfs.python.file.source.cos; | ||
|
||
import org.junit.Assert; | ||
import org.junit.Before; | ||
import org.junit.Test; | ||
import org.mockito.Mock; | ||
|
||
import java.io.IOException; | ||
import java.nio.file.Files; | ||
import java.nio.file.Paths; | ||
|
||
import static org.mockito.Mockito.when; | ||
import static org.mockito.MockitoAnnotations.initMocks; | ||
|
||
public class CosFileSourceTest { | ||
|
||
@Mock | ||
private CosClient cosClient; | ||
|
||
@Before | ||
public void setup() { | ||
initMocks(this); | ||
} | ||
|
||
@Test | ||
public void shouldGetObjectFile() throws IOException { | ||
ClassLoader classLoader = getClass().getClassLoader(); | ||
String pythonFile = classLoader.getResource("python_udf.zip").getFile(); | ||
byte[] expectedObject = Files.readAllBytes(Paths.get(pythonFile)); | ||
|
||
when(cosClient.getFile(pythonFile)).thenReturn(expectedObject); | ||
CosFileSource cosFileSource = new CosFileSource(pythonFile, cosClient); | ||
|
||
byte[] actualObject = cosFileSource.getObjectFile(); | ||
|
||
Assert.assertEquals(expectedObject, actualObject); | ||
} | ||
} |
Oops, something went wrong.