diff --git a/databricks_cli/workspace/api.py b/databricks_cli/workspace/api.py index a1813459..571abed5 100644 --- a/databricks_cli/workspace/api.py +++ b/databricks_cli/workspace/api.py @@ -37,12 +37,13 @@ class WorkspaceFileInfo(object): - def __init__(self, path, object_type, language=None, **kwargs): # noqa + def __init__(self, path, object_type, object_id, language=None, **kwargs): # noqa self.path = path self.object_type = object_type self.language = language + self.object_id = object_id - def to_row(self, is_long_form, is_absolute): + def to_row(self, is_long_form, is_absolute, with_object_id=False): path = self.path if is_absolute else self.basename if self.is_dir: stylized_path = click.style(path, 'cyan') @@ -50,10 +51,16 @@ def to_row(self, is_long_form, is_absolute): stylized_path = click.style(path, 'green') else: stylized_path = path + + result = [stylized_path] + if is_long_form: - return [self.object_type, stylized_path, self.language] - else: - return [stylized_path] + result = [self.object_type, stylized_path, self.language] + + if with_object_id: + result.append(self.object_id) + + return result @property def is_dir(self): diff --git a/databricks_cli/workspace/cli.py b/databricks_cli/workspace/cli.py index 4b95d83f..9e1b2670 100644 --- a/databricks_cli/workspace/cli.py +++ b/databricks_cli/workspace/cli.py @@ -37,14 +37,16 @@ short_help='List objects in the Databricks Workspace. ls and list are synonyms.') @click.option('--absolute', is_flag=True, default=False, help='Displays absolute paths.') -@click.option('-l', is_flag=True, default=False, +@click.option('-l', '--long', 'is_long_form', is_flag=True, default=False, help='Displays full information including ObjectType, Path, Language') +@click.option('-i', '--id', 'with_object_id', is_flag=True, default=False, + help='Displays ObjectId') @click.argument('workspace_path', type=str, nargs=-1) @debug_option @profile_option @eat_exceptions @provide_api_client -def ls_cli(api_client, l, absolute, workspace_path): # noqa +def ls_cli(api_client, is_long_form, with_object_id, absolute, workspace_path): # noqa """ List objects in the Databricks Workspace. """ @@ -53,8 +55,11 @@ def ls_cli(api_client, l, absolute, workspace_path): # noqa else: workspace_path = workspace_path[0] objects = WorkspaceApi(api_client).list_objects(workspace_path) - table = tabulate([obj.to_row(is_long_form=l, is_absolute=absolute) for obj in objects], - tablefmt='plain') + objects_table = [ + obj.to_row(is_long_form=is_long_form, is_absolute=absolute, with_object_id=with_object_id) + for obj in objects + ] + table = tabulate(objects_table, tablefmt='plain') click.echo(table) diff --git a/tests/workspace/test_api.py b/tests/workspace/test_api.py index 41bc88a2..f3c48d55 100644 --- a/tests/workspace/test_api.py +++ b/tests/workspace/test_api.py @@ -1,3 +1,4 @@ + # Databricks CLI # Copyright 2017 Databricks, Inc. # @@ -31,10 +32,12 @@ from databricks_cli.workspace.types import WorkspaceLanguage TEST_WORKSPACE_PATH = '/test/workspace/path' +TEST_WORKSPACE_OBJECT_ID = '22' TEST_JSON_RESPONSE = { 'path': TEST_WORKSPACE_PATH, 'object_type': api.DIRECTORY, - 'created_at': 124 + 'created_at': 124, + 'object_id': TEST_WORKSPACE_OBJECT_ID, } TEST_LANGUAGE = 'PYTHON' TEST_FMT = 'SOURCE' @@ -42,19 +45,22 @@ class TestWorkspaceFileInfo(object): def test_to_row_not_long_form_not_absolute(self): - file_info = api.WorkspaceFileInfo(TEST_WORKSPACE_PATH, api.NOTEBOOK) + file_info = api.WorkspaceFileInfo(TEST_WORKSPACE_PATH, api.NOTEBOOK, + TEST_WORKSPACE_OBJECT_ID) row = file_info.to_row(is_long_form=False, is_absolute=False) assert len(row) == 1 assert row[0] == file_info.basename def test_to_row_not_long_form_absolute(self): - file_info = api.WorkspaceFileInfo(TEST_WORKSPACE_PATH, api.NOTEBOOK) + file_info = api.WorkspaceFileInfo(TEST_WORKSPACE_PATH, api.NOTEBOOK, + TEST_WORKSPACE_OBJECT_ID) row = file_info.to_row(is_long_form=False, is_absolute=True) assert len(row) == 1 assert row[0] == TEST_WORKSPACE_PATH def test_to_row_long_form_absolute(self): - file_info = api.WorkspaceFileInfo(TEST_WORKSPACE_PATH, api.NOTEBOOK) + file_info = api.WorkspaceFileInfo(TEST_WORKSPACE_PATH, api.NOTEBOOK, + TEST_WORKSPACE_OBJECT_ID) row = file_info.to_row(is_long_form=True, is_absolute=True) assert len(row) == 3 assert row[0] == api.NOTEBOOK @@ -62,7 +68,8 @@ def test_to_row_long_form_absolute(self): assert row[2] is None def test_basename(self): - file_info = api.WorkspaceFileInfo(TEST_WORKSPACE_PATH, api.NOTEBOOK) + file_info = api.WorkspaceFileInfo(TEST_WORKSPACE_PATH, api.NOTEBOOK, + TEST_WORKSPACE_OBJECT_ID) assert file_info.basename == 'path' def test_from_json(self): @@ -104,8 +111,8 @@ 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_LANGUAGE, + TEST_FMT, is_overwrite=False) 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 @@ -147,18 +154,22 @@ def test_export_workspace_dir(self, workspace_api, tmpdir): def _list_objects_mock(path, headers=None): if path == '/': return [ - WorkspaceFileInfo('/a', api.DIRECTORY), - WorkspaceFileInfo('/f', api.DIRECTORY) + WorkspaceFileInfo('/a', api.DIRECTORY, TEST_WORKSPACE_OBJECT_ID), + WorkspaceFileInfo('/f', api.DIRECTORY, TEST_WORKSPACE_OBJECT_ID) ] elif path == '/a': return [ - WorkspaceFileInfo('/a/b', api.NOTEBOOK, WorkspaceLanguage.SCALA), - WorkspaceFileInfo('/a/c', api.NOTEBOOK, WorkspaceLanguage.PYTHON), - WorkspaceFileInfo('/a/d', api.NOTEBOOK, WorkspaceLanguage.R), - WorkspaceFileInfo('/a/e', api.NOTEBOOK, WorkspaceLanguage.SQL), + WorkspaceFileInfo('/a/b', api.NOTEBOOK, TEST_WORKSPACE_OBJECT_ID, + WorkspaceLanguage.SCALA), + WorkspaceFileInfo('/a/c', api.NOTEBOOK, TEST_WORKSPACE_OBJECT_ID, + WorkspaceLanguage.PYTHON), + WorkspaceFileInfo('/a/d', api.NOTEBOOK, TEST_WORKSPACE_OBJECT_ID, + WorkspaceLanguage.R), + WorkspaceFileInfo('/a/e', api.NOTEBOOK, TEST_WORKSPACE_OBJECT_ID, + WorkspaceLanguage.SQL), ] elif path == '/f': - return [WorkspaceFileInfo('/f/g', api.DIRECTORY)] + return [WorkspaceFileInfo('/f/g', api.DIRECTORY, TEST_WORKSPACE_OBJECT_ID)] elif path == '/f/g': return [] else: diff --git a/tests/workspace/test_cli.py b/tests/workspace/test_cli.py index 9c600868..aa13cf24 100644 --- a/tests/workspace/test_cli.py +++ b/tests/workspace/test_cli.py @@ -44,7 +44,7 @@ def workspace_api_mock(): def test_export_workspace_cli(workspace_api_mock, tmpdir): path = tmpdir.strpath workspace_api_mock.get_status.return_value = \ - WorkspaceFileInfo('/notebook-name', NOTEBOOK, WorkspaceLanguage.SCALA) + WorkspaceFileInfo('/notebook-name', NOTEBOOK, '22', WorkspaceLanguage.SCALA) runner = CliRunner() runner.invoke(cli.export_workspace_cli, ['--format', 'SOURCE', '/notebook-name', path]) assert workspace_api_mock.export_workspace.call_args[0][1] == os.path.join(