Auto-upgrade for colab runtimes (#12931)

[upstream:b79356dcaaa3925403eade2599558778d385097f]

Signed-off-by: Modular Magician <[email protected]>

Author: modular-magician

.changelog/12931.txt

@@ -0,0 +1,3 @@
+```release-note:enhancement
+colab: Auto-upgrade for colab runtimes
+```

google-beta/services/colab/resource_colab_runtime.go

@@ -140,6 +140,21 @@ func ResourceColabRuntime() *schema.Resource {
 					},
 				},
 			},
+			"expiration_time": {
+				Type:        schema.TypeString,
+				Computed:    true,
+				Description: `Output only. Timestamp when this NotebookRuntime will be expired.`,
+			},
+			"is_upgradable": {
+				Type:        schema.TypeBool,
+				Computed:    true,
+				Description: `Output only. Checks if the NotebookRuntime is upgradable.`,
+			},
+			"notebook_runtime_type": {
+				Type:        schema.TypeString,
+				Computed:    true,
+				Description: `Output only. The type of the notebook runtime.`,
+			},
 			"state": {
 				Type:        schema.TypeString,
 				Computed:    true,
@@ -151,6 +166,11 @@ func ResourceColabRuntime() *schema.Resource {
 				Description: `Desired state of the Colab Runtime. Set this field to 'RUNNING' to start the runtime, and 'STOPPED' to stop it.`,
 				Default:     "RUNNING",
 			},
+			"auto_upgrade": {
+				Type:        schema.TypeBool,
+				Optional:    true,
+				Description: `Triggers an upgrade anytime the runtime is started if it is upgradable.`,
+			},
 			"project": {
 				Type:     schema.TypeString,
 				Optional: true,
@@ -325,6 +345,15 @@ func resourceColabRuntimeRead(d *schema.ResourceData, meta interface{}) error {
 	if err := d.Set("state", flattenColabRuntimeState(res["state"], d, config)); err != nil {
 		return fmt.Errorf("Error reading Runtime: %s", err)
 	}
+	if err := d.Set("is_upgradable", flattenColabRuntimeIsUpgradable(res["isUpgradable"], d, config)); err != nil {
+		return fmt.Errorf("Error reading Runtime: %s", err)
+	}
+	if err := d.Set("expiration_time", flattenColabRuntimeExpirationTime(res["expirationTime"], d, config)); err != nil {
+		return fmt.Errorf("Error reading Runtime: %s", err)
+	}
+	if err := d.Set("notebook_runtime_type", flattenColabRuntimeNotebookRuntimeType(res["notebookRuntimeType"], d, config)); err != nil {
+		return fmt.Errorf("Error reading Runtime: %s", err)
+	}
 
 	return nil
 }
@@ -370,6 +399,27 @@ func resourceColabRuntimeUpdate(d *schema.ResourceData, meta interface{}) error
 		log.Printf("[DEBUG] Colab runtime %q has state %q.", name, state)
 	}
 
+	var upgrade_runtime bool
+	if d.Get("auto_upgrade").(bool) && d.Get("is_upgradable").(bool) {
+		upgrade_runtime = true
+	}
+
+	expiration_time_string := d.Get("expiration_time").(string)
+	expiration_time, err := time.Parse(time.RFC3339Nano, expiration_time_string)
+	if err != nil {
+		return err
+	}
+
+	if expiration_time.Before(time.Now()) && d.Get("notebook_runtime_type").(string) == "USER_DEFINED" {
+		upgrade_runtime = true
+	}
+
+	if upgrade_runtime {
+		if err := ModifyColabRuntime(config, d, project, billingProject, userAgent, "upgrade"); err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 
@@ -487,6 +537,18 @@ func flattenColabRuntimeState(v interface{}, d *schema.ResourceData, config *tra
 	return v
 }
 
+func flattenColabRuntimeIsUpgradable(v interface{}, d *schema.ResourceData, config *transport_tpg.Config) interface{} {
+	return v
+}
+
+func flattenColabRuntimeExpirationTime(v interface{}, d *schema.ResourceData, config *transport_tpg.Config) interface{} {
+	return v
+}
+
+func flattenColabRuntimeNotebookRuntimeType(v interface{}, d *schema.ResourceData, config *transport_tpg.Config) interface{} {
+	return v
+}
+
 func expandColabRuntimeNotebookRuntimeTemplateRef(v interface{}, d tpgresource.TerraformResourceData, config *transport_tpg.Config) (interface{}, error) {
 	l := v.([]interface{})
 	if len(l) == 0 || l[0] == nil {

google-beta/services/colab/resource_colab_runtime_generated_meta.yaml

@@ -5,14 +5,19 @@ api_service_name: 'aiplatform.googleapis.com'
 api_version: 'v1beta1'
 api_resource_type_kind: 'NotebookRuntime'
 fields:
+  - field: 'auto_upgrade'
+    provider_only: true
   - field: 'description'
   - field: 'desired_state'
     provider_only: true
   - field: 'display_name'
+  - field: 'expiration_time'
+  - field: 'is_upgradable'
   - field: 'location'
     provider_only: true
   - field: 'name'
     provider_only: true
   - field: 'notebook_runtime_template_ref.notebook_runtime_template'
+  - field: 'notebook_runtime_type'
   - field: 'runtime_user'
   - field: 'state'

google-beta/services/colab/resource_colab_runtime_generated_test.go

@@ -170,7 +170,7 @@ func TestAccColabRuntime_colabRuntimeFullExample(t *testing.T) {
 				ResourceName:            "google_colab_runtime.runtime",
 				ImportState:             true,
 				ImportStateVerify:       true,
-				ImportStateVerifyIgnore: []string{"desired_state", "location", "name"},
+				ImportStateVerifyIgnore: []string{"auto_upgrade", "desired_state", "location", "name"},
 			},
 		},
 	})
@@ -235,6 +235,8 @@ resource "google_colab_runtime" "runtime" {
 
   desired_state = "ACTIVE"
 
+  auto_upgrade = true
+
   depends_on = [
     google_colab_runtime_template.my_template
   ]

google-beta/services/colab/resource_colab_runtime_test.go

@@ -31,7 +31,7 @@ func TestAccColabRuntime_update(t *testing.T) {
 				ResourceName:            "google_colab_runtime.runtime",
 				ImportState:             true,
 				ImportStateVerify:       true,
-				ImportStateVerifyIgnore: []string{"desired_state", "location", "name"},
+				ImportStateVerifyIgnore: []string{"desired_state", "location", "name", "auto_upgrade"},
 			},
 			{
 				Config: testAccColabRuntime_no_state(context),
@@ -40,7 +40,7 @@ func TestAccColabRuntime_update(t *testing.T) {
 				ResourceName:            "google_colab_runtime.runtime",
 				ImportState:             true,
 				ImportStateVerify:       true,
-				ImportStateVerifyIgnore: []string{"desired_state", "location", "name"},
+				ImportStateVerifyIgnore: []string{"desired_state", "location", "name", "auto_upgrade"},
 			},
 			{
 				Config: testAccColabRuntime_stopped(context),
@@ -54,7 +54,21 @@ func TestAccColabRuntime_update(t *testing.T) {
 				ResourceName:            "google_colab_runtime.runtime",
 				ImportState:             true,
 				ImportStateVerify:       true,
-				ImportStateVerifyIgnore: []string{"desired_state", "location", "name"},
+				ImportStateVerifyIgnore: []string{"desired_state", "location", "name", "auto_upgrade"},
+			},
+			{
+				Config: testAccColabRuntime_full(context),
+				ConfigPlanChecks: resource.ConfigPlanChecks{
+					PreApply: []plancheck.PlanCheck{
+						plancheck.ExpectResourceAction("google_colab_runtime.runtime", plancheck.ResourceActionUpdate),
+					},
+				},
+			},
+			{
+				ResourceName:            "google_colab_runtime.runtime",
+				ImportState:             true,
+				ImportStateVerify:       true,
+				ImportStateVerifyIgnore: []string{"desired_state", "location", "name", "auto_upgrade"},
 			},
 		},
 	})
@@ -117,7 +131,9 @@ resource "google_colab_runtime" "runtime" {
   runtime_user = "[email protected]"
   description = "Full runtime"
 
-  desired_state = "ACTIVE"
+  desired_state = "RUNNING"
+
+  auto_upgrade = true
 
   depends_on = [
     google_colab_runtime_template.my_template

website/docs/r/colab_runtime.html.markdown

@@ -174,6 +174,8 @@ resource "google_colab_runtime" "runtime" {
 
   desired_state = "ACTIVE"
 
+  auto_upgrade = true
+
   depends_on = [
     google_colab_runtime_template.my_template
   ]
@@ -219,6 +221,8 @@ The following arguments are supported:
 
 * `desired_state` - (Optional) Desired state of the Colab Runtime. Set this field to `RUNNING` to start the runtime, and `STOPPED` to stop it.
 
+* `auto_upgrade` - (Optional) Triggers an upgrade anytime the runtime is started if it is upgradable.
+
 
 <a name="nested_notebook_runtime_template_ref"></a>The `notebook_runtime_template_ref` block supports:
 
@@ -235,6 +239,15 @@ In addition to the arguments listed above, the following computed attributes are
 * `state` -
   Output only. The state of the runtime.
 
+* `is_upgradable` -
+  Output only. Checks if the NotebookRuntime is upgradable.
+
+* `expiration_time` -
+  Output only. Timestamp when this NotebookRuntime will be expired.
+
+* `notebook_runtime_type` -
+  Output only. The type of the notebook runtime.
+
 
 ## Timeouts