def _configure_cli_aad_token(profile, insecure, host, jobs_api_version):
config = ProfileConfigProvider(
profile).get_config() or DatabricksConfig.empty()
if ENV_AAD_TOKEN not in os.environ:
click.echo('[ERROR] Set Environment Variable \'%s\' with your '
'AAD Token and run again.\n' % ENV_AAD_TOKEN)
click.echo(
'Commands to run to get your AAD token:\n'
'\t az login\n'
'\t token_response=$(az account get-access-token '
'--resource 2ff814a6-3304-4ab8-85cb-cd0e6f879c1d)\n'
'\t export %s=$(jq .accessToken -r <<< "$token_response")\n' %
ENV_AAD_TOKEN)
return
if not host:
host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost())
aad_token = os.environ.get(ENV_AAD_TOKEN)
new_config = DatabricksConfig.from_token(host=host,
token=aad_token,
refresh_token=None,
insecure=insecure,
jobs_api_version=jobs_api_version)
update_and_persist_config(profile, new_config)
def _configure_cli_token(profile, insecure):
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(version):
profile = get_profile_from_context()
config = ProfileConfigProvider(
profile).get_config() if profile else get_config()
new_config = config or DatabricksConfig.empty()
new_config.jobs_api_version = version
update_and_persist_config(profile, new_config)
def test_get_config_uses_default_profile():
config = DatabricksConfig.from_token("hosty", "hello")
update_and_persist_config(DEFAULT_SECTION, config)
config = get_config()
assert config.is_valid_with_token
assert config.host == "hosty"
assert config.token == "hello"
def decorator(*args, **kwargs):
ctx = click.get_current_context()
command_name = "-".join(ctx.command_path.split(" ")[1:])
command_name += "-" + str(uuid.uuid1())
profile = get_profile_from_context()
if profile:
# If we request a specific profile, only get credentials from there.
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)
# This checks if an OAuth access token has expired and will attempt to refresh it if
# a refresh token is present
if config.host and config.token and config.refresh_token:
config.token, config.refresh_token, updated = \
check_and_refresh_access_token(config.host, config.token, config.refresh_token)
if updated:
update_and_persist_config(profile, config)
kwargs['api_client'] = _get_api_client(config, command_name)
return function(*args, **kwargs)
def _configure_cli_token(profile, insecure):
PROMPT_HOST = "Databricks Host (should begin with https://)"
PROMPT_TOKEN = "Token" # NOQA
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, hide_input=True)
new_config = DatabricksConfig.from_token(host, token, insecure)
update_and_persist_config(profile, new_config)
def mock_conf_dir(request):
path = tempfile.mkdtemp()
provider._home = path
# create config
host = request.config.getoption('--host')
token = request.config.getoption('--token')
config = provider.DatabricksConfig.from_token(host, token)
provider.update_and_persist_config(provider.DEFAULT_SECTION, config)
yield
shutil.rmtree(path)
def _configure_cli_token(profile, insecure, host, jobs_api_version):
config = ProfileConfigProvider(
profile).get_config() or DatabricksConfig.empty()
if not host:
host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost())
token = click.prompt(PROMPT_TOKEN, default=config.token, hide_input=True)
new_config = DatabricksConfig.from_token(host, token, insecure,
jobs_api_version)
update_and_persist_config(profile, new_config)
def test_get_config_override_profile():
config = DatabricksConfig.from_token("yo", "lo")
update_and_persist_config(TEST_PROFILE, config)
try:
provider = ProfileConfigProvider(TEST_PROFILE)
set_config_provider(provider)
config = get_config()
assert config.host == "yo"
assert config.token == "lo"
finally:
set_config_provider(None)
def test_get_config_uses_path_env_variable(tmpdir):
cfg_file = tmpdir.join("some-cfg-path").strpath
with patch.dict('os.environ', {'DATABRICKS_CONFIG_FILE': cfg_file}):
config = DatabricksConfig.from_token("hosty", "hello")
update_and_persist_config(DEFAULT_SECTION, config)
config = get_config()
assert os.path.exists(cfg_file)
assert not os.path.exists(_get_path())
assert config.is_valid_with_token
assert config.host == "hosty"
assert config.token == "hello"
def _configure_cli_token_file(profile, token_file, host, insecure):
if not path.exists(token_file):
raise RuntimeError('Unable to read token from "{}"'.format(token_file))
with io.open(token_file, encoding='utf-8') as f:
token = f.readline().strip()
config = ProfileConfigProvider(profile).get_config() or DatabricksConfig.empty()
if not host:
host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost())
new_config = DatabricksConfig.from_token(host, token, insecure)
update_and_persist_config(profile, new_config)
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 _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 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 _configure_cli_oauth(profile, insecure, host, scope, jobs_api_version):
config = ProfileConfigProvider(
profile).get_config() or DatabricksConfig.empty()
if not host:
host = click.prompt(PROMPT_HOST, default=config.host, type=_DbfsHost())
if not scope:
scope = click.prompt(
"Pick one or more comma-separated scopes from {scopes}".format(
scopes=DEFAULT_SCOPES),
value_proc=scope_format,
type=click.Choice(DEFAULT_SCOPES),
default=None)
access_token, refresh_token = get_tokens(host, scope)
new_config = DatabricksConfig.from_token(host=host,
token=access_token,
refresh_token=refresh_token,
insecure=insecure,
jobs_api_version=jobs_api_version)
update_and_persist_config(profile, new_config)
def _configure_cli_password(profile, insecure, host, jobs_api_version):
config = ProfileConfigProvider(
profile).get_config() or DatabricksConfig.empty()
if config.password:
default_password = '******' * len(config.password)
else:
default_password = None
if not host:
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, jobs_api_version)
update_and_persist_config(profile, new_config)
def test_get_config_throw_exception_if_profile_invalid():
invalid_config = DatabricksConfig.from_token(None, None)
update_and_persist_config(DEFAULT_SECTION, invalid_config)
with pytest.raises(InvalidConfigurationError):
get_config()
def wrapper(test, *args, **kwargs):
config = DatabricksConfig.from_token('test-host', 'test-token')
update_and_persist_config(DEFAULT_SECTION, config)
return test(*args, **kwargs)
def _configure_cli_token(profile):
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)
update_and_persist_config(profile, new_config)
