Patches

databricks_cli/cli.py

@@ -35,6 +35,7 @@
 from databricks_cli.runs.cli import runs_group
 from databricks_cli.secrets.cli import secrets_group
 from databricks_cli.stack.cli import stack_group
+from databricks_cli.groups.cli import groups_group
 
 
 @click.group(context_settings=CONTEXT_SETTINGS)
@@ -55,3 +56,4 @@ def cli():
 cli.add_command(libraries_group, name='libraries')
 cli.add_command(secrets_group, name='secrets')
 cli.add_command(stack_group, name='stack')
+cli.add_command(groups_group, name='groups')

databricks_cli/click_types.py

@@ -23,8 +23,6 @@
 
 from click import ParamType, Option, MissingParameter, UsageError
 
-from databricks_cli.configure.provider import DEFAULT_SECTION
-
 
 class OutputClickType(ParamType):
     name = 'FORMAT'
@@ -117,6 +115,4 @@ def set_profile(self, profile):
         self._profile = profile
 
     def get_profile(self):
-        if self._profile is None:
-            return DEFAULT_SECTION
         return self._profile

databricks_cli/configure/cli.py

@@ -26,7 +26,7 @@
 from click import ParamType
 
 from databricks_cli.configure.provider import DatabricksConfig, update_and_persist_config, \
-    get_config_for_profile
+    ProfileConfigProvider
 from databricks_cli.utils import CONTEXT_SETTINGS
 from databricks_cli.configure.config import profile_option, get_profile_from_context, debug_option
 
@@ -37,15 +37,15 @@
 
 
 def _configure_cli_token(profile, insecure):
-    config = get_config_for_profile(profile)
+    config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty()
     host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost())
     token = click.prompt(PROMPT_TOKEN, default=config.token)
     new_config = DatabricksConfig.from_token(host, token, insecure)
     update_and_persist_config(profile, new_config)
 
 
 def _configure_cli_password(profile, insecure):
-    config = get_config_for_profile(profile)
+    config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty()
     if config.password:
         default_password = '*' * len(config.password)
     else:

databricks_cli/configure/config.py

@@ -26,7 +26,7 @@
 import six
 
 from databricks_cli.click_types import ContextObject
-from databricks_cli.configure.provider import get_config_for_profile
+from databricks_cli.configure.provider import get_config, ProfileConfigProvider
 from databricks_cli.utils import InvalidConfigurationError
 from databricks_cli.sdk import ApiClient
 
@@ -42,9 +42,14 @@ def decorator(*args, **kwargs):
         command_name = "-".join(ctx.command_path.split(" ")[1:])
         command_name += "-" + str(uuid.uuid1())
         profile = get_profile_from_context()
-        config = get_config_for_profile(profile)
-        if not config.is_valid:
-            raise InvalidConfigurationError(profile)
+        if profile:
+            # If we request a specific profile, only get credentials from tere.
+            config = ProfileConfigProvider(profile).get_config()
+        else:
+            # If unspecified, use the default provider, or allow for user overrides.
+            config = get_config()
+        if not config or not config.is_valid:
+            raise InvalidConfigurationError.for_profile(profile)
         kwargs['api_client'] = _get_api_client(config, command_name)
 
         return function(*args, **kwargs)

databricks_cli/configure/provider.py

@@ -20,10 +20,14 @@
 # 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.
+
+from abc import abstractmethod, ABCMeta
 from configparser import ConfigParser
+import os
 from os.path import expanduser, join
 
-import os
+from databricks_cli.utils import InvalidConfigurationError
+
 
 _home = expanduser('~')
 HOST = 'host'
@@ -33,6 +37,9 @@
 INSECURE = 'insecure'
 DEFAULT_SECTION = 'DEFAULT'
 
+# User-provided override for the DatabricksConfigProvider
+_config_provider = None
+
 
 def _get_path():
     return join(_home, '.databrickscfg')
