Repository: hamba/avro (go)

Issue:
Support setting OCF compression level
Probably by adding the compression level to `ocf/codec.go` `DeflateCodec` as an exported field the caller can use.

Pull request:
feat: allow deflate compression level
This allows the user to set the OCF deflate compression level.

Fixes #151

Patch:
diff --git a/ocf/codec.go b/ocf/codec.go
index d681060..a690c72 100644
--- a/ocf/codec.go
+++ b/ocf/codec.go
@@ -22,13 +22,13 @@ const (
 	Snappy  CodecName = "snappy"
 )
 
-func resolveCodec(name CodecName) (Codec, error) {
+func resolveCodec(name CodecName, lvl int) (Codec, error) {
 	switch name {
 	case Null, "":
 		return &NullCodec{}, nil
 
 	case Deflate:
-		return &DeflateCodec{}, nil
+		return &DeflateCodec{compLvl: lvl}, nil
 
 	case Snappy:
 		return &SnappyCodec{}, nil
@@ -60,10 +60,12 @@ func (*NullCodec) Encode(b []byte) []byte {
 }
 
 // DeflateCodec is a flate compression codec.
-type DeflateCodec struct{}
+type DeflateCodec struct {
+	compLvl int
+}
 
 // Decode decodes the given bytes.
-func (*DeflateCodec) Decode(b []byte) ([]byte, error) {
+func (c *DeflateCodec) Decode(b []byte) ([]byte, error) {
 	r := flate.NewReader(bytes.NewBuffer(b))
 	data, err := ioutil.ReadAll(r)
 	if err != nil {
@@ -76,10 +78,10 @@ func (*DeflateCodec) Decode(b []byte) ([]byte, error) {
 }
 
 // Encode encodes the given bytes.
-func (*DeflateCodec) Encode(b []byte) []byte {
+func (c *DeflateCodec) Encode(b []byte) []byte {
 	data := bytes.NewBuffer(make([]byte, 0, len(b)))
 
-	w, _ := flate.NewWriter(data, flate.DefaultCompression)
+	w, _ := flate.NewWriter(data, c.compLvl)
 	_, _ = w.Write(b)
 	_ = w.Close()
 
diff --git a/ocf/ocf.go b/ocf/ocf.go
index e40591d..5114a57 100644
--- a/ocf/ocf.go
+++ b/ocf/ocf.go
@@ -73,7 +73,7 @@ func NewDecoder(r io.Reader) (*Decoder, error) {
 		return nil, err
 	}
 
-	codec, err := resolveCodec(CodecName(h.Meta[codecKey]))
+	codec, err := resolveCodec(CodecName(h.Meta[codecKey]), -1)
 	if err != nil {
 		return nil, err
 	}
@@ -155,9 +155,10 @@ func (d *Decoder) readBlock() int64 {
 }
 
 type encoderConfig struct {
-	BlockLength int
-	CodecName   CodecName
-	Metadata    map[string][]byte
+	BlockLength      int
+	CodecName        CodecName
+	CodecCompression int
+	Metadata         map[string][]byte
 }
 
 // EncoderFunc represents an configuration function for Encoder.
@@ -177,6 +178,15 @@ func WithCodec(codec CodecName) EncoderFunc {
 	}
 }
 
+// WithCompressionLevel sets the compression codec to deflate and
+// the compression level on the encoder.
+func WithCompressionLevel(compLvl int) EncoderFunc {
+	return func(cfg *encoderConfig) {
+		cfg.CodecName = Deflate
+		cfg.CodecCompression = compLvl
+	}
+}
+
 // WithMetadata sets the metadata on the encoder header.
 func WithMetadata(meta map[string][]byte) EncoderFunc {
 	return func(cfg *encoderConfig) {
@@ -205,9 +215,10 @@ func NewEncoder(s string, w io.Writer, opts ...EncoderFunc) (*Encoder, error) {
 	}
 
 	cfg := encoderConfig{
-		BlockLength: 100,
-		CodecName:   Null,
-		Metadata:    map[string][]byte{},
+		BlockLength:      100,
+		CodecName:        Null,
+		CodecCompression: -1,
+		Metadata:         map[string][]byte{},
 	}
 	for _, opt := range opts {
 		opt(&cfg)
@@ -224,7 +235,7 @@ func NewEncoder(s string, w io.Writer, opts ...EncoderFunc) (*Encoder, error) {
 	_, _ = rand.Read(header.Sync[:])
 	writer.WriteVal(HeaderSchema, header)
 
-	codec, err := resolveCodec(cfg.CodecName)
+	codec, err := resolveCodec(cfg.CodecName, cfg.CodecCompression)
 	if err != nil {
 		return nil, err
 	}


Test patch:
diff --git a/ocf/ocf_test.go b/ocf/ocf_test.go
index 1431e44..b34a59c 100644
--- a/ocf/ocf_test.go
+++ b/ocf/ocf_test.go
@@ -2,6 +2,7 @@ package ocf_test
 
 import (
 	"bytes"
+	"compress/flate"
 	"errors"
 	"os"
 	"testing"
@@ -122,7 +123,7 @@ func TestDecoder(t *testing.T) {
 		t.Error(err)
 		return
 	}
-	defer f.Close()
+	t.Cleanup(func() { _ = f.Close() })
 
 	dec, err := ocf.NewDecoder(f)
 	if err != nil {
@@ -174,7 +175,7 @@ func TestDecoderDeflate(t *testing.T) {
 		t.Error(err)
 		return
 	}
-	defer f.Close()
+	t.Cleanup(func() { _ = f.Close() })
 
 	dec, err := ocf.NewDecoder(f)
 	if err != nil {
@@ -202,7 +203,7 @@ func TestDecoderDeflate_InvalidData(t *testing.T) {
 		t.Error(err)
 		return
 	}
-	defer f.Close()
+	t.Cleanup(func() { _ = f.Close() })
 
 	dec, err := ocf.NewDecoder(f)
 	if err != nil {
@@ -245,7 +246,7 @@ func TestDecoderSnappy(t *testing.T) {
 		t.Error(err)
 		return
 	}
-	defer f.Close()
+	t.Cleanup(func() { _ = f.Close() })
 
 	dec, err := ocf.NewDecoder(f)
 	if err != nil {
@@ -273,7 +274,7 @@ func TestDecoderSnappy_InvalidData(t *testing.T) {
 		t.Error(err)
 		return
 	}
-	defer f.Close()
+	t.Cleanup(func() { _ = f.Close() })
 
 	dec, err := ocf.NewDecoder(f)
 	if err != nil {
@@ -292,7 +293,7 @@ func TestDecoderSnappy_ShortCRC(t *testing.T) {
 		t.Error(err)
 		return
 	}
-	defer f.Close()
+	t.Cleanup(func() { _ = f.Close() })
 
 	dec, err := ocf.NewDecoder(f)
 	if err != nil {
@@ -311,7 +312,7 @@ func TestDecoderSnappy_InvalidCRC(t *testing.T) {
 		t.Error(err)
 		return
 	}
-	defer f.Close()
+	t.Cleanup(func() { _ = f.Close() })
 
 	dec, err := ocf.NewDecoder(f)
 	if err != nil {
@@ -504,6 +505,43 @@ func TestEncoder_EncodeCompressesDeflate(t *testing.T) {
 	assert.Equal(t, 926, buf.Len())
 }
 
+func TestEncoder_EncodeCompressesDeflateWithLevel(t *testing.T) {
+	unionStr := "union value"
+	record := FullRecord{
+		Strings: []string{"string1", "string2", "string3", "string4", "string5"},
+		Longs:   []int64{1, 2, 3, 4, 5},
+		Enum:    "C",
+		Map: map[string]int{
+			"key1": 1,
+			"key2": 2,
+			"key3": 3,
+			"key4": 4,
+			"key5": 5,
+		},
+		Nullable: &unionStr,
+		Fixed:    [16]byte{0x01, 0x02, 0x03, 0x04, 0x01, 0x02, 0x03, 0x04, 0x01, 0x02, 0x03, 0x04, 0x01, 0x02, 0x03, 0x04},
+		Record: &TestRecord{
+			Long:   1925639126735,
+			String: "I am a test record",
+			Int:    666,
+			Float:  7171.17,
+			Double: 916734926348163.01973408746523,
+			Bool:   true,
+		},
+	}
+
+	buf := &bytes.Buffer{}
+	enc, _ := ocf.NewEncoder(schema, buf, ocf.WithCompressionLevel(flate.BestCompression))
+
+	err := enc.Encode(record)
+	assert.NoError(t, err)
+
+	err = enc.Close()
+	assert.NoError(t, err)
+
+	assert.Equal(t, 926, buf.Len())
+}
+
 func TestEncoder_EncodeCompressesSnappy(t *testing.T) {
 	unionStr := "union value"
 	record := FullRecord{
@@ -553,7 +591,7 @@ func TestEncoder_EncodeError(t *testing.T) {
 func TestEncoder_EncodeWritesBlocks(t *testing.T) {
 	buf := &bytes.Buffer{}
 	enc, _ := ocf.NewEncoder(`"long"`, buf, ocf.WithBlockLength(1))
-	defer enc.Close()
+	t.Cleanup(func() { _ = enc.Close() })
 
 	err := enc.Encode(int64(1))
 
@@ -564,7 +602,7 @@ func TestEncoder_EncodeWritesBlocks(t *testing.T) {
 func TestEncoder_EncodeHandlesWriteBlockError(t *testing.T) {
 	w := &errorWriter{}
 	enc, _ := ocf.NewEncoder(`"long"`, w, ocf.WithBlockLength(1))
-	defer enc.Close()
+	t.Cleanup(func() { _ = enc.Close() })
 
 	err := enc.Encode(int64(1))
 
@@ -590,7 +628,7 @@ func TestEncodeDecodeMetadata(t *testing.T) {
 	err := enc.Encode(int64(1))
 	assert.NoError(t, err)
 
-	enc.Close()
+	_ = enc.Close()
 
 	dec, err := ocf.NewDecoder(buf)
 	assert.NoError(t, err)