Implement db restore

Author: AmrSaber

README.md

@@ -462,6 +462,11 @@ cat backup.db | kv db import -
 # Or simply:
 kv db import < backup.db
 
+# Restore from automatic backup (created during import)
+kv db restore
+# Output: Database restored from: /path/to/kv.db.backup
+# Note: This restores the backup created during the last import operation
+
 # Practical examples:
 
 # Create daily backups
@@ -487,9 +492,12 @@ gunzip -c kv-backup.db.gz | kv db import -
 
 **Safety features:**
 - Import automatically creates a backup at `<db-path>.backup` before replacing
+- You can restore from this backup anytime using `kv db restore`
 - Import validates the file is a valid database before proceeding
 - Export fails if file exists (use `--force` to override)
 - Export validates destination directory exists
+- Restore validates the backup is a valid database before restoring
+- Restore preserves the backup file after restoration
 
 ### Utility Commands
 

src/cmd/export.go

@@ -19,7 +19,9 @@ var exportCmd = &cobra.Command{
 	Long: `Export the entire database to a binary file.
 
 The exported file captures the complete state of the database and can be
-imported later using the 'import' command to restore the exact state.
+imported later using the 'kv db import' command to restore the exact state.
+When importing, a backup of the current database is automatically created,
+which can be restored using 'kv db restore' if needed.
 
 Use "-" as the file path to write to stdout (useful for piping).`,
 	Example: `  # Export database to a file

src/cmd/import.go

@@ -20,7 +20,8 @@ The imported file must be a valid database file created with the 'export' comman
 This will completely replace the current database with the imported one.
 
 WARNING: This operation is destructive. The current database will be backed up
-to <db-path>.backup before importing.
+to <db-path>.backup before importing. You can restore from this backup using
+the 'kv db restore' command if needed.
 
 Use "-" as the file path to read from stdin (useful for piping).`,
 	Example: `  # Import database from a file
@@ -87,7 +88,8 @@ Use "-" as the file path to read from stdin (useful for piping).`,
 			if err != nil {
 				common.Fail("Failed to backup current database: %v", err)
 			}
-			fmt.Printf("Current database backed up to: %s\n", backupPath)
+
+			fmt.Print("Current database backed up")
 		}
 
 		// Remove WAL files (should be empty after VACUUM, but remove just in case)

src/cmd/restore.go

@@ -0,0 +1,61 @@
+package cmd
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/AmrSaber/kv/src/common"
+	"github.com/spf13/cobra"
+)
+
+// restoreCmd represents the restore command
+var restoreCmd = &cobra.Command{
+	Use:   "restore",
+	Short: "Restore database from automatic backup",
+	Long: `Restore the database from the automatic backup file created during import.
+
+When you run 'kv db import', a backup of the current database is automatically
+created at <db-path>.backup. This command restores that backup.
+
+The backup file is preserved after restoration, allowing you to restore again
+if needed.`,
+	Args: cobra.NoArgs,
+	Run: func(cmd *cobra.Command, args []string) {
+		// Get database path
+		dbPath := common.GetDBPath()
+		backupPath := dbPath + ".backup"
+
+		// Check if backup exists
+		if _, err := os.Stat(backupPath); os.IsNotExist(err) {
+			common.Fail("No backup file found")
+		}
+
+		// Validate backup is a valid SQLite database
+		if err := common.ValidateSqliteFile(backupPath); err != nil {
+			common.Fail("Invalid backup file: %v", err)
+		}
+
+		// Close database connection
+		common.CloseDB()
+
+		// Remove current database and remove WAL files
+		_ = os.Remove(dbPath)
+		_ = os.Remove(dbPath + "-wal")
+		_ = os.Remove(dbPath + "-shm")
+
+		// Copy backup to current location
+		err := common.CopyFile(backupPath, dbPath)
+		if err != nil {
+			common.Fail("Failed to restore from backup: %v", err)
+		}
+
+		// Reopen database (migrations will run automatically)
+		common.GetDB()
+
+		fmt.Printf("Database restored from backup")
+	},
+}
+
+func init() {
+	dbCmd.AddCommand(restoreCmd)
+}

tests/export_import_test.go

