Skip to content

Commit 4b45194

Browse files
author
Chris Grass
committed
Add azure_remote_url and corresponding test.
Signed-off-by: Chris Grass <chris.grass@mcg.com>
1 parent e97adba commit 4b45194

File tree

4 files changed

+90
-1
lines changed

4 files changed

+90
-1
lines changed

flyteadmin/pkg/common/cloud.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ type CloudProvider = string
77
const (
88
AWS CloudProvider = "aws"
99
GCP CloudProvider = "gcp"
10+
Azure CloudProvider = "azure"
1011
Sandbox CloudProvider = "sandbox"
1112
Local CloudProvider = "local"
1213
None CloudProvider = "none"

flyteadmin/pkg/data/factory.go

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,11 @@ func GetRemoteDataHandler(cfg RemoteDataHandlerConfig) RemoteDataHandler {
5050
return &remoteDataHandler{
5151
remoteURL: implementations.NewGCPRemoteURL(cfg.SigningPrincipal, signedURLDuration),
5252
}
53-
53+
case common.Azure:
54+
signedURLDuration := time.Minute * time.Duration(cfg.SignedURLDurationMinutes)
55+
return &remoteDataHandler{
56+
remoteURL: implementations.NewAzureRemoteURL(*cfg.RemoteDataStoreClient, signedURLDuration),
57+
}
5458
case common.Local:
5559
logger.Infof(context.TODO(), "setting up local signer ----- ")
5660
// Since minio = aws s3, we are creating the same client but using the config primitives from aws
Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
package implementations
2+
3+
import (
4+
"context"
5+
"github.com/flyteorg/flyte/flyteadmin/pkg/data/interfaces"
6+
"github.com/flyteorg/flyte/flyteadmin/pkg/errors"
7+
"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/admin"
8+
"github.com/flyteorg/flyte/flytestdlib/storage"
9+
"github.com/flyteorg/stow"
10+
"google.golang.org/grpc/codes"
11+
"time"
12+
)
13+
14+
type AzureRemoteURL struct {
15+
remoteDataStoreClient storage.DataStore
16+
presignDuration time.Duration
17+
}
18+
19+
func (n *AzureRemoteURL) Get(ctx context.Context, uri string) (admin.UrlBlob, error) {
20+
metadata, err := n.remoteDataStoreClient.Head(ctx, storage.DataReference(uri))
21+
if err != nil {
22+
return admin.UrlBlob{}, errors.NewFlyteAdminErrorf(codes.Internal,
23+
"failed to get metadata for uri: %s with err: %v", uri, err)
24+
}
25+
26+
signedUri, err := n.remoteDataStoreClient.CreateSignedURL(ctx, storage.DataReference(uri), storage.SignedURLProperties{
27+
Scope: stow.ClientMethodGet,
28+
ExpiresIn: n.presignDuration,
29+
})
30+
if err != nil {
31+
return admin.UrlBlob{}, errors.NewFlyteAdminErrorf(codes.Internal,
32+
"failed to get metadata for uri: %s with err: %v", uri, err)
33+
}
34+
35+
return admin.UrlBlob{
36+
Url: signedUri.URL.String(),
37+
Bytes: metadata.Size(),
38+
}, nil
39+
}
40+
41+
func NewAzureRemoteURL(remoteDataStoreClient storage.DataStore, presignDuration time.Duration) interfaces.RemoteURLInterface {
42+
return &AzureRemoteURL{
43+
remoteDataStoreClient: remoteDataStoreClient,
44+
presignDuration: presignDuration,
45+
}
46+
}
Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
package implementations
2+
3+
import (
4+
"context"
5+
commonMocks "github.com/flyteorg/flyte/flyteadmin/pkg/common/mocks"
6+
"github.com/flyteorg/flyte/flytestdlib/storage"
7+
"github.com/stretchr/testify/assert"
8+
"testing"
9+
)
10+
11+
type mockMetadata struct{}
12+
13+
func (m mockMetadata) Exists() bool {
14+
return true
15+
}
16+
17+
func (m mockMetadata) Size() int64 {
18+
return 1
19+
}
20+
21+
func (m mockMetadata) Etag() string {
22+
return "etag"
23+
}
24+
25+
func TestAzureGet(t *testing.T) {
26+
inputUri := "abfs//test/data"
27+
mockStorage := commonMocks.GetMockStorageClient()
28+
mockStorage.ComposedProtobufStore.(*commonMocks.TestDataStore).HeadCb =
29+
func(ctx context.Context, reference storage.DataReference) (storage.Metadata, error) {
30+
return mockMetadata{}, nil
31+
}
32+
remoteUrl := AzureRemoteURL{
33+
remoteDataStoreClient: *mockStorage, presignDuration: 1,
34+
}
35+
36+
result, _ := remoteUrl.Get(context.TODO(), inputUri)
37+
assert.Contains(t, inputUri, result.Url)
38+
}

0 commit comments

Comments
 (0)