WIP message per mod

Author: mmatczuk

internal/impl/gcp/enterprise/changestreams/callback.go

@@ -10,6 +10,7 @@ package changestreams
 
 import (
 	"context"
+	"time"
 )
 
 // CallbackFunc is a function that is called for each change record.
@@ -24,7 +25,18 @@ import (
 type CallbackFunc func(ctx context.Context, partitionToken string, dcr *DataChangeRecord) error
 
 // UpdatePartitionWatermark updates the watermark for a partition. It's intended
-// for use by Callback function to update progress.
-func (s *Subscriber) UpdatePartitionWatermark(ctx context.Context, partitionToken string, dcr *DataChangeRecord) error {
-	return s.store.UpdateWatermark(ctx, partitionToken, dcr.CommitTimestamp)
+// for use by Callback function to update progress. If commitTimestamp is zero
+// value, the watermark is not updated.
+func (s *Subscriber) UpdatePartitionWatermark(
+	ctx context.Context,
+	partitionToken string,
+	commitTimestamp time.Time,
+) error {
+	if commitTimestamp.IsZero() {
+		return nil
+	}
+
+	s.log.Debugf("%s: updating watermark to %s", partitionToken, commitTimestamp)
+
+	return s.store.UpdateWatermark(ctx, partitionToken, commitTimestamp)
 }

internal/impl/gcp/enterprise/changestreams/subscriber_integration_test.go

@@ -362,7 +362,7 @@ func TestIntegrationSubscriberCallbackUpdatePartitionWatermark(t *testing.T) {
 			assert.Equal(t, testStartTimestamp, pm.Watermark)
 
 			// When UpdatePartitionWatermark is called
-			if err := s.UpdatePartitionWatermark(ctx, partitionToken, dcr); err != nil {
+			if err := s.UpdatePartitionWatermark(ctx, partitionToken, dcr.CommitTimestamp); err != nil {
 				require.NoError(t, err)
 			}
 		case 3:

internal/impl/gcp/enterprise/input_spanner_cdc.go

@@ -220,16 +220,16 @@ func newSpannerCDCReader(conf spannerCDCInputConfig, batching service.BatchPolic
 func (r *spannerCDCReader) emit(
 	ctx context.Context,
 	partitionToken string,
-	dcr *changestreams.DataChangeRecord,
 	msg service.MessageBatch,
+	commitTimestamp time.Time,
 ) (*ack.Once, error) {
 	if len(msg) == 0 {
 		return nil, nil
 	}
 	ackOnce := ack.NewOnce(func(ctx context.Context) error {
 		// If we processed the message and failed to update the watermark, we
 		// would try to update it on the next message, no need to return an error here.
-		if err := r.subscriber.UpdatePartitionWatermark(ctx, partitionToken, dcr); err != nil {
+		if err := r.subscriber.UpdatePartitionWatermark(ctx, partitionToken, commitTimestamp); err != nil {
 			r.log.Errorf("%s: failed to update watermark: %v", partitionToken, err)
 		}
 		return nil
@@ -259,11 +259,11 @@ func (r *spannerCDCReader) onDataChangeRecord(ctx context.Context, partitionToke
 	// On partition end, flush the remaining messages and wait for all messages
 	// to be acked before returning and marking the partition as finished.
 	if dcr == nil {
-		msg, last, err := batcher.Flush(ctx)
+		msg, ts, err := batcher.Flush(ctx)
 		if err != nil {
 			return err
 		}
-		ack, err := r.emit(ctx, partitionToken, last, msg)
+		ack, err := r.emit(ctx, partitionToken, msg, ts)
 		if err != nil {
 			return err
 		}
@@ -280,11 +280,11 @@ func (r *spannerCDCReader) onDataChangeRecord(ctx context.Context, partitionToke
 	}
 
 	if dcr == forcePeriodicFlush {
-		msg, last, err := batcher.Flush(ctx)
+		msg, ts, err := batcher.Flush(ctx)
 		if err != nil {
 			return err
 		}
-		ack, err := r.emit(ctx, partitionToken, last, msg)
+		ack, err := r.emit(ctx, partitionToken, msg, ts)
 		if err != nil {
 			return err
 		}
@@ -293,15 +293,17 @@ func (r *spannerCDCReader) onDataChangeRecord(ctx context.Context, partitionToke
 		return nil
 	}
 
-	msg, err := batcher.MaybeFlushWith(ctx, dcr)
-	if err != nil {
-		return err
+	iter := batcher.MaybeFlushWith(dcr)
+	for mb, ts := range iter.Iter(ctx) {
+		ack, err := r.emit(ctx, partitionToken, mb, ts)
+		if err != nil {
+			return err
+		}
+		batcher.AddAck(ack)
 	}
-	ack, err := r.emit(ctx, partitionToken, dcr, msg)
-	if err != nil {
+	if err := iter.Err(); err != nil {
 		return err
 	}
-	batcher.AddAck(ack)
 
 	return nil
 }

internal/impl/gcp/enterprise/input_spanner_partition_batcher.go

@@ -12,7 +12,8 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
-	"fmt"
+	"iter"
+	"strconv"
 	"sync"
 	"time"
 
@@ -21,6 +22,93 @@ import (
 	"github.com/redpanda-data/connect/v4/internal/impl/gcp/enterprise/changestreams"
 )
 
+// spannerPartitionBatchIter goes over changestreams.DataChangeRecord.Mods,
+// for every mod it creates a message and adds it to the batch, if the batch is
+// full, it yields the batch and creates a new one.
+//
+// Iff batch is returned with nonzero time, when acked the partition watermark
+// should be updated to this time.
+type spannerPartitionBatchIter struct {
+	*spannerPartitionBatcher
+	dcr *changestreams.DataChangeRecord
+	err error
+}
+
+type spannerMod struct {
+	TableName   string
+	ColumnTypes []*changestreams.ColumnType
+	Mod         *changestreams.Mod
+	ModType     string
+}
+
+func (s *spannerPartitionBatchIter) Iter(ctx context.Context) iter.Seq2[service.MessageBatch, time.Time] {
+	return func(yield func(service.MessageBatch, time.Time) bool) {
+		defer func() {
+			s.last = s.dcr
+		}()
+		if s.err != nil {
+			return
+		}
+
+		first := true
+		for i, m := range s.dcr.Mods {
+			modData := spannerMod{
+				TableName:   s.dcr.TableName,
+				ColumnTypes: s.dcr.ColumnTypes,
+				Mod:         m,
+				ModType:     s.dcr.ModType,
+			}
+
+			b, err := json.Marshal(modData)
+			if err != nil {
+				s.err = err
+				return
+			}
+
+			msg := service.NewMessage(b)
+			msg.MetaSet("commit_timestamp", s.dcr.CommitTimestamp.Format(time.RFC3339Nano))
+			msg.MetaSet("record_sequence", s.dcr.RecordSequence)
+			msg.MetaSet("server_transaction_id", s.dcr.ServerTransactionID)
+			msg.MetaSet("is_last_record_in_transaction_in_partition", strconv.FormatBool(s.dcr.IsLastRecordInTransactionInPartition))
+			msg.MetaSet("value_capture_type", s.dcr.ValueCaptureType)
+			msg.MetaSet("number_of_records_in_transaction", strconv.FormatInt(s.dcr.NumberOfRecordsInTransaction, 10))
+			msg.MetaSet("number_of_partitions_in_transaction", strconv.FormatInt(s.dcr.NumberOfPartitionsInTransaction, 10))
+			msg.MetaSet("transaction_tag", s.dcr.TransactionTag)
+			msg.MetaSet("is_system_transaction", strconv.FormatBool(s.dcr.IsSystemTransaction))
+
+			if !s.batcher.Add(msg) {
+				continue
+			}
+
+			mb, err := s.flush(ctx)
+			if err != nil {
+				s.err = err
+				return
+			}
+
+			// Return the watermark to be updated after processing the batch.
+			// Not every batch should update the watermark, we update watermark
+			// only after processing the whole DataChangeRecord.
+			var watermark time.Time
+			if first && s.last != nil {
+				watermark = s.last.CommitTimestamp
+				first = false
+			}
+			if i == len(s.dcr.Mods)-1 {
+				watermark = s.dcr.CommitTimestamp
+			}
+			if !yield(mb, watermark) {
+				return
+			}
+		}
+	}
+}
+
+// Err returns any error that occurred during iteration.
+func (s *spannerPartitionBatchIter) Err() error {
+	return s.err
+}
+
 type spannerPartitionBatcher struct {
 	batcher *service.Batcher
 	last    *changestreams.DataChangeRecord
@@ -29,29 +117,17 @@ type spannerPartitionBatcher struct {
 	rm      func()
 }
 
-func (s *spannerPartitionBatcher) MaybeFlushWith(ctx context.Context, dcr *changestreams.DataChangeRecord) (service.MessageBatch, error) {
-	b, err := json.Marshal(dcr)
-	if err != nil {
-		return nil, fmt.Errorf("marshal data change record as JSON: %w", err)
-	}
-
-	s.last = dcr
-
-	if !s.batcher.Add(service.NewMessage(b)) {
-		return nil, nil
-	}
-
-	return s.flush(ctx)
+func (s *spannerPartitionBatcher) MaybeFlushWith(dcr *changestreams.DataChangeRecord) *spannerPartitionBatchIter {
+	return &spannerPartitionBatchIter{spannerPartitionBatcher: s, dcr: dcr}
 }
 
-func (s *spannerPartitionBatcher) Flush(ctx context.Context) (service.MessageBatch, *changestreams.DataChangeRecord, error) {
+func (s *spannerPartitionBatcher) Flush(ctx context.Context) (service.MessageBatch, time.Time, error) {
 	if s.last == nil {
-		return nil, nil, nil
+		return nil, time.Time{}, nil
 	}
 
-	last := s.last
 	msg, err := s.flush(ctx)
-	return msg, last, err
+	return msg, s.last.CommitTimestamp, err
 }
 
 func (s *spannerPartitionBatcher) flush(ctx context.Context) (service.MessageBatch, error) {

internal/impl/gcp/enterprise/integration_spanner_cdc_test.go

@@ -97,11 +97,9 @@ gcp_spanner_cdc:
 	})
 	require.NoError(t, err)
 
-	expected := []changestreams.DataChangeRecord{
+	expected := []spannerMod{
 		{
-			RecordSequence:                       "00000000",
-			IsLastRecordInTransactionInPartition: false,
-			TableName:                            h.Table(),
+			TableName: h.Table(),
 			ColumnTypes: []*changestreams.ColumnType{
 				{
 					Name: "id",
@@ -122,33 +120,24 @@ gcp_spanner_cdc:
 					OrdinalPosition: 2,
 				},
 			},
-			Mods: []*changestreams.Mod{
-				{
-					Keys: spanner.NullJSON{
-						Value: map[string]interface{}{"id": "1"},
-						Valid: true,
-					},
-					NewValues: spanner.NullJSON{
-						Value: map[string]interface{}{"active": true},
-						Valid: true,
-					},
-					OldValues: spanner.NullJSON{
-						Value: map[string]interface{}{},
-						Valid: true,
-					},
+			Mod: &changestreams.Mod{
+				Keys: spanner.NullJSON{
+					Value: map[string]interface{}{"id": "1"},
+					Valid: true,
+				},
+				NewValues: spanner.NullJSON{
+					Value: map[string]interface{}{"active": true},
+					Valid: true,
+				},
+				OldValues: spanner.NullJSON{
+					Value: map[string]interface{}{},
+					Valid: true,
 				},
 			},
-			ModType:                         "INSERT",
-			ValueCaptureType:                "OLD_AND_NEW_VALUES",
-			NumberOfRecordsInTransaction:    2,
-			NumberOfPartitionsInTransaction: 1,
-			TransactionTag:                  "",
-			IsSystemTransaction:             false,
+			ModType: "INSERT",
 		},
 		{
-			RecordSequence:                       "00000001",
-			IsLastRecordInTransactionInPartition: true,
-			TableName:                            h.Table(),
+			TableName: h.Table(),
 			ColumnTypes: []*changestreams.ColumnType{
 				{
 					Name: "id",
@@ -169,40 +158,32 @@ gcp_spanner_cdc:
 					OrdinalPosition: 2,
 				},
 			},
-			Mods: []*changestreams.Mod{
-				{
-					Keys: spanner.NullJSON{
-						Value: map[string]interface{}{"id": "1"},
-						Valid: true,
-					},
-					NewValues: spanner.NullJSON{
-						Value: map[string]interface{}{},
-						Valid: true,
-					},
-					OldValues: spanner.NullJSON{
-						Value: map[string]interface{}{"active": true},
-						Valid: true,
-					},
+			Mod: &changestreams.Mod{
+				Keys: spanner.NullJSON{
+					Value: map[string]interface{}{"id": "1"},
+					Valid: true,
+				},
+				NewValues: spanner.NullJSON{
+					Value: map[string]interface{}{},
+					Valid: true,
+				},
+				OldValues: spanner.NullJSON{
+					Value: map[string]interface{}{"active": true},
+					Valid: true,
 				},
 			},
-			ModType:                         "DELETE",
-			ValueCaptureType:                "OLD_AND_NEW_VALUES",
-			NumberOfRecordsInTransaction:    2,
-			NumberOfPartitionsInTransaction: 1,
-			TransactionTag:                  "",
-			IsSystemTransaction:             false,
+			ModType: "DELETE",
 		},
 	}
 
-	var got []changestreams.DataChangeRecord
+	var got []spannerMod
 	for msg := range messages {
 		b, err := msg.AsBytes()
 		require.NoError(t, err)
 
-		var dcr changestreams.DataChangeRecord
-		require.NoError(t, json.Unmarshal(b, &dcr))
-
-		got = append(got, dcr)
+		var mod spannerMod
+		require.NoError(t, json.Unmarshal(b, &mod))
+		got = append(got, mod)
 		if len(got) == len(expected) {
 			break
 		}