Add wait_for_rollout to kubernetes_stateful_set (#914)

Co-authored-by: Stef Forrester <[email protected]>

Author: jrhouston

go.mod

@@ -25,10 +25,11 @@ require (
 	github.com/terraform-providers/terraform-provider-random v1.3.2-0.20190925210718-83518d96ae4f // indirect
 	github.com/ulikunitz/xz v0.5.6 // indirect
 	github.com/vmihailenco/msgpack v4.0.4+incompatible // indirect
-	k8s.io/api v0.0.0-20191025225708-5524a3672fbb
-	k8s.io/apimachinery v0.0.0-20191025225532-af6325b3a843
+	k8s.io/api v0.16.12
+	k8s.io/apimachinery v0.16.12
 	k8s.io/client-go v10.0.0+incompatible
 	k8s.io/kube-aggregator v0.0.0-20191025230902-aa872b06629d
+	k8s.io/kubectl v0.16.12
 )
 
 // Override invalid go-autorest pseudo-version. This can be removed once

go.sum

@@ -6,10 +6,14 @@ import (
 
 	"github.com/hashicorp/terraform-plugin-sdk/helper/resource"
 	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
-	api "k8s.io/api/apps/v1"
+	appsv1 "k8s.io/api/apps/v1"
 	"k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	pkgApi "k8s.io/apimachinery/pkg/types"
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/types"
+	kubernetes "k8s.io/client-go/kubernetes"
+	polymorphichelpers "k8s.io/kubectl/pkg/polymorphichelpers"
 )
 
 func resourceKubernetesStatefulSet() *schema.Resource {
@@ -35,6 +39,12 @@ func resourceKubernetesStatefulSet() *schema.Resource {
 					Schema: statefulSetSpecFields(false),
 				},
 			},
+			"wait_for_rollout": {
+				Type:        schema.TypeBool,
+				Description: "Wait for the rollout of the stateful set to complete. Defaults to false.",
+				Default:     false,
+				Optional:    true,
+			},
 		},
 	}
 }
