From be17ed37264ba6f5e34d53430cd24ab1b892718b Mon Sep 17 00:00:00 2001
From: ppastorf <pdrpfernandes99@gmail.com>
Date: Wed, 10 Dec 2025 14:41:51 -0300
Subject: [PATCH] feat: Add 'clean' command to remove instances on state
 ConnectionLost

---
 cmd/clean.go                 | 47 ++++++++++++++++++++++++++++++++++++
 internal/storage/instance.go | 31 ++++++++++++++++++++++++
 2 files changed, 78 insertions(+)
 create mode 100644 cmd/clean.go

diff --git a/cmd/clean.go b/cmd/clean.go
new file mode 100644
index 0000000..fc3b1d4
--- /dev/null
+++ b/cmd/clean.go
@@ -0,0 +1,47 @@
+package cmd
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/andreclaro/ssm/internal/storage"
+	"github.com/spf13/cobra"
+)
+
+// cleanCmd represents the clean command
+var cleanCmd = &cobra.Command{
+	Use:   "clean",
+	Short: "Remove instances with ConnectionLost state from the database",
+	Long: `Remove all instances from the database where the state is 'ConnectionLost'.
+
+This command is useful for cleaning up instances that have lost their connection
+and are no longer accessible.
+
+Examples:
+  ssm clean    # Remove all instances with ConnectionLost state`,
+	Run: runClean,
+}
+
+func init() {
+	rootCmd.AddCommand(cleanCmd)
+}
+
+func runClean(cmd *cobra.Command, args []string) {
+	// Initialize database
+	if err := storage.InitDB(); err != nil {
+		fmt.Fprintf(os.Stderr, "Failed to initialize database: %v\n", err)
+		os.Exit(1)
+	}
+
+	// Create repository
+	repo := storage.NewInstanceRepository()
+
+	// Delete instances with ConnectionLost state
+	count, err := repo.DeleteByState("ConnectionLost")
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Failed to clean instances: %v\n", err)
+		os.Exit(1)
+	}
+
+	fmt.Printf("Removed %d instance(s) from database\n", count)
+}
diff --git a/internal/storage/instance.go b/internal/storage/instance.go
index 961b7ec..ba7d58d 100644
--- a/internal/storage/instance.go
+++ b/internal/storage/instance.go
@@ -221,6 +221,37 @@ func (r *InstanceRepository) DeleteStale(olderThan time.Duration) error {
 	return nil
 }
 
+// DeleteByState removes instances with the specified state
+func (r *InstanceRepository) DeleteByState(state string) (int64, error) {
+	// First, get the instance IDs that will be deleted to clean up associated tags
+	var instanceIDs []string
+	if err := DB.Model(&Instance{}).Where("state = ?", state).Pluck("instance_id", &instanceIDs).Error; err != nil {
+		return 0, fmt.Errorf("failed to find instances with state %s: %w", state, err)
+	}
+
+	// Delete associated tags
+	if len(instanceIDs) > 0 {
+		if err := DB.Where("instance_id IN ?", instanceIDs).Delete(&Tag{}).Error; err != nil {
+			return 0, fmt.Errorf("failed to delete tags for instances with state %s: %w", state, err)
+		}
+	}
+
+	// Delete instances
+	result := DB.Where("state = ?", state).Delete(&Instance{})
+	if result.Error != nil {
+		return 0, fmt.Errorf("failed to delete instances with state %s: %w", state, result.Error)
+	}
+
+	if result.RowsAffected > 0 {
+		logrus.WithFields(logrus.Fields{
+			"count": result.RowsAffected,
+			"state": state,
+		}).Info("Deleted instances by state")
+	}
+
+	return result.RowsAffected, nil
+}
+
 // GetStats returns statistics about stored instances
 func (r *InstanceRepository) GetStats() (map[string]int, error) {
 	stats := make(map[string]int)