From b24a1ac9ec277098a0f7895e0238342f1edec9b2 Mon Sep 17 00:00:00 2001 From: Oleg Boiko Date: Tue, 10 Mar 2020 12:16:03 -0700 Subject: [PATCH 1/2] Adding --format option to workspace export --- databricks_cli/workspace/api.py | 4 ++-- databricks_cli/workspace/cli.py | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/databricks_cli/workspace/api.py b/databricks_cli/workspace/api.py index a1813459..43e21fdc 100644 --- a/databricks_cli/workspace/api.py +++ b/databricks_cli/workspace/api.py @@ -157,7 +157,7 @@ def import_workspace_dir(self, source_path, target_path, overwrite, exclude_hidd click.echo(('{} does not have a valid extension of {}. Skip this file and ' + 'continue.').format(cur_src, extensions)) - def export_workspace_dir(self, source_path, target_path, overwrite, headers=None): + def export_workspace_dir(self, source_path, target_path, overwrite, format=WorkspaceFormat.SOURCE, headers=None): if os.path.isfile(target_path): click.echo('{} exists as a file. Skipping this subtree {}' .format(target_path, source_path)) @@ -172,7 +172,7 @@ def export_workspace_dir(self, source_path, target_path, overwrite, headers=None elif obj.is_notebook: cur_dst = cur_dst + WorkspaceLanguage.to_extension(obj.language) try: - self.export_workspace(cur_src, cur_dst, WorkspaceFormat.SOURCE, overwrite, + self.export_workspace(cur_src, cur_dst, format, overwrite, headers=headers) click.echo('{} -> {}'.format(cur_src, cur_dst)) except LocalFileExistsException: diff --git a/databricks_cli/workspace/cli.py b/databricks_cli/workspace/cli.py index 19997e9a..d1368d20 100644 --- a/databricks_cli/workspace/cli.py +++ b/databricks_cli/workspace/cli.py @@ -147,11 +147,12 @@ def delete_cli(api_client, workspace_path, recursive): @click.argument('source_path') @click.argument('target_path') @click.option('--overwrite', '-o', is_flag=True, default=False) +@click.option('--format', '-f', default=WorkspaceFormat.SOURCE, type=FormatClickType()) @debug_option @profile_option @eat_exceptions @provide_api_client -def export_dir_cli(api_client, source_path, target_path, overwrite): +def export_dir_cli(api_client, source_path, target_path, overwrite, format): """ Recursively exports a directory from the Databricks workspace. @@ -162,7 +163,7 @@ def export_dir_cli(api_client, source_path, target_path, overwrite): workspace_api = WorkspaceApi(api_client) assert workspace_api.get_status(source_path).is_dir, 'The source path must be a directory. {}' \ .format(source_path) - workspace_api.export_workspace_dir(source_path, target_path, overwrite) + workspace_api.export_workspace_dir(source_path, target_path, overwrite, format) @click.command(context_settings=CONTEXT_SETTINGS, From 64e477cf3411b964e84fafe80623cefd367f8109 Mon Sep 17 00:00:00 2001 From: Oleg Boiko Date: Tue, 10 Mar 2020 12:27:32 -0700 Subject: [PATCH 2/2] Fixing link errors --- databricks_cli/workspace/cli.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/databricks_cli/workspace/cli.py b/databricks_cli/workspace/cli.py index d1368d20..217bb01b 100644 --- a/databricks_cli/workspace/cli.py +++ b/databricks_cli/workspace/cli.py @@ -86,7 +86,7 @@ def mkdirs_cli(api_client, workspace_path): @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, language, fmt, overwrite): # NOQA """ Imports a file from local to the Databricks workspace. @@ -94,7 +94,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, language, fmt, overwrite) # NOQA @click.command(context_settings=CONTEXT_SETTINGS, @@ -107,7 +107,7 @@ def import_workspace_cli(api_client, source_path, target_path, language, format, @profile_option @eat_exceptions @provide_api_client -def export_workspace_cli(api_client, source_path, target_path, format, overwrite): # NOQA +def export_workspace_cli(api_client, source_path, target_path, fmt, overwrite): # NOQA """ Exports a notebook from the Databricks workspace. @@ -121,7 +121,7 @@ def export_workspace_cli(api_client, source_path, target_path, format, overwrite raise RuntimeError('Export can only be called on a notebook.') extension = WorkspaceLanguage.to_extension(file_info.language) target_path = os.path.join(target_path, file_info.basename + extension) - WorkspaceApi(api_client).export_workspace(source_path, target_path, format, overwrite) # NOQA + WorkspaceApi(api_client).export_workspace(source_path, target_path, fmt, overwrite) # NOQA @click.command(context_settings=CONTEXT_SETTINGS, @@ -152,7 +152,7 @@ def delete_cli(api_client, workspace_path, recursive): @profile_option @eat_exceptions @provide_api_client -def export_dir_cli(api_client, source_path, target_path, overwrite, format): +def export_dir_cli(api_client, source_path, target_path, overwrite, fmt): """ Recursively exports a directory from the Databricks workspace. @@ -163,7 +163,7 @@ def export_dir_cli(api_client, source_path, target_path, overwrite, format): workspace_api = WorkspaceApi(api_client) assert workspace_api.get_status(source_path).is_dir, 'The source path must be a directory. {}' \ .format(source_path) - workspace_api.export_workspace_dir(source_path, target_path, overwrite, format) + workspace_api.export_workspace_dir(source_path, target_path, overwrite, fmt) @click.command(context_settings=CONTEXT_SETTINGS, @@ -200,6 +200,7 @@ def workspace_group(): """ pass + workspace_group.add_command(ls_cli, name='ls') workspace_group.add_command(ls_cli, name='list') workspace_group.add_command(mkdirs_cli, name='mkdirs')