feat : Fallback to bucketAuto Strategy In Monngodb (#158)

Co-authored-by: hash-data <[email protected]>
Co-authored-by: Ankit Sharma <[email protected]>

Author: vikash390

drivers/mongodb/go.mod

@@ -12,7 +12,6 @@ require (
 	github.com/apache/thrift v0.16.0 // indirect
 	github.com/aws/aws-sdk-go v1.43.31 // indirect
 	github.com/felixge/fgprof v0.9.5 // indirect
-	github.com/fraugster/parquet-go v0.12.0 // indirect
 	github.com/fsnotify/fsnotify v1.5.1 // indirect
 	github.com/golang/snappy v0.0.4 // indirect
 	github.com/google/pprof v0.0.0-20240227163752-401108e1b7e7 // indirect
@@ -64,14 +63,14 @@ require (
 	github.com/jackc/pgio v1.0.0 // indirect
 	github.com/jackc/pgpassfile v1.0.0 // indirect
 	github.com/jackc/pgproto3/v2 v2.3.3 // indirect
-	github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
+	github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
 	github.com/jackc/pgtype v1.14.0 // indirect
 	github.com/leodido/go-urn v1.4.0 // indirect
 	github.com/mitchellh/hashstructure v1.1.0 // indirect
 	github.com/rogpeppe/go-internal v1.13.1 // indirect
 	github.com/spf13/cobra v1.8.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
-	go.mongodb.org/mongo-driver v1.17.1
+	go.mongodb.org/mongo-driver v1.17.3
 	golang.org/x/crypto v0.31.0 // indirect
 	golang.org/x/net v0.33.0 // indirect
 	golang.org/x/sys v0.28.0 // indirect

drivers/mongodb/internal/backfill.go

@@ -124,25 +124,24 @@ func (m *Mongo) backfill(stream protocol.Stream, pool *protocol.WriterPool) erro
 
 func (m *Mongo) splitChunks(ctx context.Context, collection *mongo.Collection, stream protocol.Stream) ([]types.Chunk, error) {
 	splitVectorStrategy := func() ([]types.Chunk, error) {
-		getChunkBoundaries := func() ([]*primitive.ObjectID, error) {
-			getID := func(order int) (primitive.ObjectID, error) {
-				var doc bson.M
-				err := collection.FindOne(ctx, bson.D{}, options.FindOne().SetSort(bson.D{{Key: "_id", Value: order}})).Decode(&doc)
-				if err == mongo.ErrNoDocuments {
-					return primitive.NilObjectID, nil
-				}
-				return doc["_id"].(primitive.ObjectID), err
-			}
-
-			minID, err := getID(1)
-			if err != nil || minID == primitive.NilObjectID {
-				return nil, err
-			}
-			maxID, err := getID(-1)
-			if err != nil {
-				return nil, err
+		getID := func(order int) (primitive.ObjectID, error) {
+			var doc bson.M
+			err := collection.FindOne(ctx, bson.D{}, options.FindOne().SetSort(bson.D{{Key: "_id", Value: order}})).Decode(&doc)
+			if err == mongo.ErrNoDocuments {
+				return primitive.NilObjectID, nil
 			}
+			return doc["_id"].(primitive.ObjectID), err
+		}
 
+		minID, err := getID(1)
+		if err != nil || minID == primitive.NilObjectID {
+			return nil, err
+		}
+		maxID, err := getID(-1)
+		if err != nil {
+			return nil, err
+		}
+		getChunkBoundaries := func() ([]*primitive.ObjectID, error) {
 			var result bson.M
 			cmd := bson.D{
 				{Key: "splitVector", Value: fmt.Sprintf("%s.%s", collection.Database().Name(), collection.Name())},
@@ -173,6 +172,57 @@ func (m *Mongo) splitChunks(ctx context.Context, collection *mongo.Collection, s
 				Max: &boundaries[i+1],
 			})
 		}
+		if len(boundaries) > 0 {
+			chunks = append(chunks, types.Chunk{
+				Min: &boundaries[len(boundaries)-1],
+				Max: nil,
+			})
+		}
+		return chunks, nil
+	}
+	bucketAutoStrategy := func() ([]types.Chunk, error) {
+		logger.Info("using bucket auto strategy for stream: %s", stream.ID())
+		// Use $bucketAuto for chunking
+		pipeline := mongo.Pipeline{
+			{{Key: "$sort", Value: bson.D{{Key: "_id", Value: 1}}}},
+			{{Key: "$bucketAuto", Value: bson.D{
+				{Key: "groupBy", Value: "$_id"},
+				{Key: "buckets", Value: m.config.MaxThreads * 4},
+			}}},
+		}
+
+		cursor, err := collection.Aggregate(ctx, pipeline)
+		if err != nil {
+			return nil, fmt.Errorf("failed to execute bucketAuto aggregation: %s", err)
+		}
+		defer cursor.Close(ctx)
+
+		var buckets []struct {
+			ID struct {
+				Min primitive.ObjectID `bson:"min"`
+				Max primitive.ObjectID `bson:"max"`
+			} `bson:"_id"`
+			Count int `bson:"count"`
+		}
+
+		if err := cursor.All(ctx, &buckets); err != nil {
+			return nil, fmt.Errorf("failed to decode bucketAuto results: %s", err)
+		}
+
+		var chunks []types.Chunk
+		for _, bucket := range buckets {
+			chunks = append(chunks, types.Chunk{
+				Min: &bucket.ID.Min,
+				Max: &bucket.ID.Max,
+			})
+		}
+		if len(buckets) > 0 {
+			chunks = append(chunks, types.Chunk{
+				Min: &buckets[len(buckets)-1].ID.Max,
+				Max: nil,
+			})
+		}
+
 		return chunks, nil
 	}
 
@@ -205,14 +255,25 @@ func (m *Mongo) splitChunks(ctx context.Context, collection *mongo.Collection, s
 				Max: maxObjectID,
 			})
 		}
+		chunks = append(chunks, types.Chunk{
+			Min: generateMinObjectID(last),
+			Max: nil,
+		})
+
 		return chunks, nil
 	}
 
 	switch m.config.PartitionStrategy {
 	case "timestamp":
 		return timestampStrategy()
 	default:
-		return splitVectorStrategy()
+		chunks, err := splitVectorStrategy()
+		// check if authorization error occurs
+		if err != nil && strings.Contains(err.Error(), "not authorized") {
+			logger.Warnf("failed to get chunks via split vector strategy: %s", err)
+			return bucketAutoStrategy()
+		}
+		return chunks, err
 	}
 }
 func (m *Mongo) totalCountInCollection(ctx context.Context, collection *mongo.Collection) (int64, error) {
@@ -293,7 +354,7 @@ func generatePipeline(start, end any) mongo.Pipeline {
 		andOperation = append(andOperation, bson.D{{
 			Key: "_id",
 			Value: bson.D{{
-				Key:   "$lte",
+				Key:   "$lt",
 				Value: end,
 			}},
 		}})