From 0b264d474d5ad7966a5df7d1d42d16377ec91cda Mon Sep 17 00:00:00 2001 From: Pritish Pai Date: Sun, 11 Feb 2024 18:43:28 -0500 Subject: [PATCH 01/11] Flip filesinrepos to true if false --- src/databricks/labs/blueprint/installation.py | 11 +++++++++ tests/unit/test_installation.py | 23 +++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/src/databricks/labs/blueprint/installation.py b/src/databricks/labs/blueprint/installation.py index bf00b9b..38a5746 100644 --- a/src/databricks/labs/blueprint/installation.py +++ b/src/databricks/labs/blueprint/installation.py @@ -169,6 +169,7 @@ def load(self, type_ref: type[T], *, filename: str | None = None) -> T: the expected version using a method named `v{actual_version}_migrate` on the `type_ref` class. If the migration is successful, the method will return the migrated object. If the migration is not successful, the method will raise an `IllegalState` exception.""" + self._enable_files_in_repos() filename = self._get_filename(filename, type_ref) logger.debug(f"Loading {type_ref.__name__} from {filename}") as_dict = self._load_content(filename) @@ -669,6 +670,16 @@ def _load_csv(raw: BinaryIO) -> list[Json]: out.append(row) return out + def _enable_files_in_repos(self): + # check if "enableWorkspaceFilesystem" is set to false + workspace_file_system = self._ws.workspace_conf.get_status("enableWorkspaceFilesystem") + + logger.debug("Checking Files In Repos configuration") + + if workspace_file_system["enableWorkspaceFilesystem"] == "false": + logger.debug("enableWorkspaceFilesystem is False, enabling the config") + self._ws.workspace_conf.set_status("""{"enableWorkspaceFilesystem": "true"}""") + class MockInstallation(Installation): """Install state testing toolbelt diff --git a/tests/unit/test_installation.py b/tests/unit/test_installation.py index 6570755..fc71ee4 100644 --- a/tests/unit/test_installation.py +++ b/tests/unit/test_installation.py @@ -4,6 +4,7 @@ import pytest import yaml +import random from databricks.sdk import WorkspaceClient from databricks.sdk.core import Config from databricks.sdk.errors import NotFound @@ -327,3 +328,25 @@ def test_migrations_broken(): with pytest.raises(IllegalState): installation.load(BrokenConfig) + + +def mock_get_status(*args, **kwargs): + random_return_value = random.choice(["true", "false"]) + if args[0] == 'enableProjectTypeInWorkspace': + return {"enableProjectTypeInWorkspace": random_return_value} + elif args[0] == 'enableWorkspaceFilesystem': + return {"enableWorkspaceFilesystem": random_return_value} + + +def test_enable_files_in_repos(mocker): + ws = create_autospec(WorkspaceClient) + ws.current_user.me().user_name = "foo" + installation = Installation(ws, "ucx") + + ws.workspace_conf.get_status = mocker.patch( + "databricks.sdk.service.settings.WorkspaceConfAPI.get_status", side_effect=mock_get_status + ) + + installation._enable_files_in_repos() + assert True + From 434816c0fe9541062adaef8d525b5d367f687623 Mon Sep 17 00:00:00 2001 From: Pritish Pai Date: Sun, 11 Feb 2024 21:44:38 -0500 Subject: [PATCH 02/11] Fixing fmt errors --- tests/unit/test_installation.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_installation.py b/tests/unit/test_installation.py index fc71ee4..115d617 100644 --- a/tests/unit/test_installation.py +++ b/tests/unit/test_installation.py @@ -1,10 +1,10 @@ import io +import random from dataclasses import dataclass from unittest.mock import create_autospec import pytest import yaml -import random from databricks.sdk import WorkspaceClient from databricks.sdk.core import Config from databricks.sdk.errors import NotFound @@ -332,9 +332,9 @@ def test_migrations_broken(): def mock_get_status(*args, **kwargs): random_return_value = random.choice(["true", "false"]) - if args[0] == 'enableProjectTypeInWorkspace': + if args[0] == "enableProjectTypeInWorkspace": return {"enableProjectTypeInWorkspace": random_return_value} - elif args[0] == 'enableWorkspaceFilesystem': + elif args[0] == "enableWorkspaceFilesystem": return {"enableWorkspaceFilesystem": random_return_value} @@ -349,4 +349,3 @@ def test_enable_files_in_repos(mocker): installation._enable_files_in_repos() assert True - From 0b83d8ca581f8f444757121eddd660e907b9a42c Mon Sep 17 00:00:00 2001 From: Pritish Pai Date: Sun, 11 Feb 2024 21:56:32 -0500 Subject: [PATCH 03/11] Fixing argument for set_status --- src/databricks/labs/blueprint/installation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/labs/blueprint/installation.py b/src/databricks/labs/blueprint/installation.py index 38a5746..79e47f1 100644 --- a/src/databricks/labs/blueprint/installation.py +++ b/src/databricks/labs/blueprint/installation.py @@ -678,7 +678,7 @@ def _enable_files_in_repos(self): if workspace_file_system["enableWorkspaceFilesystem"] == "false": logger.debug("enableWorkspaceFilesystem is False, enabling the config") - self._ws.workspace_conf.set_status("""{"enableWorkspaceFilesystem": "true"}""") + self._ws.workspace_conf.set_status({"enableWorkspaceFilesystem": "true"}) class MockInstallation(Installation): From 174225c35d750a713b133edad77eacb592e71be4 Mon Sep 17 00:00:00 2001 From: Pritish Pai Date: Sun, 11 Feb 2024 22:12:39 -0500 Subject: [PATCH 04/11] Fixing tests --- tests/unit/test_installation.py | 45 ++++++++++++++++++++++++++------- 1 file changed, 36 insertions(+), 9 deletions(-) diff --git a/tests/unit/test_installation.py b/tests/unit/test_installation.py index 115d617..5555a7a 100644 --- a/tests/unit/test_installation.py +++ b/tests/unit/test_installation.py @@ -19,6 +19,15 @@ ) +@pytest.fixture() +def mock_install(mocker): + ws = create_autospec(WorkspaceClient) + ws.current_user.me().user_name = "foo" + ws.workspace_conf.get_status = mocker.patch( + "databricks.sdk.service.settings.WorkspaceConfAPI.get_status", side_effect=mock_get_status + ) + + def test_current_not_found(): ws = create_autospec(WorkspaceClient) ws.current_user.me().user_name = "foo" @@ -174,7 +183,7 @@ def test_save_typed_file_array_csv(): ) -def test_load_typed_file(): +def test_load_typed_file(mocker): ws = create_autospec(WorkspaceClient) ws.current_user.me().user_name = "foo" ws.workspace.download.return_value = io.StringIO( @@ -187,6 +196,9 @@ def test_load_typed_file(): } ) ) + ws.workspace_conf.get_status = mocker.patch( + "databricks.sdk.service.settings.WorkspaceConfAPI.get_status", side_effect=mock_get_status + ) installation = Installation(ws, "blueprint") cfg = installation.load(WorkspaceConfig) @@ -211,7 +223,9 @@ def test_load_csv_file(): @pytest.mark.parametrize("ext", ["json", "csv"]) -def test_load_typed_list_file(ext): +def test_load_typed_list_file(ext, mocker): + ws = create_autospec(WorkspaceClient) + ws.current_user.me().user_name = "foo" installation = MockInstallation( { f"workspaces.{ext}": [ @@ -220,7 +234,10 @@ def test_load_typed_list_file(ext): ] } ) - + installation._ws = ws + ws.workspace_conf.get_status = mocker.patch( + "databricks.sdk.service.settings.WorkspaceConfAPI.get_status", side_effect=mock_get_status + ) workspaces = installation.load(list[Workspace], filename=f"workspaces.{ext}") assert 2 == len(workspaces) @@ -297,8 +314,14 @@ def v2_migrate(raw: dict) -> dict: return raw -def test_migrations_on_load(): +def test_migrations_on_load(mocker): installation = MockInstallation({"config.yml": {"initial": 999}}) + ws = create_autospec(WorkspaceClient) + ws.current_user.me().user_name = "foo" + ws.workspace_conf.get_status = mocker.patch( + "databricks.sdk.service.settings.WorkspaceConfAPI.get_status", side_effect=mock_get_status + ) + installation._ws = ws cfg = installation.load(EvolvedConfig) @@ -323,18 +346,22 @@ def v1_migrate(raw: dict) -> dict: return {} -def test_migrations_broken(): +def test_migrations_broken(mocker): + ws = create_autospec(WorkspaceClient) + ws.current_user.me().user_name = "foo" installation = MockInstallation({"config.yml": {"initial": 999}}) + installation._ws = ws + ws.workspace_conf.get_status = mocker.patch( + "databricks.sdk.service.settings.WorkspaceConfAPI.get_status", side_effect=mock_get_status + ) with pytest.raises(IllegalState): installation.load(BrokenConfig) -def mock_get_status(*args, **kwargs): +def mock_get_status(*args): random_return_value = random.choice(["true", "false"]) - if args[0] == "enableProjectTypeInWorkspace": - return {"enableProjectTypeInWorkspace": random_return_value} - elif args[0] == "enableWorkspaceFilesystem": + if args[0] == "enableWorkspaceFilesystem": return {"enableWorkspaceFilesystem": random_return_value} From 39b31e090d9d51ebee9c7436bbfef9660d80784c Mon Sep 17 00:00:00 2001 From: Pritish Pai Date: Mon, 12 Feb 2024 15:57:28 -0500 Subject: [PATCH 05/11] Change function call and revert tests --- src/databricks/labs/blueprint/installation.py | 6 ++- tests/unit/test_installation.py | 38 ++----------------- 2 files changed, 8 insertions(+), 36 deletions(-) diff --git a/src/databricks/labs/blueprint/installation.py b/src/databricks/labs/blueprint/installation.py index 79e47f1..f907087 100644 --- a/src/databricks/labs/blueprint/installation.py +++ b/src/databricks/labs/blueprint/installation.py @@ -169,7 +169,6 @@ def load(self, type_ref: type[T], *, filename: str | None = None) -> T: the expected version using a method named `v{actual_version}_migrate` on the `type_ref` class. If the migration is successful, the method will return the migrated object. If the migration is not successful, the method will raise an `IllegalState` exception.""" - self._enable_files_in_repos() filename = self._get_filename(filename, type_ref) logger.debug(f"Loading {type_ref.__name__} from {filename}") as_dict = self._load_content(filename) @@ -226,7 +225,10 @@ def upload(self, filename: str, raw: bytes): try: logger.debug(f"Uploading: {dst}") attempt() - except NotFound: + except NotFound as error: + if error.error_code == "FEATURE_DISABLED": + self._enable_files_in_repos() + self.upload(filename, raw) parent_folder = os.path.dirname(dst) logger.debug(f"Creating missing folders: {parent_folder}") self._ws.workspace.mkdirs(parent_folder) diff --git a/tests/unit/test_installation.py b/tests/unit/test_installation.py index 5555a7a..e0cec7c 100644 --- a/tests/unit/test_installation.py +++ b/tests/unit/test_installation.py @@ -19,15 +19,6 @@ ) -@pytest.fixture() -def mock_install(mocker): - ws = create_autospec(WorkspaceClient) - ws.current_user.me().user_name = "foo" - ws.workspace_conf.get_status = mocker.patch( - "databricks.sdk.service.settings.WorkspaceConfAPI.get_status", side_effect=mock_get_status - ) - - def test_current_not_found(): ws = create_autospec(WorkspaceClient) ws.current_user.me().user_name = "foo" @@ -183,7 +174,7 @@ def test_save_typed_file_array_csv(): ) -def test_load_typed_file(mocker): +def test_load_typed_file(): ws = create_autospec(WorkspaceClient) ws.current_user.me().user_name = "foo" ws.workspace.download.return_value = io.StringIO( @@ -196,9 +187,6 @@ def test_load_typed_file(mocker): } ) ) - ws.workspace_conf.get_status = mocker.patch( - "databricks.sdk.service.settings.WorkspaceConfAPI.get_status", side_effect=mock_get_status - ) installation = Installation(ws, "blueprint") cfg = installation.load(WorkspaceConfig) @@ -223,9 +211,7 @@ def test_load_csv_file(): @pytest.mark.parametrize("ext", ["json", "csv"]) -def test_load_typed_list_file(ext, mocker): - ws = create_autospec(WorkspaceClient) - ws.current_user.me().user_name = "foo" +def test_load_typed_list_file(ext): installation = MockInstallation( { f"workspaces.{ext}": [ @@ -234,10 +220,6 @@ def test_load_typed_list_file(ext, mocker): ] } ) - installation._ws = ws - ws.workspace_conf.get_status = mocker.patch( - "databricks.sdk.service.settings.WorkspaceConfAPI.get_status", side_effect=mock_get_status - ) workspaces = installation.load(list[Workspace], filename=f"workspaces.{ext}") assert 2 == len(workspaces) @@ -314,14 +296,8 @@ def v2_migrate(raw: dict) -> dict: return raw -def test_migrations_on_load(mocker): +def test_migrations_on_load(): installation = MockInstallation({"config.yml": {"initial": 999}}) - ws = create_autospec(WorkspaceClient) - ws.current_user.me().user_name = "foo" - ws.workspace_conf.get_status = mocker.patch( - "databricks.sdk.service.settings.WorkspaceConfAPI.get_status", side_effect=mock_get_status - ) - installation._ws = ws cfg = installation.load(EvolvedConfig) @@ -346,14 +322,8 @@ def v1_migrate(raw: dict) -> dict: return {} -def test_migrations_broken(mocker): - ws = create_autospec(WorkspaceClient) - ws.current_user.me().user_name = "foo" +def test_migrations_broken(): installation = MockInstallation({"config.yml": {"initial": 999}}) - installation._ws = ws - ws.workspace_conf.get_status = mocker.patch( - "databricks.sdk.service.settings.WorkspaceConfAPI.get_status", side_effect=mock_get_status - ) with pytest.raises(IllegalState): installation.load(BrokenConfig) From fee849b67bc1d04fcc3b2fb39362341865601013 Mon Sep 17 00:00:00 2001 From: Pritish Pai Date: Mon, 12 Feb 2024 16:58:01 -0500 Subject: [PATCH 06/11] Unit test#2 --- src/databricks/labs/blueprint/installation.py | 1 - tests/unit/test_installation.py | 11 +++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/src/databricks/labs/blueprint/installation.py b/src/databricks/labs/blueprint/installation.py index f907087..34cf191 100644 --- a/src/databricks/labs/blueprint/installation.py +++ b/src/databricks/labs/blueprint/installation.py @@ -228,7 +228,6 @@ def upload(self, filename: str, raw: bytes): except NotFound as error: if error.error_code == "FEATURE_DISABLED": self._enable_files_in_repos() - self.upload(filename, raw) parent_folder = os.path.dirname(dst) logger.debug(f"Creating missing folders: {parent_folder}") self._ws.workspace.mkdirs(parent_folder) diff --git a/tests/unit/test_installation.py b/tests/unit/test_installation.py index e0cec7c..8b52598 100644 --- a/tests/unit/test_installation.py +++ b/tests/unit/test_installation.py @@ -346,3 +346,14 @@ def test_enable_files_in_repos(mocker): installation._enable_files_in_repos() assert True + + +def test_upload_feature_disabled_failure(): + ws = create_autospec(WorkspaceClient) + ws.current_user.me().user_name = "foo" + ws.workspace.upload.side_effect = [NotFound(error_code="FEATURE_DISABLED"), None] + installation = Installation(ws, "blueprint") + + installation.save(WorkspaceConfig(inventory_database="some_blueprint")) + + ws.workspace.mkdirs.assert_called_with("/Users/foo/.blueprint") From 817013ef84a01ab77c541685cc5c613a940737c8 Mon Sep 17 00:00:00 2001 From: Pritish Pai Date: Mon, 12 Feb 2024 23:44:12 -0500 Subject: [PATCH 07/11] Removing random and monkeypatching --- tests/integration/test_installation.py | 5 +++++ tests/unit/test_installation.py | 18 +++++++----------- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/tests/integration/test_installation.py b/tests/integration/test_installation.py index 850bc08..d02e864 100644 --- a/tests/integration/test_installation.py +++ b/tests/integration/test_installation.py @@ -83,3 +83,8 @@ def test_csv(new_installation): def test_uploading_notebooks_get_correct_urls(ext, magic, new_installation): remote_path = new_installation.upload(f"foo.{ext}", f"{magic}\nprint(1)".encode("utf8")) assert f"{new_installation.install_folder()}/foo" == remote_path + + +def test_something(ws, new_installation): + ws.workspace_conf.get_status("enableWorkspaceFilesystem") + new_installation.save(MyClass("a", "b")) diff --git a/tests/unit/test_installation.py b/tests/unit/test_installation.py index 8b52598..76720b4 100644 --- a/tests/unit/test_installation.py +++ b/tests/unit/test_installation.py @@ -1,5 +1,4 @@ import io -import random from dataclasses import dataclass from unittest.mock import create_autospec @@ -329,21 +328,18 @@ def test_migrations_broken(): installation.load(BrokenConfig) -def mock_get_status(*args): - random_return_value = random.choice(["true", "false"]) - if args[0] == "enableWorkspaceFilesystem": - return {"enableWorkspaceFilesystem": random_return_value} - - -def test_enable_files_in_repos(mocker): +def test_enable_files_in_repos(): ws = create_autospec(WorkspaceClient) ws.current_user.me().user_name = "foo" installation = Installation(ws, "ucx") - ws.workspace_conf.get_status = mocker.patch( - "databricks.sdk.service.settings.WorkspaceConfAPI.get_status", side_effect=mock_get_status - ) + # enableWorkspaceFilesystem is true + ws.workspace_conf.get_status.return_value = {"enableWorkspaceFilesystem": "true"} + installation._enable_files_in_repos() + assert True + # enableWorkspaceFilesystem is false + ws.workspace_conf.get_status.return_value = {"enableWorkspaceFilesystem": "false"} installation._enable_files_in_repos() assert True From 26f2866ae84a5b756fbdb1528a98bb5ba8d8353d Mon Sep 17 00:00:00 2001 From: Pritish Pai Date: Mon, 12 Feb 2024 23:58:44 -0500 Subject: [PATCH 08/11] Assert api call instead of private function --- tests/unit/test_installation.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_installation.py b/tests/unit/test_installation.py index 76720b4..c4d9282 100644 --- a/tests/unit/test_installation.py +++ b/tests/unit/test_installation.py @@ -1,6 +1,6 @@ import io from dataclasses import dataclass -from unittest.mock import create_autospec +from unittest.mock import create_autospec, MagicMock import pytest import yaml @@ -332,17 +332,18 @@ def test_enable_files_in_repos(): ws = create_autospec(WorkspaceClient) ws.current_user.me().user_name = "foo" installation = Installation(ws, "ucx") + ws.workspace_conf.set_status = MagicMock() # enableWorkspaceFilesystem is true ws.workspace_conf.get_status.return_value = {"enableWorkspaceFilesystem": "true"} installation._enable_files_in_repos() - assert True + ws.workspace_conf.set_status.assert_not_called() # enableWorkspaceFilesystem is false ws.workspace_conf.get_status.return_value = {"enableWorkspaceFilesystem": "false"} installation._enable_files_in_repos() - assert True - + ws.workspace_conf.set_status.assert_called_once() + ws.workspace_conf.set_status.assert_called_with({"enableWorkspaceFilesystem": "true"}) def test_upload_feature_disabled_failure(): ws = create_autospec(WorkspaceClient) From e4f806f7255aef02d4c7c46ea6792dd78d2f7a06 Mon Sep 17 00:00:00 2001 From: Pritish Pai Date: Tue, 13 Feb 2024 00:07:21 -0500 Subject: [PATCH 09/11] Remove debugging function --- tests/integration/test_installation.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/integration/test_installation.py b/tests/integration/test_installation.py index d02e864..1591942 100644 --- a/tests/integration/test_installation.py +++ b/tests/integration/test_installation.py @@ -84,7 +84,3 @@ def test_uploading_notebooks_get_correct_urls(ext, magic, new_installation): remote_path = new_installation.upload(f"foo.{ext}", f"{magic}\nprint(1)".encode("utf8")) assert f"{new_installation.install_folder()}/foo" == remote_path - -def test_something(ws, new_installation): - ws.workspace_conf.get_status("enableWorkspaceFilesystem") - new_installation.save(MyClass("a", "b")) From 905305ab3fb63dbf0ea114e2a2d8fa696b30033e Mon Sep 17 00:00:00 2001 From: Pritish Pai Date: Tue, 13 Feb 2024 00:14:26 -0500 Subject: [PATCH 10/11] Fixing small issues --- tests/integration/test_installation.py | 3 +-- tests/unit/test_installation.py | 4 +++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/integration/test_installation.py b/tests/integration/test_installation.py index 1591942..7a3f301 100644 --- a/tests/integration/test_installation.py +++ b/tests/integration/test_installation.py @@ -82,5 +82,4 @@ def test_csv(new_installation): ) def test_uploading_notebooks_get_correct_urls(ext, magic, new_installation): remote_path = new_installation.upload(f"foo.{ext}", f"{magic}\nprint(1)".encode("utf8")) - assert f"{new_installation.install_folder()}/foo" == remote_path - + assert f"{new_installation.install_folder()}/foo" == remote_path \ No newline at end of file diff --git a/tests/unit/test_installation.py b/tests/unit/test_installation.py index c4d9282..f9731ad 100644 --- a/tests/unit/test_installation.py +++ b/tests/unit/test_installation.py @@ -1,6 +1,6 @@ import io from dataclasses import dataclass -from unittest.mock import create_autospec, MagicMock +from unittest.mock import MagicMock, create_autospec import pytest import yaml @@ -219,6 +219,7 @@ def test_load_typed_list_file(ext): ] } ) + workspaces = installation.load(list[Workspace], filename=f"workspaces.{ext}") assert 2 == len(workspaces) @@ -345,6 +346,7 @@ def test_enable_files_in_repos(): ws.workspace_conf.set_status.assert_called_once() ws.workspace_conf.set_status.assert_called_with({"enableWorkspaceFilesystem": "true"}) + def test_upload_feature_disabled_failure(): ws = create_autospec(WorkspaceClient) ws.current_user.me().user_name = "foo" From eb27b59526168d1c33bbe1fce9b92a509102c9c0 Mon Sep 17 00:00:00 2001 From: Pritish Pai Date: Tue, 13 Feb 2024 00:16:39 -0500 Subject: [PATCH 11/11] EOF newline --- tests/integration/test_installation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/test_installation.py b/tests/integration/test_installation.py index 7a3f301..850bc08 100644 --- a/tests/integration/test_installation.py +++ b/tests/integration/test_installation.py @@ -82,4 +82,4 @@ def test_csv(new_installation): ) def test_uploading_notebooks_get_correct_urls(ext, magic, new_installation): remote_path = new_installation.upload(f"foo.{ext}", f"{magic}\nprint(1)".encode("utf8")) - assert f"{new_installation.install_folder()}/foo" == remote_path \ No newline at end of file + assert f"{new_installation.install_folder()}/foo" == remote_path