def initialize_databricks_cluster(install_config: InstallConfiguration,
                                  resource_group: str,
                                  artifacts_path: str,
                                  tenant_id: str = None,
                                  subscription_id: str = None):
    print("Creating Databricks cluster ... ")
    backend_keyvault_name = install_config.backend_keyvault_name
    gdc_sp_secret_value = install_config.gdc_service_principal['password']
    adb_ws_name = install_config.databricks_workspace_name
    ws_url = arm_ops.get_databricks_workspace_url(
        resource_group=resource_group, ws_name=adb_ws_name)
    if not ws_url.startswith("https://"):
        ws_url = "https://" + ws_url
    adb_access_token = ad_ops.get_databricks_access_token(
        tenant_id, subscription_id)
    managed_libraries = []
    with open("cluster_libraries.json", "r") as libs_file:
        managed_libraries = json.load(libs_file)
    provisioning_rsp = arm_ops.provision_databricks_cluster(
        install_config=install_config,
        workspace_url=ws_url,
        oauth_access_token=adb_access_token,
        gdc_sp_secret_value=gdc_sp_secret_value,
        managed_libraries=managed_libraries)
    install_config.adb_cluster_details = provisioning_rsp
    secrets_ops.set_secret(keyvault_name=backend_keyvault_name,
                           secret_name="gdc-databricks-token",
                           value=provisioning_rsp['api_token']['token_value'])
    print("Uploading artifacts onto DBFS ...")
    arm_ops.upload_artifacts(workspace_url=ws_url,
                             oauth_access_token=adb_access_token,
                             local_artifacts_path=artifacts_path,
                             dbfs_dir_path="/mnt/gdc-artifacts")
def initialize_databricks_cluster(install_config: InstallConfiguration,
                                  resource_group: str,
                                  artifacts_path: str,
                                  tenant_id: str = None,
                                  subscription_id: str = None):

    runtime_storage = install_config.runtime_storage_account_name
    storage_account_keys_list_res = az.az_cli(
        "storage account keys list --account-name " + runtime_storage)
    storage_account_access_key = storage_account_keys_list_res[0]["value"]

    print("Creating Databricks cluster ... ")
    backend_keyvault_name = install_config.backend_keyvault_name
    wc_sp_secret_value = install_config.wc_service_principal['password']
    adb_ws_name = install_config.databricks_workspace_name
    ws_url = arm_ops.get_databricks_workspace_url(
        resource_group=resource_group, ws_name=adb_ws_name)
    if not ws_url.startswith("https://"):
        ws_url = "https://" + ws_url
    adb_access_token = ad_ops.get_databricks_access_token(
        tenant_id, subscription_id)
    managed_libraries = []
    with open("cluster_libraries.json", "r") as libs_file:
        managed_libraries = json.load(libs_file)
    provisioning_rsp = arm_ops.provision_databricks_cluster(
        install_config=install_config,
        workspace_url=ws_url,
        oauth_access_token=adb_access_token,
        wc_sp_secret_value=wc_sp_secret_value,
        managed_libraries=managed_libraries)
    install_config.adb_cluster_details = provisioning_rsp

    upload_mount_storage_file(ws_url, adb_access_token)

    execute_script_mount_storage_script(
        databricks_host=ws_url,
        token=adb_access_token,
        cluster_id=provisioning_rsp['cluster_id'],
        storage_account_name=runtime_storage,
        container_name="data",
        secret_key=storage_account_access_key)

    secrets_ops.set_secret(keyvault_name=backend_keyvault_name,
                           secret_name="wc-databricks-token",
                           value=provisioning_rsp['api_token']['token_value'])

    print("Uploading artifacts onto DBFS ...")
    arm_ops.upload_artifacts(workspace_url=ws_url,
                             oauth_access_token=adb_access_token,
                             local_artifacts_path=artifacts_path,
                             dbfs_dir_path="/mnt/watercooler/scripts")
