Add deckURI to NodeExecutionData (#413)

* Add deckURI to NodeExecutionData Signed-off-by: Kevin Su <[email protected]> * Use new storage api in stdlib Signed-off-by: Kevin Su <[email protected]> * Fixed tests Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * one more test Signed-off-by: Kevin Su <[email protected]> * Add deck_uri in NodeExecutionClosure Signed-off-by: Kevin Su <[email protected]> * Fixed tests Signed-off-by: Kevin Su <[email protected]> * update tests Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * lint fix Signed-off-by: Kevin Su <[email protected]> * Updated idl Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * Add http endpoint Signed-off-by: Kevin Su <[email protected]> * few updates Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]>
flyteorg · Jun 7, 2022 · f9ba260 · f9ba260
1 parent 7cac42f
commit f9ba260
Show file tree

Hide file tree

Showing 13 changed files with 139 additions and 9 deletions.
diff --git a/flyteadmin/dataproxy/service.go b/flyteadmin/dataproxy/service.go
@@ -5,6 +5,7 @@ import (
 	"encoding/base32"
 	"encoding/base64"
 	"fmt"
+	"net/url"
 	"strings"
 	"time"
 
@@ -86,6 +87,54 @@ func (s Service) CreateUploadLocation(ctx context.Context, req *service.CreateUp
 	}, nil
 }
 
+// CreateDownloadLocation creates a temporary signed url to allow callers to download content.
+func (s Service) CreateDownloadLocation(ctx context.Context, req *service.CreateDownloadLocationRequest) (
+	*service.CreateDownloadLocationResponse, error) {
+
+	if err := s.validateCreateDownloadLocationRequest(req); err != nil {
+		return nil, err
+	}
+
+	resp, err := s.dataStore.CreateSignedURL(ctx, storage.DataReference(req.NativeUrl), storage.SignedURLProperties{
+		Scope:     stow.ClientMethodGet,
+		ExpiresIn: req.ExpiresIn.AsDuration(),
+	})
+
+	if err != nil {
+		return nil, fmt.Errorf("failed to create a signed url. Error: %w", err)
+	}
+
+	return &service.CreateDownloadLocationResponse{
+		SignedUrl: resp.URL.String(),
+		ExpiresAt: timestamppb.New(time.Now().Add(req.ExpiresIn.AsDuration())),
+	}, nil
+}
+
+func (s Service) validateCreateDownloadLocationRequest(req *service.CreateDownloadLocationRequest) error {
+	if expiresIn := req.ExpiresIn; expiresIn != nil {
+		if !expiresIn.IsValid() {
+			return fmt.Errorf("expiresIn [%v] is invalid", expiresIn)
+		}
+
+		if expiresIn.AsDuration() < 0 {
+			return fmt.Errorf("expiresIn [%v] should not less than 0",
+				expiresIn.AsDuration().String())
+		} else if expiresIn.AsDuration() > s.cfg.Download.MaxExpiresIn.Duration {
+			return fmt.Errorf("expiresIn [%v] cannot exceed max allowed expiration [%v]",
+				expiresIn.AsDuration().String(), s.cfg.Download.MaxExpiresIn.String())
+		}
+	} else {
+		req.ExpiresIn = durationpb.New(s.cfg.Download.MaxExpiresIn.Duration)
+	}
+
+	if _, err := url.Parse(req.NativeUrl); err != nil {
+		return fmt.Errorf("failed to parse native_url [%v]",
+			req.NativeUrl)
+	}
+
+	return nil
+}
+
 // createShardedStorageLocation creates a location in storage destination to maximize read/write performance in most
 // block stores. The final location should look something like: s3://<my bucket>/<shard length>/<file name>
 func createShardedStorageLocation(ctx context.Context, shardSelector ioutils.ShardSelector, store *storage.DataStore,

diff --git a/flyteadmin/dataproxy/service_test.go b/flyteadmin/dataproxy/service_test.go
@@ -5,6 +5,9 @@ import (
 	"testing"
 	"time"
 
+	commonMocks "github.com/flyteorg/flyteadmin/pkg/common/mocks"
+	stdlibConfig "github.com/flyteorg/flytestdlib/config"
+
 	"google.golang.org/protobuf/types/known/durationpb"
 
 	"github.com/flyteorg/flytestdlib/contextutils"
@@ -72,3 +75,39 @@ func TestCreateUploadLocation(t *testing.T) {
 		assert.Error(t, err)
 	})
 }
