sink: add column selector for storage sink (#5595)

api/v2/changefeed.go

@@ -1700,7 +1700,7 @@ func getVerifiedTables(
 		return nil, nil, nil, err
 	}
 
-	if err := verifyTable4MQ(replicaConfig, scheme, topic, protocol, tableInfos); err != nil {
+	if err := verifyTablesForSink(replicaConfig, scheme, topic, protocol, tableInfos); err != nil {
 		return nil, nil, nil, err
 	}
 
@@ -1715,13 +1715,21 @@ func getVerifiedTables(
 	return ineligibleTables, eligibleTables, allTables, nil
 }
 
-func verifyTable4MQ(
+func verifyTablesForSink(
 	replicaConfig *config.ReplicaConfig,
 	scheme string,
 	topic string,
 	protocol config.Protocol,
 	tableInfos []*common.TableInfo,
 ) error {
+	if config.IsStorageScheme(scheme) {
+		selectors, err := columnselector.New(replicaConfig.Sink)
+		if err != nil {
+			return err
+		}
+		return selectors.VerifyTables(tableInfos, nil)
+	}
+
 	if !config.IsMQScheme(scheme) {
 		return nil
 	}

api/v2/changefeed_test.go

@@ -20,6 +20,10 @@ import (
 	"github.com/pingcap/ticdc/pkg/config"
 	"github.com/pingcap/ticdc/pkg/errors"
 	"github.com/pingcap/ticdc/pkg/util"
+	timodel "github.com/pingcap/tidb/pkg/meta/model"
+	"github.com/pingcap/tidb/pkg/parser/ast"
+	"github.com/pingcap/tidb/pkg/parser/mysql"
+	"github.com/pingcap/tidb/pkg/types"
 	"github.com/stretchr/testify/require"
 )
 
@@ -79,3 +83,45 @@ func TestVerifyRouteConflict(t *testing.T) {
 	require.Contains(t, err.Error(), "source `db1`.`orders`")
 	require.Contains(t, err.Error(), "source `db2`.`orders`")
 }
+
+func TestVerifyTablesForSinkValidatesStorageColumnSelectors(t *testing.T) {
+	t.Parallel()
+
+	replicaCfg := config.GetDefaultReplicaConfig()
+	replicaCfg.Sink.ColumnSelectors = []*config.ColumnSelector{
+		{Matcher: []string{"test.t"}, Columns: []string{"name"}},
+	}
+	tableInfos := []*common.TableInfo{newTableInfoWithPrimaryKeyForTest()}
+
+	err := verifyTablesForSink(replicaCfg, config.FileScheme, "", config.ProtocolCanalJSON, tableInfos)
+	require.Error(t, err)
+	require.True(t, errors.ErrColumnSelectorFailed.Equal(err))
+
+	replicaCfg.Sink.ColumnSelectors[0].Columns = []string{"id", "name"}
+	require.NoError(t, verifyTablesForSink(replicaCfg, config.FileScheme, "", config.ProtocolCanalJSON, tableInfos))
+}
+
+func newTableInfoWithPrimaryKeyForTest() *common.TableInfo {
+	idFieldType := types.NewFieldType(mysql.TypeLong)
+	idFieldType.AddFlag(mysql.PriKeyFlag | mysql.NotNullFlag)
+
+	return common.WrapTableInfo("test", &timodel.TableInfo{
+		ID:         1,
+		Name:       ast.NewCIStr("t"),
+		PKIsHandle: true,
+		Columns: []*timodel.ColumnInfo{
+			{
+				ID:        1,
+				Name:      ast.NewCIStr("id"),
+				FieldType: *idFieldType,
+				State:     timodel.StatePublic,
+			},
+			{
+				ID:        2,
+				Name:      ast.NewCIStr("name"),
+				FieldType: *types.NewFieldType(mysql.TypeVarchar),
+				State:     timodel.StatePublic,
+			},
+		},
+	})
+}

cmd/storage-consumer/consumer.go

@@ -25,6 +25,7 @@ import (
 	"github.com/pingcap/log"
 	"github.com/pingcap/ticdc/cmd/util"
 	"github.com/pingcap/ticdc/downstreamadapter/sink"
+	"github.com/pingcap/ticdc/downstreamadapter/sink/columnselector"
 	"github.com/pingcap/ticdc/downstreamadapter/sink/helper"
 	commonType "github.com/pingcap/ticdc/pkg/common"
 	"github.com/pingcap/ticdc/pkg/common/event"
@@ -60,7 +61,12 @@ type indexRange struct {
 type consumer struct {
 	replicationCfg  *config.ReplicaConfig
 	codecCfg        *common.Config
+<<<<<<< HEAD
 	externalStorage storage.ExternalStorage
+=======
+	columnSelectors *columnselector.ColumnSelectors
+	externalStorage storeapi.Storage
+>>>>>>> 07e944782 (sink: add column selector for storage sink (#5595))
 	fileExtension   string
 	sink            sink.Sink
 	// tableDMLIdxMap maintains a map of <dmlPathKey, fileIndexKeyMap>
@@ -122,6 +128,10 @@ func newConsumer(ctx context.Context) (*consumer, error) {
 	if err != nil {
 		return nil, err
 	}
+	columnSelectors, err := columnselector.New(replicaConfig.Sink)
+	if err != nil {
+		return nil, err
+	}
 
 	extension := helper.GetFileExtension(protocol)
 
@@ -147,6 +157,7 @@ func newConsumer(ctx context.Context) (*consumer, error) {
 	return &consumer{
 		replicationCfg:    replicaConfig,
 		codecCfg:          codecConfig,
+		columnSelectors:   columnSelectors,
 		externalStorage:   storage,
 		fileExtension:     extension,
 		sink:              sink,
@@ -297,7 +308,19 @@ func (c *consumer) appendDMLEvents(
 
 	switch c.codecCfg.Protocol {
 	case config.ProtocolCsv:
+<<<<<<< HEAD
 		decoder, err = csv.NewDecoder(ctx, c.codecCfg, tableInfo, content)
+=======
+		tableInfo := schemaFile.TableInfo()
+		// CSV rows contain selected values without column names, so decode with the same selector.
+		decoder, err = csv.NewDecoderWithColumnSelector(
+			ctx,
+			c.codecCfg,
+			tableInfo,
+			content,
+			c.columnSelectors.GetForTableInfo(tableInfo),
+		)
+>>>>>>> 07e944782 (sink: add column selector for storage sink (#5595))
 		if err != nil {
 			return errors.Trace(err)
 		}

downstreamadapter/sink/cloudstorage/buffer_manager.go

@@ -104,12 +104,14 @@ func (c *bufferManager) run(ctx context.Context) error {
 
 func (c *bufferManager) handleDMLTask(ctx context.Context, task *task) error {
 	if len(task.encodedMsgs) == 0 {
-		task.callbacks.postEnqueue()
+		if task.postEnqueue != nil {
+			task.postEnqueue()
+		}
 		return nil
 	}
 
 	for {
-		action, entry, err := c.spool.TryEnqueue(task.encodedMsgs, task.callbacks.postEnqueue)
+		action, entry, err := c.spool.TryEnqueue(task.encodedMsgs, task.postEnqueue)
 		if err != nil {
 			return err
 		}

downstreamadapter/sink/cloudstorage/buffer_manager_test.go

@@ -171,7 +171,7 @@ func newBufferedTask(table string, dispatcherID commonType.DispatcherID, payload
 		},
 		TableInfoVersion: 1,
 		DispatcherID:     dispatcherID,
-	}, event)
+	}, event, nil)
 	msg := common.NewMsg(nil, []byte(payload))
 	msg.SetRowsCount(1)
 	t.encodedMsgs = []*common.Message{msg}

downstreamadapter/sink/cloudstorage/dml_writers.go

@@ -18,6 +18,7 @@ import (
 	"time"
 
 	"github.com/pingcap/ticdc/downstreamadapter/sink/cloudstorage/spool"
+	"github.com/pingcap/ticdc/downstreamadapter/sink/columnselector"
 	sinkmetrics "github.com/pingcap/ticdc/downstreamadapter/sink/metrics"
 	commonType "github.com/pingcap/ticdc/pkg/common"
 	commonEvent "github.com/pingcap/ticdc/pkg/common/event"
@@ -43,8 +44,9 @@ type dmlWriters struct {
 	encodeGroup *encoderGroup
 	spool       *spool.Spool
 
-	writers []*writer
-	closed  atomic.Bool
+	columnSelector *columnselector.ColumnSelectors
+	writers        []*writer
+	closed         atomic.Bool
 }
 
 func newDMLWriters(
@@ -54,6 +56,7 @@ func newDMLWriters(
 	encoderConfig *common.Config,
 	extension string,
 	statistics *metrics.Statistics,
+	columnSelector *columnselector.ColumnSelectors,
 ) (*dmlWriters, error) {
 	messageCh := chann.NewUnlimitedChannelDefault[*task]()
 	encoderGroup := newEncoderGroup(
@@ -76,12 +79,13 @@ func newDMLWriters(
 	}
 
 	return &dmlWriters{
-		changefeedID: changefeedID,
-		statistics:   statistics,
-		msgCh:        messageCh,
-		encodeGroup:  encoderGroup,
-		spool:        spool,
-		writers:      writers,
+		changefeedID:   changefeedID,
+		statistics:     statistics,
+		msgCh:          messageCh,
+		encodeGroup:    encoderGroup,
+		spool:          spool,
+		columnSelector: columnSelector,
+		writers:        writers,
 	}, nil
 }
 
@@ -168,7 +172,7 @@ func (d *dmlWriters) addDMLEvent(event *commonEvent.DMLEvent) {
 		TableInfoVersion: event.TableInfoVersion,
 		DispatcherID:     event.GetDispatcherID(),
 	}
-	d.msgCh.Push(newDMLTask(table, event))
+	d.msgCh.Push(newDMLTask(table, event, d.columnSelector.GetForTableInfo(event.TableInfo)))
 }
 
 func (d *dmlWriters) flushDMLBeforeBlock(ctx context.Context, event commonEvent.BlockEvent) error {

downstreamadapter/sink/cloudstorage/encoder_group.go

@@ -108,13 +108,12 @@ func (eg *encoderGroup) runEncoder(ctx context.Context, index int) error {
 				continue
 			}
 
-			err = encoder.AppendTxnEvent(task.event)
+			err = encoder.AppendTxnEvent(task.rowEvents)
 			if err != nil {
 				return err
 			}
 			task.encodedMsgs = encoder.Build()
-			task.replacePostFlushCallbacks()
-			task.event = nil
+			task.rowEvents = nil
 			future.Done()
 		}
 	}

downstreamadapter/sink/cloudstorage/encoder_group_test.go

@@ -16,6 +16,7 @@ package cloudstorage
 import (
 	"context"
 	"net/url"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -174,6 +175,15 @@ func TestEncodingGroupEncodeDMLTask(t *testing.T) {
 	})
 
 	dispatcherID := commonType.NewDispatcherID()
+	var flushCount atomic.Int64
+	var enqueueCount atomic.Int64
+	event := newTestDMLEvent(dispatcherID, 100)
+	event.AddPostFlushFunc(func() {
+		flushCount.Add(1)
+	})
+	event.AddPostEnqueueFunc(func() {
+		enqueueCount.Add(1)
+	})
 	taskValue := newDMLTask(
 		cloudstorage.VersionedTableName{
 			TableNameWithPhysicTableID: commonType.TableName{
@@ -184,7 +194,8 @@ func TestEncodingGroupEncodeDMLTask(t *testing.T) {
 			TableInfoVersion: 1,
 			DispatcherID:     dispatcherID,
 		},
-		newTestDMLEvent(dispatcherID, 100),
+		event,
+		nil,
 	)
 	require.NoError(t, group.add(ctx, taskValue))
 
@@ -201,7 +212,15 @@ func TestEncodingGroupEncodeDMLTask(t *testing.T) {
 				}
 				task := future.task
 				require.Equal(t, taskValue, task)
-				require.Nil(t, task.event)
+				require.Nil(t, task.rowEvents)
+				require.Len(t, task.encodedMsgs, 1)
+				require.NotNil(t, task.encodedMsgs[0].Callback)
+				task.encodedMsgs[0].Callback()
+				require.Equal(t, int64(1), flushCount.Load())
+				require.Equal(t, int64(1), enqueueCount.Load())
+				require.NotNil(t, task.postEnqueue)
+				task.postEnqueue()
+				require.Equal(t, int64(1), enqueueCount.Load())
 				done <- struct{}{}
 				return nil
 			}
@@ -250,6 +269,8 @@ func newTestDMLEvent(dispatcherID commonType.DispatcherID, tableID int64) *commo
 		PhysicalTableID:  tableID,
 		TableInfo:        tableInfo,
 		TableInfoVersion: 1,
+		Length:           1,
+		RowTypes:         []commonType.RowType{commonType.RowTypeInsert},
 		Rows:             chunk.MutRowFromValues(1, "hello world").ToRow().Chunk(),
 	}
 }

downstreamadapter/sink/cloudstorage/sink.go

@@ -21,6 +21,7 @@ import (
 	"time"
 
 	"github.com/pingcap/log"
+	"github.com/pingcap/ticdc/downstreamadapter/sink/columnselector"
 	"github.com/pingcap/ticdc/downstreamadapter/sink/helper"
 	"github.com/pingcap/ticdc/pkg/common"
 	commonEvent "github.com/pingcap/ticdc/pkg/common/event"
@@ -87,7 +88,14 @@ func Verify(ctx context.Context, changefeedID common.ChangeFeedID, sinkURI *url.
 	if err != nil {
 		return err
 	}
+<<<<<<< HEAD
 	_, err = helper.GetEncoderConfig(changefeedID, sinkURI, protocol, sinkConfig, math.MaxInt)
+=======
+	if _, err = columnselector.New(sinkConfig); err != nil {
+		return err
+	}
+	_, err = helper.GetEncoderConfig(changefeedID, sinkURI, protocol, sinkConfig, math.MaxInt, math.MaxInt)
+>>>>>>> 07e944782 (sink: add column selector for storage sink (#5595))
 	if err != nil {
 		return err
 	}
@@ -123,6 +131,10 @@ func New(
 	if err != nil {
 		return nil, err
 	}
+	columnSelectors, err := columnselector.New(sinkConfig)
+	if err != nil {
+		return nil, err
+	}
 	storage, err := util.GetExternalStorageWithDefaultTimeout(ctx, sinkURI.String())
 	if err != nil {
 		return nil, err
@@ -134,7 +146,7 @@ func New(
 			storage.Close()
 		}
 	}()
-	dmlWriters, err := newDMLWriters(changefeedID, storage, cfg, encoderConfig, ext, statistics)
+	dmlWriters, err := newDMLWriters(changefeedID, storage, cfg, encoderConfig, ext, statistics, columnSelectors)
 	if err != nil {
 		return nil, err
 	}