fix: Reuse range boundaries when resuming parallel PG migration

.trivyignore

@@ -0,0 +1,5 @@
+# Go stdlib CVEs in SUSE's Go 1.26.8 builder. Fixed in Go 1.26.9.
+# Remove once BCI ships it.
+CVE-2026-78667 exp:2026-11-30
+CVE-2026-97031 exp:2026-11-30
+CVE-2026-78669 exp:2026-11-30

cmd/migrate_from_pg.go

@@ -45,6 +45,10 @@ type pgRangeSpec struct {
 	id       int
 	startKey *string
 	endKey   *string
+	// offsetKey embeds the boundaries fingerprint so offsets only apply to the ranges they were saved for.
+	offsetKey string
+	lastKey   *string
+	count     uint64
 }
 
 func (r *MigrateFromPGCmd) Parse() error {
@@ -110,16 +114,10 @@ func (r *MigrateFromPGCmd) Run(globals *Globals) error {
 		return fmt.Errorf("failed to migrate data: %w", err)
 	}
 
-	targetPointCount, err := targetClient.Count(ctx, &qdrant.CountPoints{
-		CollectionName: r.Qdrant.Collection,
-		Exact:          qdrant.PtrOf(true),
-	})
-	if err != nil {
-		return fmt.Errorf("failed to count points in target: %w", err)
+	if err := verifyTargetCount(ctx, targetClient, r.Qdrant.Collection, sourcePointCount, r.Migration.OffsetsCollection); err != nil {
+		return err
 	}
 
-	pterm.Info.Printfln("Target collection has %d points\n", targetPointCount)
-
 	err = commons.DeleteOffsetsCollection(ctx, r.Migration.OffsetsCollection, targetClient)
 	if err != nil {
 		return fmt.Errorf("failed to delete migration marker collection: %w", err)
@@ -341,7 +339,7 @@ func (r *MigrateFromPGCmd) sampleKeyValues(ctx context.Context, conn *pgx.Conn,
 
 	// Sample random keys, but return them sorted by their native column type order.
 	query := fmt.Sprintf(
-		"SELECT %s::text FROM (SELECT %s FROM %s ORDER BY RANDOM() LIMIT $1) sub ORDER BY %s",
+		"SELECT %s::text FROM (SELECT %s FROM %s ORDER BY RANDOM() LIMIT $1) sub ORDER BY sub.%s",
 		keyColIdent, keyColIdent, tableIdent, keyColIdent,
 	)
 	rows, err := conn.Query(ctx, query, sampleSize)
@@ -401,17 +399,40 @@ func (r *MigrateFromPGCmd) migrateDataParallel(ctx context.Context, sourceConn *
 		return nil
 	}
 
-	sampleSize := r.NumWorkers * SAMPLE_SIZE_PER_WORKER
-	if uint64(sampleSize) > sourcePointCount {
-		sampleSize = int(sourcePointCount)
+	// Keys are wrapped as UUID point IDs to reuse the boundary helpers, like the per-range offsets.
+	ids, err := resolveRangeBoundaries(ctx, targetClient, r.Migration,
+		fmt.Sprintf("%s-%s-workers-%d-boundaries", r.PG.Table, r.PG.KeyColumn, r.NumWorkers),
+		func() ([]*qdrant.PointId, error) {
+			sampleSize := r.NumWorkers * SAMPLE_SIZE_PER_WORKER
+			if uint64(sampleSize) > sourcePointCount {
+				sampleSize = int(sourcePointCount)
+			}
+			keys, err := r.sampleKeyValues(ctx, sourceConn, sampleSize)
+			if err != nil {
+				return nil, fmt.Errorf("failed to sample keys: %w", err)
+			}
+			ids := make([]*qdrant.PointId, len(keys))
+			for i, key := range keys {
+				ids[i] = qdrant.NewIDUUID(key)
+			}
+			return ids, nil
+		})
+	if err != nil {
+		return err
 	}
-
-	keys, err := r.sampleKeyValues(ctx, sourceConn, sampleSize)
+	fingerprint, err := commons.BoundariesFingerprint(ids)
 	if err != nil {
-		return fmt.Errorf("failed to sample keys: %w", err)
+		return err
+	}
+	keys := make([]string, len(ids))
+	for i, id := range ids {
+		keys[i] = id.GetUuid()
 	}
 
 	ranges := r.createRangesFromSamples(keys)
+	for i := range ranges {
+		ranges[i].offsetKey = fmt.Sprintf("%s-%s-workers-%d-%s-range-%d", r.PG.Table, r.PG.KeyColumn, r.NumWorkers, fingerprint, ranges[i].id)
+	}
 
 	poolConfig, err := pgxpool.ParseConfig(r.PG.Url)
 	if err != nil {
@@ -429,15 +450,17 @@ func (r *MigrateFromPGCmd) migrateDataParallel(ctx context.Context, sourceConn *
 	defer pool.Close()
 
 	// If not restarting, load the progress for each range.
-	// The offset key includes NumWorkers, so changing num-workers will start a fresh migration.
 	var totalProcessed uint64
 	if !r.Migration.Restart {
 		for i := range ranges {
-			offsetKey := fmt.Sprintf("%s-workers-%d-range-%d", r.PG.Table, r.NumWorkers, ranges[i].id)
-			_, count, err := commons.GetStartOffset(ctx, r.Migration.OffsetsCollection, targetClient, offsetKey)
+			offsetID, count, err := commons.GetStartOffset(ctx, r.Migration.OffsetsCollection, targetClient, ranges[i].offsetKey)
 			if err != nil {
 				return fmt.Errorf("failed to get start offset: %w", err)
 			}
+			if key := offsetID.GetUuid(); key != "" {
+				ranges[i].lastKey = &key
+			}
+			ranges[i].count = count
 			totalProcessed += count
 			if count > 0 {
 				pterm.Info.Printfln("Resuming range %d (already processed %d points)", ranges[i].id, count)
@@ -479,23 +502,10 @@ func (r *MigrateFromPGCmd) migrateDataParallel(ctx context.Context, sourceConn *
 // migrateRange is the function executed by each worker in parallel migration.
 // It queries a specific key range and upserts the rows to the target.
 func (r *MigrateFromPGCmd) migrateRange(ctx context.Context, pool *pgxpool.Pool, targetClient *qdrant.Client, rg pgRangeSpec, bar *pterm.ProgressbarPrinter) error {
-	offsetKey := fmt.Sprintf("%s-workers-%d-range-%d", r.PG.Table, r.NumWorkers, rg.id)
+	offsetKey := rg.offsetKey
 	batchSize := r.Migration.BatchSize
-
-	var lastKey *string
-	var count uint64
-	if !r.Migration.Restart {
-		offsetID, c, err := commons.GetStartOffset(ctx, r.Migration.OffsetsCollection, targetClient, offsetKey)
-		if err != nil {
-			return fmt.Errorf("failed to get start offset: %w", err)
-		}
-		if offsetID != nil {
-			if uuid := offsetID.GetUuid(); uuid != "" {
-				lastKey = &uuid
-			}
-		}
-		count = c
-	}
+	lastKey := rg.lastKey
+	count := rg.count
 
 	tableIdent := pgx.Identifier{r.PG.Table}.Sanitize()
 	keyColIdent := pgx.Identifier{r.PG.KeyColumn}.Sanitize()

cmd/migrate_from_qdrant.go

@@ -129,22 +129,8 @@ func (r *MigrateFromQdrantCmd) Run(globals *Globals) error {
 		return fmt.Errorf("failed to migrate data: %w", err)
 	}
 
-	targetPointCount, err := targetClient.Count(ctx, &qdrant.CountPoints{
-		CollectionName: r.Target.Collection,
-		Exact:          qdrant.PtrOf(true),
-	})
-	if err != nil {
-		return fmt.Errorf("failed to count points in target: %w", err)
-	}
-
-	pterm.Info.Printfln("Target collection has %d points\n", targetPointCount)
-
-	// Guard against silent data loss: if the target has fewer points than the source had at the start,
-	// fail loudly and keep the offsets collection for investigation instead of reporting success.
-	if targetPointCount < sourcePointCount {
-		return fmt.Errorf("migration finished but the target has fewer points than the source: source had %d at start, target has %d "+
-			"(this can also happen if points were deleted from the source during the migration), offsets collection %q was kept",
-			sourcePointCount, targetPointCount, r.Migration.OffsetsCollection)
+	if err := verifyTargetCount(ctx, targetClient, r.Target.Collection, sourcePointCount, r.Migration.OffsetsCollection); err != nil {
+		return err
 	}
 
 	err = commons.DeleteOffsetsCollection(ctx, r.Migration.OffsetsCollection, targetClient)
@@ -490,46 +476,23 @@ func (r *MigrateFromQdrantCmd) migrateDataSequential(ctx context.Context, source
 	return nil
 }
 
-// resolveRangeBoundaries returns the sorted point IDs that split the source collection into worker ranges.
-// The boundaries are sampled randomly, so they are persisted next to the offsets on the first run and reused
-// when resuming. With --migration.restart, fresh boundaries are sampled and persisted, replacing the old ones.
-func (r *MigrateFromQdrantCmd) resolveRangeBoundaries(ctx context.Context, sourceClient *qdrant.Client, sourceCollection string, targetClient *qdrant.Client, sourcePointCount uint64) ([]*qdrant.PointId, error) {
-	boundariesKey := fmt.Sprintf("%s-workers-%d-boundaries", sourceCollection, r.NumWorkers)
-
-	if !r.Migration.Restart {
-		stored, err := commons.GetBoundaries(ctx, r.Migration.OffsetsCollection, targetClient, boundariesKey)
-		if err != nil {
-			return nil, err
-		}
-		if len(stored) > 0 {
-			pterm.Info.Printfln("Reusing %d stored range boundaries from the previous run", len(stored))
-			return stored, nil
-		}
-	}
-
-	ids, err := r.samplePointIDs(ctx, sourceClient, sourceCollection, sourcePointCount)
-	if err != nil {
-		return nil, err
-	}
-	if len(ids) == 0 {
-		return nil, fmt.Errorf("could not sample any point IDs from source collection %q", sourceCollection)
-	}
-	if err := commons.StoreBoundaries(ctx, r.Migration.OffsetsCollection, targetClient, boundariesKey, ids); err != nil {
-		return nil, err
-	}
-	return ids, nil
-}
-
 // migrateDataParallel performs the migration using multiple workers in parallel.
 func (r *MigrateFromQdrantCmd) migrateDataParallel(ctx context.Context, sourceClient *qdrant.Client, sourceCollection string, targetClient *qdrant.Client, targetCollection string, sourcePointCount uint64) error {
 	pterm.Info.Printfln("Using parallel migration with %d workers", r.NumWorkers)
 
 	// Resolve the range boundaries. On resume, the boundaries persisted by the previous run are reused.
 	// Re-sampling would produce different ranges and make the stored per-range offsets point at the wrong places.
-	ids, err := r.resolveRangeBoundaries(ctx, sourceClient, sourceCollection, targetClient, sourcePointCount)
+	ids, err := resolveRangeBoundaries(ctx, targetClient, r.Migration,
+		fmt.Sprintf("%s-workers-%d-boundaries", sourceCollection, r.NumWorkers),
+		func() ([]*qdrant.PointId, error) {
+			return r.samplePointIDs(ctx, sourceClient, sourceCollection, sourcePointCount)
+		})
 	if err != nil {
 		return err
 	}
+	if len(ids) == 0 {
+		return fmt.Errorf("could not sample any point IDs from source collection %q", sourceCollection)
+	}
 	fingerprint, err := commons.BoundariesFingerprint(ids)
 	if err != nil {
 		return err

cmd/utils.go

@@ -16,6 +16,8 @@ import (
 	"google.golang.org/grpc/keepalive"
 
 	"github.com/qdrant/go-client/qdrant"
+
+	"github.com/qdrant/migration/pkg/commons"
 )
 
 const (
@@ -213,3 +215,46 @@ func upsertWithRetry(ctx context.Context, client *qdrant.Client, req *qdrant.Ups
 	}
 	return nil
 }
+
+// resolveRangeBoundaries reuses the boundaries stored by an interrupted run, or samples and stores new ones.
+func resolveRangeBoundaries(ctx context.Context, targetClient *qdrant.Client, migration commons.MigrationConfig, boundariesKey string, sample func() ([]*qdrant.PointId, error)) ([]*qdrant.PointId, error) {
+	if !migration.Restart {
+		stored, err := commons.GetBoundaries(ctx, migration.OffsetsCollection, targetClient, boundariesKey)
+		if err != nil {
+			return nil, err
+		}
+		if len(stored) > 0 {
+			pterm.Info.Printfln("Reusing %d stored range boundaries from the previous run", len(stored))
+			return stored, nil
+		}
+	}
+
+	ids, err := sample()
+	if err != nil || len(ids) == 0 {
+		return ids, err
+	}
+	if err := commons.StoreBoundaries(ctx, migration.OffsetsCollection, targetClient, boundariesKey, ids); err != nil {
+		return nil, err
+	}
+	return ids, nil
+}
+
+// verifyTargetCount fails if the target has fewer points than the source had, keeping the offsets for a resume.
+func verifyTargetCount(ctx context.Context, targetClient *qdrant.Client, targetCollection string, sourcePointCount uint64, offsetsCollection string) error {
+	targetPointCount, err := targetClient.Count(ctx, &qdrant.CountPoints{
+		CollectionName: targetCollection,
+		Exact:          qdrant.PtrOf(true),
+	})
+	if err != nil {
+		return fmt.Errorf("failed to count points in target: %w", err)
+	}
+
+	pterm.Info.Printfln("Target collection has %d points\n", targetPointCount)
+
+	if targetPointCount < sourcePointCount {
+		return fmt.Errorf("migration finished but the target has fewer points than the source: source had %d at start, target has %d "+
+			"(this can also happen if records were deleted from the source during the migration, or the source has missing or non-unique IDs), offsets collection %q was kept",
+			sourcePointCount, targetPointCount, offsetsCollection)
+	}
+	return nil
+}