github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/pkg/sink/codec/csv/csv_encoder.go (about)

     1  // Copyright 2022 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package csv
    15  
    16  import (
    17  	"bytes"
    18  
    19  	"github.com/pingcap/tiflow/cdc/model"
    20  	"github.com/pingcap/tiflow/pkg/config"
    21  	"github.com/pingcap/tiflow/pkg/sink/codec"
    22  	"github.com/pingcap/tiflow/pkg/sink/codec/common"
    23  )
    24  
    25  // BatchEncoder encodes the events into the byte of a batch into.
    26  type BatchEncoder struct {
    27  	valueBuf  *bytes.Buffer
    28  	callback  func()
    29  	batchSize int
    30  	config    *common.Config
    31  }
    32  
    33  // AppendTxnEvent implements the TxnEventEncoder interface
    34  func (b *BatchEncoder) AppendTxnEvent(
    35  	e *model.SingleTableTxn,
    36  	callback func(),
    37  ) error {
    38  	for _, rowEvent := range e.Rows {
    39  		row, err := rowChangedEvent2CSVMsg(b.config, rowEvent)
    40  		if err != nil {
    41  			return err
    42  		}
    43  		b.valueBuf.Write(row.encode())
    44  		b.batchSize++
    45  	}
    46  	b.callback = callback
    47  	return nil
    48  }
    49  
    50  // Build implements the RowEventEncoder interface
    51  func (b *BatchEncoder) Build() (messages []*common.Message) {
    52  	if b.batchSize == 0 {
    53  		return nil
    54  	}
    55  
    56  	ret := common.NewMsg(config.ProtocolCsv, nil,
    57  		b.valueBuf.Bytes(), 0, model.MessageTypeRow, nil, nil)
    58  	ret.SetRowsCount(b.batchSize)
    59  	ret.Callback = b.callback
    60  	if b.valueBuf.Cap() > codec.MemBufShrinkThreshold {
    61  		b.valueBuf = &bytes.Buffer{}
    62  	} else {
    63  		b.valueBuf.Reset()
    64  	}
    65  	b.callback = nil
    66  	b.batchSize = 0
    67  
    68  	return []*common.Message{ret}
    69  }
    70  
    71  // newBatchEncoder creates a new csv BatchEncoder.
    72  func newBatchEncoder(config *common.Config) codec.TxnEventEncoder {
    73  	return &BatchEncoder{
    74  		config:   config,
    75  		valueBuf: &bytes.Buffer{},
    76  	}
    77  }
    78  
    79  type batchEncoderBuilder struct {
    80  	config *common.Config
    81  }
    82  
    83  // NewTxnEventEncoderBuilder creates a csv batchEncoderBuilder.
    84  func NewTxnEventEncoderBuilder(config *common.Config) codec.TxnEventEncoderBuilder {
    85  	return &batchEncoderBuilder{config: config}
    86  }
    87  
    88  // Build a csv BatchEncoder
    89  func (b *batchEncoderBuilder) Build() codec.TxnEventEncoder {
    90  	return newBatchEncoder(b.config)
    91  }