From ebf4a2a4bdc9e4b639949ae2d2542f2f1358d559 Mon Sep 17 00:00:00 2001
From: Nitin Garg <gargnitin@google.com>
Date: Tue, 17 Sep 2024 05:53:36 +0000
Subject: [PATCH] Move code to access bucket iam role to backend

---
 .../testing_on_gke/examples/dlio/run_tests.py | 12 +++-
 .../testing_on_gke/examples/fio/run_tests.py  | 14 ++++-
 .../testing_on_gke/examples/run-gke-tests.sh  | 14 +----
 .../examples/utils/run_tests_common.py        | 58 +++++++++++++++++++
 4 files changed, 83 insertions(+), 15 deletions(-)

diff --git a/perfmetrics/scripts/testing_on_gke/examples/dlio/run_tests.py b/perfmetrics/scripts/testing_on_gke/examples/dlio/run_tests.py
index 2fc859782f..7be36bcdb5 100644
--- a/perfmetrics/scripts/testing_on_gke/examples/dlio/run_tests.py
+++ b/perfmetrics/scripts/testing_on_gke/examples/dlio/run_tests.py
@@ -30,7 +30,7 @@
 
 # local imports from other directories
 sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'utils'))
-from run_tests_common import escape_commas_in_string, parse_args, run_command
+from run_tests_common import escape_commas_in_string, parse_args, run_command, add_iam_role_for_buckets
 
 # local imports from same directory
 import dlio_workload
@@ -79,6 +79,16 @@ def main(args) -> None:
       args.instance_id,
       args.machine_type,
   )
+  buckets = [dlioWorkload.bucket for dlioWorkload in dlioWorkloads]
+  role = 'roles/storage.objectUser'
+  add_iam_role_for_buckets(
+      buckets,
+      role,
+      args.project_id,
+      args.project_number,
+      args.namespace,
+      args.ksa,
+  )
   for helmInstallCommand in helmInstallCommands:
     print(f'{helmInstallCommand}')
     if not args.dry_run:
diff --git a/perfmetrics/scripts/testing_on_gke/examples/fio/run_tests.py b/perfmetrics/scripts/testing_on_gke/examples/fio/run_tests.py
index 0d946d80ab..599852ae78 100644
--- a/perfmetrics/scripts/testing_on_gke/examples/fio/run_tests.py
+++ b/perfmetrics/scripts/testing_on_gke/examples/fio/run_tests.py
@@ -29,7 +29,7 @@
 
 # local imports from other directories
 sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'utils'))
-from run_tests_common import escape_commas_in_string, parse_args, run_command
+from run_tests_common import escape_commas_in_string, parse_args, run_command, add_iam_role_for_buckets
 
 # local imports from same directory
 import fio_workload
@@ -80,6 +80,16 @@ def main(args) -> None:
       args.instance_id,
       args.machine_type,
   )
+  buckets = (fioWorkload.bucket for fioWorkload in fioWorkloads)
+  role = 'roles/storage.objectUser'
+  add_iam_role_for_buckets(
+      buckets,
+      role,
+      args.project_id,
+      args.project_number,
+      args.namespace,
+      args.ksa,
+  )
   for helmInstallCommand in helmInstallCommands:
     print(f'{helmInstallCommand}')
     if not args.dry_run:
@@ -88,4 +98,4 @@ def main(args) -> None:
 
 if __name__ == '__main__':
   args = parse_args()
-main(args)
+  main(args)
diff --git a/perfmetrics/scripts/testing_on_gke/examples/run-gke-tests.sh b/perfmetrics/scripts/testing_on_gke/examples/run-gke-tests.sh
index ec6a2b6bf4..e377fcb202 100755
--- a/perfmetrics/scripts/testing_on_gke/examples/run-gke-tests.sh
+++ b/perfmetrics/scripts/testing_on_gke/examples/run-gke-tests.sh
@@ -433,15 +433,6 @@ function createKubernetesServiceAccountForCluster() {
   kubectl config view --minify | grep namespace:
 }
 
