Add support for --cluster-name or --cluster-id to databricks cluster …

…commands (#313) * Add support for --cluster-name or --cluster-id to databricks cluster commands Add information about what is being removed to the databricks groups remove command Update the license in setup.py Add support for --cluster-name or --cluster-id to databricks libraries commands * Checkpoint work on adding the libraries cluster tests Start adding support for errors if there are multiple clusters with the same name * Remove extraneous prints * Fix a merge conflict where get_clusters_by_name was added to PolicyService instead of ClusterService * remove inadvertant groups changes * Remove get_clusters_by_name function from clusters.cli and leverage ClusterApi.get_clusters_by_name instead. Move CLUSTER_OPTIONS from clusters.cli to utils Change libraries/cli to use ClusterApi.get_clusters_by_name for getting the cluster id instead of using clusters.cli.get_clusters_by_name * Add tests for databrics libraries list --cluster-name * Exclude the .tox directory from code coverage * Add #pragma no-coverage to lines in libaries/cli that will not be hit Fix lint errors around line too long * Move from mocking the libraries API to mocking the sdk interface underneath it to get better code coverage for the libaries api * Remove unused function from clusters/cli.py Add coverage for databricks cluster edit not getting a --json or --json-file parameter * Move hardcoded test cluster data from tests/clusters/test_cli and tests/libraries/test_cli into tests/test_data Add tests for databricks cluster get --cluster-name * Cleanup lint errors * One last lint error * Add tests to databricks cluster commands that just pass the json through for output Add # pragma: no cover to all of the click group functions to clean up code-coverage * Fix overindentation of tabulate * Add no-cover to all of the _group functions to fix code coverage as they aren't tested or used * remove code from sdk/service.py * databricks cluster get --cluster-name should return the same output as databricks cluster get --cluster-id * databricks clusters get should check cluster_id as valid and not only against None Fix the tests to reflect the updated databricks cluster get --cluster-name output Fix clusters/test_api.py to not have pylint warnings * Revert changes to databricks_cli/sdk/service.py as they were whitespace only and this file is autogenerated by Databricks Co-authored-by: Allen Reese <areese999@apple.com> Co-authored-by: Andrew Chen <andrewchen@databricks.com>
databricks · Sep 17, 2020 · d602b54 · d602b54
1 parent 81c058f
commit d602b54
Show file tree

Hide file tree

Showing 23 changed files with 533 additions and 127 deletions.
diff --git a/.coveragerc b/.coveragerc
@@ -0,0 +1,2 @@
+[run]
+omit = .tox/**
diff --git a/databricks_cli/click_types.py b/databricks_cli/click_types.py
@@ -120,6 +120,18 @@ def handle_parse_result(self, ctx, opts, args):
         return super(OneOfOption, self).handle_parse_result(ctx, opts, args)
 
 
+class OptionalOneOfOption(Option):
+    def __init__(self, *args, **kwargs):
+        self.one_of = kwargs.pop('one_of')
+        super(OptionalOneOfOption, self).__init__(*args, **kwargs)
+
+    def handle_parse_result(self, ctx, opts, args):
+        cleaned_opts = set([o.replace('_', '-') for o in opts.keys()])
+        if len(cleaned_opts.intersection(set(self.one_of))) > 1:
+            raise UsageError('Only one of {} should be provided.'.format(self.one_of))
+        return super(OptionalOneOfOption, self).handle_parse_result(ctx, opts, args)
+
+
 class ContextObject(object):
     def __init__(self):
         self._profile = None
@@ -156,3 +168,17 @@ def set_profile(self, profile):
 
     def get_profile(self):
         return self._profile
+
+
+class RequiredOptions(Option):
+    def __init__(self, *args, **kwargs):
+        self.one_of = kwargs.pop('one_of')
+        super(RequiredOptions, self).__init__(*args, **kwargs)
+
+    def handle_parse_result(self, ctx, opts, args):
+        cleaned_opts = set([o.replace('_', '-') for o in opts.keys()])
+        if len(cleaned_opts.intersection(set(self.one_of))) == 0:
+            raise MissingParameter('One of {} must be provided.'.format(self.one_of))
+        if len(cleaned_opts.intersection(set(self.one_of))) > 1:
+            raise UsageError('Only one of {} should be provided.'.format(self.one_of))
+        return super(RequiredOptions, self).handle_parse_result(ctx, opts, args)
diff --git a/databricks_cli/cluster_policies/cli.py b/databricks_cli/cluster_policies/cli.py
@@ -151,7 +151,7 @@ def list_cli(api_client, output):
 @debug_option
 @profile_option
 @eat_exceptions
-def cluster_policies_group():
+def cluster_policies_group():  # pragma: no cover
     """
     Utility to interact with Databricks cluster policies.
     """

diff --git a/databricks_cli/clusters/api.py b/databricks_cli/clusters/api.py
@@ -63,6 +63,41 @@ def spark_versions(self):
     def permanent_delete(self, cluster_id):
         return self.client.permanent_delete_cluster(cluster_id)
 
+    def get_cluster_ids_by_name(self, cluster_name):
+        data = self.client.list_clusters()
+        return [c for c in data.get('clusters', []) if c.get('cluster_name') == cluster_name]
+
+    def get_cluster_id_for_name(self, cluster_name):
+        """
+        Given a cluster name, this will return a single cluster id for that name.
+        If there are multiple clusters with the same name it will raise a RuntimeError.
+        If there are no clusters with the name it will raise a RuntimeError.
+        """
+        clusters_by_name = self.get_cluster_ids_by_name(cluster_name)
+        cluster_ids = [
+            cluster['cluster_id'] for cluster in clusters_by_name if
+            cluster and 'cluster_id' in cluster
+        ]
+
+        if len(cluster_ids) == 0:
+            raise RuntimeError('No clusters with name {} were found'.format(cluster_name))
+
+        if len(cluster_ids) > 1:
+            raise RuntimeError('More than 1 cluster was named {}, '.format(cluster_name) +
+                               'please use --cluster-id.\n' +
+                               'Cluster ids found: {}'.format(', '.join(cluster_ids))
+                               )
+        return cluster_ids[0]
+
+    def get_cluster_by_name(self, cluster_name):
+        """
+        Given a cluster name, this will return the cluster config for that cluster.
+        If there are multiple clusters with the same name it will raise a RuntimeError.
+        If there are no clusters with the name it will raise a RuntimeError.
+        """
+        cluster_id = self.get_cluster_id_for_name(cluster_name)
+        return self.get_cluster(cluster_id)
+
     def get_events(self, cluster_id, start_time, end_time, order, event_types, offset, limit):
         return self.client.get_events(cluster_id, start_time, end_time, order, event_types,
                                       offset, limit)
diff --git a/databricks_cli/clusters/cli.py b/databricks_cli/clusters/cli.py
@@ -20,18 +20,20 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 import time
 from datetime import datetime
 from json import loads as json_loads
 
 import click
 from tabulate import tabulate
 
-from databricks_cli.click_types import OutputClickType, JsonClickType, ClusterIdClickType
+from databricks_cli.click_types import OutputClickType, JsonClickType, ClusterIdClickType, \
+    OneOfOption
 from databricks_cli.clusters.api import ClusterApi
-from databricks_cli.utils import eat_exceptions, CONTEXT_SETTINGS, pretty_format, json_cli_base, \
-    truncate_string
 from databricks_cli.configure.config import provide_api_client, profile_option, debug_option
+from databricks_cli.utils import eat_exceptions, CONTEXT_SETTINGS, pretty_format, json_cli_base, \
+    truncate_string, CLUSTER_OPTIONS
 from databricks_cli.version import print_version_callback, version
 
 
@@ -154,17 +156,26 @@ def delete_cli(api_client, cluster_id):
 
 
 @click.command(context_settings=CONTEXT_SETTINGS)
-@click.option('--cluster-id', required=True, type=ClusterIdClickType(),
-              help=ClusterIdClickType.help)
+@click.option('--cluster-id', cls=OneOfOption, one_of=CLUSTER_OPTIONS,
+              type=ClusterIdClickType(), default=None, help=ClusterIdClickType.help)
+@click.option('--cluster-name', cls=OneOfOption, one_of=CLUSTER_OPTIONS,
+              type=ClusterIdClickType(), default=None, help=ClusterIdClickType.help)
 @debug_option
 @profile_option
 @eat_exceptions
 @provide_api_client
-def get_cli(api_client, cluster_id):
+def get_cli(api_client, cluster_id, cluster_name):
     """
     Retrieves metadata about a cluster.
     """
-    click.echo(pretty_format(ClusterApi(api_client).get_cluster(cluster_id)))
+    if cluster_id:
+        cluster = ClusterApi(api_client).get_cluster(cluster_id)
+    elif cluster_name:
+        cluster = ClusterApi(api_client).get_cluster_by_name(cluster_name)
+    else:
+        raise RuntimeError('cluster_name and cluster_id were empty?')
+
+    click.echo(pretty_format(cluster))
 
 
 def _clusters_to_table(clusters_json):
@@ -327,7 +338,7 @@ def cluster_events_cli(api_client, cluster_id, start_time, end_time, order, even
 @debug_option
 @profile_option
 @eat_exceptions
-def clusters_group():
+def clusters_group():  # pragma: no cover
     """
     Utility to interact with Databricks clusters.
     """

diff --git a/databricks_cli/dbfs/cli.py b/databricks_cli/dbfs/cli.py
@@ -136,7 +136,7 @@ def mv_cli(api_client, src, dst):
               expose_value=False, is_eager=True, help=version)
 @debug_option
 @profile_option
-def dbfs_group():
+def dbfs_group():  # pragma: no cover
     """
     Utility to interact with DBFS.
 

diff --git a/databricks_cli/groups/cli.py b/databricks_cli/groups/cli.py
@@ -146,7 +146,7 @@ def delete_cli(api_client, group_name):
 @debug_option
 @profile_option
 @eat_exceptions
-def groups_group():
+def groups_group():  # pragma: no cover
     """Provide utility to interact with Databricks groups."""
     pass
 

diff --git a/databricks_cli/instance_pools/cli.py b/databricks_cli/instance_pools/cli.py
@@ -150,7 +150,7 @@ def list_cli(api_client, output):
 @debug_option
 @profile_option
 @eat_exceptions
-def instance_pools_group():
+def instance_pools_group():  # pragma: no cover
     """
     Utility to interact with Databricks instance pools.
     """

diff --git a/databricks_cli/jobs/cli.py b/databricks_cli/jobs/cli.py
@@ -188,7 +188,7 @@ def run_now_cli(api_client, job_id, jar_params, notebook_params, python_params,
 @debug_option
 @profile_option
 @eat_exceptions
-def jobs_group():
+def jobs_group():  # pragma: no cover
     """
     Utility to interact with jobs.
 

diff --git a/databricks_cli/libraries/cli.py b/databricks_cli/libraries/cli.py
@@ -23,22 +23,23 @@
 
 import click
 
-from databricks_cli.click_types import ClusterIdClickType, OneOfOption
+from databricks_cli.click_types import ClusterIdClickType, OneOfOption, OptionalOneOfOption
+from databricks_cli.clusters.api import ClusterApi
 from databricks_cli.configure.config import provide_api_client, profile_option, debug_option
 from databricks_cli.libraries.api import LibrariesApi
-from databricks_cli.utils import CONTEXT_SETTINGS, eat_exceptions, pretty_format
+from databricks_cli.utils import CONTEXT_SETTINGS, eat_exceptions, pretty_format, CLUSTER_OPTIONS
 from databricks_cli.version import print_version_callback, version
 
 
-def _all_cluster_statuses(config):
-    click.echo(pretty_format(LibrariesApi(config).all_cluster_statuses()))
+def _all_cluster_statuses(api_client):
+    click.echo(pretty_format(LibrariesApi(api_client).all_cluster_statuses()))
 
 
 @click.command(context_settings=CONTEXT_SETTINGS,
                short_help='Get the status of all libraries.')
 @debug_option
 @profile_option
-@eat_exceptions # noqa
+@eat_exceptions  # noqa
 @provide_api_client
 def all_cluster_statuses_cli(api_client):
     """
@@ -50,45 +51,54 @@ def all_cluster_statuses_cli(api_client):
     _all_cluster_statuses(api_client)
 
 
-def _cluster_status(api_client, cluster_id):
-    click.echo(pretty_format(LibrariesApi(api_client).cluster_status(cluster_id)))
+def _cluster_status(api_client, cluster_id, cluster_name):
+    libraries_api = LibrariesApi(api_client)
+
+    if not cluster_id:
+        cluster_id = ClusterApi(api_client).get_cluster_id_for_name(cluster_name)
+
+    click.echo(pretty_format(libraries_api.cluster_status(cluster_id)))
 
 
 @click.command(context_settings=CONTEXT_SETTINGS,
                short_help='Get the status of all libraries for a specified cluster.')
-@click.option('--cluster-id', required=True, type=ClusterIdClickType(),
-              help=ClusterIdClickType.help)
+@click.option('--cluster-id', cls=OneOfOption, one_of=CLUSTER_OPTIONS,
+              type=ClusterIdClickType(), default=None, help=ClusterIdClickType.help)
+@click.option('--cluster-name', cls=OneOfOption, one_of=CLUSTER_OPTIONS,
+              type=ClusterIdClickType(), default=None, help=ClusterIdClickType.help)
 @debug_option
 @profile_option
-@eat_exceptions # noqa
+@eat_exceptions  # noqa
 @provide_api_client
-def cluster_status_cli(api_client, cluster_id):
+def cluster_status_cli(api_client, cluster_id, cluster_name):
     """
     Get the status of all libraries for a specified cluster. A status will be available for all
     libraries installed on this cluster via the API or the libraries UI as well as libraries set to
     be installed on all clusters via the libraries UI. If a library has been set to be installed on
     all clusters, is_library_for_all_clusters will be true.
     """
-    _cluster_status(api_client, cluster_id)
+    _cluster_status(api_client, cluster_id, cluster_name)
 
 
 @click.command(context_settings=CONTEXT_SETTINGS,
                short_help='Shortcut to `all-cluster-statuses` or `cluster-status`.')
-@click.option('--cluster-id', type=ClusterIdClickType(), default=None,
-              help=ClusterIdClickType.help)
+@click.option('--cluster-id', cls=OptionalOneOfOption, one_of=CLUSTER_OPTIONS,
+              type=ClusterIdClickType(), default=None, help=ClusterIdClickType.help)
+@click.option('--cluster-name', cls=OptionalOneOfOption, one_of=CLUSTER_OPTIONS,
+              type=ClusterIdClickType(), default=None, help=ClusterIdClickType.help)
 @debug_option
 @profile_option
-@eat_exceptions # noqa
+@eat_exceptions  # noqa
 @provide_api_client
-def list_cli(api_client, cluster_id):
+def list_cli(api_client, cluster_id, cluster_name):
     """
-    Get the statsus of all libraries for all clusters or for a specified cluster.
+    Get the statuses of all libraries for all clusters or for a specified cluster.
     If the option --cluster-id is provided, then all libraries on that cluster will be listed,
     (cluster-status). If the option --cluster-id is omitted, then all libraries on all clusters
     will be listed (all-cluster-statuses).
     """
-    if cluster_id is not None:
-        _cluster_status(api_client, cluster_id)
+    if cluster_id is not None or cluster_name is not None:
+        _cluster_status(api_client, cluster_id, cluster_name)
     else:
         _all_cluster_statuses(api_client)
 
@@ -127,7 +137,7 @@ def list_cli(api_client, cluster_id):
 """
 
 
-def _get_library_from_options(jar, egg, whl, maven_coordinates, maven_repo, maven_exclusion, # noqa
+def _get_library_from_options(jar, egg, whl, maven_coordinates, maven_repo, maven_exclusion,  # noqa
                               pypi_package, pypi_repo, cran_package, cran_repo):
     maven_exclusion = list(maven_exclusion)
     if jar is not None:
@@ -153,7 +163,7 @@ def _get_library_from_options(jar, egg, whl, maven_coordinates, maven_repo, mave
         if cran_repo is not None:
             cran_library['cran']['repo'] = cran_repo
         return cran_library
-    raise AssertionError('Code not reached.')
+    raise AssertionError('Code not reached.')  # pragma: no cover
 
 
 @click.command(context_settings=CONTEXT_SETTINGS,
@@ -173,9 +183,9 @@ def _get_library_from_options(jar, egg, whl, maven_coordinates, maven_repo, mave
 @click.option('--cran-repo', help=CRAN_REPO_HELP)
 @debug_option
 @profile_option
-@eat_exceptions # noqa
+@eat_exceptions  # noqa
 @provide_api_client
-def install_cli(api_client, cluster_id, jar, egg, whl, maven_coordinates, maven_repo, # noqa
+def install_cli(api_client, cluster_id, jar, egg, whl, maven_coordinates, maven_repo,  # noqa
                 maven_exclusion, pypi_package, pypi_repo, cran_package, cran_repo):
     """
     Install a library on a cluster. Libraries must be first uploaded to dbfs or s3
@@ -227,9 +237,9 @@ def _uninstall_cli_exit_help(cluster_id):
 @click.option('--cran-repo', help=CRAN_REPO_HELP)
 @debug_option
 @profile_option
-@eat_exceptions # noqa
+@eat_exceptions  # noqa
 @provide_api_client
-def uninstall_cli(api_client, cluster_id, all, jar, egg, whl, maven_coordinates, maven_repo, # noqa
+def uninstall_cli(api_client, cluster_id, all, jar, egg, whl, maven_coordinates, maven_repo,  # noqa
                   maven_exclusion, pypi_package, pypi_repo, cran_package, cran_repo):
     """
     Mark libraries on a cluster to be uninstalled. Libraries which are marked to be uninstalled
@@ -258,7 +268,7 @@ def uninstall_cli(api_client, cluster_id, all, jar, egg, whl, maven_coordinates,
 @debug_option
 @profile_option
 @eat_exceptions
-def libraries_group():
+def libraries_group():  # pragma: no cover
     """
     Utility to interact with libraries.
 

diff --git a/databricks_cli/pipelines/cli.py b/databricks_cli/pipelines/cli.py
@@ -265,7 +265,7 @@ def _handle_duplicate_name_exception(spec, exception):
               expose_value=False, is_eager=True, help=version)
 @debug_option
 @profile_option
-def pipelines_group():
+def pipelines_group():  # pragma: no cover
     """
     Utility to interact with the Databricks pipelines.
     """

diff --git a/databricks_cli/runs/cli.py b/databricks_cli/runs/cli.py
@@ -155,7 +155,7 @@ def cancel_cli(api_client, run_id):
 @debug_option
 @profile_option
 @eat_exceptions
-def runs_group():
+def runs_group():  # pragma: no cover
     """
     Utility to interact with jobs runs.
     """

diff --git a/databricks_cli/secrets/cli.py b/databricks_cli/secrets/cli.py
@@ -316,7 +316,7 @@ def get_acl(api_client, scope, principal, output):
 @debug_option
 @profile_option
 @eat_exceptions
-def secrets_group():
+def secrets_group():  # pragma: no cover
     """
     Utility to interact with secret API.
     """

diff --git a/databricks_cli/stack/cli.py b/databricks_cli/stack/cli.py
@@ -147,7 +147,7 @@ def download(api_client, config_path, **kwargs):
               expose_value=False, is_eager=True, help=version)
 @debug_option
 @profile_option
-def stack_group():
+def stack_group():  # pragma: no cover
     """
     [Beta] Utility to deploy and download Databricks resource stacks.
     """

diff --git a/databricks_cli/tokens/cli.py b/databricks_cli/tokens/cli.py
@@ -89,7 +89,7 @@ def revoke_cli(api_client, token_id):
 @debug_option
 @profile_option
 @eat_exceptions
-def tokens_group():
+def tokens_group():  # pragma: no cover
     """Utility to interact with Databricks tokens."""
     pass
 

diff --git a/databricks_cli/utils.py b/databricks_cli/utils.py
@@ -32,6 +32,7 @@
 from databricks_cli.click_types import ContextObject
 
 CONTEXT_SETTINGS = dict(help_option_names=['-h', '--help'])
+CLUSTER_OPTIONS = ['cluster-id', 'cluster-name']
 DEBUG_MODE = False