fix: nkode round 4 — rollback safety, cleanup precision, CI-safe rollback

Findings addressed:
- Rollback gates on backup phase: refuses incomplete (phase='dump') backups
  unless --force is passed
- Rollback requires --index or --yes when multiple indexes detected (no
  interactive input() prompt that blocks CI)
- Backup cleanup uses exact .header/.data extensions and boundary check
  (prevents deleting unrelated files with shared prefix)
- Async/sync direct quantize counts len(converted) not len(batch_keys)

Tests added:
- test_rollback_skips_incomplete_backup_phase
- test_rollback_index_filter
- test_rollback_multi_index_requires_flag
- test_cleanup_only_removes_known_extensions
- test_cleanup_does_not_match_similar_prefix

Author: nkanu17

redisvl/cli/migrate.py

@@ -369,12 +369,15 @@ def estimate(self):
         disk_estimate = estimate_disk_space(plan, aof_enabled=args.aof_enabled)
         print(disk_estimate.summary())
 
+    # Phases that indicate a safe/complete backup for rollback
+    _SAFE_ROLLBACK_PHASES = frozenset({"ready", "active", "completed"})
+
     def rollback(self):
         """Restore original vectors from a backup directory (undo quantization)."""
         parser = argparse.ArgumentParser(
             usage=(
                 "rvl migrate rollback --backup-dir <dir> "
-                "[--index <name>] [--url <redis_url>]"
+                "[--index <name>] [--yes] [--force] [--url <redis_url>]"
             )
         )
         parser.add_argument(
@@ -389,6 +392,21 @@ def rollback(self):
             help="Only restore backups for this index name (filters by backup header)",
             default=None,
         )
+        parser.add_argument(
+            "--yes",
+            "-y",
+            dest="yes",
+            action="store_true",
+            help="Skip confirmation prompt for multi-index rollback",
+            default=False,
+        )
+        parser.add_argument(
+            "--force",
+            dest="force",
+            action="store_true",
+            help="Proceed even if backup phase indicates incomplete dump",
+            default=False,
+        )
         parser = add_redis_connection_options(parser)
         args = parser.parse_args(sys.argv[3:])
 
@@ -412,7 +430,7 @@ def rollback(self):
         # Derive backup base paths (strip .header suffix)
         backup_paths = [str(h.with_suffix("")) for h in header_files]
 
-        # Load and filter backups
+        # Load, filter, and validate backups
         backups_to_restore = []
         for bp in backup_paths:
             backup = VectorBackup.load(bp)
@@ -425,25 +443,36 @@ def rollback(self):
                     f"index '{backup.header.index_name}' != '{args.index_name}'"
                 )
                 continue
+            # Gate on backup phase — refuse incomplete backups unless --force
+            if backup.header.phase not in self._SAFE_ROLLBACK_PHASES:
+                if args.force:
+                    print(
+                        f"  Warning: {os.path.basename(bp)} has phase "
+                        f"'{backup.header.phase}' (incomplete dump) — "
+                        f"proceeding due to --force"
+                    )
+                else:
+                    print(
+                        f"  Skipping {os.path.basename(bp)}: backup phase "
+                        f"'{backup.header.phase}' indicates incomplete dump. "
+                        f"Use --force to restore from partial backups."
+                    )
+                    continue
             backups_to_restore.append((bp, backup))
 
         if not backups_to_restore:
             print("Error: no matching backup files found")
             sys.exit(1)
 
-        # Warn if multiple distinct indexes detected without --index filter
+        # Require --index or --yes when multiple distinct indexes detected
         distinct_indexes = {b.header.index_name for _, b in backups_to_restore}
-        if len(distinct_indexes) > 1 and not args.index_name:
+        if len(distinct_indexes) > 1 and not args.index_name and not args.yes:
             print(
-                f"Warning: found backups for {len(distinct_indexes)} distinct indexes: "
+                f"Error: found backups for {len(distinct_indexes)} distinct indexes: "
                 f"{', '.join(sorted(distinct_indexes))}. "
-                f"Use --index to filter, or press Enter to continue."
+                f"Use --index to filter or --yes to restore all."
             )
-            try:
-                input()
-            except (EOFError, KeyboardInterrupt):
-                print("\nAborted.")
-                sys.exit(1)
+            sys.exit(1)
 
         client = RedisConnectionFactory.get_redis_connection(redis_url=redis_url)
         total_restored = 0

redisvl/migration/async_executor.py