@@ -49,7 +59,7 @@ func resourceKubernetesStatefulSetCreate(d *schema.ResourceData, meta interface{
 	if err != nil {
 		return err
 	}
-	statefulSet := api.StatefulSet{
+	statefulSet := appsv1.StatefulSet{
 		ObjectMeta: metadata,
 		Spec:       *spec,
 	}
@@ -67,6 +77,14 @@ func resourceKubernetesStatefulSetCreate(d *schema.ResourceData, meta interface{
 
 	log.Printf("[INFO] StatefulSet %s created", id)
 
+	if d.Get("wait_for_rollout").(bool) {
+		log.Printf("[INFO] Waiting for StatefulSet %s to rollout", id)
+		namespace := out.ObjectMeta.Namespace
+		name := out.ObjectMeta.Name
+		return resource.Retry(d.Timeout(schema.TimeoutCreate),
+			retryUntilStatefulSetRolloutComplete(conn, namespace, name))
+	}
+
 	return resourceKubernetesStatefulSetRead(d, meta)
 }
 
@@ -156,12 +174,18 @@ func resourceKubernetesStatefulSetUpdate(d *schema.ResourceData, meta interface{
 		return fmt.Errorf("Failed to marshal update operations for StatefulSet: %s", err)
 	}
 	log.Printf("[INFO] Updating StatefulSet %q: %v", name, string(data))
-	out, err := conn.AppsV1().StatefulSets(namespace).Patch(name, pkgApi.JSONPatchType, data)
+	out, err := conn.AppsV1().StatefulSets(namespace).Patch(name, types.JSONPatchType, data)
 	if err != nil {
 		return fmt.Errorf("Failed to update StatefulSet: %s", err)
 	}
 	log.Printf("[INFO] Submitted updated StatefulSet: %#v", out)
 
+	if d.Get("wait_for_rollout").(bool) {
+		log.Printf("[INFO] Waiting for StatefulSet %s to rollout", d.Id())
+		return resource.Retry(d.Timeout(schema.TimeoutCreate),
+			retryUntilStatefulSetRolloutComplete(conn, namespace, name))
+	}
+
 	return resourceKubernetesStatefulSetRead(d, meta)
 }
 
@@ -203,3 +227,46 @@ func resourceKubernetesStatefulSetDelete(d *schema.ResourceData, meta interface{
 
 	return nil
 }
+
+// retryUntilStatefulSetRolloutComplete checks if a given job finished its execution and is either in 'Complete' or 'Failed' state.
+func retryUntilStatefulSetRolloutComplete(conn *kubernetes.Clientset, ns, name string) resource.RetryFunc {
+	return func() *resource.RetryError {
+		res, err := conn.AppsV1().StatefulSets(ns).Get(name, metav1.GetOptions{})
+		if err != nil {
+			return resource.NonRetryableError(err)
+		}
+
+		// NOTE: This is what kubectl uses to determine if a rollout is done.
+		// We are using this here because the logic for determining if a StatefulSet
+		// is done is gnarly and we don't want to duplicate it in the provider.
+		gvk := appsv1.SchemeGroupVersion.WithKind("StatefulSet")
+		gk := gvk.GroupKind()
+		statusViewer, err := polymorphichelpers.StatusViewerFor(gk)
+		if err != nil {
+			return resource.NonRetryableError(err)
+		}
+
+		obj, err := runtime.DefaultUnstructuredConverter.ToUnstructured(res)
+		if err != nil {
+			return resource.NonRetryableError(err)
+		}
+
+		// NOTE: For some reason, the Kind and apiVersion get lost when converting to unstructured.
+		obj["apiVersion"] = gvk.GroupVersion().String()
+		obj["kind"] = gvk.Kind
+		u := unstructured.Unstructured{Object: obj}
+
+		// NOTE: the revision parameter of the Status function below is not actually used.
+		// for StatefulSet so it is set to 0 here
+		_, done, err := statusViewer.Status(&u, 0)
+		if err != nil {
+			return resource.NonRetryableError(err)
+		}
+
+		if done {
+			return nil
+		}
+
+		return resource.RetryableError(fmt.Errorf("StatefulSet %s/%s is not finished rolling out", ns, name))
+	}
+}

kubernetes/resource_kubernetes_stateful_set.go

@@ -8,6 +8,7 @@ import (
 	"github.com/hashicorp/terraform-plugin-sdk/helper/resource"
 	"github.com/hashicorp/terraform-plugin-sdk/terraform"
 	api "k8s.io/api/apps/v1"
+	appsv1 "k8s.io/api/apps/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
@@ -26,6 +27,7 @@ func TestAccKubernetesStatefulSet_basic(t *testing.T) {
 				Config: testAccKubernetesStatefulSetConfigBasic(name),
 				Check: resource.ComposeAggregateTestCheckFunc(
 					testAccCheckKubernetesStatefulSetExists(statefulSetTestResourceName, &conf),
+					testAccCheckKubernetesStatefulSetRollingOut(statefulSetTestResourceName),
 					testAccKubernetesStatefulSetChecksBasic(name),
 				),
 			},
@@ -310,6 +312,26 @@ func TestAccKubernetesStatefulSet_update_pod_template(t *testing.T) {
 	})
 }
 
+func TestAccKubernetesStatefulSet_waitForRollout(t *testing.T) {
+	var conf api.StatefulSet
+	name := fmt.Sprintf("tf-acc-test-%s", acctest.RandStringFromCharSet(10, acctest.CharSetAlphaNum))
+	resource.Test(t, resource.TestCase{
+		PreCheck:      func() { testAccPreCheck(t) },
+		IDRefreshName: statefulSetTestResourceName,
+		Providers:     testAccProviders,
+		CheckDestroy:  testAccCheckKubernetesStatefulSetDestroy,
+		Steps: []resource.TestStep{
+			{
+				Config: testAccKubernetesStatefulSetConfigWaitForRollout(name),
+				Check: resource.ComposeAggregateTestCheckFunc(
+					testAccCheckKubernetesStatefulSetExists(statefulSetTestResourceName, &conf),
+					testAccCheckKubernetesStatefulSetRolledOut(statefulSetTestResourceName),
+				),
+			},
+		},
+	})
+}
+
 func testAccCheckKubernetesStatefulSetDestroy(s *terraform.State) error {
 	conn, err := testAccProvider.Meta().(KubeClientsets).MainClientset()
 	if err != nil {
@@ -337,34 +359,74 @@ func testAccCheckKubernetesStatefulSetDestroy(s *terraform.State) error {
 	return nil
 }
 
-func testAccCheckKubernetesStatefulSetExists(n string, obj *api.StatefulSet) resource.TestCheckFunc {
-	return func(s *terraform.State) error {
-		rs, ok := s.RootModule().Resources[n]
-		if !ok {
-			return fmt.Errorf("Not found: %s", n)
-		}
+func getStatefulSetFromResourceName(s *terraform.State, n string) (*appsv1.StatefulSet, error) {
+	rs, ok := s.RootModule().Resources[n]
+	if !ok {
+		return nil, fmt.Errorf("Not found: %s", n)
+	}
+
+	conn, err := testAccProvider.Meta().(KubeClientsets).MainClientset()
+	if err != nil {
+		return nil, err
+	}
 
-		conn, err := testAccProvider.Meta().(KubeClientsets).MainClientset()
+	namespace, name, err := idParts(rs.Primary.ID)
+	if err != nil {
+		return nil, err
+	}
+
+	out, err := conn.AppsV1().StatefulSets(namespace).Get(name, metav1.GetOptions{})
+	if err != nil {
+		return nil, err
+	}
+
+	return out, nil
+}
+
+func testAccCheckKubernetesStatefulSetExists(n string, obj *appsv1.StatefulSet) resource.TestCheckFunc {
+	return func(s *terraform.State) error {
+		d, err := getStatefulSetFromResourceName(s, n)
 		if err != nil {
 			return err
 		}
+		*obj = *d
+		return nil
+	}
+}
 
-		namespace, name, err := idParts(rs.Primary.ID)
+func testAccCheckKubernetesStatefulSetRollingOut(n string) resource.TestCheckFunc {
+	return func(s *terraform.State) error {
+		d, err := getStatefulSetFromResourceName(s, n)
 		if err != nil {
 			return err
 		}
 
-		out, err := conn.AppsV1().StatefulSets(namespace).Get(name, metav1.GetOptions{})
+		if d.Status.Replicas == *d.Spec.Replicas {
+			return fmt.Errorf("StatefulSet has already rolled out")
+		}
+
+		return nil
+	}
+}
+
+func testAccCheckKubernetesStatefulSetRolledOut(n string) resource.TestCheckFunc {
+	return func(s *terraform.State) error {
+		d, err := getStatefulSetFromResourceName(s, n)
 		if err != nil {
 			return err
 		}
-		*obj = *out
+
+		if d.Status.Replicas != *d.Spec.Replicas {
+			return fmt.Errorf("StatefulSet is still rolling out")
+		}
+
 		return nil
 	}
 }
 
 func testAccKubernetesStatefulSetChecksBasic(name string) resource.TestCheckFunc {
 	return resource.ComposeAggregateTestCheckFunc(
+		resource.TestCheckResourceAttr(statefulSetTestResourceName, "wait_for_rollout", "false"),
 		resource.TestCheckResourceAttrSet(statefulSetTestResourceName, "metadata.0.generation"),
 		resource.TestCheckResourceAttrSet(statefulSetTestResourceName, "metadata.0.resource_version"),
 		resource.TestCheckResourceAttrSet(statefulSetTestResourceName, "metadata.0.self_link"),
@@ -1009,3 +1071,69 @@ resource "kubernetes_stateful_set" "test" {
 }
 `, name)
 }
+
+func testAccKubernetesStatefulSetConfigWaitForRollout(name string) string {
+	return fmt.Sprintf(`
+resource "kubernetes_service" "test" {
+  metadata {
+    name = "ss-test"
+  }
+  spec {
+    port {
+      port = 80
+    }
+  }
+}
+
+resource "kubernetes_stateful_set" "test" {
+  metadata {
+    name = "%s"
+  }
+
+  spec {
+    replicas = 2
+
+    selector {
+      match_labels = {
+        app = "ss-test"
+      }
+    }
+
+    update_strategy {
+      type = "RollingUpdate"
+    }
+
+    service_name = kubernetes_service.test.metadata.0.name
+
+    template {
+      metadata {
+        labels = {
+          app = "ss-test"
+        }
+      }
+
+      spec {
+        container {
+          name  = "ss-test"
+          image = "nginx:1.19"
+
+          port {
+            container_port = 80
+          }
+
+          readiness_probe {
+            initial_delay_seconds = 5
+            http_get {
+              path = "/"
+              port = 80
+            }
+          }
+        }
+      }
+    }
+  }
+
+  wait_for_rollout = true
+}
+`, name)
+}