review

Author: jgao54

flow/connectors/mongo/qrep_partition.go

@@ -5,6 +5,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"log/slog"
+	"math"
 	"time"
 
 	"github.com/google/uuid"
@@ -27,14 +28,12 @@ func minObjectIDForTimestamp(t time.Time) bson.ObjectID {
 }
 
 // maxObjectIDForTimestamp creates an ObjectID with the given timestamp and
-// all trailing bytes set to 0xFF, producing the largest possible ObjectID
-// for that second.
+// all trailing bytes set to 0xFF (max unsigned 64-bit value), producing
+// the largest possible ObjectID for that second.
 func maxObjectIDForTimestamp(t time.Time) bson.ObjectID {
 	var oid [12]byte
 	binary.BigEndian.PutUint32(oid[0:4], uint32(t.Unix()))
-	for i := 4; i < 12; i++ {
-		oid[i] = 0xFF
-	}
+	binary.BigEndian.PutUint64(oid[4:], math.MaxUint64)
 	return oid
 }
 
@@ -53,25 +52,10 @@ func (c *MongoConnector) minMaxPartitions(
 		FullTablePartition: true,
 	}}
 
-	minRaw, err := findBoundaryID(ctx, collection, 1)
-	if err != nil {
-		c.logger.Info("[mongo] could not find min _id, collection may be empty, falling back to full table partition",
-			slog.String("error", err.Error()))
-		return fullTablePartition, nil
-	}
-	maxRaw, err := findBoundaryID(ctx, collection, -1)
+	minRaw, maxRaw, err := findMinMaxObjectIDs(ctx, collection)
 	if err != nil {
-		c.logger.Info("[mongo] could not find max _id, falling back to full table partition",
-			slog.String("error", err.Error()))
-		return fullTablePartition, nil
-	}
-
-	// Given MongoDB's type-ordered _id, if both min and max are ObjectID
-	// then every document in the collection must have an ObjectID _id.
-	if minRaw.Type != bson.TypeObjectID || maxRaw.Type != bson.TypeObjectID {
-		c.logger.Info("[mongo] _id contains non-ObjectID type, falling back to full table partition",
-			slog.String("minType", minRaw.Type.String()),
-			slog.String("maxType", maxRaw.Type.String()))
+		c.logger.Warn("[mongo] could not find min/max objectIDs, falling back to full table partition",
+			slog.String("reason", err.Error()))
 		return fullTablePartition, nil
 	}
 
@@ -115,17 +99,50 @@ func (c *MongoConnector) minMaxPartitions(
 	return partitions, nil
 }
 
-// findBoundaryID returns the raw _id value of the min (direction=1) or max (direction=-1)
-// document in the collection. Because _id is always indexed, this is an efficient O(log n) operation.
-// The caller should inspect the returned RawValue.Type to determine the BSON type.
+// findMinMaxObjectIDs returns the min and max object IDs from the collection.
+// The two boundary queries are not wrapped in a snapshot session: a concurrent
+// insert between them may shift the true min/max, but this is safe because the
+// change stream resume token is captured before either query runs. Any writes
+// that occur during or after partitioning will be replayed by the change stream.
+func findMinMaxObjectIDs(
+	ctx context.Context,
+	collection *mongo.Collection,
+) (bson.RawValue, bson.RawValue, error) {
+	minRaw, err := findBoundaryID(ctx, collection, Lower)
+	if err != nil {
+		return bson.RawValue{}, bson.RawValue{}, fmt.Errorf("failed to find min _id: %w", err)
+	}
+
+	maxRaw, err := findBoundaryID(ctx, collection, Upper)
+	if err != nil {
+		return bson.RawValue{}, bson.RawValue{}, fmt.Errorf("failed to find max _id: %w", err)
+	}
+
+	// Given MongoDB's type-ordered _id, if both min and max are ObjectID, then every
+	// document in the collection must have an ObjectID _id.
+	if minRaw.Type != bson.TypeObjectID || maxRaw.Type != bson.TypeObjectID {
+		return bson.RawValue{}, bson.RawValue{}, fmt.Errorf(
+			"_id contains non-ObjectID type (min id type=%T, max id type=%T", minRaw, maxRaw)
+	}
+	return minRaw, maxRaw, nil
+}
+
+type Boundary int
+
+const (
+	Lower Boundary = 1
+	Upper Boundary = -1
+)
+
+// findBoundaryID returns the lower- or upper-bound _id of a collection
 func findBoundaryID(
 	ctx context.Context,
 	collection *mongo.Collection,
-	direction int,
+	boundary Boundary,
 ) (bson.RawValue, error) {
 	findCmd := bson.D{
 		{Key: "find", Value: collection.Name()},
-		{Key: "sort", Value: bson.D{{Key: DefaultDocumentKeyColumnName, Value: direction}}},
+		{Key: "sort", Value: bson.D{{Key: DefaultDocumentKeyColumnName, Value: boundary}}},
 		{Key: "limit", Value: 1},
 		{Key: "projection", Value: bson.D{{Key: DefaultDocumentKeyColumnName, Value: 1}}},
 	}