Merge pull request #51 from buildkite/fix_blob_storage_buffering

fix: updated the blob storage to just pass through the reader

Author: wolfeidau

blob_storage.go

@@ -3,6 +3,7 @@ package buildkitelogs
 import (
 	"context"
 	"fmt"
+	"io"
 	"os"
 	"runtime"
 	"time"
@@ -15,7 +16,6 @@ import (
 // BlobStorage provides an abstraction over blob storage backends
 type BlobStorage struct {
 	bucket *blob.Bucket
-	ctx    context.Context
 }
 
 // BlobMetadata contains metadata for cached blobs
@@ -46,7 +46,6 @@ func NewBlobStorage(ctx context.Context, storageURL string) (*BlobStorage, error
 
 	return &BlobStorage{
 		bucket: bucket,
-		ctx:    ctx,
 	}, nil
 }
 
@@ -115,12 +114,12 @@ func GenerateBlobKey(org, pipeline, build, job string) string {
 }
 
 // Exists checks if a blob exists in storage
-func (bs *BlobStorage) Exists(key string) (bool, error) {
-	return bs.bucket.Exists(bs.ctx, key)
+func (bs *BlobStorage) Exists(ctx context.Context, key string) (bool, error) {
+	return bs.bucket.Exists(ctx, key)
 }
 
 // WriteWithMetadata writes data to blob storage with metadata
-func (bs *BlobStorage) WriteWithMetadata(key string, data []byte, metadata *BlobMetadata) error {
+func (bs *BlobStorage) WriteWithMetadata(ctx context.Context, key string, data []byte, metadata *BlobMetadata) error {
 	opts := &blob.WriterOptions{}
 
 	if metadata != nil {
@@ -136,7 +135,7 @@ func (bs *BlobStorage) WriteWithMetadata(key string, data []byte, metadata *Blob
 		}
 	}
 
-	writer, err := bs.bucket.NewWriter(bs.ctx, key, opts)
+	writer, err := bs.bucket.NewWriter(ctx, key, opts)
 	if err != nil {
 		return fmt.Errorf("failed to create blob writer: %w", err)
 	}
@@ -150,23 +149,11 @@ func (bs *BlobStorage) WriteWithMetadata(key string, data []byte, metadata *Blob
 }
 
 // ReadWithMetadata reads data from blob storage with metadata
-func (bs *BlobStorage) ReadWithMetadata(key string) ([]byte, *BlobMetadata, error) {
-	reader, err := bs.bucket.NewReader(bs.ctx, key, nil)
-	if err != nil {
-		return nil, nil, fmt.Errorf("failed to create blob reader: %w", err)
-	}
-	defer reader.Close()
-
-	// Read data
-	data := make([]byte, reader.Size())
-	if _, err := reader.Read(data); err != nil {
-		return nil, nil, fmt.Errorf("failed to read blob data: %w", err)
-	}
-
+func (bs *BlobStorage) ReadWithMetadata(ctx context.Context, key string) (*BlobMetadata, error) {
 	// Get blob attributes for metadata
-	attrs, err := bs.bucket.Attributes(bs.ctx, key)
+	attrs, err := bs.bucket.Attributes(ctx, key)
 	if err != nil {
-		return nil, nil, fmt.Errorf("failed to get blob attributes: %w", err)
+		return nil, fmt.Errorf("failed to get blob attributes: %w", err)
 	}
 
 	// Extract metadata
@@ -190,21 +177,27 @@ func (bs *BlobStorage) ReadWithMetadata(key string) ([]byte, *BlobMetadata, erro
 		}
 	}
 
-	return data, metadata, nil
+	return metadata, nil
+}
+
+// Reader returns an io.ReadCloser for streaming blob data from the specified key.
+// The caller is responsible for closing the returned reader when done.
+func (bs *BlobStorage) Reader(ctx context.Context, key string) (io.ReadCloser, error) {
+	return bs.bucket.NewReader(ctx, key, nil)
 }
 
 // GetModTime returns the modification time of a blob
-func (bs *BlobStorage) GetModTime(key string) (time.Time, error) {
-	attrs, err := bs.bucket.Attributes(bs.ctx, key)
+func (bs *BlobStorage) GetModTime(ctx context.Context, key string) (time.Time, error) {
+	attrs, err := bs.bucket.Attributes(ctx, key)
 	if err != nil {
 		return time.Time{}, fmt.Errorf("failed to get blob attributes: %w", err)
 	}
 	return attrs.ModTime, nil
 }
 
 // Delete removes a blob from storage
-func (bs *BlobStorage) Delete(key string) error {
-	return bs.bucket.Delete(bs.ctx, key)
+func (bs *BlobStorage) Delete(ctx context.Context, key string) error {
+	return bs.bucket.Delete(ctx, key)
 }
 
 // Close closes the blob storage connection

blob_storage_test.go

@@ -43,13 +43,13 @@ func TestBlobStorage(t *testing.T) {
 	}
 
 	// Test write
-	err = blobStorage.WriteWithMetadata(key, testData, metadata)
+	err = blobStorage.WriteWithMetadata(ctx, key, testData, metadata)
 	if err != nil {
 		t.Fatalf("Failed to write blob: %v", err)
 	}
 
 	// Test exists
-	exists, err := blobStorage.Exists(key)
+	exists, err := blobStorage.Exists(ctx, key)
 	if err != nil {
 		t.Fatalf("Failed to check blob existence: %v", err)
 	}
@@ -58,15 +58,11 @@ func TestBlobStorage(t *testing.T) {
 	}
 
 	// Test read
-	readData, readMetadata, err := blobStorage.ReadWithMetadata(key)
+	readMetadata, err := blobStorage.ReadWithMetadata(ctx, key)
 	if err != nil {
 		t.Fatalf("Failed to read blob: %v", err)
 	}
 
-	if string(readData) != string(testData) {
-		t.Errorf("Expected data %s, got %s", string(testData), string(readData))
-	}
-
 	if readMetadata == nil {
 		t.Fatal("Expected metadata, got nil")
 	}

cache.go

@@ -3,6 +3,7 @@ package buildkitelogs
 import (
 	"context"
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
 )
@@ -57,24 +58,25 @@ func IsCacheFileExists(org, pipeline, build, job string) (bool, string, error) {
 }
 
 // createLocalCacheFile creates a local file from blob storage for compatibility
-func createLocalCacheFile(ctx context.Context, blobStorage *BlobStorage, blobKey, org, pipeline, build, job string) (string, error) {
-	// Get local cache file path for compatibility
-	cacheFilePath, err := GetCacheFilePath(org, pipeline, build, job)
+func createLocalCacheFile(ctx context.Context, blobStorage *BlobStorage, blobKey string) (string, error) {
+	cacheFilePath, err := os.CreateTemp("", "bklog-")
 	if err != nil {
-		return "", err
+		return "", fmt.Errorf("failed to create local cache file: %w", err)
 	}
+	defer cacheFilePath.Close()
 
 	// Read from blob storage
-	data, _, err := blobStorage.ReadWithMetadata(blobKey)
+	reader, err := blobStorage.Reader(ctx, blobKey)
 	if err != nil {
 		return "", fmt.Errorf("failed to read from blob storage: %w", err)
 	}
+	defer reader.Close()
 
 	// Write to local cache file
-	err = os.WriteFile(cacheFilePath, data, 0600)
+	_, err = io.Copy(cacheFilePath, reader)
 	if err != nil {
 		return "", fmt.Errorf("failed to write local cache file: %w", err)
 	}
 
-	return cacheFilePath, nil
+	return cacheFilePath.Name(), nil
 }