feat: batch create job with secret manager sample (GoogleCloudPlatform#11911)

* feat: batch create job with secret manager sample

* fix: added name of the secret

* fix: change secret name and service account

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* Trying different arrangement

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

---------

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
Co-authored-by: Maciej Strzelczyk <[email protected]>

Author: 3 people

batch/create/create_with_secret_manager.py

@@ -0,0 +1,117 @@
+#  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
+#
+#      http://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.
+
+# [START batch_create_using_secret_manager]
+from typing import Dict, Optional
+
+from google.cloud import batch_v1
+
+
+def create_with_secret_manager(
+    project_id: str,
+    region: str,
+    job_name: str,
+    secrets: Dict[str, str],
+    service_account_email: Optional[str] = None,
+) -> batch_v1.Job:
+    """
+    This method shows how to create a sample Batch Job that will run
+    a simple command on Cloud Compute instances 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.
+
+    Args:
+        project_id: project ID or project number of the Cloud project you want to use.
+        region: name of the region you want to use to run the job. Regions that are
+            available for Batch are listed on: https://cloud.google.com/batch/docs/get-started#locations
+        job_name: the name of the job that will be created.
+            It needs to be unique for each project and region pair.
+        secrets: secrets, which should be passed to the job. Environment variables should be capitalized
+        by convention https://google.github.io/styleguide/shellguide.html#constants-and-environment-variable-names
+            The format should look like:
+                - {'SECRET_NAME': 'projects/{project_id}/secrets/{SECRET_NAME}/versions/{version}'}
+            version can be set to 'latest'.
+        service_account_email (optional): custom service account email
+
+    Returns:
+        A job object representing the job created.
+    """
+    client = batch_v1.BatchServiceClient()
+
+    # Define what will be done as part of the job.
+    task = batch_v1.TaskSpec()
+    runnable = batch_v1.Runnable()
+    runnable.script = batch_v1.Runnable.Script()
+    runnable.script.text = (
+        "echo Hello world! from task ${BATCH_TASK_INDEX}."
+        + f" ${next(iter(secrets.keys()))} is the value of the secret."
+    )
+    task.runnables = [runnable]
+    task.max_retry_count = 2
+    task.max_run_duration = "3600s"
+
+    envable = batch_v1.Environment()
+    envable.secret_variables = secrets
+    task.environment = envable
+
+    # Tasks are grouped inside a job using TaskGroups.
+    # Currently, it's possible to have only one task group.
+    group = batch_v1.TaskGroup()
+    group.task_count = 4
+    group.task_spec = task
+
+    # Policies are used to define on what kind of virtual machines the tasks will run on.
+    # Read more about local disks here: https://cloud.google.com/compute/docs/disks/persistent-disks
+    policy = batch_v1.AllocationPolicy.InstancePolicy()
+    policy.machine_type = "e2-standard-4"
+    instances = batch_v1.AllocationPolicy.InstancePolicyOrTemplate()
+    instances.policy = policy
+    allocation_policy = batch_v1.AllocationPolicy()
+    allocation_policy.instances = [instances]
+
+    service_account = batch_v1.ServiceAccount()
+    service_account.email = service_account_email
+    allocation_policy.service_account = service_account
+
+    job = batch_v1.Job()
+    job.task_groups = [group]
+    job.allocation_policy = allocation_policy
+    job.labels = {"env": "testing", "type": "script"}
+    # We use Cloud Logging as it's an out of the box available option
+    job.logs_policy = batch_v1.LogsPolicy()
+    job.logs_policy.destination = batch_v1.LogsPolicy.Destination.CLOUD_LOGGING
+
+    create_request = batch_v1.CreateJobRequest()
+    create_request.job = job
+    create_request.job_id = job_name
+    # The job's parent is the region in which the job will run
+    create_request.parent = f"projects/{project_id}/locations/{region}"
+
+    return client.create_job(create_request)
+
+
+# [END batch_create_using_secret_manager]
+
+if __name__ == "__main__":
+    import google.auth
+
+    PROJECT = google.auth.default()[1]
+    REGION = "europe-west4"
+    # Existing service account name within the project specified above.
+    name = "test-account-name"
+    secret_name = "TEST_SECRET"
+    secrets = {secret_name: f"projects/11111111/secrets/{secret_name}/versions/latest"}
+    job = create_with_secret_manager(
+        PROJECT, REGION, "secret-manager-job-batch", secrets
+    )

batch/tests/test_basics.py

@@ -21,14 +21,14 @@
 from flaky import flaky
 
 import google.auth
-from google.cloud import batch_v1
-from google.cloud import resourcemanager_v3
+from google.cloud import batch_v1, resourcemanager_v3
 import pytest
 
 from ..create.create_with_container_no_mounting import create_container_job
 from ..create.create_with_gpu_no_mounting import create_gpu_job
 from ..create.create_with_persistent_disk import create_with_pd_job
 from ..create.create_with_script_no_mounting import create_script_job
+from ..create.create_with_secret_manager import create_with_secret_manager
 from ..create.create_with_service_account import create_with_custom_service_account_job
 from ..create.create_with_ssd import create_local_ssd_job
 
@@ -42,6 +42,12 @@
 PROJECT = google.auth.default()[1]
 REGION = "europe-central2"
 ZONE = "europe-central2-b"
+SECRET_NAME = "PERMANENT_BATCH_TESTING"
+PROJECT_NUMBER = (
+    resourcemanager_v3.ProjectsClient()
+    .get_project(name=f"projects/{PROJECT}")
+    .name.split("/")[1]
+)
 
 TIMEOUT = 600  # 10 minutes
 
@@ -61,12 +67,7 @@ def job_name():
 
 @pytest.fixture()
 def service_account() -> str:
-    client = resourcemanager_v3.ProjectsClient()
-    request = resourcemanager_v3.GetProjectRequest()
-    request.name = f"projects/{PROJECT}"
-    project = client.get_project(request)
-    project_number = project.name.split("/")[-1]
-    return f"{project_number}[email protected]"
+    return f"{PROJECT_NUMBER}[email protected]"
 
 
 @pytest.fixture
@@ -132,6 +133,10 @@ def _check_service_account(job: batch_v1.Job, service_account_email: str):
     assert job.allocation_policy.service_account.email == service_account_email
 
 
+def _check_secret_set(job: batch_v1.Job, secret_name: str):
+    assert secret_name in job.task_groups[0].task_spec.environment.secret_variables
+
+
 @flaky(max_runs=3, min_passes=1)
 def test_script_job(job_name, capsys):
     job = create_script_job(PROJECT, REGION, job_name)
@@ -160,6 +165,17 @@ def test_service_account_job(job_name, service_account):
     )
 
 
+@flaky(max_runs=3, min_passes=1)
+def test_secret_manager_job(job_name, service_account):
+    secrets = {
+        SECRET_NAME: f"projects/{PROJECT_NUMBER}/secrets/{SECRET_NAME}/versions/latest"
+    }
+    job = create_with_secret_manager(
+        PROJECT, REGION, job_name, secrets, service_account
+    )
+    _test_body(job, additional_test=lambda: _check_secret_set(job, SECRET_NAME))
+
+
 @flaky(max_runs=3, min_passes=1)
 def test_ssd_job(job_name: str, disk_name: str, capsys: "pytest.CaptureFixture[str]"):
     job = create_local_ssd_job(PROJECT, REGION, job_name, disk_name)