Add deckURI to NodeExecutionData

dataproxy/service.go

@@ -86,6 +86,42 @@ 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 expiresIn := req.ExpiresIn; expiresIn != nil {
+		if !expiresIn.IsValid() {
+			fmt.Println("hello hello")
+			return nil, fmt.Errorf("expiresIn [%v] is invalid", expiresIn)
+		}
+
+		if expiresIn.AsDuration() < 0 {
+			return nil, fmt.Errorf("expiresIn [%v] should not less than 0",
+				expiresIn.AsDuration().String())
+		} else if expiresIn.AsDuration() > s.cfg.Download.MaxExpiresIn.Duration {
+			return nil, 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)
+	}
+
+	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
+}
+
 // 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,

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,32 @@ 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)
+	})
+}

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.1
 	github.com/flyteorg/flytestdlib v1.0.2

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.1 h1:z9OFS7VAsoFOyIGSfIszaMrERG8MOvS17yzpuiusb64=

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

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)

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

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

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) {

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
 }