github.com/segmentio/kafka-go@v0.4.48-0.20240318174348-3f6244eb34fd/compress/compress.go (about)

     1  package compress
     2  
     3  import (
     4  	"encoding"
     5  	"fmt"
     6  	"io"
     7  	"strconv"
     8  	"strings"
     9  
    10  	"github.com/segmentio/kafka-go/compress/gzip"
    11  	"github.com/segmentio/kafka-go/compress/lz4"
    12  	"github.com/segmentio/kafka-go/compress/snappy"
    13  	"github.com/segmentio/kafka-go/compress/zstd"
    14  )
    15  
    16  // Compression represents the compression applied to a record set.
    17  type Compression int8
    18  
    19  const (
    20  	None   Compression = 0
    21  	Gzip   Compression = 1
    22  	Snappy Compression = 2
    23  	Lz4    Compression = 3
    24  	Zstd   Compression = 4
    25  )
    26  
    27  func (c Compression) Codec() Codec {
    28  	if i := int(c); i >= 0 && i < len(Codecs) {
    29  		return Codecs[i]
    30  	}
    31  	return nil
    32  }
    33  
    34  func (c Compression) String() string {
    35  	if codec := c.Codec(); codec != nil {
    36  		return codec.Name()
    37  	}
    38  	return "uncompressed"
    39  }
    40  
    41  func (c Compression) MarshalText() ([]byte, error) {
    42  	return []byte(c.String()), nil
    43  }
    44  
    45  func (c *Compression) UnmarshalText(b []byte) error {
    46  	switch string(b) {
    47  	case "none", "uncompressed":
    48  		*c = None
    49  		return nil
    50  	}
    51  
    52  	for _, codec := range Codecs[None+1:] {
    53  		if codec.Name() == string(b) {
    54  			*c = Compression(codec.Code())
    55  			return nil
    56  		}
    57  	}
    58  
    59  	i, err := strconv.ParseInt(string(b), 10, 64)
    60  	if err == nil && i >= 0 && i < int64(len(Codecs)) {
    61  		*c = Compression(i)
    62  		return nil
    63  	}
    64  
    65  	s := &strings.Builder{}
    66  	s.WriteString("none, uncompressed")
    67  
    68  	for i, codec := range Codecs[None+1:] {
    69  		if i < (len(Codecs) - 1) {
    70  			s.WriteString(", ")
    71  		} else {
    72  			s.WriteString(", or ")
    73  		}
    74  		s.WriteString(codec.Name())
    75  	}
    76  
    77  	return fmt.Errorf("compression format must be one of %s, not %q", s, b)
    78  }
    79  
    80  var (
    81  	_ encoding.TextMarshaler   = Compression(0)
    82  	_ encoding.TextUnmarshaler = (*Compression)(nil)
    83  )
    84  
    85  // Codec represents a compression codec to encode and decode the messages.
    86  // See : https://cwiki.apache.org/confluence/display/KAFKA/Compression
    87  //
    88  // A Codec must be safe for concurrent access by multiple go routines.
    89  type Codec interface {
    90  	// Code returns the compression codec code
    91  	Code() int8
    92  
    93  	// Human-readable name for the codec.
    94  	Name() string
    95  
    96  	// Constructs a new reader which decompresses data from r.
    97  	NewReader(r io.Reader) io.ReadCloser
    98  
    99  	// Constructs a new writer which writes compressed data to w.
   100  	NewWriter(w io.Writer) io.WriteCloser
   101  }
   102  
   103  var (
   104  	// The global gzip codec installed on the Codecs table.
   105  	GzipCodec gzip.Codec
   106  
   107  	// The global snappy codec installed on the Codecs table.
   108  	SnappyCodec snappy.Codec
   109  
   110  	// The global lz4 codec installed on the Codecs table.
   111  	Lz4Codec lz4.Codec
   112  
   113  	// The global zstd codec installed on the Codecs table.
   114  	ZstdCodec zstd.Codec
   115  
   116  	// The global table of compression codecs supported by the kafka protocol.
   117  	Codecs = [...]Codec{
   118  		None:   nil,
   119  		Gzip:   &GzipCodec,
   120  		Snappy: &SnappyCodec,
   121  		Lz4:    &Lz4Codec,
   122  		Zstd:   &ZstdCodec,
   123  	}
   124  )