Options to set jobs as single user or shared to allow users to change their jobs config to UC friendly configs

dbclient/JobsClient.py

@@ -227,3 +227,46 @@ def delete_all_jobs(self):
         job_list = self.get('/jobs/list').get('jobs', [])
         for job in job_list:
             self.post('/jobs/delete', {'job_id': job['job_id']})
+
+    # MTJ Jobs not supported
+    def single_user_all_jobs(self):
+        job_list = self.get_jobs_list()
+        for job_conf in job_list:
+            job_settings = job_conf['settings']
+            job_clusters = job_settings.get('new_cluster', None)
+
+            if job_clusters:
+                job_clusters['data_security_mode'] = 'LEGACY_SINGLE_USER_STANDARD'
+                job_settings['new_cluster'] = job_clusters
+                update_job_conf = {'job_id': job_conf['job_id'],
+                                   'new_settings': job_settings}
+                self.post('/jobs/reset', update_job_conf)
+
+    # MTJ Jobs not supported
+    def shared_all_jobs(self):
+        job_list = self.get_jobs_list()
+        for job_conf in job_list:
+            job_settings = job_conf['settings']
+            job_clusters = job_settings.get('new_cluster', None)
+
+            if job_clusters:
+                job_clusters['data_security_mode'] = 'LEGACY_TABLE_ACL'
+                job_settings['new_cluster'] = job_clusters
+                update_job_conf = {'job_id': job_conf['job_id'],
+                                   'new_settings': job_settings}
+                self.post('/jobs/reset', update_job_conf)
+
+    # MTJ Jobs not supported
+    def set_policy_all_jobs(self, policy_id):
+        job_list = self.get_jobs_list()
+        for job_conf in job_list:
+            job_settings = job_conf['settings']
+            job_clusters = job_settings.get('new_cluster', None)
+
+            if job_clusters:
+                job_clusters['policy_id'] = policy_id
+                job_settings['new_cluster'] = job_clusters
+                update_job_conf = {'job_id': job_conf['job_id'],
+                                   'new_settings': job_settings}
+                self.post('/jobs/reset', update_job_conf)
+

dbclient/parser.py

@@ -391,6 +391,15 @@ def get_import_parser():
     parser.add_argument('--delete-all-jobs', action='store_true',
                         help='Delete all jobs')
 
+    parser.add_argument('--single-user-all-jobs', action='store_true',
+                        help='Set all jobs as single user to allow UC enabled clusters')
+
+    parser.add_argument('--shared-all-jobs', action='store_true',
+                        help='Set all jobs as shared to allow UC enabled clusters')
+
+    parser.add_argument('--set-policy-all-jobs', action='store',
+                        help='Set all jobs with the provided policy')
+
     parser.add_argument('--use-checkpoint', action='store_true',
                         help='use checkpointing to restart from previous state')
 

import_db.py

@@ -186,6 +186,33 @@ def main():
         end = timer()
         print("Delete all jobs time: " + str(timedelta(seconds=end - start)))
 
+    if args.single_user_all_jobs:
+        print("Setting all jobs on new clusters as single user mode to enable UC {0}".format(now))
+        start = timer()
+        jobs_c = JobsClient(client_config, checkpoint_service)
+        # log job configs
+        jobs_c.single_user_all_jobs()
+        end = timer()
+        print("Single user all jobs time: " + str(timedelta(seconds=end - start)))
+
+    if args.shared_all_jobs:
+        print("Setting all jobs on new clusters as shared mode to enable UC {0}".format(now))
+        start = timer()
+        jobs_c = JobsClient(client_config, checkpoint_service)
+        # log job configs
+        jobs_c.shared_all_jobs()
+        end = timer()
+        print("Shared access mode all jobs time: " + str(timedelta(seconds=end - start)))
+
+    if args.set_policy_all_jobs:
+        print("Setting all jobs on new clusters with the provided policy {0}".format(now))
+        start = timer()
+        jobs_c = JobsClient(client_config, checkpoint_service)
+        # log job configs
+        jobs_c.set_policy_all_jobs(args.set_policy_all_jobs)
+        end = timer()
+        print("Set policy all jobs time: " + str(timedelta(seconds=end - start)))
+
     if args.single_user:
         user_email = args.single_user
         print(f"Import user {user_email} at {now}")