@@ -81,6 +88,7 @@ def update_and_persist_config(profile, databricks_config):
     same profile.
     :param databricks_config: DatabricksConfig
     """
+    profile = profile if profile else DEFAULT_SECTION
     raw_config = _fetch_from_fs()
     _create_section_if_absent(raw_config, profile)
     _set_option(raw_config, profile, HOST, databricks_config.host)
@@ -91,35 +99,136 @@ def update_and_persist_config(profile, databricks_config):
     _overwrite_config(raw_config)
 
 
+def get_config():
+    """
+    Returns a DatabricksConfig containing the hostname and authentication used to talk to
+    the Databricks API. By default, we leverage the DefaultConfigProvider to get
+    this config, but this behavior may be overridden by calling 'set_config_provider'
+
+    If no DatabricksConfig can be found, an InvalidConfigurationError will be raised.
+    """
+    global _config_provider
+    if _config_provider:
+        config = _config_provider.get_config()
+        if config:
+            return config
+        raise InvalidConfigurationError(
+            'Custom provider returned no DatabricksConfig: %s' % _config_provider)
+
+    config = DefaultConfigProvider().get_config()
+    if config:
+        return config
+    raise InvalidConfigurationError.for_profile(None)
+
+
 def get_config_for_profile(profile):
     """
