feat: removal of heavy produce schema operation during sync for kafka and mongo

Author: vaibhav-datazip

constants/constants.go

@@ -59,3 +59,7 @@ var SkipCDCDrivers = []DriverType{Oracle, DB2}
 
 // DriversRequiringIncrementalFormatter are drivers that require special formatting for incremental value
 var DriversRequiringIncrementalFormatter = []DriverType{Oracle, DB2, MSSQL}
+
+// SyncContext is used as a context key to signal that ProduceSchema is being called during
+// sync. Drivers that perform expensive sampling (e.g. MongoDB, Kafka) can check for it to skip schema inference.
+type SyncContext struct{}

drivers/kafka/internal/kafka.go

@@ -153,6 +153,15 @@ func (k *Kafka) ProduceSchema(ctx context.Context, streamName string) (*types.St
 	stream := types.NewStream(streamName, "topics", nil)
 	stream.WithSyncMode(types.STRICTCDC)
 	stream.SyncMode = types.STRICTCDC
+	stream.SourceDefinedPrimaryKey = types.NewSet(Offset, Partition)
+
+	// During sync, skip expensive message reading for schema inference.
+	// The full schema is already available in the catalog (streams.json) from a prior discover run.
+	// Only structural metadata (sync modes, primary keys) is needed for stream validation.
+	if ctx.Value(constants.SyncContext{}) != nil {
+		logger.Infof("sync context detected, skipping schema inference for topic [%s]", streamName)
+		return stream, nil
+	}
 
 	// create reader manager for schema discovery
 	readerManager := kafkapkg.NewReaderManager(kafkapkg.ReaderConfig{
@@ -230,7 +239,6 @@ func (k *Kafka) ProduceSchema(ctx context.Context, streamName string) (*types.St
 		return nil, fmt.Errorf("failed to fetch schema for topic %s: %s", streamName, err)
 	}
 
-	stream.SourceDefinedPrimaryKey = types.NewSet(Offset, Partition)
 	return stream, nil
 }
 

drivers/mongodb/internal/mon.go

@@ -230,69 +230,76 @@ func (m *Mongo) GetStreamNames(ctx context.Context) ([]string, error) {
 	return streamNames, collections.Err()
 }
 
-// TODO: omit usage of ProduceSchema when sync command is used
 func (m *Mongo) ProduceSchema(ctx context.Context, streamName string) (*types.Stream, error) {
-	produceCollectionSchema := func(ctx context.Context, db *mongo.Database, streamName string) (*types.Stream, error) {
-		logger.Infof("producing type schema for stream [%s]", streamName)
-
-		// initialize stream
-		collection := db.Collection(streamName)
-		stream := types.NewStream(streamName, db.Name(), nil)
-		// find primary keys
-		indexesCursor, err := collection.Indexes().List(ctx, options.ListIndexes())
-		if err != nil {
+	database := m.client.Database(m.config.Database)
+	collection := database.Collection(streamName)
+	stream := types.NewStream(streamName, database.Name(), nil)
+
+	// Query indexes for primary keys (lightweight, always needed for validation)
+	indexesCursor, err := collection.Indexes().List(ctx, options.ListIndexes())
+	if err != nil {
+		return nil, err
+	}
+	defer indexesCursor.Close(ctx)
+
+	for indexesCursor.Next(ctx) {
+		var indexes bson.M
+		if err := indexesCursor.Decode(&indexes); err != nil {
 			return nil, err
 		}
-		defer indexesCursor.Close(ctx)
+		for key := range indexes["key"].(bson.M) {
+			stream.WithPrimaryKey(key)
+		}
+	}
 
-		for indexesCursor.Next(ctx) {
-			var indexes bson.M
-			if err := indexesCursor.Decode(&indexes); err != nil {
-				return nil, err
-			}
-			for key := range indexes["key"].(bson.M) {
-				stream.WithPrimaryKey(key)
-			}
+	// During sync, skip expensive document sampling for schema inference.
+	// The full schema is already available in the catalog (streams.json) from a prior discover run.
+	// Only structural metadata (primary keys, sync modes) is needed for stream validation.
+	if ctx.Value(constants.SyncContext{}) != nil {
+		logger.Infof("sync context detected, skipping schema inference for stream [%s]", streamName)
+		stream.WithSyncMode(types.FULLREFRESH, types.INCREMENTAL)
+		if m.CDCSupported() {
+			stream.UpsertField(CDCResumeToken, types.String, true, true)
+			stream.WithSyncMode(types.CDC, types.STRICTCDC)
 		}
+		return stream, nil
+	}
+
+	logger.Infof("producing type schema for stream [%s]", streamName)
+
+	// Define find options for fetching documents in ascending and descending order.
+	findOpts := []*options.FindOptions{
+		options.Find().SetLimit(10000).SetSort(bson.D{{Key: "$natural", Value: 1}}),
+		options.Find().SetLimit(10000).SetSort(bson.D{{Key: "$natural", Value: -1}}),
+	}
 
-		// Define find options for fetching documents in ascending and descending order.
-		findOpts := []*options.FindOptions{
-			options.Find().SetLimit(10000).SetSort(bson.D{{Key: "$natural", Value: 1}}),
-			options.Find().SetLimit(10000).SetSort(bson.D{{Key: "$natural", Value: -1}}),
+	if err := utils.Concurrent(ctx, findOpts, len(findOpts), func(ctx context.Context, findOpt *options.FindOptions, execNumber int) error {
+		cursor, err := collection.Find(ctx, bson.D{}, findOpt)
+		if err != nil {
+			return err
 		}
+		defer cursor.Close(ctx)
 
-		return stream, utils.Concurrent(ctx, findOpts, len(findOpts), func(ctx context.Context, findOpt *options.FindOptions, execNumber int) error {
-			cursor, err := collection.Find(ctx, bson.D{}, findOpt)
-			if err != nil {
+		for cursor.Next(ctx) {
+			var row bson.M
+			if err := cursor.Decode(&row); err != nil {
 				return err
 			}
-			defer cursor.Close(ctx)
 
-			for cursor.Next(ctx) {
-				var row bson.M
-				if err := cursor.Decode(&row); err != nil {
-					return err
-				}
-
-				filterMongoObject(row)
-				if err := typeutils.Resolve(stream, row); err != nil {
-					return err
-				}
+			filterMongoObject(row)
+			if err := typeutils.Resolve(stream, row); err != nil {
+				return err
 			}
+		}
 
-			return cursor.Err()
-		})
-	}
-	database := m.client.Database(m.config.Database)
-	// Either wait for covering 100k records from both sides for all streams
-	// Or wait till discoverCtx exits
-	stream, err := produceCollectionSchema(ctx, database, streamName)
-	if err != nil {
+		return cursor.Err()
+	}); err != nil {
 		if ctx.Err() != nil {
 			return nil, fmt.Errorf("failed to produce schema context deadline exceeded: %s", ctx.Err())
 		}
 		return nil, fmt.Errorf("failed to process collection[%s]: %s", streamName, err)
 	}
+
 	// Add all discovered fields as potential cursor fields
 	stream.Schema.Properties.Range(func(key, value interface{}) bool {
 		if fieldName, ok := key.(string); ok {
@@ -307,7 +314,7 @@ func (m *Mongo) ProduceSchema(ctx context.Context, streamName string) (*types.St
 		stream.WithSyncMode(types.CDC, types.STRICTCDC)
 	}
 
-	return stream, err
+	return stream, nil
 }
 
 func filterMongoObject(doc bson.M) {

protocol/sync.go

@@ -1,6 +1,7 @@
 package protocol
 
 import (
+	"context"
 	"fmt"
 	"strings"
 	"sync"
@@ -82,13 +83,18 @@ var syncCmd = &cobra.Command{
 		return nil
 	},
 	RunE: func(cmd *cobra.Command, _ []string) error {
+		// Mark context as sync command so drivers can skip expensive schema inference
+		// (e.g., MongoDB document sampling, Kafka message reading) since the schema
+		// is already available in the catalog (streams.json) from a prior discover run.
+		ctx := context.WithValue(cmd.Context(), constants.SyncContext{}, true)
+
 		// setup conector first
-		err := connector.Setup(cmd.Context())
+		err := connector.Setup(ctx)
 		if err != nil {
 			return err
 		}
 		// Get Source Streams, sending 0 max discover threads to discover
-		streams, err := connector.Discover(cmd.Context(), 0)
+		streams, err := connector.Discover(ctx, 0)
 		if err != nil {
 			return err
 		}

types/stream_configured.go

@@ -261,8 +261,10 @@ func (s *ConfiguredStream) Validate(source *Stream) error {
 		return fmt.Errorf("invalid sync mode[%s]; valid are %v", s.Stream.SyncMode, source.SupportedSyncModes)
 	}
 
-	// no cursor validation in cdc and backfill sync
-	if s.Stream.SyncMode == INCREMENTAL {
+	// no cursor validation in cdc and backfill sync;
+	// also skip when AvailableCursorFields is empty — the source produced a
+	// lightweight schema (e.g. sync fast-path) and we trust the streams.json.
+	if s.Stream.SyncMode == INCREMENTAL && len(source.AvailableCursorFields.Array()) > 0 {
 		primaryCursor, secondaryCursor := s.Cursor()
 		if !source.AvailableCursorFields.Exists(primaryCursor) {
 			return fmt.Errorf("invalid cursor field [%s]; valid are %v", primaryCursor, source.AvailableCursorFields)