Added a service account injection (#695)

* Added a service account injection

* Revert this

* Added service account injection

* Fixed cloud build file

* Added gcloud version debug info

* Fixed sa injection

* Removed test file

* Revert CODEOWNERS

.cloud-build/execute_changed_notebooks_cli.py

@@ -62,6 +62,12 @@ parser.add_argument(
     help="The GCP region. This is used to inject a variable value into the notebook before running.",
     required=True,
 )
+parser.add_argument(
+    "--variable_service_account",
+    type=str,
+    help="A service account. This is used to inject a variable value into the notebook before running. This is not the account that will run the notebook.",
+    required=True,
+)
 parser.add_argument(
     "--staging_bucket",
     type=str,
@@ -110,6 +116,7 @@ execute_changed_notebooks_helper.process_and_execute_notebooks(
     artifacts_bucket=args.artifacts_bucket,
     variable_project_id=args.variable_project_id,
     variable_region=args.variable_region,
+    variable_service_account=args.variable_service_account,
     private_pool_id=args.private_pool_id,
     should_parallelize=args.should_parallelize,
     timeout=args.timeout,

.cloud-build/execute_changed_notebooks_helper.py

@@ -73,6 +73,7 @@ def _process_notebook(
     notebook_path: str,
     variable_project_id: str,
     variable_region: str,
+    variable_service_account: str,
 ):
     # Read notebook
     with open(notebook_path) as f:
@@ -84,6 +85,7 @@ def _process_notebook(
         replacement_map={
             "PROJECT_ID": variable_project_id,
             "REGION": variable_region,
+            "SERVICE_ACCOUNT": variable_service_account,
         },
     )
 
@@ -118,6 +120,7 @@ def process_and_execute_notebook(
     artifacts_bucket: str,
     variable_project_id: str,
     variable_region: str,
+    variable_service_account: str,
     private_pool_id: Optional[str],
     deadline: datetime,
     notebook: str,
@@ -152,6 +155,7 @@ def process_and_execute_notebook(
             notebook_path=notebook,
             variable_project_id=variable_project_id,
             variable_region=variable_region,
+            variable_service_account=variable_service_account,
         )
 
         # Upload the pre-processed code to a GCS bucket
@@ -277,6 +281,7 @@ def process_and_execute_notebooks(
     artifacts_bucket: str,
     variable_project_id: str,
     variable_region: str,
+    variable_service_account: str,
     private_pool_id: Optional[str],
     should_parallelize: bool,
     timeout: int,
@@ -336,6 +341,7 @@ def process_and_execute_notebooks(
                             artifacts_bucket,
                             variable_project_id,
                             variable_region,
+                            variable_service_account,
                             private_pool_id,
                             deadline,
                         ),
@@ -350,6 +356,7 @@ def process_and_execute_notebooks(
                     artifacts_bucket=artifacts_bucket,
                     variable_project_id=variable_project_id,
                     variable_region=variable_region,
+                    variable_service_account=variable_service_account,
                     private_pool_id=private_pool_id,
                     deadline=deadline,
                     notebook=notebook,
@@ -406,6 +413,7 @@ def process_and_execute_notebooks(
             notebook_path=notebook,
             variable_project_id=variable_project_id,
             variable_region=variable_region,
+            variable_service_account=variable_service_account,
         )
 
         execute_notebook_helper.execute_notebook(

.cloud-build/notebook-execution-test-cloudbuild.yaml

@@ -44,7 +44,7 @@ steps:
     - -c
     - | 
       . workspace/env/bin/activate &&
-      python3 .cloud-build/execute_changed_notebooks_cli.py --test_paths_file "${_TEST_PATHS_FILE}" --base_branch "${_FORCED_BASE_BRANCH}" --container_uri ${_PYTHON_IMAGE} --staging_bucket ${_GCS_STAGING_BUCKET}  --artifacts_bucket ${_GCS_STAGING_BUCKET}/executed_notebooks/PR_${_PR_NUMBER}/BUILD_${BUILD_ID} --variable_project_id ${PROJECT_ID} --variable_region ${_GCP_REGION} `if [ ! -z "${_PRIVATE_POOL_NAME}" ]; then echo "--private_pool_id ${_PRIVATE_POOL_NAME}"; fi`
+      python3 .cloud-build/execute_changed_notebooks_cli.py --test_paths_file "${_TEST_PATHS_FILE}" --base_branch "${_FORCED_BASE_BRANCH}" --container_uri ${_PYTHON_IMAGE} --staging_bucket ${_GCS_STAGING_BUCKET}  --artifacts_bucket ${_GCS_STAGING_BUCKET}/executed_notebooks/PR_${_PR_NUMBER}/BUILD_${BUILD_ID} --variable_project_id ${PROJECT_ID} --variable_region ${_GCP_REGION} --variable_service_account ${_GCP_SERVICE_ACCOUNT} `if [ ! -z "${_PRIVATE_POOL_NAME}" ]; then echo "--private_pool_id ${_PRIVATE_POOL_NAME}"; fi`
     env:
     - 'IS_TESTING=1'
 timeout: 86400s