From 7dc14a1cb435a1c7c7cdedff569c33e9fce96392 Mon Sep 17 00:00:00 2001 From: Ziyuan Qin Date: Sun, 11 Feb 2024 14:56:00 -0800 Subject: [PATCH] move azure check and service principal confirmation prompts to for_cli function --- src/databricks/labs/ucx/cli.py | 2 +- .../labs/ucx/migration/azure_credentials.py | 23 ++++++++++--------- 2 files changed, 13 insertions(+), 12 deletions(-) diff --git a/src/databricks/labs/ucx/cli.py b/src/databricks/labs/ucx/cli.py index 0c118d146c..019e732889 100644 --- a/src/databricks/labs/ucx/cli.py +++ b/src/databricks/labs/ucx/cli.py @@ -299,7 +299,7 @@ def migrate_azure_service_principals(w: WorkspaceClient): """ logger.info("Running migrate_azure_service_principals command") prompts = Prompts() - service_principal_migration = AzureServicePrincipalMigration.for_cli(w) + service_principal_migration = AzureServicePrincipalMigration.for_cli(w, prompts) service_principal_migration.execute_migration(prompts) return diff --git a/src/databricks/labs/ucx/migration/azure_credentials.py b/src/databricks/labs/ucx/migration/azure_credentials.py index 47a458079a..e05a6047c7 100644 --- a/src/databricks/labs/ucx/migration/azure_credentials.py +++ b/src/databricks/labs/ucx/migration/azure_credentials.py @@ -79,7 +79,18 @@ def __init__( self._action_plan = 'service_principals_for_storage_credentials.csv' @classmethod - def for_cli(cls, ws: WorkspaceClient, product='ucx'): + def for_cli(cls, ws: WorkspaceClient, prompts: Prompts, product='ucx'): + if not ws.config.is_azure: + logger.error("Workspace is not on azure, please run this command on azure databricks workspaces.") + return + + csv_confirmed = prompts.confirm( + "Have you reviewed the azure_storage_account_info.csv " + "and confirm listed service principals are allowed to be checked for migration?" + ) + if csv_confirmed is not True: + return + installation = Installation.current(ws, product) config = installation.load(WorkspaceConfig) sql_backend = StatementExecutionBackend(ws, config.warehouse_id) @@ -226,16 +237,6 @@ def _validate_storage_credential(self, storage_credential, location) -> StorageC return StorageCredentialValidationResult.from_storage_credential_validation(storage_credential, validation) def execute_migration(self, prompts: Prompts): - if not self._ws.config.is_azure: - logger.error("Workspace is not on azure, please run this command on azure databricks workspaces.") - return - - csv_confirmed = prompts.confirm( - "Have you reviewed the azure_storage_account_info.csv " - "and confirm listed service principals are allowed to be checked for migration?" - ) - if csv_confirmed is not True: - return self._generate_migration_list()