aws_s3_stream: add opt-in whole-object streaming gzip (compression: gzip)

Author: Tim Riddell

internal/impl/aws/output_s3_stream.go

@@ -26,6 +26,7 @@ const (
 	ssoFieldBatching           = "batching"
 	ssoFieldContentType        = "content_type"
 	ssoFieldContentEncoding    = "content_encoding"
+	ssoFieldCompression        = "compression"
 	ssoFieldBackoff            = "backoff"
 	ssoFieldMaxRetries         = "max_retries"
 )
@@ -94,9 +95,13 @@ You can find out more [in this document](/docs/guides/cloud/aws).
 				Default("application/octet-stream").
 				Advanced(),
 			service.NewInterpolatedStringField(ssoFieldContentEncoding).
-				Description("The content encoding to set for uploaded files (e.g., gzip).").
+				Description("The content encoding to set for uploaded files (e.g., gzip). This only sets the object's `Content-Encoding` metadata and does not itself compress anything; use `compression` for that. Leave unset when using `compression`, which sets it automatically.").
 				Optional().
 				Advanced(),
+			service.NewStringEnumField(ssoFieldCompression, "none", "gzip").
+				Description("Compress the object in-stream as it is uploaded. When set to `gzip`, a single gzip stream is written across the whole object (including true multipart uploads that exceed the 5 MiB part size) and the object's `Content-Encoding` is set to `gzip` automatically. This differs from a per-message `compress` processor, which would produce a multi-member gzip file; and from batch-level `archive`+`compress`, which would merge records across partitions. When enabled, do not also set `content_encoding` or a `compress` batch processor.").
+				Default("none").
+				Advanced(),
 			service.NewIntField(ssoFieldMaxRetries).
 				Description("The maximum number of retries for each individual part upload. Set to zero to disable retries.").
 				Advanced().Default(2),
@@ -159,6 +164,7 @@ type s3StreamConfig struct {
 	MaxBufferPeriod time.Duration
 	ContentType     *service.InterpolatedString
 	ContentEncoding *service.InterpolatedString
+	Compression     string
 
 	aconf       aws.Config
 	backoffCtor func() backoff.BackOff
@@ -210,6 +216,11 @@ func s3StreamConfigFromParsed(pConf *service.ParsedConfig) (conf s3StreamConfig,
 		}
 	}
 
