feat(crds): add shard field to stage spec (#1584)

Signed-off-by: Kent Rancourt <kent.rancourt@gmail.com>
akuity · Mar 7, 2024 · 74bb750 · 74bb750
1 parent 34bcb82
commit 74bb750
Show file tree

Hide file tree

Showing 8 changed files with 349 additions and 250 deletions.
diff --git a/api/v1alpha1/stage_types.go b/api/v1alpha1/stage_types.go
@@ -138,6 +138,15 @@ func (s *Stage) GetStatus() *StageStatus {
 // StageSpec describes the sources of Freight used by a Stage and how to
 // incorporate Freight into the Stage.
 type StageSpec struct {
+	// Shard is the name of the shard that this Stage belongs to. This is an
+	// optional field. If not specified, the Stage will belong to the default
+	// shard. A defaulting webhook will sync this field with the value of the
+	// kargo.akuity.io/shard label. When the shard label is not present or differs
+	// from the value of this field, the defaulting webhook will set the label to
+	// the value of this field. If the shard label is present and this field is
+	// empty, the defaulting webhook will set the value of this field to the value
+	// of the shard label.
+	Shard string `json:"shard,omitempty"`
 	// Subscriptions describes the Stage's sources of Freight. This is a required
 	// field.
 	//

diff --git a/api/v1alpha1/types.proto b/api/v1alpha1/types.proto
@@ -232,6 +232,7 @@ message StageSpec {
   Subscriptions subscriptions = 1 [json_name = "subscriptions"];
   PromotionMechanisms promotion_mechanisms = 2 [json_name = "promotionMechanisms"];
   optional Verification verification = 3 [json_name = "verification"];
+  string shard = 4 [json_name = "shard"];
 }
 
 message Freight {

diff --git a/charts/kargo/crds/kargo.akuity.io_stages.yaml b/charts/kargo/crds/kargo.akuity.io_stages.yaml
@@ -454,6 +454,17 @@ spec:
                       type: object
                     type: array
                 type: object
+              shard:
+                description: |-
+                  Shard is the name of the shard that this Stage belongs to. This is an
+                  optional field. If not specified, the Stage will belong to the default
+                  shard. A defaulting webhook will sync this field with the value of the
+                  kargo.akuity.io/shard label. When the shard label is not present or differs
+                  from the value of this field, the defaulting webhook will set the label to
+                  the value of this field. If the shard label is present and this field is
+                  empty, the defaulting webhook will set the value of this field to the value
+                  of the shard label.
+                type: string
               subscriptions:
                 description: |-
                   Subscriptions describes the Stage's sources of Freight. This is a required

diff --git a/internal/webhook/stage/webhook.go b/internal/webhook/stage/webhook.go
@@ -62,7 +62,19 @@ func newWebhook(kubeClient client.Client) *webhook {
 
 func (w *webhook) Default(_ context.Context, obj runtime.Object) error {
 	stage := obj.(*kargoapi.Stage) // nolint: forcetypeassert
+
+	// Sync the convenience shard field with the shard label
+	if stage.Spec.Shard != "" {
+		if stage.Labels == nil {
+			stage.Labels = make(map[string]string, 1)
+		}
+		stage.Labels[kargoapi.ShardLabelKey] = stage.Spec.Shard
+	} else if stage.Labels != nil && stage.Labels[kargoapi.ShardLabelKey] != "" {
+		stage.Spec.Shard = stage.Labels[kargoapi.ShardLabelKey]
+	}
+
 	controllerutil.AddFinalizer(stage, kargoapi.FinalizerName)
+
 	return nil
 }
 

diff --git a/internal/webhook/stage/webhook_test.go b/internal/webhook/stage/webhook_test.go
@@ -6,6 +6,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/util/validation/field"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -26,14 +27,60 @@ func TestNewWebhook(t *testing.T) {
 }
 
 func TestDefault(t *testing.T) {
-	stage := &kargoapi.Stage{}
+	const testShardName = "fake-shard"
+
 	w := &webhook{}
-	err := w.Default(context.Background(), stage)
-	require.NoError(t, err)
-	require.True(
-		t,
-		controllerutil.ContainsFinalizer(stage, kargoapi.FinalizerName),
-	)
+
+	t.Run("shard stays default when not specified at all", func(t *testing.T) {
+		stage := &kargoapi.Stage{
+			Spec: &kargoapi.StageSpec{},
+		}
+		err := w.Default(context.Background(), stage)
+		require.NoError(t, err)
+		require.Empty(t, stage.Labels)
+		require.Empty(t, stage.Spec.Shard)
+	})
+
+	t.Run("sync shard label to shard field", func(t *testing.T) {
+		stage := &kargoapi.Stage{
+			Spec: &kargoapi.StageSpec{
+				Shard: testShardName,
+			},
+		}
+		err := w.Default(context.Background(), stage)
+		require.NoError(t, err)
+		require.Equal(t, testShardName, stage.Labels[kargoapi.ShardLabelKey])
+		require.Equal(t, testShardName, stage.Spec.Shard)
+	})
+
+	t.Run("sync shard field to shard label", func(t *testing.T) {
+		stage := &kargoapi.Stage{
+			ObjectMeta: metav1.ObjectMeta{
+				Labels: map[string]string{
+					kargoapi.ShardLabelKey: testShardName,
+				},
+			},
+			Spec: &kargoapi.StageSpec{
+				Shard: testShardName,
+			},
+		}
+		err := w.Default(context.Background(), stage)
+		require.NoError(t, err)
+		require.Equal(t, testShardName, stage.Labels[kargoapi.ShardLabelKey])
+		require.Equal(t, testShardName, stage.Spec.Shard)
+	})
+
+	t.Run("finalizer gets added", func(t *testing.T) {
+		stage := &kargoapi.Stage{
+			Spec: &kargoapi.StageSpec{},
+		}
+		err := w.Default(context.Background(), stage)
+		require.NoError(t, err)
+		require.True(
+			t,
+			controllerutil.ContainsFinalizer(stage, kargoapi.FinalizerName),
+		)
+	})
 }
 
 func TestValidateCreate(t *testing.T) {