@@ -1036,17 +1036,33 @@ def _index_progress(indexed: int, total: int, pct: float) -> None:
         return report
 
     def _cleanup_backup_files(self, backup_dir: str, index_name: str) -> None:
-        """Remove backup files after successful migration."""
-        import glob
+        """Remove backup files after successful migration.
 
+        Only removes files with the exact extensions produced by VectorBackup
+        (.header and .data), avoiding accidental deletion of unrelated files
+        that happen to share the same prefix.
+        """
         safe_name = index_name.replace("/", "_").replace("\\", "_").replace(":", "_")
-        pattern = str(Path(backup_dir) / f"migration_backup_{safe_name}*")
-        for f in glob.glob(pattern):
+        base_prefix = f"migration_backup_{safe_name}"
+        known_suffixes = (".header", ".data")
+        backup_dir_path = Path(backup_dir)
+
+        for entry in backup_dir_path.iterdir():
+            if not entry.is_file():
+                continue
+            name = entry.name
+            if not name.startswith(base_prefix):
+                continue
+            if not any(name.endswith(s) for s in known_suffixes):
+                continue
+            remainder = name[len(base_prefix) :]
+            if remainder and remainder[0] not in (".", "_"):
+                continue
             try:
-                Path(f).unlink()
-                logger.debug("Removed backup file: %s", f)
+                entry.unlink()
+                logger.debug("Removed backup file: %s", entry)
             except OSError as e:
-                logger.warning("Failed to remove backup file %s: %s", f, e)
+                logger.warning("Failed to remove backup file %s: %s", entry, e)
 
     # ------------------------------------------------------------------
     # Two-phase quantization: dump originals → convert from backup

redisvl/migration/executor.py

@@ -1126,17 +1126,40 @@ def _index_progress(indexed: int, total: int, pct: float) -> None:
         return report
 
     def _cleanup_backup_files(self, backup_dir: str, index_name: str) -> None:
-        """Remove backup files after successful migration."""
-        import glob
+        """Remove backup files after successful migration.
 
+        Only removes files with the exact extensions produced by VectorBackup
+        (.header and .data), avoiding accidental deletion of unrelated files
+        that happen to share the same prefix.
+        """
         safe_name = index_name.replace("/", "_").replace("\\", "_").replace(":", "_")
-        pattern = str(Path(backup_dir) / f"migration_backup_{safe_name}*")
-        for f in glob.glob(pattern):
+        base_prefix = f"migration_backup_{safe_name}"
+        # Exact suffixes written by VectorBackup
+        known_suffixes = (".header", ".data")
+        backup_dir_path = Path(backup_dir)
+
+        for entry in backup_dir_path.iterdir():
+            if not entry.is_file():
+                continue
+            name = entry.name
+            # Match: base_prefix exactly, or base_prefix + shard suffix
+            # e.g., migration_backup_myidx.header
+            #        migration_backup_myidx_shard_0.header
+            if not name.startswith(base_prefix):
+                continue
+            # Check that the file ends with a known extension
+            if not any(name.endswith(s) for s in known_suffixes):
+                continue
+            # Verify the character after the prefix is either a dot or underscore
+            # (prevents matching migration_backup_myidx2.header)
+            remainder = name[len(base_prefix) :]
+            if remainder and remainder[0] not in (".", "_"):
+                continue
             try:
-                Path(f).unlink()
-                logger.debug("Removed backup file: %s", f)
+                entry.unlink()
+                logger.debug("Removed backup file: %s", entry)
             except OSError as e:
-                logger.warning("Failed to remove backup file %s: %s", f, e)
+                logger.warning("Failed to remove backup file %s: %s", entry, e)
 
     # ------------------------------------------------------------------
     # Two-phase quantization: dump originals → convert from backup

tests/unit/test_vector_backup.py

@@ -428,3 +428,122 @@ def test_rollback_unloadable_backup_returns_none(self, tmp_path):
         bp = str(tmp_path / "bad_backup")
         result = VectorBackup.load(bp)
         assert result is None
