Reading and Writing Raw Data

The RawStore interface in Flyte provides a low-level abstraction for byte-level operations across various storage backends like S3, GCS, and local filesystems. It is the foundation for higher-level storage operations, such as reading and writing protobuf messages.

Writing Raw Data

To write bytes to storage, use the WriteRaw method. This method requires an io.Reader, the size of the data, and an Options struct which can carry additional metadata.

import (
    "bytes"
    "context"
    "github.com/flyteorg/flyte/v2/flytestdlib/storage"
)

func UploadData(ctx context.Context, store storage.RawStore, path storage.DataReference, data []byte) error {
    reader := bytes.NewReader(data)
    size := int64(len(data))
    
    // Options can be used to pass extra metadata like S3 headers
    opts := storage.Options{
        Metadata: map[string]interface{}{
            "Content-Type": "application/octet-stream",
        },
    }

    err := store.WriteRaw(ctx, path, size, opts, reader)
    if err != nil {
        return err
    }
    return nil
}

Reading Raw Data

The ReadRaw method returns an io.ReadCloser. By default, Flyte imposes a 2MB download limit to prevent accidental large memory allocations. If the file exceeds this limit, the call returns an error that you can verify using storage.IsExceedsLimit.

func DownloadData(ctx context.Context, store storage.RawStore, ref storage.DataReference) (io.ReadCloser, error) {
    reader, err := store.ReadRaw(ctx, ref)
    if err != nil {
        if storage.IsExceedsLimit(err) {
            // Handle case where file is larger than the configured limit (default 2MB)
            return nil, fmt.Errorf("file at %s exceeds download limit", ref)
        }
        return nil, err
    }
    return reader, nil
}

To disable or change this limit, you must configure storage.limits.maxDownloadMBs in the Flyte configuration.

Checking Object Metadata

Before performing operations, you can use the Head method to check if an object exists and retrieve its size or ETag. This is typically a lightweight operation compared to ReadRaw.

func GetObjectInfo(ctx context.Context, store storage.RawStore, ref storage.DataReference) error {
    metadata, err := store.Head(ctx, ref)
    if err != nil {
        return err
    }

    if !metadata.Exists() {
        return fmt.Errorf("object does not exist at %s", ref)
    }

    fmt.Printf("Size: %d, ETag: %s\n", metadata.Size(), metadata.Etag())
    return nil
}

The Metadata interface (implemented by StowMetadata for cloud stores) provides:

• Exists() bool
• Size() int64
• Etag() string
• ContentMD5() string

Listing Objects with Pagination

The List method allows you to iterate through objects under a specific prefix. It uses a Cursor to handle pagination.

func ListAllFiles(ctx context.Context, store storage.RawStore, prefix storage.DataReference) ([]storage.DataReference, error) {
    var allItems []storage.DataReference
    cursor := storage.NewCursorAtStart()
    maxResults := 100

    for {
        items, nextCursor, err := store.List(ctx, prefix, maxResults, cursor)
        if err != nil {
            return nil, err
        }

        allItems = append(allItems, items...)

        if storage.IsCursorEnd(nextCursor) {
            break
        }
        cursor = nextCursor
    }

    return allItems, nil
}

Deleting Data

To remove an object from the storage backend, use the Delete method.

func Cleanup(ctx context.Context, store storage.RawStore, ref storage.DataReference) error {
    err := store.Delete(ctx, ref)
    if err != nil {
        // In Stow-based stores (S3/GCS), this may return stow.ErrNotFound 
        // if the object does not exist.
        return err
    }
    return nil
}

Testing with dummyStore

For unit testing code that interacts with RawStore, Flyte provides a dummyStore (found in flytestdlib/storage/cached_rawstore_test.go) that allows you to mock specific behaviors using callbacks.

func TestMyStorageLogic(t *testing.T) {
    mockStore := &dummyStore{
        ReadRawCb: func(ctx context.Context, reference storage.DataReference) (io.ReadCloser, error) {
            return ioutils.NewBytesReadCloser([]byte("mock data")), nil
        },
        WriteRawCb: func(ctx context.Context, reference storage.DataReference, size int64, opts storage.Options, raw io.Reader) error {
            return nil // Simulate successful write
        },
    }

    // Use mockStore in your application logic
}

Configuration Gotchas

Multi-Container Access

By default, Flyte restricts access to the container defined in InitContainer. If you need to access multiple buckets or containers using full URIs (e.g., s3://other-bucket/key), you must enable storage.enable-multicontainer in your configuration.

Storage Type

The backend implementation is determined by the storage.type configuration. Supported types include s3, minio, local, mem, and stow. When using stow, you must also provide a StowConfig specifying the kind (e.g., google, azure, oracle).