From ea73a1d5367495070a635c9d5c47480bb11f78a3 Mon Sep 17 00:00:00 2001 From: Ziyuan Qin Date: Wed, 14 Feb 2024 15:38:26 -0800 Subject: [PATCH] rename AzureServicePrincipalMigration to ServicePrincipalMigration, rename execute_migration to run --- .../labs/ucx/azure/azure_credentials.py | 21 ++++---- src/databricks/labs/ucx/cli.py | 6 +-- .../azure/test_azure_credentials.py | 27 +++++----- tests/unit/azure/test_azure_credentials.py | 51 ++++++++++--------- 4 files changed, 54 insertions(+), 51 deletions(-) diff --git a/src/databricks/labs/ucx/azure/azure_credentials.py b/src/databricks/labs/ucx/azure/azure_credentials.py index c8334269c3..5b2ef25efe 100644 --- a/src/databricks/labs/ucx/azure/azure_credentials.py +++ b/src/databricks/labs/ucx/azure/azure_credentials.py @@ -54,7 +54,7 @@ def from_storage_credential_validation( ) -class AzureServicePrincipalMigration: +class ServicePrincipalMigration: def __init__( self, @@ -65,7 +65,6 @@ def __init__( integration_test_flag="", ): self._output_file = "azure_service_principal_migration_result.csv" - self._final_sp_list: list[ServicePrincipalMigrationInfo] = [] self._installation = installation self._ws = ws self._azure_resource_permissions = azure_resource_permissions @@ -111,7 +110,7 @@ def _list_storage_credentials(self) -> set[str]: # return the storage credential created during integration test return {storage_credential.azure_service_principal.application_id} # return no storage credential if there is none created during integration test - return {""} + return {} for storage_credential in storage_credentials: # only add service principal's application_id, ignore managed identity based storage_credential @@ -186,17 +185,17 @@ def _fetch_client_secret(self, sp_list: list[StoragePermissionMapping]) -> list[ ) return sp_list_with_secret - def _print_action_plan(self, sp_list_with_secret: list[ServicePrincipalMigrationInfo]): + def _print_action_plan(self, sp_list: list[ServicePrincipalMigrationInfo]): # print action plan to console for customer to review. - for sp in sp_list_with_secret: - print( + for sp in sp_list: + logger.info( f"Service Principal name: {sp.service_principal.principal}, " f"application_id: {sp.service_principal.client_id}, " f"privilege {sp.service_principal.privilege} " f"on location {sp.service_principal.prefix}" ) - def _generate_migration_list(self): + def _generate_migration_list(self) -> list[ServicePrincipalMigrationInfo]: """ Create the list of SP that need to be migrated, output an action plan as a csv file for users to confirm """ @@ -208,9 +207,9 @@ def _generate_migration_list(self): filtered_sp_list = [sp for sp in sp_list if sp.client_id not in sc_set] # fetch sp client_secret if any sp_list_with_secret = self._fetch_client_secret(filtered_sp_list) - self._final_sp_list = sp_list_with_secret # output the action plan for customer to confirm self._print_action_plan(sp_list_with_secret) + return sp_list_with_secret def _create_storage_credential(self, sp_migration: ServicePrincipalMigrationInfo): # prepare the storage credential properties @@ -260,9 +259,9 @@ def _validate_storage_credential( ), ) - def execute_migration(self, prompts: Prompts): + def run(self, prompts: Prompts): - self._generate_migration_list() + sp_list_with_secret = self._generate_migration_list() plan_confirmed = prompts.confirm( "Above Azure Service Principals will be migrated to UC storage credentials, please review and confirm." @@ -271,7 +270,7 @@ def execute_migration(self, prompts: Prompts): return execution_result = [] - for sp in self._final_sp_list: + for sp in sp_list_with_secret: execution_result.append(self._create_storage_credential(sp)) results_file = self._installation.save(execution_result, filename=self._output_file) diff --git a/src/databricks/labs/ucx/cli.py b/src/databricks/labs/ucx/cli.py index 5b076a0b81..dc671f09ee 100644 --- a/src/databricks/labs/ucx/cli.py +++ b/src/databricks/labs/ucx/cli.py @@ -13,7 +13,7 @@ from databricks.labs.ucx.account import AccountWorkspaces, WorkspaceInfo from databricks.labs.ucx.assessment.aws import AWSResourcePermissions from databricks.labs.ucx.azure.access import AzureResourcePermissions -from databricks.labs.ucx.azure.azure_credentials import AzureServicePrincipalMigration +from databricks.labs.ucx.azure.azure_credentials import ServicePrincipalMigration from databricks.labs.ucx.config import WorkspaceConfig from databricks.labs.ucx.framework.crawlers import StatementExecutionBackend from databricks.labs.ucx.hive_metastore import ExternalLocations, TablesCrawler @@ -297,8 +297,8 @@ def migrate_azure_service_principals(w: WorkspaceClient): """ logger.info("Running migrate_azure_service_principals command") prompts = Prompts() - service_principal_migration = AzureServicePrincipalMigration.for_cli(w, prompts) - service_principal_migration.execute_migration(prompts) + service_principal_migration = ServicePrincipalMigration.for_cli(w, prompts) + service_principal_migration.run(prompts) if __name__ == "__main__": diff --git a/tests/integration/azure/test_azure_credentials.py b/tests/integration/azure/test_azure_credentials.py index a15b7f242c..281707e673 100644 --- a/tests/integration/azure/test_azure_credentials.py +++ b/tests/integration/azure/test_azure_credentials.py @@ -1,6 +1,6 @@ import base64 import re -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest from databricks.labs.blueprint.tui import MockPrompts @@ -8,7 +8,7 @@ from databricks.labs.ucx.assessment.azure import AzureServicePrincipalInfo from databricks.labs.ucx.assessment.crawlers import _SECRET_PATTERN from databricks.labs.ucx.azure.access import StoragePermissionMapping -from databricks.labs.ucx.azure.azure_credentials import AzureServicePrincipalMigration +from databricks.labs.ucx.azure.azure_credentials import ServicePrincipalMigration @pytest.fixture @@ -75,15 +75,15 @@ def inner(read_only=False) -> dict: @pytest.fixture def execute_migration(ws): - def inner(variables: dict, integration_test_flag: str) -> AzureServicePrincipalMigration: - spn_migration = AzureServicePrincipalMigration( + def inner(variables: dict, integration_test_flag: str) -> ServicePrincipalMigration: + spn_migration = ServicePrincipalMigration( variables["installation"], ws, variables["azure_resource_permissions"], variables["azure_sp_crawler"], integration_test_flag=integration_test_flag, ) - spn_migration.execute_migration( + spn_migration.run( MockPrompts({"Above Azure Service Principals will be migrated to UC storage credentials *": "Yes"}) ) return spn_migration @@ -92,7 +92,7 @@ def inner(variables: dict, integration_test_flag: str) -> AzureServicePrincipalM def test_spn_migration_existed_storage_credential( - ws, execute_migration, make_storage_credential_from_spn, prepare_spn_migration_test + execute_migration, make_storage_credential_from_spn, prepare_spn_migration_test ): variables = prepare_spn_migration_test(read_only=False) @@ -104,11 +104,11 @@ def test_spn_migration_existed_storage_credential( directory_id=variables["directory_id"], ) - # test that the spn migration will be skipped due to above storage credential is existed - spn_migration = execute_migration(variables, integration_test_flag=variables["storage_credential_name"]) - - # because storage_credential is existing, no spn should be migrated - assert not spn_migration._final_sp_list + with patch("databricks.labs.ucx.azure.azure_credentials.ServicePrincipalMigration._create_storage_credential") as create_storage_credential: + # test that the spn migration will be skipped due to above storage credential is existed + execute_migration(variables, integration_test_flag=variables["storage_credential_name"]) + # because storage_credential is existing, no spn should be migrated + create_storage_credential.assert_not_called() @pytest.mark.parametrize("read_only", [False, True]) @@ -118,8 +118,9 @@ def test_spn_migration(ws, execute_migration, prepare_spn_migration_test, read_o try: spn_migration = execute_migration(variables, integration_test_flag="lets_migrate_the_spn") - assert spn_migration._final_sp_list[0].service_principal.principal == variables["storage_credential_name"] - assert ws.storage_credentials.get(variables["storage_credential_name"]).read_only is read_only + storage_credential = ws.storage_credentials.get(variables["storage_credential_name"]) + assert storage_credential is not None + assert storage_credential.read_only is read_only validation_result = spn_migration._installation.save.call_args.args[0][0] if read_only: diff --git a/tests/unit/azure/test_azure_credentials.py b/tests/unit/azure/test_azure_credentials.py index 08e9b9f63c..b247b85cc9 100644 --- a/tests/unit/azure/test_azure_credentials.py +++ b/tests/unit/azure/test_azure_credentials.py @@ -25,7 +25,7 @@ ) from databricks.labs.ucx.azure.access import StoragePermissionMapping from databricks.labs.ucx.azure.azure_credentials import ( - AzureServicePrincipalMigration, + ServicePrincipalMigration, ServicePrincipalMigrationInfo, ) from tests.unit.framework.mocks import MockBackend @@ -103,21 +103,21 @@ def download(path: str) -> io.BytesIO: def test_for_cli_not_azure(caplog, ws): ws.config.is_azure = False - assert AzureServicePrincipalMigration.for_cli(ws, MagicMock()) is None + assert ServicePrincipalMigration.for_cli(ws, MagicMock()) is None assert "Workspace is not on azure, please run this command on azure databricks workspaces." in caplog.text def test_for_cli_not_prompts(ws): ws.config.is_azure = True prompts = MockPrompts({"Have you reviewed the azure_storage_account_info.csv *": "No"}) - assert AzureServicePrincipalMigration.for_cli(ws, prompts) is None + assert ServicePrincipalMigration.for_cli(ws, prompts) is None def test_for_cli(ws): ws.config.is_azure = True prompts = MockPrompts({"Have you reviewed the azure_storage_account_info.csv *": "Yes"}) - assert isinstance(AzureServicePrincipalMigration.for_cli(ws, prompts), AzureServicePrincipalMigration) + assert isinstance(ServicePrincipalMigration.for_cli(ws, prompts), ServicePrincipalMigration) def test_list_storage_credentials(ws): @@ -137,7 +137,7 @@ def test_list_storage_credentials(ws): ), ] - sp_migration = AzureServicePrincipalMigration(MagicMock(), ws, MagicMock(), MagicMock()) + sp_migration = ServicePrincipalMigration(MagicMock(), ws, MagicMock(), MagicMock()) expected = {"b6420590-5e1c-4426-8950-a94cbe9b6115"} sp_migration._list_storage_credentials() @@ -165,7 +165,7 @@ def test_list_storage_credentials_for_integration_test(ws): # test storage credential: spn_for_integration_test is picked up # in integration test, we only pick up the existing storage credential created in integration test and ignore the others - sp_migration = AzureServicePrincipalMigration( + sp_migration = ServicePrincipalMigration( MagicMock(), ws, MagicMock(), MagicMock(), integration_test_flag="spn_for_integration_test" ) expected = {"b6420590-5e1c-4426-8950-a94cbe9b6115"} @@ -174,11 +174,11 @@ def test_list_storage_credentials_for_integration_test(ws): # test storage credential is not picked up # if integration test does not create storage credential, we use dummy integration_test_flag to filter out other existing storage credentials - sp_migration = AzureServicePrincipalMigration( + sp_migration = ServicePrincipalMigration( MagicMock(), ws, MagicMock(), MagicMock(), integration_test_flag="other_spn" ) sp_migration._list_storage_credentials() - assert {""} == sp_migration._list_storage_credentials() + assert {} == sp_migration._list_storage_credentials() @pytest.mark.parametrize( @@ -192,7 +192,7 @@ def test_list_storage_credentials_for_integration_test(ws): def test_read_secret_value_decode(ws, secret_bytes_value, expected_return): ws.secrets.get_secret.return_value = secret_bytes_value - sp_migration = AzureServicePrincipalMigration(MagicMock(), ws, MagicMock(), MagicMock()) + sp_migration = ServicePrincipalMigration(MagicMock(), ws, MagicMock(), MagicMock()) assert sp_migration._read_databricks_secret("test_scope", "test_key", "000") == expected_return @@ -207,7 +207,7 @@ def test_read_secret_read_exception(caplog, ws, exception, expected_log, expecte caplog.set_level(logging.INFO) ws.secrets.get_secret.side_effect = exception - sp_migration = AzureServicePrincipalMigration(MagicMock(), ws, MagicMock(), MagicMock()) + sp_migration = ServicePrincipalMigration(MagicMock(), ws, MagicMock(), MagicMock()) secret_value = sp_migration._read_databricks_secret("test_scope", "test_key", "000") assert expected_log in caplog.text @@ -281,13 +281,14 @@ def test_fetch_client_secret(ws): ), ] - sp_migration = AzureServicePrincipalMigration(MagicMock(), ws, MagicMock(), sp_crawler) + sp_migration = ServicePrincipalMigration(MagicMock(), ws, MagicMock(), sp_crawler) filtered_sp_list = sp_migration._fetch_client_secret(sp_to_be_checked) assert filtered_sp_list == expected_sp_list -def test_print_action_plan(capsys): +def test_print_action_plan(caplog): + caplog.set_level(logging.INFO) sp_list_with_secret = [ ServicePrincipalMigrationInfo( StoragePermissionMapping( @@ -300,19 +301,20 @@ def test_print_action_plan(capsys): "hello world", ) ] - sp_migration = AzureServicePrincipalMigration(MagicMock(), MagicMock(), MagicMock(), MagicMock()) + sp_migration = ServicePrincipalMigration(MagicMock(), MagicMock(), MagicMock(), MagicMock()) sp_migration._print_action_plan(sp_list_with_secret) expected_print = ( "Service Principal name: principal_1, " "application_id: app_secret1, " "privilege WRITE_FILES " - "on location prefix1\n" + "on location prefix1" ) - assert expected_print == capsys.readouterr().out + assert expected_print in caplog.text -def test_generate_migration_list(capsys, mocker, ws): +def test_generate_migration_list(caplog, mocker, ws): + caplog.set_level(logging.INFO) ws.config.is_azure = True ws.secrets.get_secret.return_value = GetSecretResponse(value="aGVsbG8gd29ybGQ=") ws.storage_credentials.list.return_value = [ @@ -335,10 +337,11 @@ def test_generate_migration_list(capsys, mocker, ws): ], ) - sp_migration = AzureServicePrincipalMigration.for_cli(ws, prompts) + sp_migration = ServicePrincipalMigration.for_cli(ws, prompts) sp_migration._generate_migration_list() - assert "app_secret2" in capsys.readouterr().out + assert "app_secret2" in caplog.text + assert "app_secret1" not in caplog.text def test_execute_migration_no_confirmation(mocker, ws): @@ -350,13 +353,13 @@ def test_execute_migration_no_confirmation(mocker, ws): } ) - mocker.patch("databricks.labs.ucx.azure.azure_credentials.AzureServicePrincipalMigration._generate_migration_list") + mocker.patch("databricks.labs.ucx.azure.azure_credentials.ServicePrincipalMigration._generate_migration_list") with patch( - "databricks.labs.ucx.azure.azure_credentials.AzureServicePrincipalMigration._create_storage_credential" + "databricks.labs.ucx.azure.azure_credentials.ServicePrincipalMigration._create_storage_credential" ) as c: - sp_migration = AzureServicePrincipalMigration.for_cli(ws, prompts) - sp_migration.execute_migration(prompts) + sp_migration = ServicePrincipalMigration.for_cli(ws, prompts) + sp_migration.run(prompts) c.assert_not_called() @@ -409,9 +412,9 @@ def test_execute_migration(caplog, capsys, mocker, ws): ], ) - sp_migration = AzureServicePrincipalMigration.for_cli(ws, prompts) + sp_migration = ServicePrincipalMigration.for_cli(ws, prompts) sp_migration._installation.save = MagicMock() - sp_migration.execute_migration(prompts) + sp_migration.run(prompts) # assert migration is complete assert "Completed migration" in capsys.readouterr().out