From c6b094e5ff5ce7fecd7c676457f91767900123a4 Mon Sep 17 00:00:00 2001 From: LuigiBrancati Date: Wed, 7 Jun 2023 23:22:22 +0200 Subject: [PATCH 1/3] Added AUTO format --- databricks_cli/workspace/types.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/databricks_cli/workspace/types.py b/databricks_cli/workspace/types.py index 8e685eac..50d1a4ee 100644 --- a/databricks_cli/workspace/types.py +++ b/databricks_cli/workspace/types.py @@ -83,11 +83,12 @@ def convert(self, value, param, ctx): class WorkspaceFormat(object): + AUTO = 'AUTO' SOURCE = 'SOURCE' HTML = 'HTML' JUPYTER = 'JUPYTER' DBC = 'DBC' - ALL = [SOURCE, HTML, JUPYTER, DBC] + ALL = [AUTO, SOURCE, HTML, JUPYTER, DBC] class FormatClickType(ParamType): From a9b18f3285734db62e477318c29892965ab4eb85 Mon Sep 17 00:00:00 2001 From: LuigiBrancati Date: Wed, 7 Jun 2023 23:23:06 +0200 Subject: [PATCH 2/3] set language as an optional parameter --- databricks_cli/sdk/service.py | 6 +++--- databricks_cli/workspace/api.py | 6 +++--- databricks_cli/workspace/cli.py | 8 ++++---- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/databricks_cli/sdk/service.py b/databricks_cli/sdk/service.py index 8083e932..2f2b86e5 100755 --- a/databricks_cli/sdk/service.py +++ b/databricks_cli/sdk/service.py @@ -1013,19 +1013,19 @@ def list(self, path, headers=None): return self.client.perform_query('GET', '/workspace/list', data=_data, headers=headers) def import_workspace( - self, path, format=None, language=None, content=None, overwrite=None, headers=None + self, path, format=None, content=None, overwrite=None, headers=None, language=None ): _data = {} if path is not None: _data['path'] = path if format is not None: _data['format'] = format - if language is not None: - _data['language'] = language if content is not None: _data['content'] = content if overwrite is not None: _data['overwrite'] = overwrite + if language is not None: + _data['language'] = language return self.client.perform_query('POST', '/workspace/import', data=_data, headers=headers) def export_workspace(self, path, format=None, direct_download=None, headers=None): diff --git a/databricks_cli/workspace/api.py b/databricks_cli/workspace/api.py index e38db7bf..88ec9ef9 100644 --- a/databricks_cli/workspace/api.py +++ b/databricks_cli/workspace/api.py @@ -103,17 +103,17 @@ def list_objects(self, workspace_path, headers=None): def mkdirs(self, workspace_path, headers=None): self.client.mkdirs(workspace_path, headers=headers) - def import_workspace(self, source_path, target_path, language, fmt, is_overwrite, headers=None): + def import_workspace(self, source_path, target_path, fmt, is_overwrite, headers=None, language=None): with open(source_path, 'rb') as f: # import_workspace must take content that is typed str. content = b64encode(f.read()).decode() self.client.import_workspace( target_path, fmt, - language, content, is_overwrite, - headers=headers) + headers=headers, + language=language) def export_workspace(self, source_path, target_path, fmt, is_overwrite, headers=None): """ diff --git a/databricks_cli/workspace/cli.py b/databricks_cli/workspace/cli.py index be03c0a8..c800569b 100644 --- a/databricks_cli/workspace/cli.py +++ b/databricks_cli/workspace/cli.py @@ -83,15 +83,15 @@ def mkdirs_cli(api_client, workspace_path): short_help='Imports a file from local to the Databricks workspace.') @click.argument('source_path') @click.argument('target_path') -@click.option('--language', '-l', required=True, type=LanguageClickType(), - help=', '.join(WorkspaceLanguage.ALL)) @click.option('--format', '-f', default=WorkspaceFormat.SOURCE, type=FormatClickType()) @click.option('--overwrite', '-o', is_flag=True, default=False) +@click.option('--language', '-l', required=False, type=LanguageClickType(), + help=', '.join(WorkspaceLanguage.ALL)) @debug_option @profile_option @eat_exceptions @provide_api_client -def import_workspace_cli(api_client, source_path, target_path, language, format, overwrite): # NOQA +def import_workspace_cli(api_client, source_path, target_path, format, overwrite, language): # NOQA """ Imports a file from local to the Databricks workspace. @@ -99,7 +99,7 @@ def import_workspace_cli(api_client, source_path, target_path, language, format, format is documented at https://docs.databricks.com/api/latest/workspace.html#notebookexportformat. """ - WorkspaceApi(api_client).import_workspace(source_path, target_path, language, format, overwrite) # NOQA + WorkspaceApi(api_client).import_workspace(source_path, target_path, format, overwrite, language) # NOQA @click.command(context_settings=CONTEXT_SETTINGS, From 4d202d0aa46b04d2e0707e0debd0486c8a5f396a Mon Sep 17 00:00:00 2001 From: LuigiBrancati Date: Wed, 7 Jun 2023 23:46:37 +0200 Subject: [PATCH 3/3] Aligned tests --- tests/workspace/test_api.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/tests/workspace/test_api.py b/tests/workspace/test_api.py index f3c48d55..ac5b37fc 100644 --- a/tests/workspace/test_api.py +++ b/tests/workspace/test_api.py @@ -111,15 +111,16 @@ def test_import_workspace(self, workspace_api, tmpdir): test_file_path = os.path.join(tmpdir.strpath, 'test') with open(test_file_path, 'w') as f: f.write('test') - workspace_api.import_workspace(test_file_path, TEST_WORKSPACE_PATH, TEST_LANGUAGE, - TEST_FMT, is_overwrite=False) + workspace_api.import_workspace(test_file_path, TEST_WORKSPACE_PATH, + TEST_FMT, is_overwrite=False, language=TEST_LANGUAGE) import_workspace_mock = workspace_api.client.import_workspace assert import_workspace_mock.call_count == 1 assert import_workspace_mock.call_args[0][0] == TEST_WORKSPACE_PATH assert import_workspace_mock.call_args[0][1] == TEST_FMT - assert import_workspace_mock.call_args[0][2] == TEST_LANGUAGE - assert import_workspace_mock.call_args[0][3] == b64encode(b'test').decode() - assert import_workspace_mock.call_args[0][4] is False + assert import_workspace_mock.call_args[0][2] == b64encode(b'test').decode() + assert import_workspace_mock.call_args[0][3] is False + assert import_workspace_mock.call_args[1]['headers'] is None + assert import_workspace_mock.call_args[1]['language'] == TEST_LANGUAGE def test_export_workspace(self, workspace_api, tmpdir): test_file_path = os.path.join(tmpdir.strpath, 'test')