[copilot][flytedirectory] multipart blob download

flyteadmin/pkg/common/mocks/storage.go

@@ -54,6 +54,11 @@
 	return storage.SignedURLResponse{URL: *signedURL}, nil
 }
 
+func (t *TestDataStore) GetItems(ctx context.Context, reference storage.DataReference) ([]string, error) {
+	var s []string
+	return s, nil
+}
+
 // Retrieves a byte array from the Blob store or an error
 func (t *TestDataStore) ReadRaw(ctx context.Context, reference storage.DataReference) (io.ReadCloser, error) {
 	return NopCloser{}, nil

flytecopilot/data/download.go

@@ -8,8 +8,10 @@
 	"io/ioutil"
 	"os"
 	"path"
+	"path/filepath"
 	"reflect"
 	"strconv"
+	"sync"
 
 	"github.com/ghodss/yaml"
 	"github.com/golang/protobuf/jsonpb"
@@ -31,24 +33,94 @@
 	mode core.IOStrategy_DownloadMode
 }
 
-// TODO add support for multipart blobs
-func (d Downloader) handleBlob(ctx context.Context, blob *core.Blob, toFilePath string) (interface{}, error) {
+func (d Downloader) handleBlob(ctx context.Context, blob *core.Blob, toPath string) (interface{}, error) {
 	ref := storage.DataReference(blob.Uri)
 	scheme, _, _, err := ref.Split()
 	if err != nil {
 		return nil, errors.Wrapf(err, "Blob uri incorrectly formatted")
 	}
+
+	if blob.GetMetadata().GetType().Dimensionality == core.BlobType_MULTIPART {
+		items, err := d.store.GetItems(ctx, ref)
+		if err != nil || len(items) == 0 {
+			logger.Errorf(ctx, "failed to collect items from multipart blob [%s]", ref)
+			return nil, err
+		}
+
+		success := 0
+		var mu sync.Mutex
+		var wg sync.WaitGroup
+		for _, absPath := range items {
+			// capture range variable
+			absPath := absPath
+
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				defer func() {
+					if err := recover(); err != nil {
+						logger.Errorf(ctx, "recover receives error: %s", err)
+					}
+				}()
+
+				ref := storage.DataReference(absPath)
+				reader, err := DownloadFileFromStorage(ctx, ref, d.store)
+				if err != nil {
+					logger.Errorf(ctx, "Failed to download from ref [%s]", ref)
+					return
+				}
+				defer func() {
+					err := reader.Close()
+					if err != nil {
+						logger.Errorf(ctx, "failed to close Blob read stream @ref [%s]. Error: %s", ref, err)
+					}
+				}()
+
+				_, _, relativePath, err := ref.Split()
+				if err != nil {
+					logger.Errorf(ctx, "Failed to parse ref [%s]", ref)
+					return
+				}
+				newPath := filepath.Join(toPath, relativePath)
+				dir := filepath.Dir(newPath)
+
+				mu.Lock()
+				// 0755: the directory can be read by anyone but can only be written by the owner
+				os.MkdirAll(dir, 0755)
+				writer, err := os.Create(newPath)
+				mu.Unlock()
+				if err != nil {
+					logger.Errorf(ctx, "failed to open file at path %s", newPath)
+					return
+				}
+				defer func() {
+					err := writer.Close()
+					if err != nil {
+						logger.Errorf(ctx, "failed to close File write stream. Error: %s", err)
+					}
+				}()
+
+				_, err = io.Copy(writer, reader)
+				if err != nil {
+					logger.Errorf(ctx, "failed to write remote data to local filesystem")
+					return
+				}
+				mu.Lock()
+				success += 1
+				mu.Unlock()
+			}()
+		}
+		wg.Wait()
+		ref = storage.DataReference(blob.Uri)
+		logger.Infof(ctx, "Successfully copied [%d] remote files from [%s] to local [%s]", success, ref, toPath)
+		return toPath, nil
+	}
+
+	// reader should be declared here (avoid being shared across all goroutines)
 	var reader io.ReadCloser
 	if scheme == "http" || scheme == "https" {
 		reader, err = DownloadFileFromHTTP(ctx, ref)
 	} else {
-		if blob.GetMetadata().GetType().Dimensionality == core.BlobType_MULTIPART {
-			logger.Warnf(ctx, "Currently only single part blobs are supported, we will force multipart to be 'path/00000'")
-			ref, err = d.store.ConstructReference(ctx, ref, "000000")
-			if err != nil {
-				return nil, err
-			}
-		}
 		reader, err = DownloadFileFromStorage(ctx, ref, d.store)
 	}
 	if err != nil {
@@ -62,9 +134,9 @@
 		}
 	}()
 