@@ -274,3 +274,123 @@ func TestExportImportRoundTrip(t *testing.T) {
 		t.Error("History should contain all versions")
 	}
 }
+
+func TestRestoreCommand(t *testing.T) {
+	cleanup := SetupTestDB(t)
+	defer cleanup()
+
+	t.Run("restore fails when no backup exists", func(t *testing.T) {
+		output := RunKVFailure(t, "db", "restore")
+		if !strings.Contains(output, "No backup file found") {
+			t.Errorf("Expected 'No backup file found' error, got: %s", output)
+		}
+	})
+
+	// Setup test data and create a backup via import
+	RunKVSuccess(t, "set", "original1", "value1")
+	RunKVSuccess(t, "set", "original2", "value2", "--password", "pass")
+	RunKVSuccess(t, "set", "original3", "value3")
+	RunKVSuccess(t, "hide", "original3")
+
+	tmpDir := t.TempDir()
+	exportPath := filepath.Join(tmpDir, "backup.db")
+	RunKVSuccess(t, "db", "export", exportPath)
+
+	// Import to create a backup
+	RunKVSuccess(t, "db", "import", exportPath)
+
+	t.Run("restore fails with invalid backup file", func(t *testing.T) {
+		// Get DB path and backup path
+		dbPath := common.GetDBPath()
+		backupPath := dbPath + ".backup"
+
+		// Save current backup
+		validBackupPath := backupPath + ".valid"
+		err := common.CopyFile(backupPath, validBackupPath)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer func() {
+			// Restore valid backup
+			_ = os.Remove(backupPath)
+			_ = common.CopyFile(validBackupPath, backupPath)
+			_ = os.Remove(validBackupPath)
+		}()
+
+		// Replace backup with invalid content
+		err = os.WriteFile(backupPath, []byte("not a database"), 0o644)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		output := RunKVFailure(t, "db", "restore")
+		if !strings.Contains(output, "Invalid backup") {
+			t.Errorf("Expected 'Invalid backup' error, got: %s", output)
+		}
+	})
+
+	t.Run("restore successfully restores from backup", func(t *testing.T) {
+		// Modify database
+		RunKVSuccess(t, "set", "new-key", "new-value")
+		RunKVSuccess(t, "delete", "original1")
+
+		// Verify changes
+		output := RunKVSuccess(t, "get", "new-key")
+		if output != "new-value" {
+			t.Error("new-key should exist before restore")
+		}
+
+		RunKVFailure(t, "get", "original1")
+
+		// Restore
+		output = RunKVSuccess(t, "db", "restore")
+		if !strings.Contains(output, "restored from") {
+			t.Errorf("Expected success message, got: %s", output)
+		}
+
+		// Verify restoration
+		output = RunKVSuccess(t, "get", "original1")
+		if output != "value1" {
+			t.Errorf("Expected 'value1', got: %s", output)
+		}
+
+		output = RunKVFailure(t, "get", "new-key")
+		if !strings.Contains(output, "does not exist") {
+			t.Error("new-key should not exist after restore")
+		}
+	})
+
+	t.Run("restore preserves backup file", func(t *testing.T) {
+		dbPath := common.GetDBPath()
+		backupPath := dbPath + ".backup"
+
+		// Verify backup still exists
+		if _, err := os.Stat(backupPath); os.IsNotExist(err) {
+			t.Error("Backup file should still exist after restore")
+		}
+
+		// Verify it's still valid
+		if err := common.ValidateSqliteFile(backupPath); err != nil {
+			t.Errorf("Backup file should still be valid: %v", err)
+		}
+	})
+
+	t.Run("restore preserves all data attributes", func(t *testing.T) {
+		// Verify locked key
+		output := RunKVSuccess(t, "list", "original2")
+		if !strings.Contains(output, "[Locked]") {
+			t.Error("original2 should be locked")
+		}
+
+		output = RunKVSuccess(t, "get", "original2", "--password", "pass")
+		if output != "value2" {
+			t.Errorf("Expected 'value2', got: %s", output)
+		}
+
+		// Verify hidden key
+		output = RunKVSuccess(t, "list", "original3")
+		if !strings.Contains(output, "[Hidden]") {
+			t.Error("original3 should be hidden")
+		}
+	})
+}