From b67dff839793c3786e98cd72f82f79cfbebe5400 Mon Sep 17 00:00:00 2001 From: Stepan Rasputny Date: Wed, 16 Oct 2024 00:37:11 +0200 Subject: [PATCH] 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 | 128 +++++++++++++++++++++++ batch/create_with_secret_manager_test.go | 58 ++++++++++ 2 files changed, 186 insertions(+) create mode 100644 batch/create_with_secret_manager.go create mode 100644 batch/create_with_secret_manager_test.go diff --git a/batch/create_with_secret_manager.go b/batch/create_with_secret_manager.go new file mode 100644 index 0000000000..c3b2fde341 --- /dev/null +++ b/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] diff --git a/batch/create_with_secret_manager_test.go b/batch/create_with_secret_manager_test.go new file mode 100644 index 0000000000..12d3374fdf --- /dev/null +++ b/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) + } +}