+
+    def test_rollback_skips_incomplete_backup_phase(self, tmp_path):
+        """Backups in 'dump' phase should be skipped without --force."""
+        from redisvl.migration.backup import VectorBackup
+
+        bp = str(tmp_path / "migration_backup_partial")
+        backup = VectorBackup.create(
+            path=bp,
+            index_name="partial_idx",
+            fields={"embedding": {"source": "float32", "target": "float16", "dims": 4}},
+            batch_size=1,
+        )
+        # Write one batch but don't mark dump complete — phase stays "dump"
+        backup.write_batch(0, ["doc:0"], {"doc:0": {"embedding": b"\x00" * 16}})
+        # Phase is "dump" — not in safe rollback phases
+        assert backup.header.phase == "dump"
+        safe_phases = frozenset({"ready", "active", "completed"})
+        assert backup.header.phase not in safe_phases
+
+    def test_rollback_index_filter(self, tmp_path):
+        """--index filter should match only backups for the specified index."""
+        self._create_backup_with_data(tmp_path, name="idx_a")
+        self._create_backup_with_data(tmp_path, name="idx_b")
+
+        from pathlib import Path
+
+        from redisvl.migration.backup import VectorBackup
+
+        header_files = sorted(Path(tmp_path).glob("*.header"))
+        assert len(header_files) == 2
+
+        # Filter for idx_a only
+        backup_paths = [str(h.with_suffix("")) for h in header_files]
+        filtered = []
+        for bp in backup_paths:
+            backup = VectorBackup.load(bp)
+            if backup and backup.header.index_name == "idx_a":
+                filtered.append(bp)
+        assert len(filtered) == 1
+        assert "idx_a" in filtered[0]
+
+    def test_rollback_multi_index_requires_flag(self, tmp_path):
+        """Multiple distinct indexes should require --index or --yes."""
+        self._create_backup_with_data(tmp_path, name="idx_a")
+        self._create_backup_with_data(tmp_path, name="idx_b")
+
+        from pathlib import Path
+
+        from redisvl.migration.backup import VectorBackup
+
+        header_files = sorted(Path(tmp_path).glob("*.header"))
+        backup_paths = [str(h.with_suffix("")) for h in header_files]
+        backups = []
+        for bp in backup_paths:
+            backup = VectorBackup.load(bp)
+            if backup:
+                backups.append(backup)
+        distinct = {b.header.index_name for b in backups}
+        assert len(distinct) > 1  # Multi-index — should require --index or --yes
+
+
+class TestBackupCleanup:
+    """Tests for tightened backup file cleanup."""
+
+    def test_cleanup_only_removes_known_extensions(self, tmp_path):
+        """Cleanup should only remove .header and .data files."""
+        # Create files with various extensions
+        (tmp_path / "migration_backup_test.header").touch()
+        (tmp_path / "migration_backup_test.data").touch()
+        (tmp_path / "migration_backup_test.log").touch()  # should NOT be deleted
+        (tmp_path / "migration_backup_test_shard_0.header").touch()
+        (tmp_path / "migration_backup_test_shard_0.data").touch()
+        (tmp_path / "unrelated_file.txt").touch()  # should NOT be deleted
+
+        # Simulate the cleanup logic
+        base_prefix = "migration_backup_test"
+        known_suffixes = (".header", ".data")
+        deleted = []
+        for entry in tmp_path.iterdir():
+            if not entry.is_file():
+                continue
+            name = entry.name
+            if not name.startswith(base_prefix):
+                continue
+            if not any(name.endswith(s) for s in known_suffixes):
+                continue
+            remainder = name[len(base_prefix) :]
+            if remainder and remainder[0] not in (".", "_"):
+                continue
+            deleted.append(name)
+
+        assert "migration_backup_test.header" in deleted
+        assert "migration_backup_test.data" in deleted
+        assert "migration_backup_test_shard_0.header" in deleted
+        assert "migration_backup_test_shard_0.data" in deleted
+        assert "migration_backup_test.log" not in deleted
+        assert "unrelated_file.txt" not in deleted
+
+    def test_cleanup_does_not_match_similar_prefix(self, tmp_path):
+        """migration_backup_foo should not match migration_backup_foobar."""
+        (tmp_path / "migration_backup_foo.header").touch()
+        (tmp_path / "migration_backup_foobar.header").touch()
+
+        base_prefix = "migration_backup_foo"
+        known_suffixes = (".header", ".data")
+        deleted = []
+        for entry in tmp_path.iterdir():
+            name = entry.name
+            if not name.startswith(base_prefix):
+                continue
+            if not any(name.endswith(s) for s in known_suffixes):
+                continue
+            remainder = name[len(base_prefix) :]
+            if remainder and remainder[0] not in (".", "_"):
+                continue
+            deleted.append(name)
+
+        assert "migration_backup_foo.header" in deleted
+        assert "migration_backup_foobar.header" not in deleted