Merge #150501

150501: changefeedccl: add protobuf encoder support with envelope=enriched r=asg0451 a=elizaMkraule

This change adds support for enriched envelopes when using protobuf format

Fixes: #150499

Release note (general change): The protobuf format for changefeeds
now support enriched envelopes.

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

Author: craig[bot]

pkg/ccl/changefeedccl/changefeed_test.go

@@ -12218,13 +12218,13 @@ func TestChangefeedProtobuf(t *testing.T) {
 	type testCase struct {
 		envelope     string
 		withDiff     bool
+		withSource   bool
 		expectedRows []string
 	}
 
 	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"}}`,
@@ -12238,12 +12238,48 @@ func TestChangefeedProtobuf(t *testing.T) {
 			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":1}->{"after":{"values":{"discount":15.75,"id":1,"name":"Chair","options":["Brown","Black"],"tax":"2.500"}},"key":{"id":1},"topic":"pricing"}`,
+				`pricing: {"id":2}->{"after":{"values":{"discount":20,"id":2,"name":"Table","options":["Brown","Black"],"tax":"1.23456789"}},"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"}`,
+				`pricing: {"id":3}->{"after":{"values":{"discount":50,"id":3,"name":"Sofa","options":["Gray"],"tax":"4.250"}},"key":{"id":3},"topic":"pricing"}`,
+				`pricing: {"id":2}->{"before":{"values":{"discount":25.5,"id":2,"name":"Table","options":["Brown","Black"],"tax":"1.23456789"}},"key":{"id":2},"topic":"pricing"}`,
+			},
+		},
+		{
+			envelope: "wrapped",
+			expectedRows: []string{
+				`pricing: {"id":1}->{"after":{"values":{"discount":15.75,"id":1,"name":"Chair","options":["Brown","Black"],"tax":"2.500"}},"key":{"id":1},"topic":"pricing"}`,
+				`pricing: {"id":2}->{"after":{"values":{"discount":20,"id":2,"name":"Table","options":["Brown","Black"],"tax":"1.23456789"}},"key":{"id":2},"topic":"pricing"}`,
+				`pricing: {"id":2}->{"after":{"values":{"discount":25.5,"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"}},"key":{"id":1},"topic":"pricing"}`,
+				`pricing: {"id":3}->{"after":{"values":{"discount":50,"id":3,"name":"Sofa","options":["Gray"],"tax":"4.250"}},"key":{"id":3},"topic":"pricing"}`,
+				`pricing: {"id":2}->{"key":{"id":2},"topic":"pricing"}`,
+			},
+		},
+		{
+			envelope:   "enriched",
+			withDiff:   true,
+			withSource: true,
+			expectedRows: []string{
+				`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}, "op": 2}`,
+				`pricing: {"id":1}->{"after": {"values": {"discount": 15.75, "id": 1, "name": "Chair", "options": ["Brown", "Black"], "tax": "2.500"}}, "key": {"id": 1}, "op": 1}`,
+				`pricing: {"id":2}->{"after": {"values": {"discount": 20, "id": 2, "name": "Table", "options": ["Brown", "Black"], "tax": "1.23456789"}}, "key": {"id": 2}, "op": 1}`,
+				`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}, "op": 2}`,
+				`pricing: {"id":2}->{"before": {"values": {"discount": 25.5, "id": 2, "name": "Table", "options": ["Brown", "Black"], "tax": "1.23456789"}}, "key": {"id": 2}, "op": 3}`,
+				`pricing: {"id":3}->{"after": {"values": {"discount": 50, "id": 3, "name": "Sofa", "options": ["Gray"], "tax": "4.250"}}, "key": {"id": 3}, "op": 1}`,
+			},
+		},
+		{
+			envelope: "enriched",
+			withDiff: true,
+			expectedRows: []string{
+				`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}, "op": 2}`,
+				`pricing: {"id":1}->{"after": {"values": {"discount": 15.75, "id": 1, "name": "Chair", "options": ["Brown", "Black"], "tax": "2.500"}}, "key": {"id": 1}, "op": 1}`,
+				`pricing: {"id":2}->{"after": {"values": {"discount": 20, "id": 2, "name": "Table", "options": ["Brown", "Black"], "tax": "1.23456789"}}, "key": {"id": 2}, "op": 1}`,
+				`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}, "op": 2}`,
+				`pricing: {"id":2}->{"before": {"values": {"discount": 25.5, "id": 2, "name": "Table", "options": ["Brown", "Black"], "tax": "1.23456789"}}, "key": {"id": 2}, "op": 3}`,
+				`pricing: {"id":3}->{"after": {"values": {"discount": 50, "id": 3, "name": "Sofa", "options": ["Gray"], "tax": "4.250"}}, "key": {"id": 3}, "op": 1}`,
 			},
 		},
 	}
