changefeedccl: optimize enriched source provider

Optimize enriched source provider by avoiding unnecessary
allocations and json construction.

Fixes: #141798

Release note: None

Author: asg0451

pkg/ccl/changefeedccl/avro/avro.go

@@ -166,17 +166,20 @@ type DataRecord struct {
 	native map[string]interface{}
 }
 
+type functionalNativeFromRowFn func(row cdcevent.Row, dest map[string]any)
+
 type FunctionalRecord struct {
-	// Use the Record just for schema info
+	// Use the Record just for schema info.
 	Record
 
-	nativeFromRowFn func(row cdcevent.Row) (map[string]any, error)
+	// nativeFromRowFn populates the native map from a row.
+	nativeFromRowFn functionalNativeFromRowFn
+	// Reuse this map to avoid repeated map allocation when encoding.
+	native map[string]any
 }
 
 func NewFunctionalRecord(
-	name, namespace string,
-	fields []*SchemaField,
-	nativeFromRowFn func(row cdcevent.Row) (map[string]any, error),
+	name, namespace string, fields []*SchemaField, nativeFromRowFn functionalNativeFromRowFn,
 ) (*FunctionalRecord, error) {
 	schema := &FunctionalRecord{
 		Record: Record{
@@ -186,6 +189,7 @@ func NewFunctionalRecord(
 			Fields:     fields,
 		},
 		nativeFromRowFn: nativeFromRowFn,
+		native:          make(map[string]any),
 	}
 
 	schemaJSON, err := json.Marshal(schema)
@@ -199,6 +203,16 @@ func NewFunctionalRecord(
 	return schema, nil
 }
 
+// nativeFromRow populates the native map from a row and returns it. The
+// returned map will be reused in subsequent calls.
+func (r *FunctionalRecord) nativeFromRow(row cdcevent.Row) map[string]any {
+	if r.native == nil {
+		r.native = make(map[string]any, len(r.Fields))
+	}
+	r.nativeFromRowFn(row, r.native)
+	return r.native
+}
+
 // Metadata is the `EnvelopeRecord` metadata.
 type Metadata map[string]interface{}
 
@@ -1158,11 +1172,7 @@ func (r *EnvelopeRecord) BinaryFromRow(
 	}
 
 	if r.Opts.SourceField {
-		sourceNative, err := r.Source.nativeFromRowFn(recordRow)
-		if err != nil {
-			return nil, err
-		}
-		native[`source`] = goavro.Union(unionKey(&r.Source.Record), sourceNative)
+		native[`source`] = goavro.Union(unionKey(&r.Source.Record), r.Source.nativeFromRow(recordRow))
 	}
 	if r.Opts.TsField {
 		native[`ts_ns`] = nil

pkg/ccl/changefeedccl/changefeed_processors.go

@@ -1250,7 +1250,10 @@ func newChangeFrontierProcessor(
 
 	// This changeFrontier's encoder will only be used for resolved events which
 	// never have a source field, so we pass an empty enriched source provider.
-	sourceProvider := newEnrichedSourceProvider(encodingOpts, enrichedSourceData{})
+	sourceProvider, err := newEnrichedSourceProvider(encodingOpts, enrichedSourceData{})
+	if err != nil {
+		return nil, err
+	}
 	if cf.encoder, err = getEncoder(
 		ctx, encodingOpts, AllTargets(spec.Feed), spec.Feed.Select != "",
 		makeExternalConnectionProvider(ctx, flowCtx.Cfg.DB), sliMetrics,

pkg/ccl/changefeedccl/changefeed_stmt.go

@@ -604,7 +604,10 @@ func createChangefeedJobRecord(
 	if err != nil {
 		return nil, err
 	}
-	sourceProvider := newEnrichedSourceProvider(encodingOpts, enrichedSourceData{})
+	sourceProvider, err := newEnrichedSourceProvider(encodingOpts, enrichedSourceData{})
+	if err != nil {
+		return nil, err
+	}
 	if _, err := getEncoder(ctx, encodingOpts, AllTargets(details), details.Select != "",
 		makeExternalConnectionProvider(ctx, p.ExecCfg().InternalDB), nil, sourceProvider); err != nil {
 		return nil, err

pkg/ccl/changefeedccl/changefeed_test.go

@@ -3938,7 +3938,8 @@ func TestChangefeedEnriched(t *testing.T) {
 	// Create an enriched source provider with no data. The contents of source
 	// will be tested in another test, we just want to make sure the structure &
 	// schema is right here.
-	esp := newEnrichedSourceProvider(changefeedbase.EncodingOptions{}, enrichedSourceData{})
+	esp, err := newEnrichedSourceProvider(changefeedbase.EncodingOptions{}, enrichedSourceData{})
+	require.NoError(t, err)
 	source, err := esp.GetJSON(cdcevent.Row{})
 	require.NoError(t, err)
 
@@ -4236,7 +4237,7 @@ func TestChangefeedEnrichedSourceWithData(t *testing.T) {
 	}
 
 	for _, testCase := range cases {
-		t.Run(testCase.name, func(t *testing.T) {
+		testutils.RunTrueAndFalse(t, "mvcc_ts", func(t *testing.T, withMVCCTS bool) {
 			clusterName := "clusterName123"
 			dbVersion := "v999.0.0"
 			defer build.TestingOverrideVersion(dbVersion)()
@@ -4248,7 +4249,11 @@ func TestChangefeedEnrichedSourceWithData(t *testing.T) {
 
 					sqlDB.Exec(t, `CREATE TABLE foo (i INT PRIMARY KEY)`)
 					sqlDB.Exec(t, `INSERT INTO foo values (0)`)
-					testFeed := feed(t, f, fmt.Sprintf(`CREATE CHANGEFEED FOR foo WITH envelope=enriched, enriched_properties='source', format=%s`, testCase.format))
+					stmt := fmt.Sprintf(`CREATE CHANGEFEED FOR foo WITH envelope=enriched, enriched_properties='source', format=%s`, testCase.format)
+					if withMVCCTS {
+						stmt += ", mvcc_timestamp"
+					}
+					testFeed := feed(t, f, stmt)
 					defer closeFeed(t, testFeed)
 
 					var jobID int64
@@ -4276,35 +4281,50 @@ func TestChangefeedEnrichedSourceWithData(t *testing.T) {
 							sink = sinkTypeSinklessBuffer.String()
 						}
 
+						const dummyMvccTimestamp = "1234567890.0001"
+						jobIDStr := strconv.FormatInt(jobID, 10)
+
 						var assertion string
 						if testCase.format == "avro" {
-							assertion = fmt.Sprintf(
-								`{
-								"source": {
-									"cluster_id": {"string": "%s"},
-									"cluster_name": {"string": "%s"},
-									"db_version": {"string": "%s"},
-									"job_id": {"string": "%d"},
-									"node_id": {"string": "%s"},
-									"node_name": {"string": "%s"},
-									"changefeed_sink": {"string": "%s"},
-									"source_node_locality": {"string": "%s"}
-								}
-							}`,
-								clusterID, clusterName, dbVersion, jobID, nodeID, nodeName, sink, sourceNodeLocality)
+							assertionMap := map[string]any{
+								"source": map[string]any{
+									"cluster_id":   map[string]any{"string": clusterID},
+									"cluster_name": map[string]any{"string": clusterName},
+									"db_version":   map[string]any{"string": dbVersion},
+									"job_id":       map[string]any{"string": jobIDStr},
+									// Note that the field is still present in the avro schema, so it appears here as nil.
+									"mvcc_timestamp":       nil,
+									"node_id":              map[string]any{"string": nodeID},
+									"node_name":            map[string]any{"string": nodeName},
+									"changefeed_sink":      map[string]any{"string": sink},
+									"source_node_locality": map[string]any{"string": sourceNodeLocality},
+								},
+							}
+							if withMVCCTS {
+								mvccTsMap := actualSource["source"].(map[string]any)["mvcc_timestamp"].(map[string]any)
+								assertReasonableMVCCTimestamp(t, mvccTsMap["string"].(string))
+
+								mvccTsMap["string"] = dummyMvccTimestamp
+								assertionMap["source"].(map[string]any)["mvcc_timestamp"] = map[string]any{"string": dummyMvccTimestamp}
+							}
+							assertion = toJSON(t, assertionMap)
 						} else {
-							assertion = fmt.Sprintf(
-								`{
-								"cluster_id": "%s",
-								"cluster_name": "%s",
-								"db_version": "%s",
-								"job_id": "%d",
-								"node_id": "%s",
-								"node_name": "%s",
-								"changefeed_sink": "%s",
-								"source_node_locality": "%s"
-							}`,
-								clusterID, clusterName, dbVersion, jobID, nodeID, nodeName, sink, sourceNodeLocality)
+							assertionMap := map[string]any{
+								"cluster_id":           clusterID,
+								"cluster_name":         clusterName,
+								"db_version":           dbVersion,
+								"job_id":               jobIDStr,
+								"node_id":              nodeID,
+								"node_name":            nodeName,
+								"changefeed_sink":      sink,
+								"source_node_locality": sourceNodeLocality,
+							}
+							if withMVCCTS {
+								assertReasonableMVCCTimestamp(t, actualSource["mvcc_timestamp"].(string))
+								actualSource["mvcc_timestamp"] = dummyMvccTimestamp
+								assertionMap["mvcc_timestamp"] = dummyMvccTimestamp
+							}
+							assertion = toJSON(t, assertionMap)
 						}
 
 						value, err := reformatJSON(actualSource)
@@ -10748,3 +10768,9 @@ func TestChangefeedProtectedTimestampUpdate(t *testing.T) {
 
 	cdcTest(t, testFn, feedTestForceSink("kafka"), withTxnRetries)
 }
+
+func assertReasonableMVCCTimestamp(t *testing.T, ts string) {
+	epochNanos := parseTimeToHLC(t, ts).WallTime
+	now := timeutil.Now()
+	require.GreaterOrEqual(t, epochNanos, now.Add(-1*time.Hour).UnixNano())
+}

pkg/ccl/changefeedccl/changefeedbase/options.go

@@ -943,21 +943,26 @@ func (e EncodingOptions) Validate() error {
 		return errors.Errorf(`%s is only usable with %s=%s/%s`, OptHeadersJSONColumnName, OptFormat, OptFormatJSON, OptFormatAvro)
 	}
 
-	if e.Envelope != OptEnvelopeWrapped && e.Format != OptFormatJSON && e.Format != OptFormatParquet {
+	// TODO(#140110): refactor this logic.
+	if (e.Envelope != OptEnvelopeWrapped && e.Envelope != OptEnvelopeEnriched) && e.Format != OptFormatJSON && e.Format != OptFormatParquet {
 		requiresWrap := []struct {
 			k string
 			b bool
 		}{
 			{OptKeyInValue, e.KeyInValue},
+			// NOTE: topic_in_value is allowed for envelope=enriched, but has no
+			// effect. This is because the enriched envelope already has much of
+			// the information contained in the topic (ie table name), but this
+			// option is required for the webhook sink so we must permit it.
 			{OptTopicInValue, e.TopicInValue},
 			{OptUpdatedTimestamps, e.UpdatedTimestamps},
 			{OptMVCCTimestamps, e.MVCCTimestamps},
 			{OptDiff, e.Diff},
 		}
 		for _, v := range requiresWrap {
 			if v.b {
-				return errors.Errorf(`%s is only usable with %s=%s`,
-					v.k, OptEnvelope, OptEnvelopeWrapped)
+				return errors.Errorf(`%[1]s is only usable with %[2]s=%[3]s or %[2]s=%[4]s`,
+					v.k, OptEnvelope, OptEnvelopeWrapped, OptEnvelopeEnriched)
 			}
 		}
 	}

pkg/ccl/changefeedccl/encoder_avro.go

@@ -254,8 +254,7 @@ func (e *confluentAvroEncoder) EncodeValue(
 			recordDataSchema = currentSchema
 		case changefeedbase.OptEnvelopeEnriched:
 			afterDataSchema = currentSchema
-			opts = avro.EnvelopeOpts{AfterField: true, BeforeField: e.beforeField, MVCCTimestampField: e.mvccTimestampField, UpdatedField: e.updatedField,
-				OpField: true, TsField: true, SourceField: e.sourceField}
+			opts = avro.EnvelopeOpts{AfterField: true, BeforeField: e.beforeField, OpField: true, TsField: true, SourceField: e.sourceField}
 			if e.sourceField {
 				if sourceDataSchema, err = e.enrichedSourceProvider.GetAvro(updatedRow, e.schemaPrefix); err != nil {
 					return nil, err

pkg/ccl/changefeedccl/encoder_json.go

@@ -784,7 +784,10 @@ func EncodeAsJSONChangefeedWithFlags(
 	}
 	// This encoder is not used in the context of a real changefeed so make an empty
 	// source provider.
-	sourceProvider := newEnrichedSourceProvider(opts, enrichedSourceData{})
+	sourceProvider, err := newEnrichedSourceProvider(opts, enrichedSourceData{})
+	if err != nil {
+		return nil, err
+	}
 
 	// If this function ends up needing to be optimized, cache or pool these.
 	// Nontrivial to do as an encoder generally isn't safe to call on different

pkg/ccl/changefeedccl/encoder_json_test.go

@@ -156,7 +156,7 @@ func TestJSONEncoderJSONNullAsObject(t *testing.T) {
 		// NOTE: This is no longer required in go 1.22+, but bazel still requires it. See https://github.com/bazelbuild/rules_go/issues/3924
 		c := c
 		t.Run(c.name, func(t *testing.T) {
-			e, err := getEncoder(ctx, opts, targets, false, nil, nil, getTestingEnrichedSourceProvider(opts))
+			e, err := getEncoder(ctx, opts, targets, false, nil, nil, getTestingEnrichedSourceProvider(t, opts))
 			require.NoError(t, err)
 
 			row := cdcevent.TestingMakeEventRow(tableDesc, 0, c.row, false)
@@ -201,7 +201,7 @@ func TestJSONEncoderJSONNullAsObjectEdgeCases(t *testing.T) {
 			rowenc.EncDatum{Datum: tree.NewDJSON(json.NullJSONValue)},
 		}
 		e, err := getEncoder(ctx, opts, targets, false, nil, nil,
-			getTestingEnrichedSourceProvider(opts))
+			getTestingEnrichedSourceProvider(t, opts))
 		require.NoError(t, err)
 
 		row := cdcevent.TestingMakeEventRow(tableDesc, 0, eRow, false)
@@ -225,7 +225,7 @@ func TestJSONEncoderJSONNullAsObjectEdgeCases(t *testing.T) {
 			rowenc.EncDatum{Datum: tree.DNull},
 		}
 		e, err := getEncoder(ctx, opts, twoJSONsTargets, false, nil, nil,
-			getTestingEnrichedSourceProvider(opts))
+			getTestingEnrichedSourceProvider(t, opts))
 		require.NoError(t, err)
 
 		row := cdcevent.TestingMakeEventRow(twoJSONsTableDesc, 0, eRow, false)
@@ -243,7 +243,7 @@ func TestJSONEncoderJSONNullAsObjectEdgeCases(t *testing.T) {
 			rowenc.EncDatum{Datum: tree.NewDJSON(json.NullJSONValue)},
 		}
 		e, err := getEncoder(ctx, opts, twoJSONsTargets, false, nil, nil,
-			getTestingEnrichedSourceProvider(opts))
+			getTestingEnrichedSourceProvider(t, opts))
 		require.NoError(t, err)
 
 		row := cdcevent.TestingMakeEventRow(twoJSONsTableDesc, 0, eRow, false)
@@ -265,7 +265,7 @@ func TestJSONEncoderJSONNullAsObjectEdgeCases(t *testing.T) {
 			rowenc.EncDatum{Datum: tree.DNull},
 		}
 		e, err := getEncoder(ctx, disabledOpts, twoJSONsTargets, false, nil, nil,
-			getTestingEnrichedSourceProvider(disabledOpts))
+			getTestingEnrichedSourceProvider(t, disabledOpts))
 		require.NoError(t, err)
 
 		row := cdcevent.TestingMakeEventRow(twoJSONsTableDesc, 0, eRow, false)
@@ -287,7 +287,7 @@ func TestJSONEncoderJSONNullAsObjectEdgeCases(t *testing.T) {
 			rowenc.EncDatum{Datum: tree.NewDJSON(obj)},
 		}
 		e, err := getEncoder(ctx, opts, twoJSONsTargets, false, nil, nil,
-			getTestingEnrichedSourceProvider(opts))
+			getTestingEnrichedSourceProvider(t, opts))
 		require.NoError(t, err)
 
 		row := cdcevent.TestingMakeEventRow(twoJSONsTableDesc, 0, eRow, false)

pkg/ccl/changefeedccl/encoder_test.go

@@ -142,13 +142,13 @@ func TestEncoders(t *testing.T) {
 			resolved: `{"resolved":{"string":"1.0000000002"}}`,
 		},
 		`format=avro,envelope=key_only,updated`: {
-			err: `updated is only usable with envelope=wrapped`,
+			err: `updated is only usable with envelope=wrapped or envelope=enriched`,
 		},
 		`format=avro,envelope=key_only,diff`: {
-			err: `diff is only usable with envelope=wrapped`,
+			err: `diff is only usable with envelope=wrapped or envelope=enriched`,
 		},
 		`format=avro,envelope=key_only,updated,diff`: {
-			err: `updated is only usable with envelope=wrapped`,
+			err: `updated is only usable with envelope=wrapped or envelope=enriched`,
 		},
 		`format=avro,envelope=row`: {
 			err: `envelope=row is not supported with format=avro`,
@@ -241,7 +241,7 @@ func TestEncoders(t *testing.T) {
 				return
 			}
 			require.NoError(t, o.Validate())
-			e, err := getEncoder(context.Background(), o, targets, false, nil, nil, getTestingEnrichedSourceProvider(o))
+			e, err := getEncoder(context.Background(), o, targets, false, nil, nil, getTestingEnrichedSourceProvider(t, o))
 			require.NoError(t, err)
 
 			rowInsert := cdcevent.TestingMakeEventRow(tableDesc, 0, row, false)
@@ -387,7 +387,7 @@ func TestAvroEncoderWithTLS(t *testing.T) {
 				StatementTimeName: changefeedbase.StatementTimeName(tableDesc.GetName()),
 			})
 
-			e, err := getEncoder(context.Background(), opts, targets, false, nil, nil, getTestingEnrichedSourceProvider(opts))
+			e, err := getEncoder(context.Background(), opts, targets, false, nil, nil, getTestingEnrichedSourceProvider(t, opts))
 			require.NoError(t, err)
 
 			rowInsert := cdcevent.TestingMakeEventRow(tableDesc, 0, row, false)
@@ -419,7 +419,7 @@ func TestAvroEncoderWithTLS(t *testing.T) {
 			defer noCertReg.Close()
 			opts.SchemaRegistryURI = noCertReg.URL()
 
-			enc, err := getEncoder(context.Background(), opts, targets, false, nil, nil, getTestingEnrichedSourceProvider(opts))
+			enc, err := getEncoder(context.Background(), opts, targets, false, nil, nil, getTestingEnrichedSourceProvider(t, opts))
 			require.NoError(t, err)
 			_, err = enc.EncodeKey(context.Background(), rowInsert)
 			require.Regexp(t, "x509", err)
@@ -432,7 +432,7 @@ func TestAvroEncoderWithTLS(t *testing.T) {
 			defer wrongCertReg.Close()
 			opts.SchemaRegistryURI = wrongCertReg.URL()
 
-			enc, err = getEncoder(context.Background(), opts, targets, false, nil, nil, getTestingEnrichedSourceProvider(opts))
+			enc, err = getEncoder(context.Background(), opts, targets, false, nil, nil, getTestingEnrichedSourceProvider(t, opts))
 			require.NoError(t, err)
 			_, err = enc.EncodeKey(context.Background(), rowInsert)
 			require.Regexp(t, `contacting confluent schema registry.*: x509`, err)
@@ -974,7 +974,7 @@ func BenchmarkEncoders(b *testing.B) {
 		b.StopTimer()
 
 		encoder, err := getEncoder(context.Background(), opts, targets, false, nil, nil,
-			getTestingEnrichedSourceProvider(opts))
+			getTestingEnrichedSourceProvider(b, opts))
 		if err != nil {
 			b.Fatal(err)
 		}
@@ -1248,7 +1248,7 @@ func TestJsonRountrip(t *testing.T) {
 
 			// TODO(#139660): test this with other envelopes.
 			opts := jsonEncoderOptions{EncodingOptions: changefeedbase.EncodingOptions{Envelope: changefeedbase.OptEnvelopeBare}}
-			encoder, err := makeJSONEncoder(context.Background(), opts, getTestingEnrichedSourceProvider(opts.EncodingOptions), makeChangefeedTargets(test.name))
+			encoder, err := makeJSONEncoder(context.Background(), opts, getTestingEnrichedSourceProvider(t, opts.EncodingOptions), makeChangefeedTargets(test.name))
 			require.NoError(t, err)
 
 			// Encode the value to a string and parse it. Assert that the parsed json matches the