feat: create batch using secret manager (#4382)

* feat: create batch using secret manager

* extend error trace for failed job

* Revert extending error trace

batch/create_with_secret_manager.go

@@ -0,0 +1,128 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package snippets
+
+// [START batch_create_using_secret_manager]
+import (
+	"context"
+	"fmt"
+	"io"
+
+	batch "cloud.google.com/go/batch/apiv1"
+	"cloud.google.com/go/batch/apiv1/batchpb"
+	durationpb "google.golang.org/protobuf/types/known/durationpb"
+)
+
+// Creates and runs a job with passing secrets from secret manager.
+// Note: Job's service account should have the permissions to access secrets.
+//   - Secret Manager Secret Accessor (roles/secretmanager.secretAccessor) IAM role.
+func createJobWithSecretManager(w io.Writer, projectID, jobName string, secrets map[string]string) (*batchpb.Job, error) {
+	// projectID := "your_project_id"
+	// region := "us-central1"
+	// jobName := "some-job"
+	// secrets := map[string]string{"SECRET_NAME": "projects/{projectID}/secrets/{SECRET_NAME}/versions/{version}"}
+
+	// Note: Environment variables should be capitalized by convention
+	// https://google.github.io/styleguide/shellguide.html#constants-and-environment-variable-names
+	// Version of the secret can be set to "latest"
+
+	region := "us-central1"
+	displayName1 := "script 1"
+
+	ctx := context.Background()
+	batchClient, err := batch.NewClient(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("batchClient error: %w", err)
+	}
+	defer batchClient.Close()
+
+	runn1 := &batchpb.Runnable{
+		Executable: &batchpb.Runnable_Script_{
+			Script: &batchpb.Runnable_Script{
+				Command: &batchpb.Runnable_Script_Text{
+					Text: "echo Hello world from script 1 for task ${BATCH_TASK_INDEX}",
+				},
+			},
+		},
+		DisplayName: displayName1,
+	}
+
+	taskSpec := &batchpb.TaskSpec{
+		ComputeResource: &batchpb.ComputeResource{
+			// CpuMilli is milliseconds per cpu-second. This means the task requires 2 whole CPUs.
+			CpuMilli:  2000,
+			MemoryMib: 16,
+		},
+		MaxRunDuration: &durationpb.Duration{
+			Seconds: 3600,
+		},
+		MaxRetryCount: 2,
+		Runnables:     []*batchpb.Runnable{runn1},
+		Environment: &batchpb.Environment{
+			SecretVariables: secrets,
+		},
+	}
+
+	taskGroups := []*batchpb.TaskGroup{
+		{
+			TaskCount: 4,
+			TaskSpec:  taskSpec,
+		},
+	}
+
+	labels := map[string]string{"env": "testing", "type": "container"}
+
+	// Policies are used to define on what kind of virtual machines the tasks will run on.
+	// In this case, we tell the system to use "e2-standard-4" machine type.
+	// Read more about machine types here: https://cloud.google.com/compute/docs/machine-types
+	allocationPolicy := &batchpb.AllocationPolicy{
+		Instances: []*batchpb.AllocationPolicy_InstancePolicyOrTemplate{{
+			PolicyTemplate: &batchpb.AllocationPolicy_InstancePolicyOrTemplate_Policy{
+				Policy: &batchpb.AllocationPolicy_InstancePolicy{
+					MachineType: "e2-standard-4",
+				},
+			},
+		}},
+	}
+
+	// We use Cloud Logging as it's an out of the box available option
+	logsPolicy := &batchpb.LogsPolicy{
+		Destination: batchpb.LogsPolicy_CLOUD_LOGGING,
+	}
+
+	job := &batchpb.Job{
+		Name:             jobName,
+		TaskGroups:       taskGroups,
+		AllocationPolicy: allocationPolicy,
+		Labels:           labels,
+		LogsPolicy:       logsPolicy,
+	}
+
+	request := &batchpb.CreateJobRequest{
+		Parent: fmt.Sprintf("projects/%s/locations/%s", projectID, region),
+		JobId:  jobName,
+		Job:    job,
+	}
+
+	created_job, err := batchClient.CreateJob(ctx, request)
+	if err != nil {
+		return nil, fmt.Errorf("unable to create job: %w", err)
+	}
+
+	fmt.Fprintf(w, "Job created: %v\n", created_job)
+	return created_job, nil
+}
+
+// [END batch_create_using_secret_manager]

batch/create_with_secret_manager_test.go

@@ -0,0 +1,58 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package snippets
+
+import (
+	"bytes"
+	"fmt"
+	"math/rand"
+	"reflect"
+	"testing"
+	"time"
+
+	"github.com/GoogleCloudPlatform/golang-samples/internal/testutil"
+)
+
+func TestCreateJobWithSecretManager(t *testing.T) {
+	buf := &bytes.Buffer{}
+	var r *rand.Rand = rand.New(
+		rand.NewSource(time.Now().UnixNano()))
+	tc := testutil.SystemTest(t)
+	jobName := fmt.Sprintf("test-job-go-%v-%v", time.Now().Format("2006-01-02"), r.Int())
+	region := "us-central1"
+	secretName := "PERMANENT_BATCH_TESTING"
+	wantSecrets := map[string]string{
+		secretName: fmt.Sprintf("projects/%s/secrets/%s/versions/latest", tc.ProjectID, secretName),
+	}
+
+	job, err := createJobWithSecretManager(buf, tc.ProjectID, jobName, wantSecrets)
+	if err != nil {
+		t.Errorf("createJobWithPD got err: %v", err)
+	}
+
+	_, err = jobSucceeded(tc.ProjectID, region, jobName)
+	if err != nil {
+		t.Errorf("Could not verify job completion: %v", err)
+	}
+
+	gotSecrets := job.GetTaskGroups()[0].GetTaskSpec().GetEnvironment().GetSecretVariables()
+	if !reflect.DeepEqual(gotSecrets, wantSecrets) {
+		t.Errorf("No secrets set")
+	}
+
+	if err := deleteJob(buf, tc.ProjectID, region, jobName); err != nil {
+		t.Errorf("deleteJob got err: %v", err)
+	}
+}