From 169eaf1921157b0f1c8b3917551b8cebe9377bbf Mon Sep 17 00:00:00 2001
From: Martin Sucha <martin.sucha@kiwi.com>
Date: Fri, 22 Sep 2023 13:11:42 +0200
Subject: [PATCH] Expose whether an outgoing frame was compressed or not

Clients might want to track this info.
---
 conn.go       |  1 +
 frame.go      |  6 +++---
 frame_test.go | 27 +++++++++++++++++++++++++--
 3 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/conn.go b/conn.go
index 09a8bb74e..73137ad52 100644
--- a/conn.go
+++ b/conn.go
@@ -1254,6 +1254,7 @@ type ObservedStream struct {
 	FramePayloadUncompressedSize int
 	// FramePayloadCompressedSize is the compressed size of the frame payload (without frame header).
 	// This field is only available in StreamStarted.
+	// FramePayloadCompressedSize is zero if the frame was not compressed.
 	FramePayloadCompressedSize int
 	// QueryValuesSize is the total uncompressed size of query values in the frame (without other query options).
 	// For a batch, it is the sum for all queries in the batch.
diff --git a/frame.go b/frame.go
index ce5d3403c..8fd7258cc 100644
--- a/frame.go
+++ b/frame.go
@@ -864,10 +864,10 @@ func (f *framer) finish() (outFrameInfo, error) {
 		}
 
 		f.buf = append(f.buf[:f.headSize], compressed...)
+
+		info.compressedSize = len(f.buf) - f.headSize
 	}
-	length := len(f.buf) - f.headSize
-	info.compressedSize = length
-	f.setLength(length)
+	f.setLength(len(f.buf) - f.headSize)
 
 	return info, nil
 }
diff --git a/frame_test.go b/frame_test.go
index 7a36d23e3..299ffbcc4 100644
--- a/frame_test.go
+++ b/frame_test.go
@@ -107,6 +107,7 @@ func TestFrameReadTooLong(t *testing.T) {
 func TestOutFrameInfo(t *testing.T) {
 	tests := map[string]struct {
 		frame        frameBuilder
+		compress     bool
 		expectedInfo outFrameInfo
 	}{
 		"query": {
@@ -132,6 +133,7 @@ func TestOutFrameInfo(t *testing.T) {
 				},
 				customPayload: nil,
 			},
+			compress: true,
 			expectedInfo: outFrameInfo{
 				op:               opQuery,
 				uncompressedSize: 81,
@@ -155,6 +157,7 @@ func TestOutFrameInfo(t *testing.T) {
 				},
 				customPayload: nil,
 			},
+			compress: true,
 			expectedInfo: outFrameInfo{
 				op:               opExecute,
 				compressedSize:   50,
@@ -198,6 +201,7 @@ func TestOutFrameInfo(t *testing.T) {
 				defaultTimestampValue: 0,
 				customPayload:         nil,
 			},
+			compress: true,
 			expectedInfo: outFrameInfo{
 				op:               opBatch,
 				compressedSize:   96,
@@ -207,7 +211,8 @@ func TestOutFrameInfo(t *testing.T) {
 			},
 		},
 		"options": {
-			frame: &writeOptionsFrame{},
+			frame:    &writeOptionsFrame{},
+			compress: true,
 			expectedInfo: outFrameInfo{
 				op:               opOptions,
 				compressedSize:   0,
@@ -220,6 +225,7 @@ func TestOutFrameInfo(t *testing.T) {
 			frame: &writeRegisterFrame{
 				events: []string{"event1", "event2"},
 			},
+			compress: true,
 			expectedInfo: outFrameInfo{
 				op:               opRegister,
 				compressedSize:   20,
@@ -228,10 +234,27 @@ func TestOutFrameInfo(t *testing.T) {
 				queryCount:       0,
 			},
 		},
+		"register uncompressed": {
+			frame: &writeRegisterFrame{
+				events: []string{"event1", "event2"},
+			},
+			compress: false,
+			expectedInfo: outFrameInfo{
+				op:               opRegister,
+				compressedSize:   0,
+				uncompressedSize: 18,
+				queryValuesSize:  0,
+				queryCount:       0,
+			},
+		},
 	}
 	for name, test := range tests {
 		t.Run(name, func(t *testing.T) {
-			fr := newFramer(SnappyCompressor{}, 4)
+			var compressor Compressor
+			if test.compress {
+				compressor = SnappyCompressor{}
+			}
+			fr := newFramer(compressor, 4)
 			ofi, err := test.frame.buildFrame(fr, 42)
 			if err != nil {
 				t.Fatal(err)