def provision_databricks_cluster(install_config: InstallConfiguration,
                                 workspace_url: str, oauth_access_token: str,
                                 gdc_sp_secret_value: str,
                                 managed_libraries: list = None,
                                 gdc_sp_secret_name: str = "gdc-service-principal-secret",
                                 gdc_graph_api_sp_secret_name = "graph-api-service-principal-secret",
                                 secret_scope_name: str = "gdc",
                                 adb_cluster_name: str = "default-gdc-cluster",
                                 max_worker: int = 2,
                                 node_type_id: str = "Standard_DS3_v2",
                                 autotermination_minutes: int = 60):
    """

    :param managed_libraries: list of json object in format https://docs.databricks.com/dev-tools/api/latest/libraries.html#example-request
    :param workspace_url:
    :param oauth_access_token:
    :param gdc_sp_secret_value:
    :param gdc_sp_secret_name:
    :param secret_scope_name:
    :param adb_cluster_name:
    :param max_worker:
    :param node_type_id:
    :param autotermination_minutes:
    :return:  dict {
        "cluster_id": cluster_id,
        "api_token": adb_api_token
    }
    """
    print("Provisioning ADB cluster ...")
    assert oauth_access_token is not None
    adb_client = DatabricksAPI(host=workspace_url, token=oauth_access_token)
    scopes = adb_client.secret.list_scopes().get("scopes", [])
    if not any(x for x in scopes if x.get("name") == secret_scope_name):
        adb_client.secret.create_scope(scope=secret_scope_name,
                                       initial_manage_principal="users")

    adb_client.secret.put_secret(scope=secret_scope_name, key=gdc_sp_secret_name, string_value=gdc_sp_secret_value)
    # both databricks jobs use gdc-service service principal to access Graph API and other component
    # but we've introduce two secrets for flexibility even thought they have same value for now
    adb_client.secret.put_secret(scope=secret_scope_name, key=gdc_graph_api_sp_secret_name,
                                 string_value=gdc_sp_secret_value)

    adb_api_token = adb_client.token.create_token(comment="GDC Pipeline API token")
    cluster_id = None
    clusters = adb_client.cluster.list_clusters().get("clusters", [])
    cluster_rsp = list([x for x in clusters if x.get("cluster_name") == adb_cluster_name])
    if not cluster_rsp:
        print("Creating a new cluster %s" % adb_cluster_name)
        cluster_rsp = adb_client.cluster.create_cluster(cluster_name=adb_cluster_name, autoscale={
            "min_workers": 1,
            "max_workers": max_worker
        }, node_type_id=node_type_id, driver_node_type_id=node_type_id,
                                                        autotermination_minutes=autotermination_minutes,
                                                        enable_elastic_disk=True,
                                                        spark_version="6.6.x-scala2.11")
    else:
        print("Cluster %s exists at %s" % (adb_cluster_name, workspace_url))
        cluster_rsp = cluster_rsp[0]

    # capture cluster details as soon as it's available
    install_config.adb_cluster_details = {
        "cluster_id": cluster_rsp['cluster_id'],
        "api_token": adb_api_token
    }
    cluster_id = cluster_rsp['cluster_id']
    if managed_libraries:
        cluster_info = adb_client.cluster.get_cluster(cluster_id=cluster_id)
        cluster_state = cluster_info['state']
        # possible values PENDING, TERMINATED and RUNNING
        if cluster_state == "TERMINATED":
            print("Starting cluster %s " % cluster_id)
            adb_client.cluster.start_cluster(cluster_id=cluster_id)

        cluster_state = "PENDING"
        while cluster_state == "PENDING" or cluster_state == "RESTARTING" or cluster_state == "RESIZING":
            print("Waiting cluster %s " % cluster_id)
            sleep(5)
            cluster_info = adb_client.cluster.get_cluster(cluster_id=cluster_id)
            cluster_state = cluster_info['state']
            print("Cluster is now in state %s " % cluster_state)

        if cluster_state == "TERMINATING" or cluster_state == "TERMINATED" or cluster_state == "ERROR":
            print("Can't install managed libraries, cluster %s is not running" % cluster_id)
            raise RuntimeError("Can't install managed libraries, cluster %s is not running. Check Databricks Workspace Portal for details and  try again later" % cluster_id)
        else:
            try:
                print("Installing managed libraries on cluster %s " % cluster_id)
                install_managed_libraries(adb_client, cluster_id, managed_libraries)
            except BaseException as e:
                print("Failed to install libraries into cluster %s " % cluster_id)
                print(e)

    return {
        "cluster_id": cluster_id,
        "api_token": adb_api_token
    }