Support ClickHouse having extra columns compared to source for initial load (#2794)

Currently to construct the column selectors for the Insert into Select
to ClickHouse in initial load, we get the columns from ClickHouse.
This results in initial load failing if there is an extra column on
target vs source.

The reason we were using the CH schema and not the source schema here
was so that column exclusion was automatically taken care of. But this
caused an unideal situation where we support source having extra columns
for initial load and ClickHouse having extra columns for CDC normalize.

This PR re-wires things to make initial load use the source schema and
gives it column exclusion info

Author: Amogh-Bharadwaj

flow/connectors/clickhouse/qrep.go

@@ -6,7 +6,6 @@ import (
 	"log/slog"
 	"strings"
 
-	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/s3"
 
@@ -33,26 +32,11 @@ func (c *ClickHouseConnector) SyncQRepRecords(
 		slog.String("destinationTable", destTable),
 	)
 
-	tblSchema, err := c.getTableSchema(ctx, destTable)
-	if err != nil {
-		return 0, fmt.Errorf("failed to get schema of table %s: %w", destTable, err)
-	}
-	c.logger.Info("Called QRep sync function and obtained table schema", flowLog)
+	c.logger.Info("Called QRep sync function", flowLog)
 
 	avroSync := NewClickHouseAvroSyncMethod(config, c)
 
-	return avroSync.SyncQRepRecords(ctx, config, partition, tblSchema, stream)
-}
-
-func (c *ClickHouseConnector) getTableSchema(ctx context.Context, tableName string) ([]driver.ColumnType, error) {
-	queryString := fmt.Sprintf("SELECT * FROM `%s` LIMIT 0", tableName)
-	rows, err := c.query(ctx, queryString)
-	if err != nil {
-		return nil, fmt.Errorf("failed to execute query: %w", err)
-	}
-	defer rows.Close()
-
-	return rows.ColumnTypes(), nil
+	return avroSync.SyncQRepRecords(ctx, config, partition, stream)
 }
 
 func (c *ClickHouseConnector) ConsolidateQRepPartitions(_ context.Context, config *protos.QRepConfig) error {

flow/connectors/clickhouse/qrep_avro_sync.go

@@ -7,8 +7,6 @@ import (
 	"strings"
 	"time"
 
-	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
-
 	"github.com/PeerDB-io/peerdb/flow/connectors/utils"
 	avro "github.com/PeerDB-io/peerdb/flow/connectors/utils/avro"
 	"github.com/PeerDB-io/peerdb/flow/generated/protos"
@@ -105,7 +103,6 @@ func (s *ClickHouseAvroSyncMethod) SyncQRepRecords(
 	ctx context.Context,
 	config *protos.QRepConfig,
 	partition *protos.QRepPartition,
-	dstTableSchema []driver.ColumnType,
 	stream *model.QRecordStream,
 ) (int, error) {
 	dstTableName := config.DestinationTableIdentifier
@@ -146,17 +143,14 @@ func (s *ClickHouseAvroSyncMethod) SyncQRepRecords(
 	endpoint := s.credsProvider.Provider.GetEndpointURL()
 	region := s.credsProvider.Provider.GetRegion()
 	avroFileUrl := utils.FileURLForS3Service(endpoint, region, s3o.Bucket, avroFile.FilePath)
-	selectedColumnNames := make([]string, 0, len(dstTableSchema))
-	insertedColumnNames := make([]string, 0, len(dstTableSchema))
-	for _, col := range dstTableSchema {
-		colName := col.Name()
-		if strings.EqualFold(colName, signColName) ||
-			strings.EqualFold(colName, config.SyncedAtColName) ||
-			strings.EqualFold(colName, versionColName) ||
-			(sourceSchemaAsDestinationColumn && strings.EqualFold(colName, sourceSchemaColName)) {
-			continue
+	selectedColumnNames := make([]string, 0, len(schema.Fields))
+	insertedColumnNames := make([]string, 0, len(schema.Fields))
+	for _, colName := range schema.GetColumnNames() {
+		for _, excludedColumn := range config.Exclude {
+			if colName == excludedColumn {
+				continue
+			}
 		}
-
 		avroColName, ok := columnNameAvroFieldMap[colName]
 		if !ok {
 			s.logger.Error("destination column not found in avro schema",
@@ -184,7 +178,7 @@ func (s *ClickHouseAvroSyncMethod) SyncQRepRecords(
 		sessionTokenPart = fmt.Sprintf(", '%s'", creds.AWS.SessionToken)
 	}
 
-	hashColName := columnNameAvroFieldMap[dstTableSchema[0].Name()]
+	hashColName := columnNameAvroFieldMap[schema.Fields[0].Name]
 	numParts, err := internal.PeerDBClickHouseInitialLoadPartsPerPartition(ctx, s.config.Env)
 	if err != nil {
 		s.logger.Warn("failed to get chunking parts, proceeding without chunking", slog.Any("error", err))

flow/e2e/clickhouse/clickhouse.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"reflect"
+	"strings"
 	"testing"
 	"time"
 
@@ -99,6 +100,32 @@ func (s ClickHouseSuite) Teardown(ctx context.Context) {
 	require.NoError(s.t, s.connector.Close())
 }
 
+type TestClickHouseColumn struct {
+	Name string
+	Type string
+}
+
+// CreateRMTTable creates a ReplacingMergeTree table with the given name and columns.
+func (s ClickHouseSuite) CreateRMTTable(tableName string, columns []TestClickHouseColumn, orderingKey string) error {
+	ch, err := connclickhouse.Connect(s.t.Context(), nil, s.Peer().GetClickhouseConfig())
+	if err != nil {
+		return err
+	}
+	defer ch.Close()
+
+	columnStrings := make([]string, len(columns))
+	for i, col := range columns {
+		columnStrings[i] = fmt.Sprintf("`%s` %s", col.Name, col.Type)
+	}
+
+	// Join the column definitions into a single string
+	columnStr := strings.Join(columnStrings, ", ")
+
+	// Create the table with ReplacingMergeTree engine
+	createTableQuery := fmt.Sprintf("CREATE TABLE `%s` (%s) ENGINE = ReplacingMergeTree() ORDER BY `%s`", tableName, columnStr, orderingKey)
+	return ch.Exec(s.t.Context(), createTableQuery)
+}
+
 func (s ClickHouseSuite) GetRows(table string, cols string) (*model.QRecordBatch, error) {
 	ch, err := connclickhouse.Connect(s.t.Context(), nil, s.Peer().GetClickhouseConfig())
 	if err != nil {

flow/e2e/clickhouse/peer_flow_ch_test.go

@@ -1250,6 +1250,47 @@ func (s ClickHouseSuite) Test_SchemaAsColumn() {
 	e2e.RequireEnvCanceled(s.t, env)
 }
 
+func (s ClickHouseSuite) Test_Extra_CH_Columns() {
+	srcTableName := "test_extra_ch_cols"
+	srcFullName := s.attachSchemaSuffix("test_extra_ch_cols")
+	dstTableName := "test_extra_ch_cols"
+
+	require.NoError(s.t, s.source.Exec(s.t.Context(),
+		fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s (id INT PRIMARY KEY, "key" TEXT NOT NULL)`, srcFullName)))
+
+	require.NoError(s.t, s.CreateRMTTable(dstTableName, []TestClickHouseColumn{
+		{Name: "id", Type: "Int32"},
+		{Name: "key", Type: "String"},
+		{Name: "updatedAt", Type: "String"},
+		{Name: "_peerdb_is_deleted", Type: "Int8"},
+		{Name: "_peerdb_synced_at", Type: "DateTime"},
+		{Name: "_peerdb_version", Type: "Int64"},
+	}, "id"),
+	)
+
+	require.NoError(s.t, s.source.Exec(s.t.Context(), fmt.Sprintf(`INSERT INTO %s (id,"key") VALUES (1,'init')`, srcFullName)))
+
+	connectionGen := e2e.FlowConnectionGenerationConfig{
+		FlowJobName:      s.attachSuffix("test_extra_ch_cols"),
+		TableNameMapping: map[string]string{srcFullName: dstTableName},
+		Destination:      s.Peer().Name,
+	}
+	flowConnConfig := connectionGen.GenerateFlowConnectionConfigs(s)
+	flowConnConfig.DoInitialSnapshot = true
+
+	tc := e2e.NewTemporalClient(s.t)
+	env := e2e.ExecutePeerflow(s.t.Context(), tc, peerflow.CDCFlowWorkflow, flowConnConfig, nil)
+	e2e.SetupCDCFlowStatusQuery(s.t, env, flowConnConfig)
+
+	e2e.EnvWaitForEqualTablesWithNames(env, s, "waiting on initial", srcTableName, dstTableName, "id,\"key\"")
+
+	require.NoError(s.t, s.source.Exec(s.t.Context(), fmt.Sprintf(`INSERT INTO %s (id,"key") VALUES (2,'cdc')`, srcFullName)))
+	e2e.EnvWaitForEqualTablesWithNames(env, s, "waiting on cdc", srcTableName, dstTableName, "id,\"key\"")
+
+	env.Cancel(s.t.Context())
+	e2e.RequireEnvCanceled(s.t, env)
+}
+
 func (s ClickHouseSuite) Test_NullEngine() {
 	srcTableName := "test_nullengine"
 	srcFullName := s.attachSchemaSuffix(srcTableName)

flow/workflows/snapshot_flow.go

@@ -216,6 +216,7 @@ func (s *SnapshotFlowExecution) cloneTable(
 		Script:                     s.config.Script,
 		Env:                        s.config.Env,
 		ParentMirrorName:           flowName,
+		Exclude:                    mapping.Exclude,
 	}
 
 	boundSelector.SpawnChild(childCtx, QRepFlowWorkflow, nil, config, nil)

protos/flow.proto

@@ -317,6 +317,7 @@ message QRepConfig {
   map<string, string> env = 24;
 
   string parent_mirror_name = 25;
+  repeated string exclude = 26;
 }
 
 message QRepPartition {

ui/app/mirrors/create/helpers/common.ts

@@ -69,4 +69,5 @@ export const blankQRepSetting: QRepConfig = {
   system: TypeSystem.Q,
   env: {},
   parentMirrorName: '',
+  exclude: [],
 };