+
+func TestCreateDownloadLocation(t *testing.T) {
+	dataStore := commonMocks.GetMockStorageClient()
+	s, err := NewService(config.DataProxyConfig{Download: config.DataProxyDownloadConfig{MaxExpiresIn: stdlibConfig.Duration{Duration: time.Hour}}}, dataStore)
+	assert.NoError(t, err)
+
+	t.Run("Invalid expiry", func(t *testing.T) {
+		_, err = s.CreateDownloadLocation(context.Background(), &service.CreateDownloadLocationRequest{
+			NativeUrl: "s3://bucket/key",
+			ExpiresIn: durationpb.New(-time.Hour),
+		})
+		assert.Error(t, err)
+	})
+
+	t.Run("valid config", func(t *testing.T) {
+		_, err = s.CreateDownloadLocation(context.Background(), &service.CreateDownloadLocationRequest{
+			NativeUrl: "s3://bucket/key",
+			ExpiresIn: durationpb.New(time.Hour),
+		})
+		assert.NoError(t, err)
+	})
+
+	t.Run("use default ExpiresIn", func(t *testing.T) {
+		_, err = s.CreateDownloadLocation(context.Background(), &service.CreateDownloadLocationRequest{
+			NativeUrl: "s3://bucket/key",
+		})
+		assert.NoError(t, err)
+	})
+
+	t.Run("invalid URL", func(t *testing.T) {
+		_, err = s.CreateDownloadLocation(context.Background(), &service.CreateDownloadLocationRequest{
+			NativeUrl: "bucket/key",
+		})
+		assert.NoError(t, err)
+	})
+}
diff --git a/flyteadmin/go.mod b/flyteadmin/go.mod
@@ -13,7 +13,7 @@ require (
 	github.com/cloudevents/sdk-go/v2 v2.8.0
 	github.com/coreos/go-oidc v2.2.1+incompatible
 	github.com/evanphx/json-patch v4.9.0+incompatible
-	github.com/flyteorg/flyteidl v1.1.0
+	github.com/flyteorg/flyteidl v1.1.4
 	github.com/flyteorg/flyteplugins v1.0.0
 	github.com/flyteorg/flytepropeller v1.1.3
 	github.com/flyteorg/flytestdlib v1.0.2

diff --git a/flyteadmin/go.sum b/flyteadmin/go.sum
@@ -380,8 +380,8 @@ github.com/fatih/structtag v1.2.0/go.mod h1:mBJUNpUnHmRKrKlQQlmCrh5PuhftFbNv8Ys4
 github.com/felixge/httpsnoop v1.0.1 h1:lvB5Jl89CsZtGIWuTcDM1E/vkVs49/Ml7JJe07l8SPQ=
 github.com/felixge/httpsnoop v1.0.1/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
 github.com/flyteorg/flyteidl v1.0.0/go.mod h1:JW0z1ZaHS9zWvDAwSMIyGhsf+V4zrzBBgh5IuqzMFCM=
-github.com/flyteorg/flyteidl v1.1.0 h1:f8tdMXOuorS/d+4Ut2QarfDbdCOriK0S+EnlQzrwz9E=
-github.com/flyteorg/flyteidl v1.1.0/go.mod h1:JW0z1ZaHS9zWvDAwSMIyGhsf+V4zrzBBgh5IuqzMFCM=
+github.com/flyteorg/flyteidl v1.1.4 h1:P6YgFYcmBxoLcTegv301i5oYKBCvjHGW3ujRT9s4dvI=
+github.com/flyteorg/flyteidl v1.1.4/go.mod h1:f1tvw5CDjqmrzNxKpRYr6BdAhHL8f7Wp1Duxl0ZOV4g=
 github.com/flyteorg/flyteplugins v1.0.0 h1:77hUJjiIxBmQ9rd3+cXjSGnzOVAFrSzCd59aIaYFB/8=
 github.com/flyteorg/flyteplugins v1.0.0/go.mod h1:4Cpn+9RfanIieTTh2XsuL6zPYXtsR5UDe8YaEmXONT4=
 github.com/flyteorg/flytepropeller v1.1.3 h1:RuS/mkbEhjGyUy2XIs7sHOaio1BK8TUZMGKiIN0/pqE=

diff --git a/flyteadmin/pkg/common/mocks/storage.go b/flyteadmin/pkg/common/mocks/storage.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"net/url"
 	"strings"
 
 	"github.com/flyteorg/flytestdlib/storage"
@@ -44,7 +45,11 @@ func (t *TestDataStore) GetBaseContainerFQN(ctx context.Context) storage.DataRef
 }
 
 func (t *TestDataStore) CreateSignedURL(ctx context.Context, reference storage.DataReference, properties storage.SignedURLProperties) (storage.SignedURLResponse, error) {
-	return storage.SignedURLResponse{}, fmt.Errorf("unsupported")
+	signedURL, err := url.Parse(reference.String())
+	if err != nil {
+		return storage.SignedURLResponse{}, err
+	}
+	return storage.SignedURLResponse{URL: *signedURL}, nil
 }
 
 // Retrieves a byte array from the Blob store or an error

