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.