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 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().host == TEST_HOST
assert get_config().token == TEST_TOKEN
assert ProfileConfigProvider(TEST_PROFILE).get_config().host == TEST_HOST_2
assert ProfileConfigProvider(TEST_PROFILE).get_config().token == TEST_TOKEN
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 connect(profile):
"""Initialize Databricks API client
Args:
profile (str): Databricks CLI profile string
Returns:
ApiClient: Databricks ApiClient object
"""
config = ProfileConfigProvider(profile).get_config()
if config is None:
print_error("Cannot initialize ApiClient")
bye(1)
verify = config.insecure is None
if config.is_valid_with_token:
api_client = ApiClient(host=config.host,
token=config.token,
verify=verify)
api_client.default_headers[
"user-agent"] = "databrickslabs-jupyterlab-%s" % __version__
return api_client
else:
print_error(
"No token found for profile '%s'.\nUsername/password in .databrickscfg is not supported by databrickslabs-jupyterlab"
% profile)
bye(1)
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_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 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 _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 _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 pick_config(environment_data: Dict[str, Any]) -> Tuple[str, DatabricksConfig]:
config = EnvironmentVariableConfigProvider().get_config()
if config:
config_type = "ENV"
dbx_echo("Using configuration from the environment variables")
else:
dbx_echo("No environment variables provided, using the ~/.databrickscfg")
config = ProfileConfigProvider(environment_data["profile"]).get_config()
config_type = "PROFILE"
if not config:
raise Exception(
f"""Couldn't get profile with name: {environment_data["profile"]}. Please check the config settings"""
)
return config_type, config
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 tere.
config = ProfileConfigProvider(profile).get_config()
else:
raise ValueError("Please provide profile field")
if not config or not config.is_valid:
raise InvalidConfigurationError.for_profile(profile)
os.environ["DATABRICKS_HOST"] = config.host
os.environ["DATABRICKS_TOKEN"] = config.token
return function(*args, **kwargs)
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 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)
def _configure_cli_password(profile, insecure, host):
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)
update_and_persist_config(profile, new_config)
def prepare_environment(environment: str):
environment_data = InfoFile.get("environments").get(environment)
if not environment_data:
raise Exception(
f"No environment {environment} provided in the project file")
config = EnvironmentVariableConfigProvider().get_config()
if config:
config_type = "ENV"
dbx_echo("Using configuration from the environment variables")
else:
dbx_echo(
"No environment variables provided, using the ~/.databrickscfg")
config = ProfileConfigProvider(
environment_data["profile"]).get_config()
config_type = "PROFILE"
if not config:
raise Exception(
f"""Couldn't get profile with name: {environment_data["profile"]}. Please check the config settings"""
)
api_client = _get_api_client(config, command_name="cicdtemplates-")
_prepare_workspace_dir(api_client, environment_data["workspace_dir"])
if config_type == "ENV":
mlflow.set_tracking_uri(DATABRICKS_MLFLOW_URI)
elif config_type == "PROFILE":
mlflow.set_tracking_uri(
f'{DATABRICKS_MLFLOW_URI}://{environment_data["profile"]}')
else:
raise NotImplementedError(
f"Config type: {config_type} is not implemented")
experiment = mlflow.get_experiment_by_name(
environment_data["workspace_dir"])
if not experiment:
mlflow.create_experiment(environment_data["workspace_dir"],
environment_data["artifact_location"])
mlflow.set_experiment(environment_data["workspace_dir"])
return api_client
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)
# Create the egg:
#################
print("Preparing Recommenders library file ({})...".format(args.eggname))
myegg = create_egg(args.path_to_recommenders,
local_eggname=args.eggname,
overwrite=args.overwrite)
print("Created: {}".format(myegg))
############################
# Interact with Databricks:
############################
# first make sure you are using the correct profile and connecting to the intended workspace
my_api_client = _get_api_client(
ProfileConfigProvider(args.profile).get_config())
# Create a cluster if flagged
if args.create_cluster:
# treat args.cluster_id as the name, because if you create a cluster, you do not know its id yet.
DEFAULT_CLUSTER_CONFIG["cluster_name"] = args.cluster_id
cluster_info = ClusterApi(my_api_client).create_cluster(
DEFAULT_CLUSTER_CONFIG)
args.cluster_id = cluster_info["cluster_id"]
print("Creating a new cluster with name {}. New cluster_id={}".format(
DEFAULT_CLUSTER_CONFIG["cluster_name"], args.cluster_id))
# Upload the egg:
upload_path = Path(args.dbfs_path, args.eggname).as_posix()
# Check if file exists to alert user.