Add support for DBFS partial deletes (#280)

* Implemented retry logic for partial deletes

* Added unit tests for partial delete handling

* Changed expected partial delete request response. Added log to indicate progress on partial deletes.

* Added progress message with total number of files deleted in partial deletes.

* Fixed lint error: import order.

* Add retries for other 503 errors in delete.

* Added retry delay for delete.

* Added more tests for dbfs delete.

Author: bogdanghita-db

databricks_cli/dbfs/api.py

@@ -27,6 +27,8 @@
 import shutil
 import tempfile
 
+import re
+import time
 import click
 
 from requests.exceptions import HTTPError
@@ -37,6 +39,12 @@
 from databricks_cli.dbfs.exceptions import LocalFileExistsException
 
 BUFFER_SIZE_BYTES = 2**20
+DELETE_MAX_CONSECUTIVE_503_RETRIES = 3
+DELETE_503_RETRY_DELAY_MILLIS = 30 * 1000
+
+
+class ParseException(Exception):
+    pass
 
 
 class FileInfo(object):
@@ -69,11 +77,14 @@ def __eq__(self, other):
 class DbfsErrorCodes(object):
     RESOURCE_DOES_NOT_EXIST = 'RESOURCE_DOES_NOT_EXIST'
     RESOURCE_ALREADY_EXISTS = 'RESOURCE_ALREADY_EXISTS'
+    TEMPORARILY_UNAVAILABLE = 'TEMPORARILY_UNAVAILABLE'
+    PARTIAL_DELETE = 'PARTIAL_DELETE'
 
 
 class DbfsApi(object):
-    def __init__(self, api_client):
+    def __init__(self, api_client, delete_retry_delay_millis=DELETE_503_RETRY_DELAY_MILLIS):
         self.client = DbfsService(api_client)
+        self.delete_retry_delay_millis = delete_retry_delay_millis
 
     def list_files(self, dbfs_path, headers=None):
         list_response = self.client.list(dbfs_path.absolute_path, headers=headers)
@@ -123,8 +134,52 @@ def get_file(self, dbfs_path, dst_path, overwrite, headers=None):
                 offset += bytes_read
                 local_file.write(b64decode(data))
 
+    @staticmethod
+    def get_num_files_deleted(partial_delete_error):
+        try:
+            message = partial_delete_error.response.json()['message']
+        except (AttributeError, KeyError):
+            raise ParseException("Unable to retrieve the number of deleted files.")
+        m = re.compile(r".*operation has deleted (\d+) files.*").match(message)
+        if not m:
+            raise ParseException(
+                "Unable to retrieve the number of deleted files from the error message: {}".format(
+                    message))
+        return int(m.group(1))
+
     def delete(self, dbfs_path, recursive, headers=None):
-        self.client.delete(dbfs_path.absolute_path, recursive=recursive, headers=headers)
+        num_consecutive_503_retries = 0
+        num_files_deleted = 0
+        while True:
+            try:
+                self.client.delete(dbfs_path.absolute_path, recursive=recursive, headers=headers)
+            except HTTPError as e:
+                if e.response.status_code == 503:
+                    try:
+                        error_code = e.response.json()['error_code']
+                    except (AttributeError, KeyError):
+                        error_code = None
+                    # Handle partial delete exceptions: retry until all the files have been deleted
+                    if error_code == DbfsErrorCodes.PARTIAL_DELETE:
+                        try:
+                            num_files_deleted += DbfsApi.get_num_files_deleted(e)
+                            click.echo("\rDeleted {} files. Delete in progress...\033[K".format(
+                                num_files_deleted), nl=False)
+                        except ParseException:
+                            click.echo("\rDelete in progress...\033[K", nl=False)
+                        num_consecutive_503_retries = 0
+                        continue
+                    # Retry at most DELETE_MAX_CONSECUTIVE_503_ERRORS times for other 503 errors
+                    elif num_consecutive_503_retries < DELETE_MAX_CONSECUTIVE_503_RETRIES:
+                        num_consecutive_503_retries += 1
+                        time.sleep(float(self.delete_retry_delay_millis) / 1000)
+                        continue
+                    else:
+                        raise e
+                else:
+                    raise e
+            break
+        click.echo("\rDelete finished successfully.\033[K")
 
     def mkdirs(self, dbfs_path, headers=None):
         self.client.mkdirs(dbfs_path.absolute_path, headers=headers)

tests/dbfs/test_api.py

@@ -48,6 +48,20 @@ def get_resource_does_not_exist_exception():
     return requests.exceptions.HTTPError(response=response)
 
 
+def get_temporarily_unavailable_exception():
+    response = requests.Response()
+    response.status_code = 503
+    response._content = ('{{"error_code": "{}"}}'.format(api.DbfsErrorCodes.TEMPORARILY_UNAVAILABLE)).encode() #  NOQA
+    return requests.exceptions.HTTPError(response=response)
+
+
+def get_partial_delete_exception(message="[...] operation has deleted 10 files [...]"):
+    response = requests.Response()
+    response.status_code = 503
+    response._content = ('{{"error_code": "{}","message": "{}"}}'.format(api.DbfsErrorCodes.PARTIAL_DELETE, message)).encode() #  NOQA
+    return requests.exceptions.HTTPError(response=response)
+
+
 class TestFileInfo(object):
     def test_to_row_not_long_form_not_absolute(self):
         file_info = api.FileInfo(TEST_DBFS_PATH, False, 1)
@@ -165,3 +179,52 @@ def test_cat(self, dbfs_api):
         with mock.patch('databricks_cli.dbfs.api.click') as click_mock:
             dbfs_api.cat('dbfs:/whatever-doesnt-matter')
             click_mock.echo.assert_called_with('a', nl=False)
+
+    def test_partial_delete(self, dbfs_api):
+        e_partial_delete = get_partial_delete_exception()
+        e_temporarily_unavailable = get_temporarily_unavailable_exception()
+        # Simulate partial deletes and 503 exceptions followed by a full successful delete
+        exception_sequence = \
+            [e_temporarily_unavailable, e_partial_delete, e_partial_delete] + \
+            [e_temporarily_unavailable] * api.DELETE_MAX_CONSECUTIVE_503_RETRIES + \
+            [e_partial_delete, None]
+        dbfs_api.client.delete = mock.Mock(side_effect=exception_sequence)
+        dbfs_api.delete_retry_delay_millis = 1
+        # Should succeed
+        dbfs_api.delete(DbfsPath('dbfs:/whatever-doesnt-matter'), recursive=True)
+
+    def test_partial_delete_service_unavailable(self, dbfs_api):
+        e_partial_delete = get_partial_delete_exception()
+        e_temporarily_unavailable = get_temporarily_unavailable_exception()
+        # Simulate more than api.DELETE_MAX_CONSECUTIVE_503_ERRORS 503 errors that are not partial
+        # deletes (error_code != PARTIAL_DELETE)
+        exception_sequence = \
+            [e_partial_delete] + \
+            [e_temporarily_unavailable] * (api.DELETE_MAX_CONSECUTIVE_503_RETRIES + 1) + \
+            [e_partial_delete, None]
+        dbfs_api.client.delete = mock.Mock(side_effect=exception_sequence)
+        dbfs_api.delete_retry_delay_millis = 1
+        with pytest.raises(e_temporarily_unavailable.__class__) as thrown:
+            dbfs_api.delete(DbfsPath('dbfs:/whatever-doesnt-matter'), recursive=True)
+        # Should raise the same e_temporarily_unavailable exception instance
+        assert thrown.value == e_temporarily_unavailable
+
+    def test_partial_delete_exception_message_parse_error(self, dbfs_api):
+        message = "unexpected partial delete exception message"
+        e_partial_delete = get_partial_delete_exception(message)
+        dbfs_api.client.delete = mock.Mock(side_effect=[e_partial_delete, None])
+        dbfs_api.delete_retry_delay_millis = 1
+        # Should succeed
+        dbfs_api.delete(DbfsPath('dbfs:/whatever-doesnt-matter'), recursive=True)
+
+    def test_get_num_files_deleted(self):
+        e_partial_delete = get_partial_delete_exception()
+        # Should succeed
+        api.DbfsApi.get_num_files_deleted(e_partial_delete)
+
+    def test_get_num_files_deleted_parse_error(self):
+        message = "unexpected partial delete exception message"
+        e_partial_delete = get_partial_delete_exception(message)
+        # Should raise api.ParseException
+        with pytest.raises(api.ParseException):
+            api.DbfsApi.get_num_files_deleted(e_partial_delete)