From 6252f8051ce65ba80af00e1e3a369dda53682e7f Mon Sep 17 00:00:00 2001
From: Ljubisa Gacevic <ljubisa.rs@gmail.com>
Date: Thu, 23 Oct 2025 16:20:03 +0200
Subject: [PATCH] feat(nuke): update statefulset image with new flag

---
 cmd/beekeeper/cmd/nuke.go | 3 +++
 pkg/nuker/nuker.go        | 6 ++++++
 2 files changed, 9 insertions(+)

diff --git a/cmd/beekeeper/cmd/nuke.go b/cmd/beekeeper/cmd/nuke.go
index f5446c42..50a17991 100644
--- a/cmd/beekeeper/cmd/nuke.go
+++ b/cmd/beekeeper/cmd/nuke.go
@@ -15,6 +15,7 @@ const (
 	optionNameUseRandomNeighboorhood = "use-random-neighborhood"
 	optionNameDeploymentType         = "deployment-type"
 	optionNameStatefulSets           = "stateful-sets"
+	optionNameImage                  = "image"
 	beeLabelSelector                 = "app.kubernetes.io/name=bee"
 )
 
@@ -47,6 +48,7 @@ The command supports two modes:
 					K8sClient:             c.k8sClient,
 					NodeProvider:          nodeClient,
 					UseRandomNeighborhood: c.globalConfig.GetBool(optionNameUseRandomNeighboorhood),
+					Image:                 c.globalConfig.GetString(optionNameImage),
 				})
 
 				statefulSetNames := c.globalConfig.GetStringSlice(optionNameStatefulSets)
@@ -79,6 +81,7 @@ The command supports two modes:
 	cmd.Flags().Bool(optionNameUseRandomNeighboorhood, false, "Use random neighborhood for Bee nodes (default: false)")
 	cmd.Flags().String(optionNameDeploymentType, string(node.DeploymentTypeBeekeeper), "Indicates how the cluster was deployed: 'beekeeper' or 'helm'.")
 	cmd.Flags().StringSlice(optionNameStatefulSets, nil, "List of StatefulSet names to target for nuke (e.g., 'bootnode-0,bootnode-1'). When provided, uses direct StatefulSet targeting instead of NodeProvider.")
+	cmd.Flags().String(optionNameImage, "", "Container image to use when restarting pods (defaults to current image if not set).")
 
 	c.root.AddCommand(cmd)
 
diff --git a/pkg/nuker/nuker.go b/pkg/nuker/nuker.go
index 5019e6fb..331886b2 100644
--- a/pkg/nuker/nuker.go
+++ b/pkg/nuker/nuker.go
@@ -20,6 +20,7 @@ type ClientConfig struct {
 	NodeProvider          node.NodeProvider
 	K8sClient             *k8s.Client
 	UseRandomNeighborhood bool
+	Image                 string
 }
 
 type Client struct {
@@ -27,6 +28,7 @@ type Client struct {
 	nodeProvider          node.NodeProvider
 	k8sClient             *k8s.Client
 	useRandomNeighborhood bool
+	image                 string
 }
 
 func New(cfg *ClientConfig) *Client {
@@ -42,6 +44,7 @@ func New(cfg *ClientConfig) *Client {
 		log:          cfg.Log,
 		k8sClient:    cfg.K8sClient,
 		nodeProvider: cfg.NodeProvider,
+		image:        cfg.Image,
 	}
 }
 
@@ -197,6 +200,9 @@ func (c *Client) updateAndRollbackStatefulSet(ctx context.Context, namespace str
 			latestSS.Spec.Replicas = originalSS.Spec.Replicas
 			latestSS.Spec.Template.Spec.Containers[0].Command = restartArgs
 			latestSS.Spec.Template.Spec.Containers[0].ReadinessProbe = originalSS.Spec.Template.Spec.Containers[0].ReadinessProbe
+			if c.image != "" {
+				latestSS.Spec.Template.Spec.Containers[0].Image = c.image
+			}
 
 			return c.k8sClient.StatefulSet.Update(ctx, namespace, latestSS)
 		}); err != nil {