Auto-upgrade for colab runtimes

mmv1/products/colab/Runtime.yaml

@@ -60,12 +60,17 @@ examples:
       key_name: 'acctest.BootstrapKMSKeyInLocation(t, "us-central1").CryptoKey.Name'
     ignore_read_extra:
       - 'desired_state'
+      - 'auto_upgrade'
 virtual_fields:
   - name: 'desired_state'
     description: |
       Desired state of the Colab Runtime. Set this field to `RUNNING` to start the runtime, and `STOPPED` to stop it.
     type: String
     default_value: "RUNNING"
+  - name: 'auto_upgrade'
+    description: |
+      Triggers an upgrade anytime the runtime is started if it is upgradable.
+    type: Boolean
 parameters:
   - name: 'location'
     type: String
@@ -110,3 +115,18 @@ properties:
     description: |
       Output only. The state of the runtime.
     output: true
+  - name: isUpgradable
+    type: Boolean
+    description: |
+      Output only. Checks if the NotebookRuntime is upgradable.
+    output: true
+  - name: expirationTime
+    type: String
+    description: |
+      Output only. Timestamp when this NotebookRuntime will be expired.
+    output: true
+  - name: notebookRuntimeType
+    type: String
+    description: |
+      Output only. The type of the notebook runtime.
+    output: true

mmv1/templates/terraform/custom_update/colab_runtime.go.tmpl

@@ -37,4 +37,25 @@ if desired_state != "" && state != desired_state {
   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

mmv1/templates/terraform/examples/colab_runtime_full.tf.tmpl

@@ -55,6 +55,8 @@ resource "google_colab_runtime" "{{$.PrimaryResourceId}}" {
 
   desired_state = "ACTIVE"
 
+  auto_upgrade = true
+
   depends_on = [
     google_colab_runtime_template.my_template
   ]

mmv1/third_party/terraform/services/colab/resource_colab_runtime_test.go

@@ -29,7 +29,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),
@@ -38,7 +38,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),
@@ -52,7 +52,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"},
 			},
 		},
 	})
@@ -115,7 +129,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