diff --git a/flyteadmin/pkg/config/config.go b/flyteadmin/pkg/config/config.go
@@ -29,7 +29,12 @@ type ServerConfig struct {
 }
 
 type DataProxyConfig struct {
-	Upload DataProxyUploadConfig `json:"upload" pflag:",Defines data proxy upload configuration."`
+	Upload   DataProxyUploadConfig   `json:"upload" pflag:",Defines data proxy upload configuration."`
+	Download DataProxyDownloadConfig `json:"download" pflag:",Defines data proxy download configuration."`
+}
+
+type DataProxyDownloadConfig struct {
+	MaxExpiresIn config.Duration `json:"maxExpiresIn" pflag:",Maximum allowed expiration duration."`
 }
 
 type DataProxyUploadConfig struct {
@@ -86,6 +91,9 @@ var defaultServerConfig = &ServerConfig{
 			MaxExpiresIn:          config.Duration{Duration: time.Hour},
 			DefaultFileNameLength: 20,
 		},
+		Download: DataProxyDownloadConfig{
+			MaxExpiresIn: config.Duration{Duration: time.Hour},
+		},
 	},
 }
 var serverConfig = config.MustRegisterSection(SectionKey, defaultServerConfig)

diff --git a/flyteadmin/pkg/config/serverconfig_flags.go b/flyteadmin/pkg/config/serverconfig_flags.go
diff --git a/flyteadmin/pkg/config/serverconfig_flags_test.go b/flyteadmin/pkg/config/serverconfig_flags_test.go
diff --git a/flyteadmin/pkg/manager/impl/node_execution_manager_test.go b/flyteadmin/pkg/manager/impl/node_execution_manager_test.go
@@ -7,6 +7,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/flyteorg/flyteadmin/pkg/manager/impl/util"
+
 	genModel "github.com/flyteorg/flyteadmin/pkg/repositories/gen/models"
 
 	eventWriterMocks "github.com/flyteorg/flyteadmin/pkg/async/events/mocks"
