Make colab runtime startable/stoppable (#12904)

mmv1/products/colab/Runtime.yaml

@@ -31,13 +31,24 @@ async:
     full_url: 'https://{{location}}-aiplatform.googleapis.com/v1/{{op_id}}'
 custom_code:
   encoder: 'templates/terraform/encoders/colab_runtime.go.tmpl'
+  post_create: 'templates/terraform/post_create/colab_runtime.go.tmpl'
+  custom_update: 'templates/terraform/custom_update/colab_runtime.go.tmpl'
+  constants: 'templates/terraform/constants/colab_runtime.go.tmpl'
 examples:
   - name: 'colab_runtime_basic'
     primary_resource_id: 'runtime'
     primary_resource_name: 'fmt.Sprintf("tf-test-colab-runtime%s", context["random_suffix"])'
     region_override: 'us-central1'
     vars:
       runtime_name: 'colab-runtime'
+  - name: 'colab_runtime_stopped'
+    primary_resource_id: 'runtime'
+    primary_resource_name: 'fmt.Sprintf("tf-test-colab-runtime%s", context["random_suffix"])'
+    region_override: 'us-central1'
+    vars:
+      runtime_name: 'colab-runtime'
+    ignore_read_extra:
+      - 'desired_state'
   - name: 'colab_runtime_full'
     primary_resource_id: 'runtime'
     primary_resource_name: 'fmt.Sprintf("tf-test-colab-runtime%s", context["random_suffix"])'
@@ -47,16 +58,26 @@ examples:
       key_name: 'my-crypto-key'
     test_vars_overrides:
       key_name: 'acctest.BootstrapKMSKeyInLocation(t, "us-central1").CryptoKey.Name'
+    ignore_read_extra:
+      - 'desired_state'
+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"
 parameters:
   - name: 'location'
     type: String
     required: true
     url_param_only: true
     description: 'The location for the resource: https://cloud.google.com/colab/docs/locations'
+    immutable: true
   - name: 'name'
     type: String
     url_param_only: true
     description: 'The resource name of the Runtime'
+    immutable: true
 properties:
   - name: notebookRuntimeTemplateRef
     type: NestedObject
@@ -65,18 +86,27 @@ properties:
     properties:
       - name: 'notebookRuntimeTemplate'
         type: String
+        immutable: true
         required: true
         description: 'The resource name of the NotebookRuntimeTemplate based on which a NotebookRuntime will be created.'
         diff_suppress_func: 'tpgresource.ProjectNumberDiffSuppress'
   - name: 'runtimeUser'
     type: String
     required: true
+    immutable: true
     description: 'The user email of the NotebookRuntime.'
   - name: 'displayName'
     type: String
+    immutable: true
     description:
       Required. The display name of the Runtime.
     required: true
   - name: description
     type: String
+    immutable: true
     description: 'The description of the Runtime.'
+  - name: state
+    type: String
+    description: |
+      Output only. The state of the runtime.
+    output: true

mmv1/products/colab/RuntimeTemplate.yaml

@@ -41,7 +41,7 @@ iam_policy:
     - 'projects/{{project}}/locations/{{location}}/notebookRuntimeTemplates/{{runtime_template}}'
     - '{{runtime_template}}'
 custom_code:
-  post_create: 'templates/terraform/post_create/colab_runtime_template.tmpl'
+  post_create: 'templates/terraform/post_create/colab_runtime_template.go.tmpl'
 examples:
   - name: 'colab_runtime_template_basic'
     primary_resource_id: 'runtime-template'

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

@@ -0,0 +1,42 @@
+func ModifyColabRuntimeOperation(config *transport_tpg.Config, d *schema.ResourceData, project string, billingProject string, userAgent string, method string) (map[string]interface{}, error) {
+	url, err := tpgresource.ReplaceVars(d, config, "{{"{{"}}ColabBasePath{{"}}"}}projects/{{"{{"}}project{{"}}"}}/locations/{{"{{"}}location{{"}}"}}/notebookRuntimes/{{"{{"}}name{{"}}"}}:"+method)
+	if err != nil {
+		return nil, err
+	}
+
+	res, err := transport_tpg.SendRequest(transport_tpg.SendRequestOptions{
+		Config: config,
+		Method: "POST",
+		Project: billingProject,
+		RawURL: url,
+		UserAgent: userAgent,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("Unable to %q google_colab_runtime %q: %s", method, d.Id(), err)
+	}
+	return res, nil
+}
+
+{{- if ne $.Compiler "terraformgoogleconversion-codegen" }}
+func waitForColabOperation(config *transport_tpg.Config, d *schema.ResourceData, project string, billingProject string, userAgent string, response map[string]interface{}) error {
+	var opRes map[string]interface{}
+	err := ColabOperationWaitTimeWithResponse(
+		config, response, &opRes, project, "Waiting for Colab Runtime Operation", userAgent,
+		d.Timeout(schema.TimeoutUpdate))
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func ModifyColabRuntime(config *transport_tpg.Config, d *schema.ResourceData, project string, billingProject string, userAgent string, method string) error {
+	dRes, err := ModifyColabRuntimeOperation(config, d, project, billingProject, userAgent, method)
+	if err != nil {
+		return err
+	}
+	if err := waitForColabOperation(config, d, project, billingProject, userAgent, dRes); err != nil {
+		return fmt.Errorf("Error with Colab runtime method: %s", err)
+	}
+	return nil
+}
+{{- end }}

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

@@ -0,0 +1,40 @@
+name := d.Get("name").(string)
+state := d.Get("state").(string)
+desired_state := d.Get("desired_state").(string)
+
+project, err := tpgresource.GetProject(d, config)
+if err != nil {
+ 	return err
+}
+
+userAgent, err := tpgresource.GenerateUserAgentString(d, config.UserAgent)
+if err != nil {
+  return err
+}
+
+billingProject := ""
+if bp, err := tpgresource.GetBillingProject(d, config); err == nil {
+  billingProject = bp
+}
+
+if desired_state != "" && state != desired_state {
+  var verb string
+
+  switch desired_state {
+    case "STOPPED":
+      verb = "stop"
+    case "RUNNING":
+      verb = "start"
+    default:
+      return fmt.Errorf("desired_state has to be RUNNING or STOPPED")
+  }
+
+  if err := ModifyColabRuntime(config, d, project, billingProject, userAgent, verb); err != nil {
+    return err
+  }
+
+} else {
+  log.Printf("[DEBUG] Colab runtime %q has state %q.", name, state)
+}
+
+return nil

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

@@ -53,6 +53,8 @@ resource "google_colab_runtime" "{{$.PrimaryResourceId}}" {
   runtime_user = "gterraformtestuser@gmail.com"
   description = "Full runtime"
 
+  desired_state = "ACTIVE"
+
   depends_on = [
     google_colab_runtime_template.my_template
   ]

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

@@ -0,0 +1,31 @@
+resource "google_colab_runtime_template" "my_template" {
+  name = "{{index $.Vars "runtime_name"}}"
+  display_name = "Runtime template basic"
+  location = "us-central1"
+
+  machine_spec {
+    machine_type     = "e2-standard-4"
+  }
+
+  network_spec {
+    enable_internet_access = true
+  }
+}
+
+resource "google_colab_runtime" "{{$.PrimaryResourceId}}" {
+  name = "{{index $.Vars "runtime_name"}}"
+  location = "us-central1" 
+
+  notebook_runtime_template_ref {
+    notebook_runtime_template = google_colab_runtime_template.my_template.id
+  }
+
+  desired_state = "STOPPED"
+
+  display_name = "Runtime stopped"
+  runtime_user = "gterraformtestuser@gmail.com"
+
+  depends_on = [
+    google_colab_runtime_template.my_template,
+  ]
+}

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

@@ -0,0 +1,5 @@
+if p, ok := d.GetOk("desired_state"); ok && p.(string) == "STOPPED" {
+  if err := ModifyColabRuntime(config, d, project, billingProject, userAgent, "stop"); err != nil {
+    return err
+  }
+}