+	// In-writer compression (defaults to "none").
+	if conf.Compression, err = pConf.FieldString(ssoFieldCompression); err != nil {
+		return
+	}
+
 	// AWS config
 	if conf.aconf, err = GetSession(context.TODO(), pConf); err != nil {
 		return
@@ -388,6 +399,7 @@ func (s *s3StreamOutput) writeToPartition(ctx context.Context, partitionKey stri
 				MaxBufferPeriod: s.conf.MaxBufferPeriod,
 				ContentType:     contentType,
 				ContentEncoding: contentEncoding,
+				Compression:     s.conf.Compression,
 				BackoffCtor:     s.conf.backoffCtor,
 			})
 			if err != nil {

internal/impl/aws/output_s3_stream_integration_test.go

@@ -1,6 +1,8 @@
 package aws
 
 import (
+	"bytes"
+	"compress/gzip"
 	"context"
 	"fmt"
 	"io"
@@ -364,3 +366,96 @@ max_buffer_period: 1s
 
 	assert.Equal(t, "single message\n", string(content))
 }
+
+// TestS3StreamOutput_IntegrationGzipCompression tests the compression: gzip
+// option end-to-end: the object must download as a single valid gzip stream
+// (Content-Encoding gzip) that decompresses to the exact concatenation of the
+// records — for both a small PutObject-path object and a large multipart one.
+func TestS3StreamOutput_IntegrationGzipCompression(t *testing.T) {
+	integration.CheckSkip(t)
+
+	servicePort := GetLocalStack(t, nil)
+	bucketName := "test-stream-gzip"
+
+	s3Client := getTestS3Client(context.Background(), t, servicePort)
+	_, err := s3Client.CreateBucket(context.Background(), &s3.CreateBucketInput{
+		Bucket: aws.String(bucketName),
+	})
+	require.NoError(t, err)
+
+	cases := []struct {
+		name    string
+		key     string
+		records int
+		// bufBytes seals parts at this compressed size; small value + many
+		// records forces a true multipart upload.
+		bufBytes int
+	}{
+		{name: "small_putobject", key: "gz/small.jsonl.gz", records: 20, bufBytes: 10 * 1024 * 1024},
+		{name: "large_multipart", key: "gz/large.jsonl.gz", records: 400000, bufBytes: 5 * 1024 * 1024},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			configYAML := fmt.Sprintf(`
+bucket: %s
+path: '%s'
+region: us-east-1
+force_path_style_urls: true
+endpoint: http://localhost:%s
+content_type: 'application/json'
+compression: gzip
+max_buffer_bytes: %d
+max_buffer_count: 1000000
+max_buffer_period: 1s
+`, bucketName, tc.key, servicePort, tc.bufBytes)
+
+			parsedConf, err := s3StreamOutputSpec().ParseYAML(configYAML, nil)
+			require.NoError(t, err)
+
+			wConf, err := s3StreamConfigFromParsed(parsedConf)
+			require.NoError(t, err)
+			require.Equal(t, "gzip", wConf.Compression)
+
+			output, err := newS3StreamOutput(wConf, service.MockResources())
+			require.NoError(t, err)
+
+			ctx := context.Background()
+			require.NoError(t, output.Connect(ctx))
+
+			var expected bytes.Buffer
+			batch := service.MessageBatch{}
+			for i := range tc.records {
+				line := fmt.Appendf(nil, `{"id":%d,"msg":"gzip stream test record"}%s`, i, "\n")
+				expected.Write(line)
+				batch = append(batch, service.NewMessage(line))
+			}
+
+			require.NoError(t, output.WriteBatch(ctx, batch))
+			require.NoError(t, output.Close(ctx))
+
+			getResp, err := s3Client.GetObject(ctx, &s3.GetObjectInput{
+				Bucket: aws.String(bucketName),
+				Key:    aws.String(tc.key),
+			})
+			require.NoError(t, err)
+			defer getResp.Body.Close()
+
+			assert.Equal(t, "gzip", aws.ToString(getResp.ContentEncoding), "Content-Encoding should be gzip")
+
+			raw, err := io.ReadAll(getResp.Body)
+			require.NoError(t, err)
+
+			// The stored object must be ONE valid gzip stream (not multi-member
+			// per-record) that decompresses to the exact input.
+			zr, err := gzip.NewReader(bytes.NewReader(raw))
+			require.NoError(t, err, "object is not a valid gzip stream")
+			got, err := io.ReadAll(zr)
+			require.NoError(t, err)
+			require.NoError(t, zr.Close())
+
+			assert.Equal(t, expected.Bytes(), got)
+			assert.Less(t, len(raw), expected.Len(), "gzip object should be smaller than the raw input")
+		})
+	}
+}

internal/impl/aws/output_s3_stream_writer.go