@@ -12272,7 +12308,9 @@ func TestChangefeedProtobuf(t *testing.T) {
 				if tc.withDiff {
 					opts = append(opts, "diff")
 				}
-
+				if tc.withSource {
+					opts = append(opts, "enriched_properties='source'")
+				}
 				feed := feed(t, f, fmt.Sprintf("CREATE CHANGEFEED FOR pricing WITH %s", strings.Join(opts, ", ")))
 				defer closeFeed(t, feed)
 
@@ -12281,7 +12319,26 @@ func TestChangefeedProtobuf(t *testing.T) {
 				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)
+				if tc.envelope == "enriched" {
+					sourceAssertion := func(source map[string]any) {
+						if tc.withSource {
+							require.NotNil(t, source)
+							require.Equal(t, "kafka", source["changefeed_sink"])
+							require.Equal(t, "d", source["database_name"])
+							require.Equal(t, "public", source["schema_name"])
+							require.Equal(t, "pricing", source["table_name"])
+							require.Equal(t, "cockroachdb", source["origin"])
+							require.ElementsMatch(t, []any{"id"}, source["primary_keys"].([]any))
+						} else {
+							require.Nil(t, source)
+						}
+					}
+					assertPayloadsEnriched(t, feed, tc.expectedRows, sourceAssertion)
+				} else {
+
+					assertPayloads(t, feed, tc.expectedRows)
+				}
+
 			}
 			cdcTest(t, testFn, feedTestForceSink("kafka"))
 		})

pkg/ccl/changefeedccl/changefeedbase/options.go

