def test_configure_cli_insecure(): runner = CliRunner() runner.invoke(cli.configure_cli, ['--token', '--insecure'], input=(TEST_HOST + '\n' + TEST_TOKEN + '\n')) assert get_config_for_profile(DEFAULT_SECTION).host == TEST_HOST assert get_config_for_profile(DEFAULT_SECTION).token == TEST_TOKEN assert get_config_for_profile(DEFAULT_SECTION).insecure == 'True'
def test_configure_cli(): runner = CliRunner() runner.invoke(cli.configure_cli, input=(TEST_HOST + '\n' + TEST_USER + '\n' + TEST_PASSWORD + '\n' + TEST_PASSWORD + '\n')) assert get_config_for_profile(DEFAULT_SECTION).host == TEST_HOST assert get_config_for_profile(DEFAULT_SECTION).username == TEST_USER assert get_config_for_profile(DEFAULT_SECTION).password == TEST_PASSWORD
def test_configure_two_sections(): runner = CliRunner() runner.invoke(cli.configure_cli, ['--token'], input=(TEST_HOST + '\n' + TEST_TOKEN + '\n')) runner.invoke(cli.configure_cli, ['--token', '--profile', TEST_PROFILE], input=(TEST_HOST_2 + '\n' + TEST_TOKEN + '\n')) assert get_config_for_profile(DEFAULT_SECTION).host == TEST_HOST assert get_config_for_profile(DEFAULT_SECTION).token == TEST_TOKEN assert get_config_for_profile(TEST_PROFILE).host == TEST_HOST_2 assert get_config_for_profile(TEST_PROFILE).token == TEST_TOKEN
def test_update_and_persist_config_case_sensitive(): config = DatabricksConfig.from_token(TEST_HOST, TEST_TOKEN) update_and_persist_config(TEST_PROFILE, config) config_2 = DatabricksConfig.from_password(TEST_HOST, TEST_USER, TEST_PASSWORD) update_and_persist_config(TEST_PROFILE.upper(), config_2) config = get_config_for_profile(TEST_PROFILE) assert config.is_valid_with_token assert not config.is_valid_with_password config_2 = get_config_for_profile(TEST_PROFILE.upper()) assert config_2.is_valid_with_password assert not config_2.is_valid_with_token
def test_get_config_if_token_environment_set(): with patch.dict('os.environ', {'DATABRICKS_HOST': TEST_HOST, 'DATABRICKS_TOKEN': TEST_TOKEN}): config = get_config_for_profile(TEST_PROFILE) assert config.host == TEST_HOST assert config.token == TEST_TOKEN
def get_databricks_host_creds(profile=None): """ Reads in configuration necessary to make HTTP requests to a Databricks server. This uses the Databricks CLI's ConfigProvider interface to load the DatabricksConfig object. This method will throw an exception if sufficient auth cannot be found. :param profile: Databricks CLI profile. If not provided, we will read the default profile. :return: :py:class:`mlflow.rest_utils.MlflowHostCreds` which includes the hostname and authentication information necessary to talk to the Databricks server. """ if not hasattr(provider, 'get_config'): eprint( "Warning: support for databricks-cli<0.8.0 is deprecated and will be removed" " in a future version.") config = provider.get_config_for_profile(profile) elif profile: config = provider.ProfileConfigProvider(profile).get_config() else: config = provider.get_config() if not config or not config.host: _fail_malformed_databricks_auth(profile) insecure = hasattr(config, 'insecure') and config.insecure if config.username is not None and config.password is not None: return MlflowHostCreds(config.host, username=config.username, password=config.password, ignore_tls_verification=insecure) elif config.token: return MlflowHostCreds(config.host, token=config.token, ignore_tls_verification=insecure) _fail_malformed_databricks_auth(profile)
def decorator(*args, **kwargs): profile = get_profile_from_context() config = get_config_for_profile(profile) if not config.is_valid: raise InvalidConfigurationError(profile) kwargs['api_client'] = _get_api_client(config) return function(*args, **kwargs)
def test_get_config_if_password_environment_set(): with patch.dict('os.environ', {'DATABRICKS_HOST': TEST_HOST, 'DATABRICKS_USERNAME': TEST_USER, 'DATABRICKS_PASSWORD': TEST_PASSWORD}): config = get_config_for_profile(TEST_PROFILE) assert config.host == TEST_HOST assert config.username == TEST_USER assert config.password == TEST_PASSWORD
def test_update_and_persist_config(): config = DatabricksConfig.from_token(TEST_HOST, TEST_TOKEN) update_and_persist_config(DEFAULT_SECTION, config) config = get_config_for_profile(DEFAULT_SECTION) assert config.is_valid_with_token assert config.host == TEST_HOST assert config.token == TEST_TOKEN # Overwrite conf for same section. config = DatabricksConfig.from_password(TEST_HOST, TEST_USER, TEST_PASSWORD) update_and_persist_config(DEFAULT_SECTION, config) config = get_config_for_profile(DEFAULT_SECTION) assert config.is_valid_with_password assert not config.is_valid_with_token assert config.host == TEST_HOST assert config.username == TEST_USER assert config.password == TEST_PASSWORD
def test_get_config_for_profile_empty(): config = get_config_for_profile(TEST_PROFILE) assert not config.is_valid_with_password assert not config.is_valid_with_token assert config.host is None assert config.username is None assert config.password is None assert config.token is None
def get_databricks_host_creds(server_uri=None): """ Reads in configuration necessary to make HTTP requests to a Databricks server. This uses the Databricks CLI's ConfigProvider interface to load the DatabricksConfig object. If no Databricks CLI profile is found corresponding to the server URI, this function will attempt to retrieve these credentials from the Databricks Secret Manager. For that to work, the server URI will need to be of the following format: "databricks://scope:prefix". In the Databricks Secret Manager, we will query for a secret in the scope "<scope>" for secrets with keys of the form "<prefix>-host" and "<prefix>-token". Note that this prefix *cannot* be empty if trying to authenticate with this method. If found, those host credentials will be used. This method will throw an exception if sufficient auth cannot be found. :param server_uri: A URI that specifies the Databricks profile you want to use for making requests. :return: :py:class:`mlflow.rest_utils.MlflowHostCreds` which includes the hostname and authentication information necessary to talk to the Databricks server. """ profile, path = get_db_info_from_uri(server_uri) if not hasattr(provider, "get_config"): _logger.warning( "Support for databricks-cli<0.8.0 is deprecated and will be removed" " in a future version.") config = provider.get_config_for_profile(profile) elif profile: config = provider.ProfileConfigProvider(profile).get_config() else: config = provider.get_config() # if a path is specified, that implies a Databricks tracking URI of the form: # databricks://profile-name/path-specifier if (not config or not config.host) and path: dbutils = _get_dbutils() if dbutils: # Prefix differentiates users and is provided as path information in the URI key_prefix = path host = dbutils.secrets.get(scope=profile, key=key_prefix + "-host") token = dbutils.secrets.get(scope=profile, key=key_prefix + "-token") if host and token: config = provider.DatabricksConfig.from_token(host=host, token=token, insecure=False) if not config or not config.host: _fail_malformed_databricks_auth(profile) insecure = hasattr(config, "insecure") and config.insecure if config.username is not None and config.password is not None: return MlflowHostCreds( config.host, username=config.username, password=config.password, ignore_tls_verification=insecure, ) elif config.token: return MlflowHostCreds(config.host, token=config.token, ignore_tls_verification=insecure) _fail_malformed_databricks_auth(profile)
def _configure_cli_password(profile, insecure): config = get_config_for_profile(profile) if config.password: default_password = '******' * len(config.password) else: default_password = None host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost()) username = click.prompt(PROMPT_USERNAME, default=config.username) password = click.prompt(PROMPT_PASSWORD, default=default_password, hide_input=True, confirmation_prompt=True) if password == default_password: password = config.password new_config = DatabricksConfig.from_password(host, username, password, insecure) update_and_persist_config(profile, new_config)
def _get_db_hostname_and_auth(): """ Reads the hostname & auth token to use for running on Databricks from the config file created by the Databricks CLI. """ home_dir = os.path.expanduser("~") cfg_file = os.path.join(home_dir, ".databrickscfg") if not os.path.exists(cfg_file): raise ExecutionException("Could not find profile for Databricks CLI in %s. Make sure the " "the Databricks CLI is installed and that credentials have been " "configured as described in " "https://github.com/databricks/databricks-cli" % cfg_file) else: config = provider.get_config_for_profile(provider.DEFAULT_SECTION) return config.host, config.token, config.username, config.password
def get_databricks_http_request_kwargs_or_fail(profile=None): """ Reads in configuration necessary to make HTTP requests to a Databricks server. This uses the Databricks CLI's ConfigProvider interface to load the DatabricksConfig object. This method will throw an exception if sufficient auth cannot be found. :param profile: Databricks CLI profile. If not provided, we will read the default profile. :return: Dictionary with parameters that can be passed to http_request(). This will at least include the hostname and headers sufficient to authenticate to Databricks. """ if not hasattr(provider, 'get_config'): eprint( "Warning: support for databricks-cli<0.8.0 is deprecated and will be removed" " in a future version.") config = provider.get_config_for_profile(profile) elif profile: config = provider.ProfileConfigProvider(profile).get_config() else: config = provider.get_config() hostname = config.host if not hostname: _fail_malformed_databricks_auth(profile) auth_str = None if config.username is not None and config.password is not None: basic_auth_str = ("%s:%s" % (config.username, config.password)).encode("utf-8") auth_str = "Basic " + base64.standard_b64encode(basic_auth_str).decode( "utf-8") elif config.token: auth_str = "Bearer %s" % config.token else: _fail_malformed_databricks_auth(profile) headers = { "Authorization": auth_str, } verify = True if hasattr(config, 'insecure') and config.insecure: verify = False return { 'hostname': hostname, 'headers': headers, 'verify': verify, }
def get_databricks_hostname_and_auth(): """ Reads the hostname & auth token to use for running on Databricks from the config file created by the Databricks CLI. Returns a tuple of (hostname, auth, token) to use when making API requests. """ profile = provider.DEFAULT_SECTION config = provider.get_config_for_profile(profile) if config.username is not None and config.password is not None: return config.host, (config.username, config.password), config.token if config.token: return config.host, None, config.token raise Exception( "Got malformed Databricks CLI profile '%s'. Please make sure the Databricks " "CLI is properly configured as described at " "https://github.com/databricks/databricks-cli." % profile)
def get_databricks_http_request_kwargs_or_fail(profile=None): """ Reads in configuration necessary to make HTTP requests to a Databricks server. This uses the Databricks CLI's ConfigProvider interface to load the DatabricksConfig object. This method will throw an exception if sufficient auth cannot be found. :param profile: Databricks CLI profile. If not provided, we will read the default profile. :return: Dictionary with parameters that can be passed to http_request(). This will at least include the hostname and headers sufficient to authenticate to Databricks. """ if not profile: profile = provider.DEFAULT_SECTION config = provider.get_config_for_profile(profile) hostname = config.host if not hostname: _fail_malformed_databricks_auth(profile) basic_auth_str = None if config.username is not None and config.password is not None: basic_auth_str = ("%s:%s" % (config.username, config.password)).encode("utf-8") elif config.token: basic_auth_str = ("token:%s" % config.token).encode("utf-8") if not basic_auth_str: _fail_malformed_databricks_auth(profile) headers = { "Authorization": "Basic " + base64.standard_b64encode(basic_auth_str).decode("utf-8") } secure_verify = True if hasattr(config, 'insecure') and config.insecure: secure_verify = False return { 'hostname': hostname, 'headers': headers, 'secure_verify': secure_verify, }
def get_host_token(profile=None): """ Get the host and token for a profile from ~/.databrickscfg. """ cfg = provider.get_config( ) if profile is None else provider.get_config_for_profile(profile) return (cfg.host, cfg.token)
def get_credentials(profile): cfg = provider.get_config( ) if profile is None else provider.get_config_for_profile(profile) return (cfg.host, cfg.token)
def get_credentials(profile): from databricks_cli.configure import provider cfg = provider.get_config( ) if profile is None else provider.get_config_for_profile(profile) return (cfg.host, cfg.token)
def get_host_token(profile=None): cfg = provider.get_config( ) if profile is None else provider.get_config_for_profile(profile) return (cfg.host, cfg.token)
def _configure_cli_token(profile, insecure): config = get_config_for_profile(profile) 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 tgt_api_client(load_env): target_profile = os.environ.get("AZURE_TARGET_WORKSPACE") config = get_config_for_profile(target_profile) return ApiClient(host=config.host, token=config.token)