@@ -2,6 +2,7 @@ package aws
 
 import (
 	"bytes"
+	"compress/gzip"
 	"context"
 	"errors"
 	"fmt"
@@ -45,10 +46,17 @@ type S3StreamingWriter struct {
 	completedParts []types.CompletedPart
 	backoffCtor    func() backoff.BackOff
 
-	// Buffering
-	messageBuffer [][]byte
+	// Buffering.
+	//
+	// uploadBuffer holds the bytes pending upload as the next S3 part. When
+	// compression is enabled, gzipWriter streams compressed bytes into
+	// uploadBuffer, so uploadSize (and therefore all part-size decisions) is
+	// measured in COMPRESSED bytes; the single gzip stream is simply split
+	// across parts at arbitrary byte boundaries and reassembled by S3.
+	gzipWriter    *gzip.Writer
 	uploadBuffer  *bytes.Buffer
 	uploadSize    int64
+	bufferedCount int
 
 	// Statistics tracking
 	totalMessages int64
@@ -75,6 +83,7 @@ type S3StreamingWriterConfig struct {
 	MaxBufferPeriod time.Duration // Maximum time to buffer before flushing (default: 10s)
 	ContentType     string        // Content type for S3 object
 	ContentEncoding string        // Content encoding for S3 object (optional)
+	Compression     string        // In-writer compression over the whole object: "" / "none" / "gzip"
 	BackoffCtor     func() backoff.BackOff
 }
 
@@ -118,12 +127,20 @@ func NewS3StreamingWriter(config S3StreamingWriterConfig) (*S3StreamingWriter, e
 		key:             config.Key,
 		backoffCtor:     config.BackoffCtor,
 		uploadBuffer:    bytes.NewBuffer(nil),
-		messageBuffer:   make([][]byte, 0, config.MaxBufferCount),
 		created:         time.Now(),
 		lastWrite:       time.Now(),
 		lastFlush:       time.Now(),
 	}
 
+	// When gzip compression is enabled the writer owns the compression: a single
+	// gzip stream is written across all parts and the object is tagged with the
+	// gzip Content-Encoding. This is the only way to get whole-object (not
+	// per-record) gzip that can still exceed the 5 MiB multipart threshold.
+	if config.Compression == "gzip" {
+		w.gzipWriter = gzip.NewWriter(w.uploadBuffer)
+		w.contentEncoding = "gzip"
+	}
+
 	return w, nil
 }
 
@@ -177,16 +194,24 @@ func (w *S3StreamingWriter) WriteBytes(ctx context.Context, data []byte) error {
 		return errors.New("writer not initialized")
 	}
 
-	// Add to message buffer
-	w.messageBuffer = append(w.messageBuffer, data)
-	w.uploadBuffer.Write(data)
-	w.uploadSize += int64(len(data))
+	if w.gzipWriter != nil {
+		// Stream the record through gzip into uploadBuffer. Part-size decisions
+		// are made on the compressed bytes accumulated so far.
+		if _, err := w.gzipWriter.Write(data); err != nil {
+			return fmt.Errorf("failed to write to gzip stream: %w", err)
+		}
+		w.uploadSize = int64(w.uploadBuffer.Len())
+	} else {
+		w.uploadBuffer.Write(data)
+		w.uploadSize += int64(len(data))
+	}
+	w.bufferedCount++
 	w.totalMessages++
 	w.totalBytes += int64(len(data))
 	w.lastWrite = time.Now()
 
 	// Check if we should flush
-	if w.uploadSize >= w.maxBufferBytes || len(w.messageBuffer) >= w.maxBufferCount {
+	if w.uploadSize >= w.maxBufferBytes || w.bufferedCount >= w.maxBufferCount {
 		return w.flush(ctx)
 	}
 
@@ -228,10 +253,11 @@ retryLoop:
 				PartNumber: aws.Int32(w.partNumber),
 			})
 
-			// Clear buffers
+			// Clear buffers (gzipWriter keeps streaming into the now-empty
+			// uploadBuffer; the uploaded part is a prefix of the gzip stream).
 			w.uploadBuffer.Reset()
 			w.uploadSize = 0
-			w.messageBuffer = w.messageBuffer[:0]
+			w.bufferedCount = 0
 			w.lastFlush = time.Now()
 
 			// Reset flush timer
@@ -323,7 +349,7 @@ func (w *S3StreamingWriter) forceFlush(ctx context.Context) error {
 	// Clear buffers
 	w.uploadBuffer.Reset()
 	w.uploadSize = 0
-	w.messageBuffer = w.messageBuffer[:0]
+	w.bufferedCount = 0
 	w.lastFlush = time.Now()
 
 	return nil
@@ -370,6 +396,18 @@ func (w *S3StreamingWriter) Close(ctx context.Context) error {
 		w.flushTimer.Stop()
 	}
 
+	// Finalize the gzip stream: this flushes any bytes still held inside the
+	// compressor and writes the gzip footer into uploadBuffer, after which
+	// uploadSize reflects the final compressed tail. Must happen before the
+	// size checks below so small/empty objects are sized correctly.
+	if w.gzipWriter != nil {
+		if err := w.gzipWriter.Close(); err != nil {
+			w.abortMultipartUpload(ctx)
+			return fmt.Errorf("failed to finalize gzip stream: %w", err)
+		}
+		w.uploadSize = int64(w.uploadBuffer.Len())
+	}
+
 	// A small or empty object (no parts yet and < 5 MiB buffered) cannot be a
 	// standalone multipart part, so finalize it with a single PutObject.
 	if len(w.completedParts) == 0 && w.uploadSize < s3MinMultipartPartSize {