🚀 Update Workspace SSH key reconciliation logic

.changes/unreleased/ENHANCEMENTS-478-20240822-123848.yaml

@@ -0,0 +1,5 @@
+kind: ENHANCEMENTS
+body: '`Workspace`: Update SSH key reconciliation logic to reduce the number of API calls.'
+time: 2024-08-22T12:38:48.871032+02:00
+custom:
+    PR: "478"

.changes/unreleased/NOTES-478-20240822-123914.yaml

@@ -0,0 +1,5 @@
+kind: NOTES
+body: The `Workspace` CRD has been changed. Please follow the Helm chart instructions on how to upgrade it.
+time: 2024-08-22T12:39:14.407611+02:00
+custom:
+    PR: "478"

api/v1alpha2/workspace_types.go

@@ -674,6 +674,10 @@ type WorkspaceStatus struct {
 	//
 	//+optional
 	DefaultProjectID string `json:"defaultProjectID,omitempty"`
+	// SSH Key ID.
+	//
+	//+optional
+	SSHKeyID string `json:"sshKeyID,omitempty"`
 	// Workspace Destroy Run ID.
 	//
 	//+optional

charts/hcp-terraform-operator/crds/app.terraform.io_workspaces.yaml

@@ -796,6 +796,9 @@ spec:
                       status.
                     type: string
                 type: object
+              sshKeyID:
+                description: SSH Key ID.
+                type: string
               terraformVersion:
                 description: Workspace Terraform version.
                 pattern: ^\d{1}\.\d{1,2}\.\d{1,2}$

config/crd/bases/app.terraform.io_workspaces.yaml

@@ -793,6 +793,9 @@ spec:
                       status.
                     type: string
                 type: object
+              sshKeyID:
+                description: SSH Key ID.
+                type: string
               terraformVersion:
                 description: Workspace Terraform version.
                 pattern: ^\d{1}\.\d{1,2}\.\d{1,2}$

controllers/workspace_controller.go

@@ -527,7 +527,7 @@ func (r *WorkspaceReconciler) reconcileWorkspace(ctx context.Context, w *workspa
 	}
 
 	// reconcile SSH key
-	err = r.reconcileSSHKey(ctx, w, workspace)
+	err = w.reconcileSSHKey(ctx, workspace)
 	if err != nil {
 		w.log.Error(err, "Reconcile SSH Key", "msg", "failed to assign ssh key ID")
 		r.Recorder.Eventf(&w.instance, corev1.EventTypeWarning, "ReconcileSSHKey", "Failed to assign SSH Key ID")

controllers/workspace_controller_sshkey.go

@@ -5,83 +5,76 @@ package controllers
 
 import (
 	"context"
+	"errors"
 	"fmt"
 
 	tfc "github.com/hashicorp/go-tfe"
 )
 
-func (r *WorkspaceReconciler) getSSHKeyIDByName(ctx context.Context, w *workspaceInstance) (string, error) {
-	SSHKeyName := w.instance.Spec.SSHKey.Name
-
-	listOpts := &tfc.SSHKeyListOptions{
-		ListOptions: tfc.ListOptions{
-			PageSize: maxPageSize,
-		},
+func (w *workspaceInstance) getSSHKeyID(ctx context.Context) (string, error) {
+	if w.instance.Spec.SSHKey == nil {
+		return "", errors.New("instance spec.SSHKey is nil")
 	}
-	for {
-		SSHKeys, err := w.tfClient.Client.SSHKeys.List(ctx, w.instance.Spec.Organization, listOpts)
-		if err != nil {
-			return "", err
+
+	if w.instance.Spec.SSHKey.Name != "" {
+		w.log.Info("Reconcile SSH Key", "msg", "getting SSH key ID by name")
+		listOpts := &tfc.SSHKeyListOptions{
+			ListOptions: tfc.ListOptions{
+				PageNumber: 1,
+				PageSize:   maxPageSize,
+			},
 		}
-		for _, s := range SSHKeys.Items {
-			if s.Name == SSHKeyName {
-				return s.ID, nil
+		for {
+			sshKeyList, err := w.tfClient.Client.SSHKeys.List(ctx, w.instance.Spec.Organization, listOpts)
+			if err != nil {
+				return "", err
 			}
-		}
-		if SSHKeys.NextPage == 0 {
-			break
-		}
-		listOpts.PageNumber = SSHKeys.NextPage
-	}
 
-	return "", fmt.Errorf("ssh key ID was not found for ssh key name %q", SSHKeyName)
-}
-
-func (r *WorkspaceReconciler) getSSHKeyID(ctx context.Context, w *workspaceInstance) (string, error) {
-	specSSHKey := w.instance.Spec.SSHKey
+			for _, s := range sshKeyList.Items {
+				if s.Name == w.instance.Spec.SSHKey.Name {
+					return s.ID, nil
+				}
+			}
 
-	if specSSHKey.Name != "" {
-		w.log.Info("Reconcile SSH Key", "msg", "getting ssh key ID by name")
-		return r.getSSHKeyIDByName(ctx, w)
+			if sshKeyList.NextPage == 0 {
+				break
+			}
+			listOpts.PageNumber = sshKeyList.NextPage
+		}
+		return "", fmt.Errorf("ssh key ID was not found for SSH key name %q", w.instance.Spec.SSHKey.Name)
 	}
 
-	w.log.Info("Reconcile SSH Key", "msg", "getting ssh key ID from the spec.sshKey.ID")
-	return specSSHKey.ID, nil
+	w.log.Info("Reconcile SSH Key", "msg", "getting SSH key ID from the spec.sshKey.ID")
+	return w.instance.Spec.SSHKey.ID, nil
 }
 
-func (r *WorkspaceReconciler) reconcileSSHKey(ctx context.Context, w *workspaceInstance, workspace *tfc.Workspace) error {
-	if w.instance.Spec.SSHKey == nil {
-		// Verify whether a Workspace has an SSH key and unassign it if so
-		if workspace.SSHKey != nil {
-			w.log.Info("Reconcile SSH Key", "msg", "unassigning the ssh key")
-			_, err := w.tfClient.Client.Workspaces.UnassignSSHKey(ctx, workspace.ID)
+func (w *workspaceInstance) reconcileSSHKey(ctx context.Context, workspace *tfc.Workspace) error {
+	w.log.Info("Reconcile SSH Key", "msg", "new reconciliation event")
+
+	if w.instance.Spec.SSHKey == nil && workspace.SSHKey != nil {
+		w.log.Info("Reconcile SSH Key", "msg", "removing the SSH key")
+		if _, err := w.tfClient.Client.Workspaces.UnassignSSHKey(ctx, workspace.ID); err != nil {
 			return err
 		}
-
-		return nil
+		w.instance.Status.SSHKeyID = ""
 	}
 
-	SSHKeyID, err := r.getSSHKeyID(ctx, w)
-	if err != nil {
-		return err
-	}
-
-	// Assign an SSH key to a workspace if nothing is assigned
-	if workspace.SSHKey == nil {
-		w.log.Info("Reconcile SSH Key", "msg", "assigning the ssh key")
-		_, err := w.tfClient.Client.Workspaces.AssignSSHKey(ctx, workspace.ID, tfc.WorkspaceAssignSSHKeyOptions{
-			SSHKeyID: tfc.String(SSHKeyID),
-		})
-		return err
-	}
-
-	// Assign an SSH key to a workspace if it is different from the one in spec
-	if workspace.SSHKey.ID != SSHKeyID {
-		w.log.Info("Reconcile SSH Key", "msg", "assigning the ssh key")
-		_, err := w.tfClient.Client.Workspaces.AssignSSHKey(ctx, workspace.ID, tfc.WorkspaceAssignSSHKeyOptions{
-			SSHKeyID: tfc.String(SSHKeyID),
-		})
-		return err
+	if w.instance.Spec.SSHKey != nil {
+		if workspace.SSHKey == nil || workspace.SSHKey.ID != w.instance.Status.SSHKeyID {
+			sshKeyID, err := w.getSSHKeyID(ctx)
+			if err != nil {
+				return err
+			}
+			w.log.Info("Reconcile SSH Key", "msg", "assigning the SSH key")
+			opt := tfc.WorkspaceAssignSSHKeyOptions{
+				SSHKeyID: &sshKeyID,
+			}
+			workspace, err = w.tfClient.Client.Workspaces.AssignSSHKey(ctx, workspace.ID, opt)
+			if err != nil {
+				return err
+			}
+			w.instance.Status.SSHKeyID = workspace.SSHKey.ID
+		}
 	}
 
 	return nil

controllers/workspace_controller_sshkey_test.go

@@ -11,14 +11,13 @@ import (
 	"fmt"
 	"time"
 
+	tfc "github.com/hashicorp/go-tfe"
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
-
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
 
-	tfc "github.com/hashicorp/go-tfe"
 	appv1alpha2 "github.com/hashicorp/hcp-terraform-operator/api/v1alpha2"
 )
 
@@ -69,7 +68,8 @@ var _ = Describe("Workspace controller", Ordered, func() {
 						Key: secretKey,
 					},
 				},
-				Name: workspace,
+				Name:        workspace,
+				Description: "SSH Key",
 			},
 			Status: appv1alpha2.WorkspaceStatus{},
 		}
@@ -82,57 +82,101 @@ var _ = Describe("Workspace controller", Ordered, func() {
 	})
 
 	Context("SSH Key", func() {
-		It("can handle SSH Key by name", func() {
+		It("can be created by name", func() {
 			instance.Spec.SSHKey = &appv1alpha2.SSHKey{
 				Name: sshKeyName,
 			}
 			// Create a new Kubernetes workspace object and wait until the controller finishes the reconciliation
 			createWorkspace(instance)
-			isReconciledSSHKeyByName(instance)
+			isReconciledSSHKey(instance)
+		})
+		It("can be created by ID", func() {
+			instance.Spec.SSHKey = &appv1alpha2.SSHKey{
+				ID: sshKeyID,
+			}
+			// Create a new Kubernetes workspace object and wait until the controller finishes the reconciliation
+			createWorkspace(instance)
+			isReconciledSSHKey(instance)
+		})
 
+		It("can be restored by name", func() {
+			instance.Spec.SSHKey = &appv1alpha2.SSHKey{
+				Name: sshKeyName,
+			}
+			// Create a new Kubernetes workspace object and wait until the controller finishes the reconciliation
+			createWorkspace(instance)
+			isReconciledSSHKey(instance)
 			// Delete the SSH key manually and wait until the controller revert this change
 			ws, err := tfClient.Workspaces.UnassignSSHKey(ctx, instance.Status.WorkspaceID)
 			Expect(err).Should(Succeed())
 			Expect(ws).ShouldNot(BeNil())
-			isReconciledSSHKeyByName(instance)
+			isReconciledSSHKey(instance)
+		})
+		It("can be restored by ID", func() {
+			instance.Spec.SSHKey = &appv1alpha2.SSHKey{
+				ID: sshKeyID,
+			}
+			// Create a new Kubernetes workspace object and wait until the controller finishes the reconciliation
+			createWorkspace(instance)
+			isReconciledSSHKey(instance)
+			// Delete the SSH key manually and wait until the controller revert this change
+			ws, err := tfClient.Workspaces.UnassignSSHKey(ctx, instance.Status.WorkspaceID)
+			Expect(err).Should(Succeed())
+			Expect(ws).ShouldNot(BeNil())
+			isReconciledSSHKey(instance)
+		})
 
+		It("can be updated by name", func() {
+			instance.Spec.SSHKey = &appv1alpha2.SSHKey{
+				Name: sshKeyName,
+			}
+			// Create a new Kubernetes workspace object and wait until the controller finishes the reconciliation
+			createWorkspace(instance)
+			isReconciledSSHKey(instance)
 			// Update the SSH key
 			Expect(k8sClient.Get(ctx, namespacedName, instance)).Should(Succeed())
 			instance.Spec.SSHKey = &appv1alpha2.SSHKey{
 				Name: sshKeyName2,
 			}
 			Expect(k8sClient.Update(ctx, instance)).Should(Succeed())
-			isReconciledSSHKeyByName(instance)
-
-			// Detach the SSH key
-			Expect(k8sClient.Get(ctx, namespacedName, instance))
-			instance.Spec.SSHKey = nil
-			Expect(k8sClient.Update(ctx, instance)).Should(Succeed())
-			isSSHKeyEmpty(instance)
+			isReconciledSSHKey(instance)
 		})
-
-		It("can handle SSH Key by id", func() {
+		It("can be updated by ID", func() {
 			instance.Spec.SSHKey = &appv1alpha2.SSHKey{
 				ID: sshKeyID,
 			}
 			// Create a new Kubernetes workspace object and wait until the controller finishes the reconciliation
 			createWorkspace(instance)
-			isReconciledSSHKeyByID(instance)
-
-			// Delete the SSH key manually and wait until the controller revert this change
-			ws, err := tfClient.Workspaces.UnassignSSHKey(ctx, instance.Status.WorkspaceID)
-			Expect(err).Should(Succeed())
-			Expect(ws).ShouldNot(BeNil())
-			isReconciledSSHKeyByID(instance)
-
+			isReconciledSSHKey(instance)
 			// Update the SSH key
 			Expect(k8sClient.Get(ctx, namespacedName, instance))
 			instance.Spec.SSHKey = &appv1alpha2.SSHKey{
 				ID: sshKeyID2,
 			}
 			Expect(k8sClient.Update(ctx, instance)).Should(Succeed())
-			isReconciledSSHKeyByID(instance)
+			isReconciledSSHKey(instance)
+		})
 
+		It("can be removed by name", func() {
+			instance.Spec.SSHKey = &appv1alpha2.SSHKey{
+				Name: sshKeyName,
+			}
+			// Create a new Kubernetes workspace object and wait until the controller finishes the reconciliation
+			createWorkspace(instance)
+			isReconciledSSHKey(instance)
+			// Detach the SSH key
+			Expect(k8sClient.Get(ctx, namespacedName, instance))
+			instance.Spec.SSHKey = nil
+			Expect(k8sClient.Update(ctx, instance)).Should(Succeed())
+			isSSHKeyEmpty(instance)
+		})
+		It("can be removed by ID", func() {
+			instance.Spec.SSHKey = &appv1alpha2.SSHKey{
+				ID: sshKeyID,
+			}
+			// Create a new Kubernetes workspace object and wait until the controller finishes the reconciliation
+			createWorkspace(instance)
+			isReconciledSSHKey(instance)
 			// Delete the SSH key
 			Expect(k8sClient.Get(ctx, namespacedName, instance))
 			instance.Spec.SSHKey = nil
@@ -142,38 +186,21 @@ var _ = Describe("Workspace controller", Ordered, func() {
 	})
 })
 
-func isReconciledSSHKeyByName(instance *appv1alpha2.Workspace) {
+func isReconciledSSHKey(instance *appv1alpha2.Workspace) {
 	namespacedName := getNamespacedName(instance)
 
 	Eventually(func() bool {
 		Expect(k8sClient.Get(ctx, namespacedName, instance)).Should(Succeed())
-		ws, err := tfClient.Workspaces.ReadByID(ctx, instance.Status.WorkspaceID)
-		Expect(err).Should(Succeed())
-		Expect(ws).ShouldNot(BeNil())
-		if ws.SSHKey == nil {
+		if instance.Status.SSHKeyID == "" {
 			return false
-		} else {
-			s, err := tfClient.SSHKeys.Read(ctx, ws.SSHKey.ID)
-			Expect(err).Should(Succeed())
-			Expect(s).ShouldNot(BeNil())
-			return s.Name == instance.Spec.SSHKey.Name
 		}
-	}).Should(BeTrue())
-}
-
-func isReconciledSSHKeyByID(instance *appv1alpha2.Workspace) {
-	namespacedName := getNamespacedName(instance)
-
-	Eventually(func() bool {
-		Expect(k8sClient.Get(ctx, namespacedName, instance)).Should(Succeed())
 		ws, err := tfClient.Workspaces.ReadByID(ctx, instance.Status.WorkspaceID)
 		Expect(err).Should(Succeed())
 		Expect(ws).ShouldNot(BeNil())
 		if ws.SSHKey == nil {
 			return false
-		} else {
-			return ws.SSHKey.ID == instance.Spec.SSHKey.ID
 		}
+		return ws.SSHKey.ID == instance.Status.SSHKeyID
 	}).Should(BeTrue())
 }
 
@@ -182,6 +209,9 @@ func isSSHKeyEmpty(instance *appv1alpha2.Workspace) {
 
 	Eventually(func() bool {
 		Expect(k8sClient.Get(ctx, namespacedName, instance)).Should(Succeed())
+		if instance.Status.SSHKeyID != "" {
+			return false
+		}
 		ws, err := tfClient.Workspaces.ReadByID(ctx, instance.Status.WorkspaceID)
 		Expect(err).Should(Succeed())
 		Expect(ws).ShouldNot(BeNil())
@@ -191,8 +221,8 @@ func isSSHKeyEmpty(instance *appv1alpha2.Workspace) {
 
 func createSSHKey(sshKeyName string) string {
 	sk, err := rsa.GenerateKey(rand.Reader, 2048)
-	Expect(sk).ShouldNot(BeNil())
 	Expect(err).Should(Succeed())
+	Expect(sk).ShouldNot(BeNil())
 	var privateKeyBytes []byte = x509.MarshalPKCS1PrivateKey(sk)
 	privateKeyBlock := &pem.Block{
 		Type:  "RSA PRIVATE KEY",