From bef22a79c938d8178c24da30569b7273926c727a Mon Sep 17 00:00:00 2001 From: Andrew Snare Date: Mon, 15 Jul 2024 16:08:27 +0200 Subject: [PATCH] `WorkspacePath`: Fix the .rename() and .replace() implementations to return the target path (#130) This PR updates the `.rename()` and `.replace()` implementations on `WorkspacePath` so that they return the target path. (This is part of the public API.) In addition, the `.rename()` method no longer extends the public API with the additional `overwrite` keyword argument. The contract for `.rename()` is that it is allowed to always fail if the target exists, which is how it now behaves. Out of scope for this PR is defining the exact error that is raised if the target exists. This PR is part of the preparation for factoring out common code to implement support for DBFS paths. --- src/databricks/labs/blueprint/paths.py | 11 ++++++++--- tests/unit/test_paths.py | 8 ++++---- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/src/databricks/labs/blueprint/paths.py b/src/databricks/labs/blueprint/paths.py index 44b04a1..f31cec8 100644 --- a/src/databricks/labs/blueprint/paths.py +++ b/src/databricks/labs/blueprint/paths.py @@ -482,16 +482,21 @@ def rmdir(self, recursive=False): """Remove a directory in Databricks Workspace""" self._ws.workspace.delete(self.as_posix(), recursive=recursive) - def rename(self, target, overwrite=False): + def _rename(self, target, overwrite: bool): """Rename a file or directory in Databricks Workspace""" - dst = WorkspacePath(self._ws, target) + dst = self.with_segments(target) with self._ws.workspace.download(self.as_posix(), format=ExportFormat.AUTO) as f: self._ws.workspace.upload(dst.as_posix(), f.read(), format=ImportFormat.AUTO, overwrite=overwrite) self.unlink() + return dst + + def rename(self, target): + """Rename a file or directory in Databricks Workspace, failing if the target already exists.""" + return self._rename(target, overwrite=False) def replace(self, target): """Rename a file or directory in Databricks Workspace, overwriting the target if it exists.""" - return self.rename(target, overwrite=True) + return self._rename(target, overwrite=True) def unlink(self, missing_ok=False): """Remove a file in Databricks Workspace.""" diff --git a/tests/unit/test_paths.py b/tests/unit/test_paths.py index feee490..f618064 100644 --- a/tests/unit/test_paths.py +++ b/tests/unit/test_paths.py @@ -584,19 +584,19 @@ def test_rmdir_removes_directory_recursive() -> None: ws.workspace.delete.assert_called_once_with("/test/path", recursive=True) -def test_rename_file_without_overwrite() -> None: +def test_rename_file() -> None: ws = create_autospec(WorkspaceClient) workspace_path = WorkspacePath(ws, "/test/path") ws.workspace.download.return_value.__enter__.return_value.read.return_value = b"test" - workspace_path.rename("/new/path") + assert workspace_path.rename("/new/path") == WorkspacePath(ws, "/new/path") ws.workspace.upload.assert_called_once_with("/new/path", b"test", format=ImportFormat.AUTO, overwrite=False) -def test_rename_file_with_overwrite() -> None: +def test_replace_file() -> None: ws = create_autospec(WorkspaceClient) workspace_path = WorkspacePath(ws, "/test/path") ws.workspace.download.return_value.__enter__.return_value.read.return_value = b"test" - workspace_path.rename("/new/path", overwrite=True) + assert workspace_path.replace("/new/path") == WorkspacePath(ws, "/new/path") ws.workspace.upload.assert_called_once_with("/new/path", b"test", format=ImportFormat.AUTO, overwrite=True)