diff --git a/.trivyignore b/.trivyignore new file mode 100644 index 0000000..d2b4469 --- /dev/null +++ b/.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 \ No newline at end of file diff --git a/cmd/migrate_from_pg.go b/cmd/migrate_from_pg.go index 3c3f4a9..1a1d871 100644 --- a/cmd/migrate_from_pg.go +++ b/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() diff --git a/cmd/migrate_from_qdrant.go b/cmd/migrate_from_qdrant.go index f4c117b..af8d3b1 100644 --- a/cmd/migrate_from_qdrant.go +++ b/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 diff --git a/cmd/utils.go b/cmd/utils.go index e51d26c..90f10fe 100644 --- a/cmd/utils.go +++ b/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 +}