Add --long option to workspace ls cli (#310)

* Add --long option to workspace ls cli
Add --id option to workspace ls cli to display the object id as well

* Add missing object_id argument to tests
Make with_object_id default to False in to_row
Fix lint over 100 characters error

Co-authored-by: Allen Reese <areese999@apple.com>

Author: areese

databricks_cli/workspace/api.py

@@ -37,23 +37,30 @@
 
 
 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')
         elif self.is_library:
             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):

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)
 
 

tests/workspace/test_api.py

@@ -1,3 +1,4 @@
+
 # Databricks CLI
 # Copyright 2017 Databricks, Inc.
 #
@@ -31,38 +32,44 @@
 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'
 
 
 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
         assert row[1] == TEST_WORKSPACE_PATH
         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:

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(