Add file_extension field to BlobType

flytecopilot/cmd/download.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/golang/protobuf/proto"
 	"github.com/spf13/cobra"
 
 	"github.com/flyteorg/flyte/flytecopilot/data"
@@ -49,6 +50,11 @@ func GetUploadModeVals() []string {
 }
 
 func (d *DownloadOptions) Download(ctx context.Context) error {
+	variableMap := &core.VariableMap{}
+	if err := proto.Unmarshal(d.inputInterface, variableMap); err != nil {
+		logger.Warnf(ctx, "Bad input interface passed, failed to unmarshal err: %s", err)
+	}
+
 	if d.remoteOutputsPrefix == "" {
 		return fmt.Errorf("to-output-prefix is required")
 	}
@@ -77,7 +83,7 @@ func (d *DownloadOptions) Download(ctx context.Context) error {
 			childCtx, cancelFn = context.WithTimeout(ctx, d.timeout)
 		}
 		defer cancelFn()
-		err := dl.DownloadInputs(childCtx, storage.DataReference(d.remoteInputsPath), d.localDirectoryPath)
+		err := dl.DownloadInputs(childCtx, variableMap, storage.DataReference(d.remoteInputsPath), d.localDirectoryPath)
 		if err != nil {
 			logger.Errorf(ctx, "Downloading failed, err %s", err)
 			return err
@@ -116,6 +122,6 @@ func NewDownloadCommand(opts *RootOptions) *cobra.Command {
 	downloadCmd.Flags().StringVarP(&downloadOpts.metadataFormat, "format", "m", core.DataLoadingConfig_JSON.String(), fmt.Sprintf("What should be the output format for the primitive and structured types. Options [%v]", GetFormatVals()))
 	downloadCmd.Flags().StringVarP(&downloadOpts.downloadMode, "download-mode", "d", core.IOStrategy_DOWNLOAD_EAGER.String(), fmt.Sprintf("Download mode to use. Options [%v]", GetDownloadModeVals()))
 	downloadCmd.Flags().DurationVarP(&downloadOpts.timeout, "timeout", "t", time.Hour*1, "Max time to allow for downloads to complete, default is 1H")
-	downloadCmd.Flags().BytesBase64VarP(&downloadOpts.inputInterface, "input-interface", "i", nil, "Input interface proto message - core.VariableMap, base64 encoced string")
+	downloadCmd.Flags().BytesBase64VarP(&downloadOpts.inputInterface, "input-interface", "i", nil, "Input interface proto message - core.VariableMap, base64 encoded string")
 	return downloadCmd
 }

flytecopilot/cmd/download_test.go

@@ -10,6 +10,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/golang/protobuf/proto"
 	"github.com/stretchr/testify/assert"
 
 	"github.com/flyteorg/flyte/flyteidl/clients/go/coreutils"
@@ -137,6 +138,100 @@ func TestDownloadOptions_Download(t *testing.T) {
 		assert.ElementsMatch(t, []string{"inputs.json", "inputs.pb", "x", "y", "blob"}, collectFile(tmpDir))
 	})
 
+	t.Run("primitiveAndBlobInputsWithFileExtension", func(t *testing.T) {
+		tmpDir, err := ioutil.TempDir(tmpFolderLocation, tmpPrefix)
+		assert.NoError(t, err)
+		defer func() {
+			assert.NoError(t, os.RemoveAll(tmpDir))
+		}()
+		dopts.localDirectoryPath = tmpDir
+
+		s := promutils.NewTestScope()
+		store, err := storage.NewDataStore(&storage.Config{Type: storage.TypeMemory}, s.NewSubScope("storage"))
+		assert.NoError(t, err)
+		dopts.RootOptions = &RootOptions{
+			Scope: s,
+			Store: store,
+		}
+
+		iface := &core.VariableMap{
+			Variables: map[string]*core.Variable{
+				"blob": {
+					Type: &core.LiteralType{Type: &core.LiteralType_Blob{Blob: &core.BlobType{Dimensionality: core.BlobType_SINGLE, Format: "xyz", FileExtension: "xyz"}}},
+				},
+				"csv_blob": {
+					Type: &core.LiteralType{Type: &core.LiteralType_Blob{Blob: &core.BlobType{Dimensionality: core.BlobType_SINGLE, Format: "csv", FileExtension: "csv"}}},
+				},
+				"invalid_extension_blob": {
+					Type: &core.LiteralType{Type: &core.LiteralType_Blob{Blob: &core.BlobType{Dimensionality: core.BlobType_SINGLE, Format: "xyz", FileExtension: "/../../../invalid"}}},
+				},
+			},
+		}
+		d, err := proto.Marshal(iface)
+		assert.NoError(t, err)
+		dopts.inputInterface = d
+
+		blobLoc := storage.DataReference("blob-loc")
+		br := bytes.NewBuffer([]byte("Hello World!"))
+		assert.NoError(t, store.WriteRaw(ctx, blobLoc, int64(br.Len()), storage.Options{}, br))
+		assert.NoError(t, store.WriteProtobuf(ctx, storage.DataReference(inputPath), storage.Options{}, &core.LiteralMap{
+			Literals: map[string]*core.Literal{
+				"x": coreutils.MustMakePrimitiveLiteral(1),
+				"y": coreutils.MustMakePrimitiveLiteral("hello"),
+				"blob": {Value: &core.Literal_Scalar{
+					Scalar: &core.Scalar{
+						Value: &core.Scalar_Blob{
+							Blob: &core.Blob{
+								Uri: blobLoc.String(),
+								Metadata: &core.BlobMetadata{
+									Type: &core.BlobType{
+										Dimensionality: core.BlobType_SINGLE,
+										Format:         "xyz",
+										FileExtension:  "xyz",
+									},
+								},
+							},
+						},
+					},
+				}},
+				"csv_blob": {Value: &core.Literal_Scalar{
+					Scalar: &core.Scalar{
+						Value: &core.Scalar_Blob{
+							Blob: &core.Blob{
+								Uri: blobLoc.String(),
+								Metadata: &core.BlobMetadata{
+									Type: &core.BlobType{
+										Dimensionality: core.BlobType_SINGLE,
+										Format:         "xyz",
+										FileExtension:  "xyz",
+									},
+								},
+							},
+						},
+					},
+				}},
+				"invalid_extension_blob": {Value: &core.Literal_Scalar{
+					Scalar: &core.Scalar{
+						Value: &core.Scalar_Blob{
+							Blob: &core.Blob{
+								Uri: blobLoc.String(),
+								Metadata: &core.BlobMetadata{
+									Type: &core.BlobType{
+										Dimensionality: core.BlobType_SINGLE,
+										Format:         "xyz",
+										FileExtension:  "/../../../invalid",
+									},
+								},
+							},
+						},
+					},
+				}},
+			},
+		}))
+		assert.NoError(t, dopts.Download(ctx), "Download Operation failed")
+		assert.ElementsMatch(t, []string{"inputs.json", "inputs.pb", "x", "y", "blob.xyz", "csv_blob.csv", "invalid_extension_blob"}, collectFile(tmpDir))
+	})
+
 	t.Run("primitiveAndMissingBlobInputs", func(t *testing.T) {
 		tmpDir, err := ioutil.TempDir(tmpFolderLocation, tmpPrefix)
 		assert.NoError(t, err)

flytecopilot/data/download.go

@@ -9,6 +9,7 @@ import (
 	"path"
 	"path/filepath"
 	"reflect"
+	"regexp"
 	"strconv"
 	"strings"
 	"sync"
@@ -26,13 +27,46 @@ import (
 	"github.com/flyteorg/flyte/flytestdlib/storage"
 )
 
+var validFileExtensionRe = regexp.MustCompile(`^[a-zA-Z0-9]+(\.[a-zA-Z0-9]+)*$`)
+
 type Downloader struct {
 	format core.DataLoadingConfig_LiteralMapFormat
 	store  *storage.DataStore
 	// TODO support download mode
 	mode core.IOStrategy_DownloadMode
 }
 
+// By default, blobs (FlyteFiles) were not and still are not written with a
+// file extension. For example, a data: FlyteFile["csv"] blob is written
+// to "inputs/data", even though Format="csv".
+//
+// When FileExtension="" (the default), this old behavior is preserved.
+//
+// However, an input blob
+// `data: Annotated[FlyteFile["csv"], FileDownloadConfig(file_extension="csv")]`
+// should be written to "inputs/data.csv" (when FileExtension="csv" - new behavior).
+func resolveVarFilenames(ctx context.Context, vars *core.VariableMap) map[string]string {
+	varFilenames := make(map[string]string, len(vars.GetVariables()))
+	for varName, variable := range vars.GetVariables() {
+		varType := variable.GetType()
+		switch varType.GetType().(type) {
+		case *core.LiteralType_Blob:
+			ext := varType.GetBlob().GetFileExtension()
+			if ext == "" {
+				varFilenames[varName] = varName
+			} else if !validFileExtensionRe.MatchString(ext) {
+				logger.Warnf(ctx, "invalid file extension for variable %q [%q], ignoring...", varName, ext)
+				varFilenames[varName] = varName
+			} else {
+				varFilenames[varName] = varName + "." + ext
+			}
+		default:
+			varFilenames[varName] = varName
+		}
+	}
+	return varFilenames
+}
+
 // TODO add timeout and rate limit
 // TODO use chunk to download
 func (d Downloader) handleBlob(ctx context.Context, blob *core.Blob, toPath string) (interface{}, error) {
@@ -432,7 +466,7 @@ func (d Downloader) handleLiteral(ctx context.Context, lit *core.Literal, filePa
 		if err != nil {
 			return nil, nil, errors.Wrapf(err, "failed to create directory [%s]", filePath)
 		}
-		v, m, err := d.RecursiveDownload(ctx, lit.GetMap(), filePath, writeToFile)
+		v, m, err := d.RecursiveDownload(ctx, lit.GetMap(), filePath, make(map[string]string), writeToFile)
 		if err != nil {
 			return nil, nil, err
 		}
@@ -468,7 +502,7 @@ type downloadedResult struct {
 	v   interface{}
 }
 
-func (d Downloader) RecursiveDownload(ctx context.Context, inputs *core.LiteralMap, dir string, writePrimitiveToFile bool) (VarMap, *core.LiteralMap, error) {
+func (d Downloader) RecursiveDownload(ctx context.Context, inputs *core.LiteralMap, dir string, varFilenames map[string]string, writePrimitiveToFile bool) (VarMap, *core.LiteralMap, error) {
 	childCtx, cancel := context.WithCancel(ctx)
 	defer cancel()
 	if inputs == nil || len(inputs.GetLiterals()) == 0 {
@@ -486,7 +520,11 @@ func (d Downloader) RecursiveDownload(ctx context.Context, inputs *core.LiteralM
 			}
 			logger.Infof(ctx, "read object at location [%s]", offloadedMetadataURI)
 		}
-		varPath := path.Join(dir, variable)
+		filename := variable
+		if varFilename, ok := varFilenames[variable]; ok {
+			filename = varFilename
+		}
+		varPath := path.Join(dir, filename)
 		lit := literal
 		f[variable] = futures.NewAsyncFuture(childCtx, func(ctx2 context.Context) (interface{}, error) {
 			v, lit, err := d.handleLiteral(ctx2, lit, varPath, writePrimitiveToFile)
@@ -520,7 +558,7 @@ func (d Downloader) RecursiveDownload(ctx context.Context, inputs *core.LiteralM
 	return vmap, m, nil
 }
 
-func (d Downloader) DownloadInputs(ctx context.Context, inputRef storage.DataReference, outputDir string) error {
+func (d Downloader) DownloadInputs(ctx context.Context, vars *core.VariableMap, inputRef storage.DataReference, outputDir string) error {
 	logger.Infof(ctx, "Downloading inputs from [%s]", inputRef)
 	defer logger.Infof(ctx, "Exited downloading inputs from [%s]", inputRef)
 	if err := os.MkdirAll(outputDir, os.ModePerm); err != nil {
@@ -533,7 +571,9 @@ func (d Downloader) DownloadInputs(ctx context.Context, inputRef storage.DataRef
 		logger.Errorf(ctx, "Failed to download inputs from [%s], err [%s]", inputRef, err)
 		return errors.Wrapf(err, "failed to download input metadata message from remote store")
 	}
-	varMap, lMap, err := d.RecursiveDownload(ctx, inputs, outputDir, true)
+
+	varFilenames := resolveVarFilenames(ctx, vars)
+	varMap, lMap, err := d.RecursiveDownload(ctx, inputs, outputDir, varFilenames, true)
 	if err != nil {
 		return errors.Wrapf(err, "failed to download input variable from remote store")
 	}

flytecopilot/data/download_test.go

@@ -231,7 +231,7 @@ func TestRecursiveDownload(t *testing.T) {
 			}
 		}()
 
-		varMap, lMap, err := d.RecursiveDownload(context.Background(), inputs, toPath, true)
+		varMap, lMap, err := d.RecursiveDownload(context.Background(), inputs, toPath, make(map[string]string), true)
 		assert.NoError(t, err)
 		assert.NotNil(t, varMap)
 		assert.NotNil(t, lMap)
@@ -309,7 +309,7 @@ func TestRecursiveDownload(t *testing.T) {
 			}
 		}()
 
-		varMap, lMap, err := d.RecursiveDownload(context.Background(), inputs, toPath, true)
+		varMap, lMap, err := d.RecursiveDownload(context.Background(), inputs, toPath, make(map[string]string), true)
 		assert.NoError(t, err)
 		assert.NotNil(t, varMap)
 		assert.NotNil(t, lMap)

flytecopilot/data/upload.go

@@ -107,11 +107,11 @@ func (u Uploader) handleBlobType(ctx context.Context, localPath string, toPath s
 			}
 		}
 
-		return coreutils.MakeLiteralForBlob(toPath, false, ""), nil
+		return coreutils.MakeLiteralForBlob(toPath, false, "", ""), nil
 	}
 	size := info.Size()
 	// Should we make this a go routine as well, so that we can introduce timeouts
-	return coreutils.MakeLiteralForBlob(toPath, false, ""), UploadFileToStorage(ctx, fpath, toPath, size, u.store)
+	return coreutils.MakeLiteralForBlob(toPath, false, "", ""), UploadFileToStorage(ctx, fpath, toPath, size, u.store)
 }
 
 func (u Uploader) RecursiveUpload(ctx context.Context, vars *core.VariableMap, fromPath string, metaOutputPath, dataRawPath storage.DataReference) error {

flyteidl/.gitignore

@@ -8,6 +8,9 @@ vendor
 
 dist
 gen/pb_python/flyteidl.egg-info/
+gen/pb_python/flyteidl/**/__pycache__/
+pip-wheel-metadata/*.dist-info/
+*.egg-info/
 
 .virtualgo
 docs/build/

flyteidl/clients/go/coreutils/literals.go

@@ -485,7 +485,7 @@ func MakeLiteralForStructuredDataSet(path storage.DataReference, columns []*core
 	}
 }
 
-func MakeLiteralForBlob(path storage.DataReference, isDir bool, format string) *core.Literal {
+func MakeLiteralForBlob(path storage.DataReference, isDir bool, format string, fileExtension string) *core.Literal {
 	dim := core.BlobType_SINGLE
 	if isDir {
 		dim = core.BlobType_MULTIPART
@@ -500,6 +500,7 @@ func MakeLiteralForBlob(path storage.DataReference, isDir bool, format string) *
 							Type: &core.BlobType{
 								Dimensionality: dim,
 								Format:         format,
+								FileExtension:  fileExtension,
 							},
 						},
 					},
@@ -601,7 +602,7 @@ func MakeLiteralForType(t *core.LiteralType, v interface{}) (*core.Literal, erro
 
 	case *core.LiteralType_Blob:
 		isDir := newT.Blob.GetDimensionality() == core.BlobType_MULTIPART
-		lv := MakeLiteralForBlob(storage.DataReference(fmt.Sprintf("%v", v)), isDir, newT.Blob.GetFormat())
+		lv := MakeLiteralForBlob(storage.DataReference(fmt.Sprintf("%v", v)), isDir, newT.Blob.GetFormat(), newT.Blob.GetFileExtension())
 		return lv, nil
 
 	case *core.LiteralType_Schema:

flyteidl/clients/go/coreutils/literals_test.go

@@ -388,9 +388,10 @@ func TestMakeLiteralForSimpleType(t *testing.T) {
 
 func TestMakeLiteralForBlob(t *testing.T) {
 	type args struct {
-		path   storage.DataReference
-		isDir  bool
-		format string
+		path          storage.DataReference
+		isDir         bool
+		format        string
+		fileExtension string
 	}
 	tests := []struct {
 		name string
@@ -399,10 +400,11 @@ func TestMakeLiteralForBlob(t *testing.T) {
 	}{
 		{"simple-key", args{path: "/key", isDir: false, format: "xyz"}, &core.Blob{Uri: "/key", Metadata: &core.BlobMetadata{Type: &core.BlobType{Format: "xyz", Dimensionality: core.BlobType_SINGLE}}}},
 		{"simple-dir", args{path: "/key", isDir: true, format: "xyz"}, &core.Blob{Uri: "/key", Metadata: &core.BlobMetadata{Type: &core.BlobType{Format: "xyz", Dimensionality: core.BlobType_MULTIPART}}}},
+		{"simple-key-with-extension", args{path: "/key", isDir: false, format: "xyz", fileExtension: "xyz"}, &core.Blob{Uri: "/key", Metadata: &core.BlobMetadata{Type: &core.BlobType{Format: "xyz", Dimensionality: core.BlobType_SINGLE, FileExtension: "xyz"}}}},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			if got := MakeLiteralForBlob(tt.args.path, tt.args.isDir, tt.args.format); !reflect.DeepEqual(got.GetScalar().GetBlob(), tt.want) {
+			if got := MakeLiteralForBlob(tt.args.path, tt.args.isDir, tt.args.format, tt.args.fileExtension); !reflect.DeepEqual(got.GetScalar().GetBlob(), tt.want) {
 				t.Errorf("MakeLiteralForBlob() = %v, want %v", got, tt.want)
 			}
 		})