diff --git a/dbclient/JobsClient.py b/dbclient/JobsClient.py index e4445820..50debda8 100644 --- a/dbclient/JobsClient.py +++ b/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) + diff --git a/dbclient/parser.py b/dbclient/parser.py index 8cbbb786..11620d8e 100644 --- a/dbclient/parser.py +++ b/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') diff --git a/import_db.py b/import_db.py index 508830a0..a94a5d01 100644 --- a/import_db.py +++ b/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}")