pkg/resource: Add migration label support to CrdbCluster CR

Makefile

@@ -373,4 +373,3 @@ publish-operator-openshift:
 .PHONY: release/publish-openshift-bundle
 release/publish-openshift-bundle:
 	./build/release/teamcity-publish-openshift-bundle.sh
-

WORKSPACE

@@ -123,7 +123,7 @@ oci_pull(
     name = "redhat_ubi_minimal",
     platforms = [
         "linux/amd64",
-        "linux/arm64/v8"
+        "linux/arm64/v8",
     ],
     registry = "registry.access.redhat.com",
     repository = "ubi9/ubi-minimal",

examples/smoketest.yaml

@@ -40,4 +40,4 @@ spec:
   tlsEnabled: true
   image:
     name: cockroachdb/cockroach:v25.2.4
-  nodes: 3
+  nodes: 3

hack/build/kustomize.bzl

@@ -57,7 +57,7 @@ def kustomization(
   # bazel build <name> to generate manifest.yaml
   native.genrule(
       name = name,
-      exec_tools = tools,
+      tools = tools,
       srcs = [":{}".format(rule_srcs), ":{}".format(rule_main)],
       outs = [out],
       visibility = visibility,

pkg/resource/statefulset.go

@@ -51,6 +51,8 @@ const (
 	DbInitContainerName = "db-init"
 
 	terminationGracePeriodSecs = 300
+
+	migrationLabel = "crdb.io/migrating"
 )
 
 type StatefulSetBuilder struct {
@@ -69,6 +71,16 @@ func (b StatefulSetBuilder) Build(obj client.Object) error {
 		ss.ObjectMeta.Name = b.StatefulSetName()
 	}
 
+	// This is used to prevent scaling up if the custom migration label is set
+	currentReplicas := ss.Spec.Replicas
+	desiredReplicas := b.Spec().Nodes
+
+	if val, ok := b.Cluster.Unwrap().Labels[migrationLabel]; ok && val == "true" {
+		if currentReplicas != nil && *currentReplicas < desiredReplicas {
+			desiredReplicas = *currentReplicas
+		}
+	}
+
 	ss.Annotations = b.Spec().AdditionalAnnotations
 
 	if ss.Annotations == nil {
@@ -78,7 +90,7 @@ func (b StatefulSetBuilder) Build(obj client.Object) error {
 	ss.Annotations[CrdbContainerImageAnnotation] = b.Cluster.GetAnnotationContainerImage()
 	ss.Spec = appsv1.StatefulSetSpec{
 		ServiceName: b.Cluster.DiscoveryServiceName(),
-		Replicas:    ptr.Int32(b.Spec().Nodes),
+		Replicas:    ptr.Int32(desiredReplicas),
 		UpdateStrategy: appsv1.StatefulSetUpdateStrategy{
 			RollingUpdate: &appsv1.RollingUpdateStatefulSetStrategy{},
 		},

pkg/resource/statefulset_test.go

@@ -22,22 +22,26 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"testing"
 
 	api "github.com/cockroachdb/cockroach-operator/apis/v1alpha1"
 	"github.com/cockroachdb/cockroach-operator/pkg/labels"
+	"github.com/cockroachdb/cockroach-operator/pkg/ptr"
 	"github.com/cockroachdb/cockroach-operator/pkg/resource"
 	"github.com/cockroachdb/cockroach-operator/pkg/testutil"
 	"github.com/cockroachdb/cockroach-operator/pkg/utilfeature"
 	"github.com/google/go-cmp/cmp"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	appsv1 "k8s.io/api/apps/v1"
-
-	"testing"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
 var update = flag.Bool("update", false, "update the golden files of this test")
 
+const migrationLabel = "crdb.io/migrating"
+
 func TestStatefulSetBuilder(t *testing.T) {
 	// turn on featuregate to test rules
 	require.NoError(t, utilfeature.DefaultMutableFeatureGate.Set("AffinityRules=true,TolerationRules=true"))
@@ -118,3 +122,97 @@ func load(t *testing.T, file string) []byte {
 
 	return content
 }
+
+func TestStatefulSetStopScaleUpWithMigrationLabel(t *testing.T) {
+	tests := []struct {
+		name             string
+		crNodes          int32
+		currentReplicas  int32
+		labels           map[string]string
+		expectedReplicas int32
+	}{
+		{
+			name:             "Normal Scale Up - No Label",
+			crNodes:          3,
+			currentReplicas:  2,
+			labels:           nil,
+			expectedReplicas: 3,
+		},
+		{
+			name:            "Stop Scale Up - Label Present",
+			crNodes:         3,
+			currentReplicas: 2,
+			labels: map[string]string{
+				migrationLabel: "true",
+			},
+			expectedReplicas: 2,
+		},
+		{
+			name:            "Allow Scale Down - Label Present",
+			crNodes:         3,
+			currentReplicas: 4,
+			labels: map[string]string{
+				migrationLabel: "true",
+			},
+			expectedReplicas: 3,
+		},
+		{
+			name:            "Stop Scale Up - False Label Value",
+			crNodes:         3,
+			currentReplicas: 2,
+			labels: map[string]string{
+				migrationLabel: "false",
+			},
+			expectedReplicas: 3,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cr := &api.CrdbCluster{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-cluster",
+					Namespace: "default",
+					Labels:    tt.labels,
+				},
+				Spec: api.CrdbClusterSpec{
+					Nodes:              tt.crNodes,
+					CockroachDBVersion: "v20.2.5",
+					Image: &api.PodImage{
+						Name: "cockroachdb/cockroach:v20.2.5",
+					},
+					DataStore: api.Volume{
+						VolumeClaim: &api.VolumeClaim{
+							PersistentVolumeClaimSpec: corev1.PersistentVolumeClaimSpec{},
+						},
+					},
+				},
+			}
+
+			cluster := resource.NewCluster(cr)
+			commonLabels := labels.Common(cr)
+
+			// Pre-populate the existing StatefulSet with current replicas
+			actual := &appsv1.StatefulSet{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-cluster",
+					Namespace: "default",
+				},
+				Spec: appsv1.StatefulSetSpec{
+					Replicas: ptr.Int32(tt.currentReplicas),
+				},
+			}
+
+			builder := resource.StatefulSetBuilder{
+				Cluster:   &cluster,
+				Selector:  commonLabels.Selector(cluster.Spec().AdditionalLabels),
+				Telemetry: "kubernetes-operator-test",
+			}
+
+			err := builder.Build(actual)
+			require.NoError(t, err)
+
+			assert.Equal(t, tt.expectedReplicas, *actual.Spec.Replicas)
+		})
+	}
+}