Core Storage Abstraction Overview

Flyte components often need to interact with various cloud storage providers (AWS S3, Google Cloud Storage, Azure Blob Storage) or local filesystems without being tied to a specific provider's SDK. If you hardcode S3-specific logic into a plugin, it becomes unusable for users on GCP.

The DataStore abstraction in flytestdlib/storage provides a unified interface that handles these differences, allowing you to write storage-agnostic code that works across all supported environments.

Addressing Data with DataReference

In Flyte, every piece of data is addressed using a DataReference. This is a string-based type that represents a URI to a storage location.

// From flytestdlib/storage/storage.go
type DataReference string

A DataReference typically follows the format scheme://container/key (e.g., s3://my-bucket/path/to/file). However, Flyte handles cloud-specific nuances internally. For example, when working with Azure ADLS Gen2, the Split() method correctly extracts the container from the userinfo part of the URL (container@account) rather than the host.

The DataStore Interface

The DataStore class is the primary entry point for all storage operations. It is a composed interface that combines three distinct functional areas:

1. RawStore: Low-level byte-level access (Read/Write/Copy/Delete).
2. ProtobufStore: High-level structured data access for Protobuf messages.
3. ReferenceConstructor: Utilities for building portable, storage-agnostic paths.

Byte-Level Operations (RawStore)

When you need to move raw files or byte streams, use the RawStore methods. A common scenario is uploading a local file to a remote location during a task execution.

// Example from flytecopilot/data/utils.go
func UploadFileToStorage(ctx context.Context, filePath string, toPath storage.DataReference, size int64, store *storage.DataStore) error {
	f, err := os.Open(filePath)
	if err != nil {
		return err
	}
	defer f.Close()
	
	// WriteRaw handles the stream and creates the container/bucket if it doesn't exist
	return store.WriteRaw(ctx, toPath, size, storage.Options{}, f)
}

The RawStore interface (defined in flytestdlib/storage/storage.go) includes methods like Head for metadata, List for directory-like exploration, and CopyRaw for server-side copies.

Structured Data (ProtobufStore)

Flyte relies heavily on Protobuf for its internal data types, such as LiteralMap. The ProtobufStore facet of DataStore automates the serialization and deserialization process.

// Example adapted from flyteplugins/go/tasks/pluginmachinery/catalog/cache_service/client.go
func readCachedOutput(ctx context.Context, store *storage.DataStore, outputURI string) (*corepb.LiteralMap, error) {
	outputs := &corepb.LiteralMap{}
	
	// ReadProtobuf automatically fetches the bytes and unmarshals them into the provided message
	if err := store.ReadProtobuf(ctx, storage.DataReference(outputURI), outputs); err != nil {
		return nil, fmt.Errorf("failed to read output data from %q: %w", outputURI, err)
	}
	
	return outputs, nil
}

Portable Path Construction

To avoid hardcoding URI schemes or path separators, use the ReferenceConstructor. This ensures that your code remains portable across different storage backends.

// Example from flytestdlib/storage/storage_test.go
func ExampleNewDataStore() {
	ctx := context.Background()
	// ... store initialization ...

	// ConstructReference joins parts using the correct separator for the underlying store
	ref, _ := store.ConstructReference(ctx, storage.DataReference("s3://my-bucket"), "metadata", "inputs.pb")
	// Result: "s3://my-bucket/metadata/inputs.pb"
}

Configuration and Limits

Flyte's storage behavior is controlled via the Config struct. You can switch between storage types by setting storage.type to s3, gcs, local, mem, or stow.

Download Limits

To prevent memory exhaustion, Flyte enforces a download limit on ReadRaw calls. This is configured via storage.limits.maxDownloadMBs (defaulting to 2MB). If you attempt to read a file larger than this limit, the call will return an ErrExceedsLimit error.

Multi-container Access

By default, Flyte may be restricted to a specific base container. If your application needs to access multiple buckets or containers, ensure storage.enable-multicontainer is set to true. This allows the DataStore to resolve any valid DataReference regardless of the configured base container.

Metadata Checks

When using the Head() method to check for a file's existence, always verify the Exists() method on the returned Metadata object. Some backend implementations might return a non-nil metadata object even if the file is missing, rather than returning a 404-style error immediately.