fix: handle aws integration attach/detach on stack update (#47)

internal/logging/keys.go

@@ -6,8 +6,9 @@ const (
 	RunId    = "run.id"
 	RunState = "run.state"
 
-	StackName = "stack.name"
-	StackId   = "stack.id"
+	StackName             = "stack.name"
+	StackId               = "stack.id"
+	StackAWSIntegrationId = "stack.aws_integration_id"
 
 	SpaceId   = "space.id"
 	SpaceName = "space.name"

internal/spacelift/repository/stack.go

@@ -2,12 +2,15 @@ package repository
 
 import (
 	"context"
+	"slices"
 
 	"github.com/pkg/errors"
 	"github.com/shurcooL/graphql"
 	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/log"
 
 	"github.com/spacelift-io/spacelift-operator/api/v1beta1"
+	"github.com/spacelift-io/spacelift-operator/internal/logging"
 	spaceliftclient "github.com/spacelift-io/spacelift-operator/internal/spacelift/client"
 	"github.com/spacelift-io/spacelift-operator/internal/spacelift/models"
 	"github.com/spacelift-io/spacelift-operator/internal/spacelift/repository/slug"
@@ -102,17 +105,49 @@ func (r *stackRepository) attachAWSIntegration(ctx context.Context, stack *v1bet
 		"read":  graphql.Boolean(stack.Spec.AWSIntegration.Read),
 		"write": graphql.Boolean(stack.Spec.AWSIntegration.Write),
 	}
+
 	if err := c.Mutate(ctx, &mutation, awsIntegrationAttachVars); err != nil {
 		return err
 	}
 
 	return nil
 }
 
+type awsIntegrationDetachMutation struct {
+	AWSIntegrationDetach struct {
+		ID string `graphql:"id"`
+	} `graphql:"awsIntegrationDetach(id: $id)"`
+}
+
+func (r *stackRepository) detachAWSIntegration(ctx context.Context, stack *v1beta1.Stack, id string) error {
+	c, err := spaceliftclient.DefaultClient(ctx, r.client, stack.Namespace)
+	if err != nil {
+		return errors.Wrap(err, "unable to fetch spacelift client while detaching AWS integration")
+	}
+	var mutation awsIntegrationDetachMutation
+	vars := map[string]any{
+		"id": graphql.ID(id),
+	}
+
+	if err := c.Mutate(ctx, &mutation, vars); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+type stackUpdateMutationAWSIntegration struct {
+	ID            string `graphql:"id"`
+	IntegrationID string `graphql:"integrationId"`
+	Read          bool   `graphql:"read"`
+	Write         bool   `graphql:"write"`
+}
+
 type stackUpdateMutation struct {
 	StackUpdate struct {
-		ID    string `graphql:"id"`
-		State string `graphql:"state"`
+		ID                      string                              `graphql:"id"`
+		State                   string                              `graphql:"state"`
+		AttachedAWSIntegrations []stackUpdateMutationAWSIntegration `graphql:"attachedAwsIntegrations"`
 	} `graphql:"stackUpdate(id: $id, input: $input)"`
 }
 
@@ -134,6 +169,39 @@ func (r *stackRepository) Update(ctx context.Context, stack *v1beta1.Stack) (*mo
 		return nil, errors.Wrap(err, "unable to create stack")
 	}
 
+	logger := log.FromContext(ctx).WithValues(logging.StackId, mutation.StackUpdate.ID)
+
+	// First we check if there are any integrations to detach
+	// Any existing attachment that does not match spec.AWSIntegration will be detached.
+	attachedIntegrations := mutation.StackUpdate.AttachedAWSIntegrations
+	for i, integration := range attachedIntegrations {
+		if stack.Spec.AWSIntegration == nil ||
+			stack.Spec.AWSIntegration.Id != integration.IntegrationID ||
+			(stack.Spec.AWSIntegration.Id == integration.IntegrationID &&
+				(stack.Spec.AWSIntegration.Read != integration.Read ||
+					stack.Spec.AWSIntegration.Write != integration.Write)) {
+			if err := r.detachAWSIntegration(ctx, stack, integration.ID); err != nil {
+				return nil, errors.Wrap(err, "unable to detach AWS integration from stack")
+			}
+			logger.Info("Detached AWS integration from stack", logging.StackAWSIntegrationId, integration.IntegrationID)
+			// If we are detaching an integration, we also reflect this change to the mutation array.
+			// This allows an integration to be detached and reattached in a row if a read or write attribute has been changed.
+			mutation.StackUpdate.AttachedAWSIntegrations = append(mutation.StackUpdate.AttachedAWSIntegrations[:i], mutation.StackUpdate.AttachedAWSIntegrations[i+1:]...)
+		}
+	}
+
+	if stack.Spec.AWSIntegration != nil {
+		shouldAttachInteration := !slices.ContainsFunc(mutation.StackUpdate.AttachedAWSIntegrations, func(s stackUpdateMutationAWSIntegration) bool {
+			return s.IntegrationID == stack.Spec.AWSIntegration.Id
+		})
+		if shouldAttachInteration {
+			if err := r.attachAWSIntegration(ctx, stack); err != nil {
+				return nil, errors.Wrap(err, "unable to attach AWS integration to stack")
+			}
+			logger.Info("Attached AWS integration to stack", logging.StackAWSIntegrationId, stack.Spec.AWSIntegration.Id)
+		}
+	}
+
 	// TODO(michalg): URL can never change here, should we still generate it for k8s api?
 	url := c.URL("/stack/%s", mutation.StackUpdate.ID)
 	return &models.Stack{

internal/spacelift/repository/stack_test.go

@@ -191,3 +191,151 @@ func Test_stackRepository_Update(t *testing.T) {
 	assert.Equal(t, "stack-name", actualVars["id"])
 	assert.IsType(t, structs.StackInput{}, actualVars["input"])
 }
+
+func Test_stackRepository_Update_WithAWSIntegration(t *testing.T) {
+	originalClient := spaceliftclient.DefaultClient
+	defer func() { spaceliftclient.DefaultClient = originalClient }()
+	fakeClient := mocks.NewClient(t)
+	spaceliftclient.DefaultClient = func(_ context.Context, _ client.Client, _ string) (spaceliftclient.Client, error) {
+		return fakeClient, nil
+	}
+
+	fakeStackId := "stack-id"
+	var actualVars map[string]any
+	fakeClient.EXPECT().
+		Mutate(mock.Anything, mock.AnythingOfType("*repository.stackUpdateMutation"), mock.Anything).
+		Run(func(_ context.Context, mutation any, vars map[string]interface{}, _ ...graphql.RequestOption) {
+			actualVars = vars
+			updateMutation := mutation.(*stackUpdateMutation)
+			updateMutation.StackUpdate.ID = fakeStackId
+			updateMutation.StackUpdate.AttachedAWSIntegrations = []stackUpdateMutationAWSIntegration{
+				{
+					ID:            "attachment-id",
+					IntegrationID: "another-integration-id",
+					Read:          true,
+					Write:         true,
+				},
+			}
+		}).Return(nil)
+	fakeClient.EXPECT().URL("/stack/%s", fakeStackId).Return("")
+
+	var detachVars map[string]any
+	fakeClient.EXPECT().Mutate(mock.Anything, mock.AnythingOfType("*repository.awsIntegrationDetachMutation"), mock.Anything).
+		Run(func(_ context.Context, _ any, vars map[string]any, _ ...graphql.RequestOption) {
+			detachVars = vars
+		}).
+		Return(nil)
+	var attachVars map[string]any
+	fakeClient.EXPECT().Mutate(mock.Anything, mock.AnythingOfType("*repository.awsIntegrationAttachMutation"), mock.Anything).
+		Run(func(_ context.Context, _ any, vars map[string]any, _ ...graphql.RequestOption) {
+			attachVars = vars
+		}).
+		Return(nil)
+
+	repo := NewStackRepository(nil)
+
+	fakeStack := &v1beta1.Stack{
+		ObjectMeta: v1.ObjectMeta{
+			Name: "stack-name",
+		},
+		Spec: v1beta1.StackSpec{
+			SpaceId: utils.AddressOf("space-id"),
+			AWSIntegration: &v1beta1.AWSIntegration{
+				Id:    "integration-id",
+				Read:  true,
+				Write: true,
+			},
+		},
+		Status: v1beta1.StackStatus{
+			Id: fakeStackId,
+		},
+	}
+	_, err := repo.Update(context.Background(), fakeStack)
+	require.NoError(t, err)
+	assert.Equal(t, "stack-name", actualVars["id"])
+	assert.IsType(t, structs.StackInput{}, actualVars["input"])
+	assert.Equal(t, map[string]any{
+		"id": graphql.ID("attachment-id"),
+	}, detachVars)
+	assert.Equal(t, map[string]any{
+		"id":    "integration-id",
+		"stack": fakeStackId,
+		"read":  graphql.Boolean(true),
+		"write": graphql.Boolean(true),
+	}, attachVars)
+}
+
+func Test_stackRepository_Update_WithAWSIntegration_UpdateExistingIntegration(t *testing.T) {
+	originalClient := spaceliftclient.DefaultClient
+	defer func() { spaceliftclient.DefaultClient = originalClient }()
+	fakeClient := mocks.NewClient(t)
+	spaceliftclient.DefaultClient = func(_ context.Context, _ client.Client, _ string) (spaceliftclient.Client, error) {
+		return fakeClient, nil
+	}
+
+	fakeStackId := "stack-id"
+	var actualVars map[string]any
+	fakeClient.EXPECT().
+		Mutate(mock.Anything, mock.AnythingOfType("*repository.stackUpdateMutation"), mock.Anything).
+		Run(func(_ context.Context, mutation any, vars map[string]interface{}, _ ...graphql.RequestOption) {
+			actualVars = vars
+			updateMutation := mutation.(*stackUpdateMutation)
+			updateMutation.StackUpdate.ID = fakeStackId
+			updateMutation.StackUpdate.AttachedAWSIntegrations = []stackUpdateMutationAWSIntegration{
+				{
+					ID:            "attachment-id",
+					IntegrationID: "integration-id",
+					Read:          true,
+					Write:         true,
+				},
+			}
+		}).Return(nil)
+	fakeClient.EXPECT().URL("/stack/%s", fakeStackId).Return("")
+
+	var detachVars map[string]any
+	fakeClient.EXPECT().Mutate(mock.Anything, mock.AnythingOfType("*repository.awsIntegrationDetachMutation"), mock.Anything).
+		Run(func(_ context.Context, _ any, vars map[string]any, _ ...graphql.RequestOption) {
+			detachVars = vars
+		}).
+		Return(nil)
+	var attachVars map[string]any
+	fakeClient.EXPECT().Mutate(mock.Anything, mock.AnythingOfType("*repository.awsIntegrationAttachMutation"), mock.Anything).
+		Run(func(_ context.Context, _ any, vars map[string]any, _ ...graphql.RequestOption) {
+			attachVars = vars
+		}).
+		Return(nil)
+
+	repo := NewStackRepository(nil)
+
+	fakeStack := &v1beta1.Stack{
+		ObjectMeta: v1.ObjectMeta{
+			Name: "stack-name",
+		},
+		Spec: v1beta1.StackSpec{
+			SpaceId: utils.AddressOf("space-id"),
+			// Because Write has been changed from true to false
+			// The integration should be detached and reattached
+			AWSIntegration: &v1beta1.AWSIntegration{
+				Id:    "integration-id",
+				Read:  true,
+				Write: false,
+			},
+		},
+		Status: v1beta1.StackStatus{
+			Id: fakeStackId,
+		},
+	}
+	_, err := repo.Update(context.Background(), fakeStack)
+	require.NoError(t, err)
+	assert.Equal(t, "stack-name", actualVars["id"])
+	assert.IsType(t, structs.StackInput{}, actualVars["input"])
+	assert.Equal(t, map[string]any{
+		"id": graphql.ID("attachment-id"),
+	}, detachVars)
+	assert.Equal(t, map[string]any{
+		"id":    "integration-id",
+		"stack": fakeStackId,
+		"read":  graphql.Boolean(true),
+		"write": graphql.Boolean(false),
+	}, attachVars)
+}