@@ -935,8 +935,8 @@ func (e EncodingOptions) Validate() error {
 	}
 
 	if e.Envelope == OptEnvelopeEnriched {
-		if e.Format != OptFormatJSON && e.Format != OptFormatAvro {
-			return errors.Errorf(`%s=%s is only usable with %s=%s/%s`, OptEnvelope, OptEnvelopeEnriched, OptFormat, OptFormatJSON, OptFormatAvro)
+		if e.Format != OptFormatJSON && e.Format != OptFormatAvro && e.Format != OptFormatProtobuf {
+			return errors.Errorf(`%s=%s is only usable with %s=%s/%s/%s`, OptEnvelope, OptEnvelopeEnriched, OptFormat, OptFormatJSON, OptFormatAvro, OptFormatProtobuf)
 		}
 	} else {
 		if len(e.EnrichedProperties) > 0 {

pkg/ccl/changefeedccl/changefeedpb/changefeed.proto

@@ -64,8 +64,9 @@ message EnrichedEnvelope {
     Record after = 1;
     Record before = 2;
     Op op = 3;
-    int64 ts_ns = 4;
-    EnrichedSource source = 5;
+    Key key = 4;
+    int64 ts_ns = 5;
+    EnrichedSource source = 6;
 }
 
 // Resolved carries resolved timestamp information for a changefeed span.

pkg/ccl/changefeedccl/encoder.go

@@ -63,7 +63,7 @@ func getEncoder(
 		//information on why this was needed.
 		return nil, nil
 	case changefeedbase.OptFormatProtobuf:
-		return newProtobufEncoder(ctx, protobufEncoderOptions{EncodingOptions: opts}, targets), nil
+		return newProtobufEncoder(ctx, protobufEncoderOptions{EncodingOptions: opts}, targets, sourceProvider), nil
 	default:
 		return nil, errors.AssertionFailedf(`unknown format: %s`, opts.Format)
 	}

pkg/ccl/changefeedccl/encoder_protobuf.go

@@ -17,18 +17,21 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondatapb"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
+	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/errors"
 	"github.com/gogo/protobuf/types"
 )
 
 type protobufEncoder struct {
-	updatedField       bool
-	mvccTimestampField bool
-	beforeField        bool
-	keyInValue         bool
-	topicInValue       bool
-	envelopeType       changefeedbase.EnvelopeType
-	targets            changefeedbase.Targets
+	updatedField                   bool
+	mvccTimestampField             bool
+	beforeField                    bool
+	keyInValue                     bool
+	topicInValue                   bool
+	sourceField                    bool
+	envelopeType                   changefeedbase.EnvelopeType
+	targets                        changefeedbase.Targets
+	enrichedEnvelopeSourceProvider *enrichedSourceProvider
 }
 
 // protobufEncoderOptions wraps EncodingOptions for initializing a protobufEncoder.
@@ -40,16 +43,21 @@ var _ Encoder = &protobufEncoder{}
 
 // newProtobufEncoder constructs a new protobufEncoder from the given options and targets.
 func newProtobufEncoder(
-	ctx context.Context, opts protobufEncoderOptions, targets changefeedbase.Targets,
+	ctx context.Context,
+	opts protobufEncoderOptions,
+	targets changefeedbase.Targets,
+	sourceProvider *enrichedSourceProvider,
 ) Encoder {
 	return &protobufEncoder{
-		envelopeType:       opts.Envelope,
-		keyInValue:         opts.KeyInValue,
-		topicInValue:       opts.TopicInValue,
-		beforeField:        opts.Diff,
-		updatedField:       opts.UpdatedTimestamps,
-		mvccTimestampField: opts.MVCCTimestamps,
-		targets:            targets,
+		envelopeType:                   opts.Envelope,
+		keyInValue:                     opts.KeyInValue,
+		topicInValue:                   opts.TopicInValue,
+		beforeField:                    opts.Diff,
+		updatedField:                   opts.UpdatedTimestamps,
+		mvccTimestampField:             opts.MVCCTimestamps,
+		sourceField:                    inSet(changefeedbase.EnrichedPropertySource, opts.EnrichedProperties),
+		targets:                        targets,
+		enrichedEnvelopeSourceProvider: sourceProvider,
 	}
 }
 
@@ -71,11 +79,64 @@ func (e *protobufEncoder) EncodeValue(
 		return e.buildBare(evCtx, updatedRow, prevRow)
 	case changefeedbase.OptEnvelopeWrapped:
 		return e.buildWrapped(ctx, evCtx, updatedRow, prevRow)
+	case changefeedbase.OptEnvelopeEnriched:
+		return e.buildEnriched(ctx, evCtx, updatedRow, prevRow)
 	default:
 		return nil, errors.AssertionFailedf("envelope format not supported: %s", e.envelopeType)
 	}
 }
 
+func (e *protobufEncoder) buildEnriched(
+	ctx context.Context, evCtx eventContext, updatedRow cdcevent.Row, prevRow cdcevent.Row,
+) ([]byte, error) {
+	var after *changefeedpb.Record
+	var err error
+	if updatedRow.IsInitialized() && !updatedRow.IsDeleted() {
+		after, err = encodeRowToRecord(updatedRow)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	var before *changefeedpb.Record
+	if e.beforeField {
+		if prevRow.IsInitialized() && !prevRow.IsDeleted() {
+			before, err = encodeRowToRecord(prevRow)
+			if err != nil {
+				return nil, err
+			}
+		}
+	}
+	var keyMsg *changefeedpb.Key
+	if e.keyInValue {
+		keyMsg, err = buildKeyMessage(updatedRow)
+		if err != nil {
+			return nil, err
+		}
+	}
+	var src *changefeedpb.EnrichedSource
+	if e.sourceField {
+		src, err = e.enrichedEnvelopeSourceProvider.GetProtobuf(evCtx, updatedRow, prevRow)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	enriched := &changefeedpb.EnrichedEnvelope{
+		After:  after,
+		Before: before,
+		Key:    keyMsg,
+		TsNs:   timeutil.Now().UnixNano(),
+		Op:     inferOp(updatedRow, prevRow),
+		Source: src,
+	}
+
+	env := &changefeedpb.Message{
+		Data: &changefeedpb.Message_Enriched{Enriched: enriched},
+	}
+	return protoutil.Marshal(env)
+}
+
 // EncodeResolvedTimestamp encodes a resolved timestamp message for the specified topic.
 func (e *protobufEncoder) EncodeResolvedTimestamp(
 	ctx context.Context, topic string, ts hlc.Timestamp,
@@ -140,8 +201,6 @@ func (e *protobufEncoder) buildWrapped(
 		if err != nil {
 			return nil, err
 		}
-	} else {
-		after = &changefeedpb.Record{}
 	}
 
 	var before *changefeedpb.Record
@@ -151,8 +210,6 @@ func (e *protobufEncoder) buildWrapped(
 			if err != nil {
 				return nil, err
 			}
-		} else {
-			before = &changefeedpb.Record{}
 		}
 	}
 	var keyMsg *changefeedpb.Key
@@ -256,6 +313,19 @@ func buildKeyMessage(row cdcevent.Row) (*changefeedpb.Key, error) {
 	return &changefeedpb.Key{Key: keyMap}, nil
 }
 
+func inferOp(updated, prev cdcevent.Row) changefeedpb.Op {
+	switch deduceOp(updated, prev) {
+	case eventTypeCreate:
+		return changefeedpb.Op_OP_CREATE
+	case eventTypeUpdate:
+		return changefeedpb.Op_OP_UPDATE
+	case eventTypeDelete:
+		return changefeedpb.Op_OP_DELETE
+	default:
+		return changefeedpb.Op_OP_UNSPECIFIED
+	}
+}
+
 // datumToProtoValue converts a tree.Datum into a changefeedpb.Value.
 // It handles all common CockroachDB datum types and maps them to their
 // corresponding protobuf representation.

pkg/ccl/changefeedccl/encoder_protobuf_test.go

@@ -163,7 +163,7 @@ func Test_ProtoEncoderAllTypes(t *testing.T) {
 			}, false)
 
 			opts := changefeedbase.EncodingOptions{Envelope: changefeedbase.OptEnvelopeBare}
-			enc := newProtobufEncoder(ctx, protobufEncoderOptions{EncodingOptions: opts}, targets)
+			enc := newProtobufEncoder(ctx, protobufEncoderOptions{EncodingOptions: opts}, targets, nil)
 
 			evCtx := eventContext{updated: hlc.Timestamp{WallTime: 42}}
 			valBytes, err := enc.EncodeValue(ctx, evCtx, row, cdcevent.Row{})
@@ -440,7 +440,7 @@ func Test_ProtoEncoder_Escaping(t *testing.T) {
 	}, false)
 
 	opts := changefeedbase.EncodingOptions{Envelope: changefeedbase.OptEnvelopeBare}
-	enc := newProtobufEncoder(ctx, protobufEncoderOptions{EncodingOptions: opts}, targets)
+	enc := newProtobufEncoder(ctx, protobufEncoderOptions{EncodingOptions: opts}, targets, nil)
 
 	valBytes, err := enc.EncodeValue(ctx, eventContext{updated: hlc.Timestamp{WallTime: 42}}, row, cdcevent.Row{})
 	require.NoError(t, err)
@@ -485,7 +485,7 @@ func TestProtoEncoder_BareEnvelope_WithMetadata(t *testing.T) {
 		},
 	}
 