-    Reads from the filesystem and gets a DatabricksConfig for the specified profile. If it does not
-    exist, then return a DatabricksConfig with fields set.
+    [Deprecated] Reads from the filesystem and gets a DatabricksConfig for the
+    specified profile. If it does not exist, then return a DatabricksConfig with fields set
+    to None.
+
+    Internal callers should prefer get_config() to use user-specified overrides, and
+    to return appropriate error messages as opposited to invalid configurations.
+
+    If you want to read from a specific profile, please instead use
+    'ProfileConfigProvider(profile).get_config()'.
+
+    This method is maintained for backwards-compatibility. It may be removed in future versions.
+
     :return: DatabricksConfig
     """
-    if is_environment_set():
+    profile = profile if profile else DEFAULT_SECTION
+    config = EnvironmentVariableConfigProvider().get_config()
+    if config and config.is_valid:
+        return config
+
+    config = ProfileConfigProvider(profile).get_config()
+    if config:
+        return config
+    return DatabricksConfig(None, None, None, None, None)
+
+
+def set_config_provider(provider):
+    """
+    Sets a DatabricksConfigProvider that will be used for all future calls to get_config(),
+    used by the Databricks CLI code to discover the user's credentials.
+    """
+    global _config_provider
+    if provider and not isinstance(provider, DatabrickConfigProvider):
+        raise Exception('Must be instance of DatabrickConfigProvider: %s' % _config_provider)
+    _config_provider = provider
+
+
+def get_config_provider():
+    """
+    Returns the current DatabricksConfigProvider.
+    If None, the DefaultConfigProvider will be used.
+    """
+    global _config_provider
+    return _config_provider
+
+
+class DatabrickConfigProvider(object):
+    """
+    Responsible for providing hostname and authentication information to make
+    API requests against the Databricks REST API.
+    This method should generally return None if it cannot provide credentials, in order
+    to facilitate chanining of providers.
+    """
+
+    __metaclass__ = ABCMeta
+
+    @abstractmethod
+    def get_config(self):
+        pass
+
+
+class DefaultConfigProvider(DatabrickConfigProvider):
+    """Prefers environment variables, and then the default profile."""
+    def __init__(self):
+        self.env_provider = EnvironmentVariableConfigProvider()
+        self.default_profile_provider = ProfileConfigProvider()
+
+    def get_config(self):
+        env_config = self.env_provider.get_config()
+        if env_config:
+            return env_config
+
+        profile_config = self.default_profile_provider.get_config()
+        if profile_config:
+            return profile_config
+
+
+class EnvironmentVariableConfigProvider(DatabrickConfigProvider):
+    """Loads from system environment variables."""
+    def get_config(self):
         host = os.environ.get('DATABRICKS_HOST')
         username = os.environ.get('DATABRICKS_USERNAME')
         password = os.environ.get('DATABRICKS_PASSWORD')
         token = os.environ.get('DATABRICKS_TOKEN')
         insecure = os.environ.get('DATABRICKS_INSECURE')
-        return DatabricksConfig(host, username, password, token, insecure)
-    raw_config = _fetch_from_fs()
-    host = _get_option_if_exists(raw_config, profile, HOST)
-    username = _get_option_if_exists(raw_config, profile, USERNAME)
-    password = _get_option_if_exists(raw_config, profile, PASSWORD)
-    token = _get_option_if_exists(raw_config, profile, TOKEN)
-    insecure = _get_option_if_exists(raw_config, profile, INSECURE)
-    return DatabricksConfig(host, username, password, token, insecure)
+        config = DatabricksConfig(host, username, password, token, insecure)
+        if config.is_valid:
+            return config
+        return None
+
 
+class ProfileConfigProvider(DatabrickConfigProvider):
+    """Loads from the databrickscfg file."""
+    def __init__(self, profile=DEFAULT_SECTION):
+        self.profile = profile
 
-def is_environment_set():
-    token_exists = (os.environ.get('DATABRICKS_HOST')
-                    and os.environ.get('DATABRICKS_TOKEN'))
-    password_exists = (os.environ.get('DATABRICKS_HOST')
-                       and os.environ.get('DATABRICKS_USERNAME')
-                       and os.environ.get('DATABRICKS_PASSWORD'))
-    return token_exists or password_exists
+    def get_config(self):
+        raw_config = _fetch_from_fs()
+        host = _get_option_if_exists(raw_config, self.profile, HOST)
+        username = _get_option_if_exists(raw_config, self.profile, USERNAME)
+        password = _get_option_if_exists(raw_config, self.profile, PASSWORD)
+        token = _get_option_if_exists(raw_config, self.profile, TOKEN)
+        insecure = _get_option_if_exists(raw_config, self.profile, INSECURE)
+        config = DatabricksConfig(host, username, password, token, insecure)
+        if config.is_valid:
+            return config
+        return None
 
 
 class DatabricksConfig(object):
@@ -138,6 +247,10 @@ def from_token(cls, host, token, insecure=None):
     def from_password(cls, host, username, password, insecure=None):
         return DatabricksConfig(host, username, password, None, insecure)
 
+    @classmethod
+    def empty(cls):
+        return DatabricksConfig(None, None, None, None, None)
+
     @property
     def is_valid_with_token(self):
         return self.host is not None and self.token is not None

databricks_cli/groups/api.py

@@ -0,0 +1,101 @@
+"""Implement Databricks Groups API, interfacing with the GroupsService."""
+# Databricks CLI
+# Copyright 2017 Databricks, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"), except
+# that the use of services to which certain application programming
+# interfaces (each, an "API") connect requires that the user first obtain
+# a license for the use of the APIs from Databricks, Inc. ("Databricks"),
+# by creating an account at www.databricks.com and agreeing to either (a)
+# the Community Edition Terms of Service, (b) the Databricks Terms of
+# Service, or (c) another written agreement between Licensee and Databricks
+# for the use of the APIs.
+#
+# You may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+from databricks_cli.sdk import GroupsService
+from databricks_cli.groups.exceptions import GroupsException
+
+
+class GroupsApi(object):
+    """Implement the databricks '2.0/groups' API Interface."""
+
+    def __init__(self, api_client):
+        self.client = GroupsService(api_client)
+
+    def add_member(self, parent_name, member_type, member_name):
+        """Add a user or group to a group.
+
+        member_type is either 'group' or 'user'.
+        member_name is the name of the member.
+        """
+        if member_type == "group":
+            return self.client.add_to_group(parent_name=parent_name,
+                                            group_name=member_name)
+        elif member_type == "user":
+            return self.client.add_to_group(parent_name=parent_name,
+                                            user_name=member_name)
+        else:
+            raise GroupsException("Invalid 'member_type' {}".format(
+                member_type
+            ))
+
+    def create(self, group_name):
+        """Create a new group with the given name."""
+        return self.client.create_group(group_name)
+
+    def list_members(self, group_name):
+        """Return all of the members of a particular group."""
+        return self.client.get_group_members(group_name)
+
+    def list_all(self):
+        """Return all of the groups in an organization."""
+        return self.client.get_groups()
+
+    def list_parents(self, member_type, member_name):
+        """Retrieve all groups in which a given user or group is a member.
+
+        member_type is either 'group' or 'user'.
+        member_name is the name of the member.
+
+        Note: this method is non-recursive - it will return all groups in
+        which the given user or group is a member but not the groups in which
+        those groups are members).
+        """
+        if member_type == "group":
+            return self.client.get_groups_for_principal(group_name=member_name)
+        elif member_type == "user":
+            return self.client.get_groups_for_principal(user_name=member_name)
+        else:
+            raise GroupsException("Invalid 'member_type' {}".format(
+                member_type
+            ))
+
+    def remove_member(self, parent_name, member_type, member_name):
+        """Remove a user or group from a group.
+
+        member_type is either 'group' or 'user'.
+        member_name is the name of the member.
+        """
+        if member_type == "group":
+            return self.client.remove_from_group(parent_name=parent_name,
+                                                 group_name=member_name)
+        elif member_type == "user":
+            return self.client.remove_from_group(parent_name=parent_name,
+                                                 user_name=member_name)
+        else:
+            raise GroupsException("Invalid 'member_type' {}".format(
+                member_type
+            ))
+
+    def delete(self, group_name):
+        """Remove a group from this organization."""
+        return self.client.remove_group(group_name)