-function addGCSAccessPermissions() {
-  test -f "${workload_config}"
-  grep -wh '\"bucket\"' "${workload_config}" | cut -d: -f2 | cut -d, -f1 | cut -d \" -f2 | sort | uniq | grep -v ' ' | while read workload_bucket; do
-    gcloud storage buckets add-iam-policy-binding gs://${workload_bucket} \
-      --member "principal://iam.googleapis.com/projects/${project_number}/locations/global/workloadIdentityPools/${project_id}.svc.id.goog/subject/ns/${appnamespace}/sa/${ksa}" \
-      --role "roles/storage.objectUser"
-  done
-}
-
 function ensureGcsfuseCode() {
   echo "Ensuring we have gcsfuse code ..."
   # clone gcsfuse code if needed
@@ -519,12 +510,12 @@ function deleteAllPods() {
 
 function deployAllFioHelmCharts() {
   echo "Deploying all fio helm charts ..."
-  cd "${gke_testing_dir}"/examples/fio && python3 ./run_tests.py --workload-config "${workload_config}" --instance-id ${instance_id} --machine-type="${machine_type}" && cd -
+  cd "${gke_testing_dir}"/examples/fio && python3 ./run_tests.py --workload-config "${workload_config}" --instance-id ${instance_id} --machine-type="${machine_type}" --project-id=${project_id} --project-number=${project_number} --namespace=${appnamespace} --ksa=${ksa} && cd -
 }
 
 function deployAllDlioHelmCharts() {
   echo "Deploying all dlio helm charts ..."
-  cd "${gke_testing_dir}"/examples/dlio && python3 ./run_tests.py --workload-config "${workload_config}" --instance-id ${instance_id} --machine-type="${machine_type}" && cd -
+  cd "${gke_testing_dir}"/examples/dlio && python3 ./run_tests.py --workload-config "${workload_config}" --instance-id ${instance_id} --machine-type="${machine_type}" --project-id=${project_id} --project-number=${project_number} --namespace=${appnamespace} --ksa=${ksa} && cd -
 }
 
 function listAllHelmCharts() {
@@ -620,7 +611,6 @@ createKubernetesServiceAccountForCluster
 ensureGcsfuseCode
 
 # GCP/GKE configuration dependent on GCSFuse/CSI driver source code
-addGCSAccessPermissions
 createCustomCsiDriverIfNeeded
 
 # Run latest workload configuration
diff --git a/perfmetrics/scripts/testing_on_gke/examples/utils/run_tests_common.py b/perfmetrics/scripts/testing_on_gke/examples/utils/run_tests_common.py
index 13d56f7b85..e2c2aa63e2 100644
--- a/perfmetrics/scripts/testing_on_gke/examples/utils/run_tests_common.py
+++ b/perfmetrics/scripts/testing_on_gke/examples/utils/run_tests_common.py
@@ -70,6 +70,33 @@ def parse_args():
       help='Machine-type of the GCE VM or GKE cluster node e.g. n2-standard-32',
       required=True,
   )
+  parser.add_argument(
+      '--project-id',
+      metavar='project-id of the user gke cluster',
+      help='project-id of the user gke cluster e.g. gcs-fuse-test',
+      required=True,
+  )
+  parser.add_argument(
+      '--project-number',
+      metavar='project-number of the user gke cluster',
+      help='project-number of the user gke cluster e.g. 927584127901',
+      required=True,
+      type=int,
+  )
+  parser.add_argument(
+      '--namespace',
+      metavar='kubectl namespace of the user',
+      help='kubectl namespace of the user e.g. default',
+      required=False,
+      default='default',
+  )
+  parser.add_argument(
+      '--ksa',
+      metavar='kubernetes service account of the user',
+      help='kubernetest service account of the user e.g. default',
+      required=False,
+      default='default',
+  )
   parser.add_argument(
       '-n',
       '--dry-run',
@@ -84,6 +111,9 @@ def parse_args():
   for argument in [
       'instance_id',
       'machine_type',
+      'project_id',
+      'namespace',
+      'ksa',
   ]:
     value = getattr(args, argument)
     if len(value) == 0 or str.isspace(value):
@@ -98,3 +128,31 @@ def parse_args():
       )
 
   return args
+
+
+def add_iam_role_for_buckets(
+    buckets: set,
+    role: str,
+    project_id: str,
+    project_number: str,
+    namespace: str,
+    ksa: str,
+):
+  print(
+      f'Adding role {role} to all the relevant buckets to'
+      f' ksa={ksa} in namespace={namespace} ...\n\n'
+  )
+  for bucket in buckets:
+    command = (
+        f'gcloud storage buckets add-iam-policy-binding gs://{bucket} --member'
+        f' principal://iam.googleapis.com/projects/{project_number}/locations/global/workloadIdentityPools/{project_id}.svc.id.goog/subject/ns/{namespace}/sa/{ksa} --role'
+        f' {role}'
+    )
+    print(command)
+    ret = run_command(command)
+    if ret != 0:
+      raise Exception(
+          f'Failed to add role {role} for {bucket}: exit-code={ret}'
+      )
+
+  pass