-	writer, err := os.Create(toFilePath)
+	writer, err := os.Create(toPath)
 	if err != nil {
-		return nil, errors.Wrapf(err, "failed to open file at path %s", toFilePath)
+		return nil, errors.Wrapf(err, "failed to open file at path %s", toPath)
 	}
 	defer func() {
 		err := writer.Close()
@@ -76,12 +148,11 @@
 	if err != nil {
 		return nil, errors.Wrapf(err, "failed to write remote data to local filesystem")
 	}
-	logger.Infof(ctx, "Successfully copied [%d] bytes remote data from [%s] to local [%s]", v, ref, toFilePath)
-	return toFilePath, nil
+	logger.Infof(ctx, "Successfully copied [%d] bytes remote data from [%s] to local [%s]", v, ref, toPath)
+	return toPath, nil
 }
 
 func (d Downloader) handleSchema(ctx context.Context, schema *core.Schema, toFilePath string) (interface{}, error) {
-	// TODO Handle schema type
 	return d.handleBlob(ctx, &core.Blob{Uri: schema.Uri, Metadata: &core.BlobMetadata{Type: &core.BlobType{Dimensionality: core.BlobType_MULTIPART}}}, toFilePath)
 }
 

flytecopilot/data/download_test.go

@@ -0,0 +1,151 @@
+package data
+
+import (
+	"bytes"
+	"context"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core"
+	"github.com/flyteorg/flyte/flytestdlib/promutils"
+	"github.com/flyteorg/flyte/flytestdlib/storage"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHandleBlobMultipart(t *testing.T) {
+	t.Run("Successful Query", func(t *testing.T) {
+		s, err := storage.NewDataStore(&storage.Config{Type: storage.TypeMemory}, promutils.NewTestScope())
+		assert.NoError(t, err)
+		ref := storage.DataReference("s3://container/folder/file1")
+		s.WriteRaw(context.Background(), ref, 0, storage.Options{}, bytes.NewReader([]byte{}))
+		ref = storage.DataReference("s3://container/folder/file2")
+		s.WriteRaw(context.Background(), ref, 0, storage.Options{}, bytes.NewReader([]byte{}))
+
+		d := Downloader{store: s}
+
+		blob := &core.Blob{
+			Uri: "s3://container/folder",
+			Metadata: &core.BlobMetadata{
+				Type: &core.BlobType{
+					Dimensionality: core.BlobType_MULTIPART,
+				},
+			},
+		}
+
+		toPath := "./inputs"
+		defer func() {
+			err := os.RemoveAll(toPath)
+			if err != nil {
+				t.Errorf("Failed to delete directory: %v", err)
+			}
+		}()
+
+		result, err := d.handleBlob(context.Background(), blob, toPath)
+		assert.NoError(t, err)
+		assert.Equal(t, toPath, result)
+
+		// Check if files were created and data written
+		for _, file := range []string{"file1", "file2"} {
+			if _, err := os.Stat(filepath.Join(toPath, "folder", file)); os.IsNotExist(err) {
+				t.Errorf("expected file %s to exist", file)
+			}
+		}
+	})
+
+	t.Run("No Items", func(t *testing.T) {
+		s, err := storage.NewDataStore(&storage.Config{Type: storage.TypeMemory}, promutils.NewTestScope())
+		assert.NoError(t, err)
+
+		d := Downloader{store: s}
+
+		blob := &core.Blob{
+			Uri: "s3://container/folder",
+			Metadata: &core.BlobMetadata{
+				Type: &core.BlobType{
+					Dimensionality: core.BlobType_MULTIPART,
+				},
+			},
+		}
+
+		toPath := "./inputs"
+		defer func() {
+			err := os.RemoveAll(toPath)
+			if err != nil {
+				t.Errorf("Failed to delete directory: %v", err)
+			}
+		}()
+
+		result, err := d.handleBlob(context.Background(), blob, toPath)
+		assert.Error(t, err)
+		assert.Nil(t, result)
+	})
+}
+
+func TestHandleBlobSinglePart(t *testing.T) {
+	s, err := storage.NewDataStore(&storage.Config{Type: storage.TypeMemory}, promutils.NewTestScope())
+	assert.NoError(t, err)
+	ref := storage.DataReference("s3://container/file")
+	s.WriteRaw(context.Background(), ref, 0, storage.Options{}, bytes.NewReader([]byte{}))
+
+	d := Downloader{store: s}
+
+	blob := &core.Blob{
+		Uri: "s3://container/file",
+		Metadata: &core.BlobMetadata{
+			Type: &core.BlobType{
+				Dimensionality: core.BlobType_SINGLE,
+			},
+		},
+	}
+
+	toPath := "./input"
+	defer func() {
+		err := os.RemoveAll(toPath)
+		if err != nil {
+			t.Errorf("Failed to delete file: %v", err)
+		}
+	}()
+
+	result, err := d.handleBlob(context.Background(), blob, toPath)
+	assert.NoError(t, err)
+	assert.Equal(t, toPath, result)
+
+	// Check if files were created and data written
+	if _, err := os.Stat(toPath); os.IsNotExist(err) {
+		t.Errorf("expected file %s to exist", toPath)
+	}
+}
+
+func TestHandleBlobHTTP(t *testing.T) {
+	s, err := storage.NewDataStore(&storage.Config{Type: storage.TypeMemory}, promutils.NewTestScope())
+	assert.NoError(t, err)
+	d := Downloader{store: s}
+
+	blob := &core.Blob{
+		Uri: "https://raw.githubusercontent.com/flyteorg/flyte/master/README.md",
+		Metadata: &core.BlobMetadata{
+			Type: &core.BlobType{
+				Dimensionality: core.BlobType_SINGLE,
+			},
+		},
+	}
+
+	toPath := "./input"
+	defer func() {
+		err := os.RemoveAll(toPath)
+		if err != nil {
+			t.Errorf("Failed to delete file: %v", err)
+		}
+	}()
+
+	result, err := d.handleBlob(context.Background(), blob, toPath)
+	assert.NoError(t, err)
+	assert.Equal(t, toPath, result)
+
+	// Check if files were created and data written
+	if _, err := os.Stat(toPath); os.IsNotExist(err) {
+		t.Errorf("expected file %s to exist", toPath)
+	}
+}

flytepropeller/pkg/utils/failing_datastore.go

@@ -27,6 +27,10 @@ func (FailingRawStore) Head(ctx context.Context, reference storage.DataReference
 	return nil, fmt.Errorf("failed metadata fetch")
 }
 
+func (FailingRawStore) GetItems(ctx context.Context, reference storage.DataReference) ([]string, error) {
+	return nil, fmt.Errorf("failed get items")
+}
+
 func (FailingRawStore) ReadRaw(ctx context.Context, reference storage.DataReference) (io.ReadCloser, error) {
 	return nil, fmt.Errorf("failed read raw")
 }

flytepropeller/pkg/utils/failing_datastore_test.go

@@ -19,6 +19,9 @@ func TestFailingRawStore(t *testing.T) {
 	c := f.GetBaseContainerFQN(ctx)
 	assert.Equal(t, storage.DataReference(""), c)
 
+	_, err = f.GetItems(ctx, "")
+	assert.Error(t, err)
+
 	_, err = f.ReadRaw(ctx, "")
 	assert.Error(t, err)
 

flytestdlib/storage/cached_rawstore.go

@@ -50,6 +50,15 @@ func (s *cachedRawStore) Head(ctx context.Context, reference DataReference) (Met
 	return s.RawStore.Head(ctx, reference)
 }
 
+// GetItems retrieves the paths of all items from the Blob store or an error
+func (s *cachedRawStore) GetItems(ctx context.Context, reference DataReference) ([]string, error) {
+	_, span := otelutils.NewSpan(ctx, otelutils.BlobstoreClientTracer, "flytestdlib.storage.cachedRawStore/GetItems")
+	defer span.End()
+
+	// freecache does not support full cache scanning
+	return s.RawStore.GetItems(ctx, reference)
+}
+
 // ReadRaw retrieves a byte array from the Blob store or an error
 func (s *cachedRawStore) ReadRaw(ctx context.Context, reference DataReference) (io.ReadCloser, error) {
 	ctx, span := otelutils.NewSpan(ctx, otelutils.BlobstoreClientTracer, "flytestdlib.storage.cachedRawStore/ReadRaw")

flytestdlib/storage/cached_rawstore_test.go

@@ -55,6 +55,7 @@ func dummyCacheStore(t *testing.T, store RawStore, metrics *cacheMetrics) *cache
 type dummyStore struct {
 	copyImpl
 	HeadCb     func(ctx context.Context, reference DataReference) (Metadata, error)
+	GetItemsCb func(ctx context.Context, reference DataReference) ([]string, error)
 	ReadRawCb  func(ctx context.Context, reference DataReference) (io.ReadCloser, error)
 	WriteRawCb func(ctx context.Context, reference DataReference, size int64, opts Options, raw io.Reader) error
 	DeleteCb   func(ctx context.Context, reference DataReference) error
@@ -73,6 +74,10 @@ func (d *dummyStore) Head(ctx context.Context, reference DataReference) (Metadat
 	return d.HeadCb(ctx, reference)
 }
 
+func (d *dummyStore) GetItems(ctx context.Context, reference DataReference) ([]string, error) {
+	return d.GetItemsCb(ctx, reference)
+}
+
 func (d *dummyStore) ReadRaw(ctx context.Context, reference DataReference) (io.ReadCloser, error) {
 	return d.ReadRawCb(ctx, reference)
 }
@@ -105,6 +110,14 @@ func TestCachedRawStore(t *testing.T) {
 			}
 			return MemoryMetadata{}, fmt.Errorf("err")
 		},
+		GetItemsCb: func(ctx context.Context, reference DataReference) ([]string, error) {
+			var s []string
+			if reference == "k1" {
+				s = append(s, "item")
+				return s, nil
+			}
+			return s, fmt.Errorf("err")
+		},
 		WriteRawCb: func(ctx context.Context, reference DataReference, size int64, opts Options, raw io.Reader) error {
 			if writeCalled {
 				assert.FailNow(t, "Should not be writeCalled")
@@ -160,6 +173,18 @@ func TestCachedRawStore(t *testing.T) {
 		assert.False(t, m.Exists())
 	})
 
+	t.Run("Get Items", func(t *testing.T) {
+		items, err := cStore.GetItems(ctx, k1)
+		assert.NoError(t, err)
+		assert.Equal(t, 1, len(items))
+	})
+
+	t.Run("No Items", func(t *testing.T) {
+		items, err := cStore.GetItems(ctx, k2)
+		assert.Error(t, err)
+		assert.Equal(t, 0, len(items))
+	})
+
 	t.Run("ReadCachePopulate", func(t *testing.T) {
 		o, err := cStore.ReadRaw(ctx, k1)
 		assert.NoError(t, err)