Auto merge apply with multi workspace support

pkg/mocks/helpers.go

@@ -1,6 +1,7 @@
 package mocks
 
 import (
+	"errors"
 	"fmt"
 	"path/filepath"
 	"regexp"
@@ -175,6 +176,8 @@ func (ts *TestSuite) InitTestSuite() {
 	ts.MockApiClient.EXPECT().AddTags(gomock.Any(), gomock.Any(), "tfbuddylock", "101").AnyTimes()
 
 	ts.MockStreamClient.EXPECT().AddRunMeta(gomock.Any()).AnyTimes()
+	ts.MockStreamClient.EXPECT().AddWorkspaceMeta(gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes()
+	ts.MockStreamClient.EXPECT().GetWorkspaceMeta(gomock.Any(), gomock.Any()).Return(nil, errors.New("not found")).AnyTimes()
 
 	ts.MockProject.EXPECT().GetPathWithNamespace().Return(ts.MetaData.ProjectNameNS).AnyTimes()
 

pkg/runstream/interfaces.go

@@ -12,6 +12,8 @@ type StreamClient interface {
 	PublishTFRunEvent(ctx context.Context, re RunEvent) error
 	AddRunMeta(rmd RunMetadata) error
 	GetRunMeta(runID string) (RunMetadata, error)
+	AddWorkspaceMeta(rmd WorkspaceMetadata, mrID, workspace string) error
+	GetWorkspaceMeta(mrID, workspace string) (*TFCWorkspacesMetadata, error)
 	NewTFRunPollingTask(meta RunMetadata, delay time.Duration) RunPollingTask
 	SubscribeTFRunPollingTasks(cb func(task RunPollingTask) bool) (closer func(), err error)
 	SubscribeTFRunEvents(queue string, cb func(run RunEvent) bool) (closer func(), err error)

pkg/runstream/stream.go

@@ -12,9 +12,10 @@ const RunMetadataKvBucket = "RUN_METADATA"
 
 type Stream struct {
 	//nc         *nats.Conn
-	js         nats.JetStreamContext
-	metadataKV nats.KeyValue
-	pollingKV  nats.KeyValue
+	js          nats.JetStreamContext
+	metadataKV  nats.KeyValue
+	pollingKV   nats.KeyValue
+	workspaceKV nats.KeyValue
 }
 
 func NewStream(js nats.JetStreamContext) StreamClient {
@@ -23,11 +24,13 @@ func NewStream(js nats.JetStreamContext) StreamClient {
 	configureTFRunPollingTaskStream(js)
 	kv, _ := configureTFRunMetadataKVStore(js)
 	pollingKV, _ := configureRunPollingKVStore(js)
+	workspaceKV, _ := configureWorkspaceMetadataKVStore(js)
 
 	s := &Stream{
 		js,
 		kv,
 		pollingKV,
+		workspaceKV,
 	}
 
 	s.startPollingTaskDispatcher()

pkg/runstream/workspaces_metadata.go

@@ -0,0 +1,73 @@
+package runstream
+
+import (
+	"encoding/json"
+	"fmt"
+	"time"
+
+	"github.com/nats-io/nats.go"
+)
+
+const WorkspaceMetadataKvBucket = "WORKSPACE_METADATA"
+
+func configureWorkspaceMetadataKVStore(js nats.JetStreamContext) (nats.KeyValue, error) {
+	cfg := &nats.KeyValueConfig{
+		Bucket:      WorkspaceMetadataKvBucket,
+		Description: "KV store for Workspace Metadata",
+		TTL:         time.Hour * 720,
+		Storage:     nats.FileStorage,
+		Replicas:    1,
+	}
+
+	for store := range js.KeyValueStores() {
+		if store.Bucket() == cfg.Bucket {
+			return js.KeyValue(cfg.Bucket)
+		}
+	}
+
+	return js.CreateKeyValue(cfg)
+}
+
+type WorkspaceMetadata interface {
+	GetCountExecutedWorkspaces() int
+	GetCountTotalWorkspaces() int
+}
+type TFCWorkspacesMetadata struct {
+	CountExecutedWorkspaces int `json:"count_executed_workspaces"`
+	CountTotalWorkspaces    int `json:"count_total_workspaces"`
+}
+
+func (t *TFCWorkspacesMetadata) GetCountExecutedWorkspaces() int {
+	return t.CountExecutedWorkspaces
+}
+func (t *TFCWorkspacesMetadata) GetCountTotalWorkspaces() int {
+	return t.CountTotalWorkspaces
+}
+func encodeWorkspaceMetadata(run WorkspaceMetadata) ([]byte, error) {
+	return json.Marshal(run)
+}
+
+func decodeWorkspaceMetadata(b []byte) (*TFCWorkspacesMetadata, error) {
+	rmd := &TFCWorkspacesMetadata{}
+	err := json.Unmarshal(b, &rmd)
+
+	return rmd, err
+}
+func (s *Stream) AddWorkspaceMeta(rmd WorkspaceMetadata, mrID, workspace string) error {
+	b, err := encodeWorkspaceMetadata(rmd)
+	if err != nil {
+		return err
+	}
+	_, err = s.metadataKV.Put(getKey(mrID, workspace), b)
+	return err
+}
+func getKey(mrID, workspace string) string {
+	return fmt.Sprintf("%s-%s", mrID, workspace)
+}
+func (s *Stream) GetWorkspaceMeta(mrID, workspace string) (*TFCWorkspacesMetadata, error) {
+	entry, err := s.metadataKV.Get(getKey(mrID, workspace))
+	if err != nil {
+		return nil, err
+	}
+	return decodeWorkspaceMetadata(entry.Value())
+}

pkg/tfc_trigger/tfc_trigger.go

@@ -441,7 +441,14 @@ func (t *TFCTrigger) TriggerTFCEvents(ctx context.Context) (*TriggeredTFCWorkspa
 		log.Debug().Msg("No Terraform changes found in changeset.")
 		return nil, nil
 	}
-
+	//only set workspace metadata for a MR when the first run is triggered. This prevents the count of executed workspaces from being reset
+	wsMeta, err := t.runstream.GetWorkspaceMeta(fmt.Sprintf("%d", t.GetMergeRequestIID()), t.GetProjectNameWithNamespace())
+	if err != nil || wsMeta == nil {
+		t.runstream.AddWorkspaceMeta(&runstream.TFCWorkspacesMetadata{
+			CountTotalWorkspaces:    len(triggeredWorkspaces),
+			CountExecutedWorkspaces: 0,
+		}, fmt.Sprintf("%d", t.GetMergeRequestIID()), t.GetProjectNameWithNamespace())
+	}
 	return workspaceStatus, nil
 }
 

pkg/tfc_trigger/tfc_trigger_test.go

@@ -2,6 +2,7 @@ package tfc_trigger_test
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"os"
 	"testing"
@@ -296,6 +297,10 @@ func TestTFCEvents_MultiWorkspaceApply(t *testing.T) {
 	}).Times(2)
 
 	testSuite.MockStreamClient.EXPECT().AddRunMeta(gomock.Any()).Times(2)
+	testSuite.MockStreamClient.EXPECT().AddWorkspaceMeta(&runstream.TFCWorkspacesMetadata{
+		CountTotalWorkspaces:    2,
+		CountExecutedWorkspaces: 0,
+	}, fmt.Sprintf("%d", testSuite.MetaData.MRIID), testSuite.MetaData.ProjectNameNS)
 	testSuite.InitTestSuite()
 	testLogger := zltest.New(t)
 	log.Logger = log.Logger.Output(testLogger)
@@ -474,7 +479,11 @@ func TestTFCEvents_WorkspaceApplyModifiedBothSrcDstBranches(t *testing.T) {
 	testSuite.MockGitClient.EXPECT().GetMergeRequestModifiedFiles(gomock.Any(), testSuite.MetaData.MRIID, testSuite.MetaData.ProjectNameNS).Return([]string{"main.tf"}, nil)
 
 	mockStreamClient := mocks.NewMockStreamClient(mockCtrl)
-
+	mockStreamClient.EXPECT().GetWorkspaceMeta(gomock.Any(), gomock.Any()).Return(nil, errors.New("no record"))
+	mockStreamClient.EXPECT().AddWorkspaceMeta(&runstream.TFCWorkspacesMetadata{
+		CountTotalWorkspaces:    1,
+		CountExecutedWorkspaces: 0,
+	}, fmt.Sprintf("%d", testSuite.MetaData.MRIID), testSuite.MetaData.ProjectNameNS)
 	testSuite.InitTestSuite()
 
 	testLogger := zltest.New(t)

pkg/vcs/gitlab/mr_status_updater.go

@@ -208,8 +208,29 @@ func (p *RunStatusUpdater) mergeMRIfPossible(ctx context.Context, rmd runstream.
 	if !rmd.GetAutoMerge() {
 		return
 	}
+	//check that all triggered workspaces have been executed or increment
+	wsMeta, err := p.rs.GetWorkspaceMeta(fmt.Sprintf("%d", rmd.GetMRInternalID()), rmd.GetMRProjectNameWithNamespace())
+	if err != nil {
+		log.Debug().AnErr("err", err).Msg("get workspace metadata")
+		return
+	}
+	wsMeta.CountExecutedWorkspaces++
+
+	if wsMeta.CountTotalWorkspaces > wsMeta.CountExecutedWorkspaces {
+		err = p.rs.AddWorkspaceMeta(wsMeta, fmt.Sprintf("%d", rmd.GetMRInternalID()), rmd.GetMRProjectNameWithNamespace())
+		if err != nil {
+			log.Debug().AnErr("err", err).Msg("add workspace metadata")
+			span.RecordError(err)
+		}
+		return
+	}
+	if wsMeta.CountExecutedWorkspaces > wsMeta.CountTotalWorkspaces {
+		log.Debug().Msg("count executed workspaces is greater than total workspaces")
+		span.RecordError(errors.New("count executed workspaces is greater than total workspaces"))
+		return
+	}
 
-	err := p.client.MergeMR(ctx, rmd.GetMRInternalID(), rmd.GetMRProjectNameWithNamespace())
+	err = p.client.MergeMR(ctx, rmd.GetMRInternalID(), rmd.GetMRProjectNameWithNamespace())
 	if err != nil {
 		span.RecordError(err)
 	}

pkg/vcs/gitlab/mr_status_updater_test.go

@@ -22,6 +22,10 @@ func TestAutoMergeNoChangesApply(t *testing.T) {
 	testSuite.MockGitClient.EXPECT().MergeMR(gomock.Any(), gomock.Any(), gomock.Any())
 	testSuite.MockGitClient.EXPECT().GetPipelinesForCommit(gomock.Any(), gomock.Any(), gomock.Any()).Return([]vcs.ProjectPipeline{&GitlabPipeline{&gogitlab.PipelineInfo{ID: 1}}}, nil).AnyTimes()
 	testSuite.MockGitClient.EXPECT().SetCommitStatus(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, errors.New("could not commit status")).AnyTimes()
+	testSuite.MockStreamClient.EXPECT().GetWorkspaceMeta(gomock.Any(), gomock.Any()).Return(&runstream.TFCWorkspacesMetadata{
+		CountExecutedWorkspaces: 0,
+		CountTotalWorkspaces:    1,
+	}, nil)
 	testSuite.InitTestSuite()
 	r := &RunStatusUpdater{
 		tfc:    testSuite.MockApiClient,
@@ -71,6 +75,12 @@ func TestAutoMergeApply(t *testing.T) {
 	testSuite.MockGitClient.EXPECT().MergeMR(gomock.Any(), gomock.Any(), gomock.Any())
 	testSuite.MockGitClient.EXPECT().GetPipelinesForCommit(gomock.Any(), gomock.Any(), gomock.Any()).Return([]vcs.ProjectPipeline{&GitlabPipeline{&gogitlab.PipelineInfo{ID: 1}}}, nil).AnyTimes()
 	testSuite.MockGitClient.EXPECT().SetCommitStatus(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, errors.New("could not commit status")).AnyTimes()
+
+	testSuite.MockStreamClient.EXPECT().GetWorkspaceMeta(gomock.Any(), gomock.Any()).Return(&runstream.TFCWorkspacesMetadata{
+		CountExecutedWorkspaces: 0,
+		CountTotalWorkspaces:    1,
+	}, nil)
+
 	testSuite.InitTestSuite()
 	r := &RunStatusUpdater{
 		tfc:    testSuite.MockApiClient,
@@ -86,6 +96,54 @@ func TestAutoMergeApply(t *testing.T) {
 	})
 }
 
+func TestAutoMergeApplyMultiWorkspace(t *testing.T) {
+
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	testSuite := mocks.CreateTestSuite(mockCtrl, mocks.TestOverrides{}, t)
+
+	testSuite.MockGitClient.EXPECT().MergeMR(gomock.Any(), gomock.Any(), gomock.Any())
+	testSuite.MockGitClient.EXPECT().GetPipelinesForCommit(gomock.Any(), gomock.Any(), gomock.Any()).Return([]vcs.ProjectPipeline{&GitlabPipeline{&gogitlab.PipelineInfo{ID: 1}}}, nil).AnyTimes()
+	testSuite.MockGitClient.EXPECT().SetCommitStatus(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, errors.New("could not commit status")).AnyTimes()
+
+	//workspace 1 in same mr
+	testSuite.MockStreamClient.EXPECT().GetWorkspaceMeta("101", "zapier/test").Return(&runstream.TFCWorkspacesMetadata{
+		CountExecutedWorkspaces: 0,
+		CountTotalWorkspaces:    2,
+	}, nil)
+	//workspace 2 in same mr
+	testSuite.MockStreamClient.EXPECT().GetWorkspaceMeta("101", "zapier/test").Return(&runstream.TFCWorkspacesMetadata{
+		CountExecutedWorkspaces: 1,
+		CountTotalWorkspaces:    2,
+	}, nil)
+
+	testSuite.InitTestSuite()
+	r := &RunStatusUpdater{
+		tfc:    testSuite.MockApiClient,
+		client: testSuite.MockGitClient,
+		rs:     testSuite.MockStreamClient,
+	}
+	r.updateCommitStatusForRun(context.Background(), &tfe.Run{
+		Status:     tfe.RunApplied,
+		HasChanges: true,
+	}, &runstream.TFRunMetadata{
+		Action:                               "apply",
+		AutoMerge:                            true,
+		MergeRequestIID:                      101,
+		MergeRequestProjectNameWithNamespace: "zapier/test",
+	})
+
+	r.updateCommitStatusForRun(context.Background(), &tfe.Run{
+		Status:     tfe.RunApplied,
+		HasChanges: true,
+	}, &runstream.TFRunMetadata{
+		Action:                               "apply",
+		AutoMerge:                            true,
+		MergeRequestIID:                      101,
+		MergeRequestProjectNameWithNamespace: "zapier/test",
+	})
+}
 func TestAutoMergeTargetedApply(t *testing.T) {
 
 	mockCtrl := gomock.NewController(t)