Merge #149696

149696:  changefeedccl: protobuf support envelope=wrapped with kafka r=asg0451 a=elizaMkraule

This change adds support for wrapped envelopes for changefeeds
in protobuf format.

Release note (general change): Changefeeds using the protobuf format
now support wrapped envelopes in kafka sinks.

Fixes #148936

Co-authored-by: Eliza Kraule <[email protected]>

Author: craig[bot]

pkg/ccl/changefeedccl/changefeed_test.go

@@ -12105,37 +12105,79 @@ func TestDatabaseLevelChangefeed(t *testing.T) {
 	cdcTest(t, testFn)
 }
 
-func TestChangefeedBareFullProtobuf(t *testing.T) {
+func TestChangefeedProtobuf(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
 
-	testFn := func(t *testing.T, s TestServer, f cdctest.TestFeedFactory) {
-		sqlDB := sqlutils.MakeSQLRunner(s.DB)
+	type testCase struct {
+		envelope     string
+		withDiff     bool
+		expectedRows []string
+	}
 
-		sqlDB.Exec(t, `
-			CREATE TABLE pricing (
-				id INT PRIMARY KEY,
-				name STRING,
-				discount FLOAT,
-				tax DECIMAL,
-				options STRING[]
-			)`)
-		sqlDB.Exec(t, `
-			INSERT INTO pricing VALUES
-				(1, 'Chair', 15.75, 2.500, ARRAY['Brown', 'Black']), 
-				(2, 'Table', 20.00, 1.23456789, ARRAY['Brown', 'Black'])
-		`)
-		pricingFeed := feed(t, f,
-			`CREATE CHANGEFEED FOR pricing WITH envelope='bare', format='protobuf', key_in_value, topic_in_value`)
-		defer closeFeed(t, pricingFeed)
+	tests := []testCase{
+		{
+			envelope: "bare",
+			withDiff: false,
+			expectedRows: []string{
+				`pricing: {"id":1}->{"values":{"discount":15.75,"id":1,"name":"Chair","options":["Brown","Black"],"tax":"2.500"},"__crdb__":{"key":{"id":1},"topic":"pricing"}}`,
+				`pricing: {"id":2}->{"values":{"discount":20,"id":2,"name":"Table","options":["Brown","Black"],"tax":"1.23456789"},"__crdb__":{"key":{"id":2},"topic":"pricing"}}`,
+				`pricing: {"id":2}->{"values":{"discount":25.5,"id":2,"name":"Table","options":["Brown","Black"],"tax":"1.23456789"},"__crdb__":{"key":{"id":2},"topic":"pricing"}}`,
+				`pricing: {"id":1}->{"values":{"discount":10,"id":1,"name":"Armchair","options":["Red"],"tax":"1.000"},"__crdb__":{"key":{"id":1},"topic":"pricing"}}`,
+				`pricing: {"id":3}->{"values":{"discount":50,"id":3,"name":"Sofa","options":["Gray"],"tax":"4.250"},"__crdb__":{"key":{"id":3},"topic":"pricing"}}`,
+				`pricing: {"id":2}->{"values":{"discount":null,"id":2,"name":null,"options":null,"tax":null},"__crdb__":{"key":{"id":2},"topic":"pricing"}}`,
+			},
+		},
+		{
+			envelope: "wrapped",
+			withDiff: true,
+			expectedRows: []string{
+				`pricing: {"id":1}->{"after":{"values":{"discount":15.75,"id":1,"name":"Chair","options":["Brown","Black"],"tax":"2.500"}},"before":{},"key":{"id":1},"topic":"pricing"}`,
+				`pricing: {"id":2}->{"after":{"values":{"discount":20,"id":2,"name":"Table","options":["Brown","Black"],"tax":"1.23456789"}},"before":{},"key":{"id":2},"topic":"pricing"}`,
+				`pricing: {"id":2}->{"after":{"values":{"discount":25.5,"id":2,"name":"Table","options":["Brown","Black"],"tax":"1.23456789"}},"before":{"values":{"discount":20,"id":2,"name":"Table","options":["Brown","Black"],"tax":"1.23456789"}},"key":{"id":2},"topic":"pricing"}`,
+				`pricing: {"id":1}->{"after":{"values":{"discount":10,"id":1,"name":"Armchair","options":["Red"],"tax":"1.000"}},"before":{"values":{"discount":15.75,"id":1,"name":"Chair","options":["Brown","Black"],"tax":"2.500"}},"key":{"id":1},"topic":"pricing"}`,
+				`pricing: {"id":3}->{"after":{"values":{"discount":50,"id":3,"name":"Sofa","options":["Gray"],"tax":"4.250"}},"before":{},"key":{"id":3},"topic":"pricing"}`,
+				`pricing: {"id":2}->{"after":{},"before":{"values":{"discount":25.5,"id":2,"name":"Table","options":["Brown","Black"],"tax":"1.23456789"}},"key":{"id":2},"topic":"pricing"}`,
+			},
+		},
+	}
 
-		expected := []string{
-			`pricing: {"id":1}->{"values":{"discount":15.75,"id":1,"name":"Chair","options":["Brown","Black"],"tax":"2.500"},"__crdb__":{"key":{"id":1},"topic":"pricing"}}`,
-			`pricing: {"id":2}->{"values":{"discount":20,"id":2,"name":"Table","options":["Brown","Black"],"tax":"1.23456789"},"__crdb__":{"key":{"id":2},"topic":"pricing"}}`,
-		}
+	for _, tc := range tests {
+		t.Run(fmt.Sprintf("envelope=%s", tc.envelope), func(t *testing.T) {
+			testFn := func(t *testing.T, s TestServer, f cdctest.TestFeedFactory) {
+				sqlDB := sqlutils.MakeSQLRunner(s.DB)
 
-		assertPayloads(t, pricingFeed, expected)
-	}
+				sqlDB.Exec(t, `
+					CREATE TABLE pricing (
+						id INT PRIMARY KEY,
+						name STRING,
+						discount FLOAT,
+						tax DECIMAL,
+						options STRING[]
+					)`)
+				sqlDB.Exec(t, `
+					INSERT INTO pricing VALUES
+						(1, 'Chair', 15.75, 2.500, ARRAY['Brown', 'Black']), 
+						(2, 'Table', 20.00, 1.23456789, ARRAY['Brown', 'Black'])`)
+
+				var opts []string
+				opts = append(opts, fmt.Sprintf("envelope='%s'", tc.envelope))
+				opts = append(opts, "format='protobuf'", "key_in_value", "topic_in_value")
+				if tc.withDiff {
+					opts = append(opts, "diff")
+				}
 
-	cdcTest(t, testFn, feedTestForceSink("kafka"))
+				feed := feed(t, f, fmt.Sprintf("CREATE CHANGEFEED FOR pricing WITH %s", strings.Join(opts, ", ")))
+				defer closeFeed(t, feed)
+
+				sqlDB.Exec(t, `UPDATE pricing SET discount = 25.50 WHERE id = 2`)
+				sqlDB.Exec(t, `UPSERT INTO pricing (id, name, discount, tax, options) VALUES (1, 'Armchair', 10.00, 1.000, ARRAY['Red'])`)
+				sqlDB.Exec(t, `INSERT INTO pricing VALUES (3, 'Sofa', 50.00, 4.250, ARRAY['Gray'])`)
+				sqlDB.Exec(t, `DELETE FROM pricing WHERE id = 2`)
+
+				assertPayloads(t, feed, tc.expectedRows)
+			}
+			cdcTest(t, testFn, feedTestForceSink("kafka"))
+		})
+	}
 }

pkg/ccl/changefeedccl/encoder_protobuf.go

@@ -69,6 +69,8 @@ func (e *protobufEncoder) EncodeValue(
 	switch e.envelopeType {
 	case changefeedbase.OptEnvelopeBare:
 		return e.buildBare(evCtx, updatedRow, prevRow)
+	case changefeedbase.OptEnvelopeWrapped:
+		return e.buildWrapped(ctx, evCtx, updatedRow, prevRow)
 	default:
 		return nil, errors.AssertionFailedf("envelope format not supported: %s", e.envelopeType)
 	}
@@ -126,6 +128,71 @@ func (e *protobufEncoder) buildBare(
 	return protoutil.Marshal(env)
 }
 
+// buildWrapped constructs a WrappedEnvelope serializes it.
+func (e *protobufEncoder) buildWrapped(
+	ctx context.Context, evCtx eventContext, updatedRow, prevRow cdcevent.Row,
+) ([]byte, error) {
+
+	var after *changefeedpb.Record
+	var err error
+	if !updatedRow.IsDeleted() {
+		after, err = encodeRowToRecord(updatedRow)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		after = &changefeedpb.Record{}
+	}
+
+	var before *changefeedpb.Record
+	if e.beforeField {
+		if prevRow.IsInitialized() && !prevRow.IsDeleted() {
+			before, err = encodeRowToRecord(prevRow)
+			if err != nil {
+				return nil, err
+			}
+		} else {
+			before = &changefeedpb.Record{}
+		}
+	}
+	var keyMsg *changefeedpb.Key
+	if e.keyInValue {
+		keyMsg, err = buildKeyMessage(updatedRow)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	var topicStr string
+	if e.topicInValue {
+		topicStr = evCtx.topic
+	}
+
+	var updatedStr string
+	if e.updatedField {
+		updatedStr = evCtx.updated.AsOfSystemTime()
+	}
+
+	var mvccStr string
+	if e.mvccTimestampField {
+		mvccStr = evCtx.mvcc.AsOfSystemTime()
+	}
+
+	wrapped := &changefeedpb.WrappedEnvelope{
+		After:         after,
+		Before:        before,
+		Key:           keyMsg,
+		Topic:         topicStr,
+		Updated:       updatedStr,
+		MvccTimestamp: mvccStr,
+	}
+
+	env := &changefeedpb.Message{
+		Data: &changefeedpb.Message_Wrapped{Wrapped: wrapped},
+	}
+	return protoutil.Marshal(env)
+}
+
 // buildMetadata returns metadata to include in the BareEnvelope.
 func (e *protobufEncoder) buildMetadata(
 	evCtx eventContext, row cdcevent.Row,
@@ -155,6 +222,9 @@ func (e *protobufEncoder) buildMetadata(
 
 // encodeRowToRecord converts a Row into a Record proto.
 func encodeRowToRecord(row cdcevent.Row) (*changefeedpb.Record, error) {
+	if !row.HasValues() {
+		return nil, nil
+	}
 	record := &changefeedpb.Record{Values: make(map[string]*changefeedpb.Value, row.NumValueColumns())}
 	if err := row.ForEachColumn().Datum(func(d tree.Datum, col cdcevent.ResultColumn) error {
 		val, err := datumToProtoValue(d)

pkg/ccl/changefeedccl/helpers_test.go

@@ -391,11 +391,22 @@ func assertPayloadsBaseErr(
 			if err := protoutil.Unmarshal(m.Value, &msg); err != nil {
 				return err
 			}
-			m.Value, err = gojson.Marshal(msg.GetBare())
-			if err != nil {
-				return err
-			}
 
+			switch env := msg.GetData().(type) {
+			case *changefeedpb.Message_Bare:
+				m.Value, err = gojson.Marshal(env.Bare)
+				if err != nil {
+					return err
+				}
+			case *changefeedpb.Message_Wrapped:
+				m.Value, err = gojson.Marshal(env.Wrapped)
+				if err != nil {
+					return err
+				}
+
+			default:
+				return errors.Newf("unexpected message type: %T", env)
+			}
 			var key changefeedpb.Key
 			if err := protoutil.Unmarshal(m.Key, &key); err != nil {
 				return err