separate out network error from invalid object id

Author: jgao54

flow/connectors/mongo/qrep_partition.go

@@ -33,21 +33,19 @@ func (c *MongoConnector) minMaxPartitions(
 		FullTablePartition: true,
 	}}
 
-	minRaw, maxRaw, err := findMinMaxObjectIDs(ctx, collection, c.config.ReadPreference)
+	minID, maxID, found, err := findMinMaxObjectIDs(ctx, collection, c.config.ReadPreference)
 	if err != nil {
-		c.logger.Warn("[mongo] could not find min/max objectIDs, falling back to full table partition",
-			slog.String("reason", err.Error()))
+		return nil, err
+	}
+	if !found {
+		c.logger.Info("[mongo] no valid min/max ObjectIDs found, falling back to full table partition")
 		return fullTablePartition, nil
 	}
-
-	minID := minRaw.ObjectID()
-	maxID := maxRaw.ObjectID()
-
 	minInt := new(big.Int).SetBytes(minID[:])
 	maxInt := new(big.Int).SetBytes(maxID[:])
 	intRange := new(big.Int).Sub(maxInt, minInt)
 	if intRange.Sign() <= 0 {
-		c.logger.Info("[mongo] min/max ObjectIDs are equal, falling back to full table partition")
+		c.logger.Info("[mongo] min/max ObjectID range is non-positive, falling back to full table partition")
 		return fullTablePartition, nil
 	}
 
@@ -109,25 +107,24 @@ func findMinMaxObjectIDs(
 	ctx context.Context,
 	collection *mongo.Collection,
 	readPreference protos.ReadPreference,
-) (bson.RawValue, bson.RawValue, error) {
+) (bson.ObjectID, bson.ObjectID, bool, error) {
 	minRaw, err := findBoundaryID(ctx, collection, Lower, readPreference)
 	if err != nil {
-		return bson.RawValue{}, bson.RawValue{}, fmt.Errorf("failed to find min _id: %w", err)
+		return bson.ObjectID{}, bson.ObjectID{}, false, fmt.Errorf("failed to find min _id: %w", err)
+	}
+	if minRaw.Type != bson.TypeObjectID {
+		return bson.ObjectID{}, bson.ObjectID{}, false, nil
 	}
 
 	maxRaw, err := findBoundaryID(ctx, collection, Upper, readPreference)
 	if err != nil {
-		return bson.RawValue{}, bson.RawValue{}, fmt.Errorf("failed to find max _id: %w", err)
+		return bson.ObjectID{}, bson.ObjectID{}, false, 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=%s, max id type=%s)",
-			minRaw.Type.String(), maxRaw.Type.String())
+	if maxRaw.Type != bson.TypeObjectID {
+		return bson.ObjectID{}, bson.ObjectID{}, false, nil
 	}
-	return minRaw, maxRaw, nil
+
+	return minRaw.ObjectID(), maxRaw.ObjectID(), true, nil
 }
 
 type Boundary int
@@ -158,11 +155,11 @@ func findBoundaryID(
 	defer cursor.Close(ctx)
 
 	if !cursor.Next(ctx) {
-		return bson.RawValue{}, fmt.Errorf("collection is empty")
-	}
-	rv := cursor.Current.Lookup(DefaultDocumentKeyColumnName)
-	if rv.IsZero() {
-		return bson.RawValue{}, fmt.Errorf("document _id is nil")
+		if err := cursor.Err(); err != nil {
+			return bson.RawValue{}, fmt.Errorf("cursor error: %w", err)
+		}
+		// no results
+		return bson.RawValue{}, nil
 	}
-	return rv, nil
+	return cursor.Current.Lookup(DefaultDocumentKeyColumnName), nil
 }