-	encoder := newProtobufEncoder(context.Background(), encOpts, mkTargets(tableDesc))
+	encoder := newProtobufEncoder(context.Background(), encOpts, mkTargets(tableDesc), nil)
 
 	valueBytes, err := encoder.EncodeValue(context.Background(), evCtx, row, cdcevent.Row{})
 	require.NoError(t, err)

pkg/ccl/changefeedccl/enriched_source_provider.go

@@ -15,6 +15,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/ccl/changefeedccl/avro"
 	"github.com/cockroachdb/cockroach/pkg/ccl/changefeedccl/cdcevent"
 	"github.com/cockroachdb/cockroach/pkg/ccl/changefeedccl/changefeedbase"
+	"github.com/cockroachdb/cockroach/pkg/ccl/changefeedccl/changefeedpb"
 	"github.com/cockroachdb/cockroach/pkg/ccl/changefeedccl/kcjsonschema"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
 	"github.com/cockroachdb/cockroach/pkg/sql"
@@ -192,6 +193,41 @@ func (p *enrichedSourceProvider) KafkaConnectJSONSchema() kcjsonschema.Schema {
 	return kafkaConnectJSONSchema
 }
 
+func (p *enrichedSourceProvider) GetProtobuf(
+	evCtx eventContext, updated, prev cdcevent.Row,
+) (*changefeedpb.EnrichedSource, error) {
+	md := updated.Metadata
+	tableInfo, ok := p.sourceData.tableSchemaInfo[md.TableID]
+	if !ok {
+		return nil, errors.AssertionFailedf("table %d not found in tableSchemaInfo", md.TableID)
+	}
+
+	src := &changefeedpb.EnrichedSource{
+		JobId:              p.sourceData.jobID,
+		ChangefeedSink:     p.sourceData.sink,
+		DbVersion:          p.sourceData.dbVersion,
+		ClusterName:        p.sourceData.clusterName,
+		ClusterId:          p.sourceData.clusterID,
+		SourceNodeLocality: p.sourceData.sourceNodeLocality,
+		NodeName:           p.sourceData.nodeName,
+		NodeId:             p.sourceData.nodeID,
+		Origin:             originCockroachDB,
+		DatabaseName:       tableInfo.dbName,
+		SchemaName:         tableInfo.schemaName,
+		TableName:          tableInfo.tableName,
+		PrimaryKeys:        tableInfo.primaryKeys,
+	}
+
+	if p.opts.mvccTimestamp {
+		src.MvccTimestamp = evCtx.mvcc.AsOfSystemTime()
+	}
+	if p.opts.updated {
+		src.TsNs = evCtx.updated.WallTime
+		src.TsHlc = evCtx.updated.AsOfSystemTime()
+	}
+	return src, nil
+}
+
 // GetJSON returns a json object for the source data.
 func (p *enrichedSourceProvider) GetJSON(
 	updated cdcevent.Row, evCtx eventContext,