@@ -1192,8 +1194,9 @@ func TestGetNodeExecutionData(t *testing.T) {
 	expectedClosure := admin.NodeExecutionClosure{
 		Phase: core.NodeExecution_SUCCEEDED,
 		OutputResult: &admin.NodeExecutionClosure_OutputUri{
-			OutputUri: "output uri",
+			OutputUri: util.OutputsFile,
 		},
+		DeckUri: util.DeckFile,
 	}
 	dynamicWorkflowClosureRef := "s3://my-s3-bucket/foo/bar/dynamic.pb"
 
@@ -1233,7 +1236,7 @@ func TestGetNodeExecutionData(t *testing.T) {
 				Url:   "inputs",
 				Bytes: 100,
 			}, nil
-		} else if uri == "output uri" {
+		} else if uri == util.OutputsFile {
 			return admin.UrlBlob{
 				Url:   "outputs",
 				Bytes: 200,
@@ -1260,7 +1263,7 @@ func TestGetNodeExecutionData(t *testing.T) {
 			marshalled, _ := proto.Marshal(fullInputs)
 			_ = proto.Unmarshal(marshalled, msg)
 			return nil
-		} else if reference.String() == "output uri" {
+		} else if reference.String() == util.OutputsFile {
 			marshalled, _ := proto.Marshal(fullOutputs)
 			_ = proto.Unmarshal(marshalled, msg)
 			return nil

diff --git a/flyteadmin/pkg/manager/impl/util/data.go b/flyteadmin/pkg/manager/impl/util/data.go
@@ -13,6 +13,11 @@ import (
 	"github.com/golang/protobuf/proto"
 )
 
+const (
+	OutputsFile = "outputs.pb"
+	DeckFile    = "deck.html"
+)
+
 func shouldFetchData(config *runtimeInterfaces.RemoteDataConfig, urlBlob admin.UrlBlob) bool {
 	return config.Scheme == common.Local || config.Scheme == common.None || config.MaxSizeInBytes == 0 ||
 		urlBlob.Bytes < config.MaxSizeInBytes

diff --git a/flyteadmin/pkg/manager/impl/util/data_test.go b/flyteadmin/pkg/manager/impl/util/data_test.go
@@ -154,7 +154,6 @@ func TestGetInputs(t *testing.T) {
 		assert.True(t, proto.Equal(fullInputs, testLiteralMap))
 		assert.Empty(t, inputURLBlob)
 	})
-
 }
 
 func TestGetOutputs(t *testing.T) {

diff --git a/flyteadmin/pkg/repositories/transformers/node_execution.go b/flyteadmin/pkg/repositories/transformers/node_execution.go
@@ -95,6 +95,8 @@ func addTerminalState(
 		nodeExecutionModel.ErrorKind = &k
 		nodeExecutionModel.ErrorCode = &request.Event.GetError().Code
 	}
+	closure.DeckUri = request.Event.DeckUri
+
 	return nil
 }
 

diff --git a/flyteadmin/pkg/server/service.go b/flyteadmin/pkg/server/service.go
@@ -210,6 +210,11 @@ func newHTTPServer(ctx context.Context, cfg *config.ServerConfig, _ *authConfig.
 		return nil, errors.Wrap(err, "error registering identity service")
 	}
 
+	err = service.RegisterDataProxyServiceHandlerFromEndpoint(ctx, gwmux, grpcAddress, grpcConnectionOpts)
+	if err != nil {
+		return nil, errors.Wrap(err, "error registering data proxy service")
+	}
+
 	mux.Handle("/", gwmux)
 
 	return mux, nil
-Original file line number
+Diff line change
@@ Expand Up / @@ -154,7 +154,6 @@ func TestGetInputs(t *testing.T) { @@
     		assert.True(t, proto.Equal(fullInputs, testLiteralMap))
     		assert.Empty(t, inputURLBlob)
     	})
     }
     func TestGetOutputs(t *testing.T) {
@@ Expand Down @@