SVR-387: Add org to CreateUploadLocation (#134)

flyteadmin/dataproxy/service.go

@@ -12,6 +12,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/samber/lo"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/protobuf/types/known/durationpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
@@ -61,7 +62,7 @@ func (s Service) CreateUploadLocation(ctx context.Context, req *service.CreateUp
 	// If we fall in here, that means that the full path is deterministic and we should check for existence.
 	if len(req.Filename) > 0 && len(req.FilenameRoot) > 0 {
 		knownLocation, err := createStorageLocation(ctx, s.dataStore, s.cfg.Upload,
-			req.Project, req.Domain, req.FilenameRoot, req.Filename)
+			req.Org, req.Project, req.Domain, req.FilenameRoot, req.Filename)
 		if err != nil {
 			return nil, errors.NewFlyteAdminErrorf(codes.Internal, "failed to create storage location, Error: %v", err)
 		}
@@ -115,7 +116,7 @@ func (s Service) CreateUploadLocation(ctx context.Context, req *service.CreateUp
 		prefix = base32.StdEncoding.EncodeToString(req.ContentMd5)
 	}
 	storagePath, err := createStorageLocation(ctx, s.dataStore, s.cfg.Upload,
-		req.Project, req.Domain, prefix, req.Filename)
+		req.Org, req.Project, req.Domain, prefix, req.Filename)
 	if err != nil {
 		return nil, errors.NewFlyteAdminErrorf(codes.Internal, "failed to create shardedStorageLocation, Error: %v", err)
 	}
@@ -273,6 +274,9 @@ func (s Service) validateCreateDownloadLinkRequest(req *service.CreateDownloadLi
 func createStorageLocation(ctx context.Context, store *storage.DataStore,
 	cfg config.DataProxyUploadConfig, keyParts ...string) (storage.DataReference, error) {
 
+	keyParts = lo.Filter(keyParts, func(key string, _ int) bool {
+		return key != ""
+	})
 	storagePath, err := store.ConstructReference(ctx, store.GetBaseContainerFQN(ctx),
 		append([]string{cfg.StoragePrefix}, keyParts...)...)
 	if err != nil {

flyteadmin/dataproxy/service_test.go

@@ -422,3 +422,23 @@ func TestService_Error(t *testing.T) {
 		assert.Error(t, err, "no task executions")
 	})
 }
+
+func TestCreateStorageLocation(t *testing.T) {
+	ctx := context.TODO()
+	dataStore := commonMocks.GetMockStorageClient()
+	expectedStoragePath := storage.DataReference("s3://bucket/prefix/foo/bar/baz")
+	t.Run("no empty parts", func(t *testing.T) {
+		storagePath, err := createStorageLocation(ctx, dataStore, config.DataProxyUploadConfig{
+			StoragePrefix: "prefix",
+		}, "foo", "bar", "baz")
+		assert.NoError(t, err)
+		assert.Equal(t, expectedStoragePath, storagePath)
+	})
+	t.Run("with empty parts", func(t *testing.T) {
+		storagePath, err := createStorageLocation(ctx, dataStore, config.DataProxyUploadConfig{
+			StoragePrefix: "prefix",
+		}, "foo", "bar", "", "baz")
+		assert.NoError(t, err)
+		assert.Equal(t, expectedStoragePath, storagePath)
+	})
+}

flyteadmin/pkg/manager/impl/launch_plan_manager_test.go

@@ -7,6 +7,11 @@ import (
 	"testing"
 	"time"
 
+	"github.com/golang/protobuf/proto"
+	"github.com/golang/protobuf/ptypes"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/protobuf/types/known/anypb"
 
 	"github.com/flyteorg/flyte/flyteadmin/pkg/artifacts"
@@ -27,11 +32,6 @@ import (
 	artifactsIdl "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/artifacts"
 	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core"
 	mockScope "github.com/flyteorg/flyte/flytestdlib/promutils"
-	"github.com/golang/protobuf/proto"
-	"github.com/golang/protobuf/ptypes"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/mock"
-	"google.golang.org/grpc/codes"
 )
 
 var active = int32(admin.LaunchPlanState_ACTIVE)