import sys import os import requests import json import time import argparse import logging from pprint import pprint def parse_arguments(): parser = argparse.ArgumentParser() parser.add_argument('--dbxApiVersion', help='Databricks api version', default='2.0') parser.add_argument('--loglevel', help='Databricks api version', default='INFO') requiredNamed = parser.add_argument_group('required named arguments') requiredNamed.add_argument('--dbxWorkspace', help='Databrics workspace name', required=True) requiredNamed.add_argument('--dbxClusterId', help='Databrics cluster id', required=True) requiredNamed.add_argument('--dbxAuthToken', help='Databrics auth token', required=True) return parser.parse_args() class BearerAuth(requests.auth.AuthBase): def __init__(self, token): self.token = token def __call__(self, r): r.headers["authorization"] = "Bearer " + self.token return r def dbx_cluster_status_check(databricks_url, cluster_id, auth, logger, timeout_sec=300): #wait for cluster to start timeout = time.time() + timeout_sec # 2 minutes from now running_cluster = requests.get(f"{databricks_url}/clusters/get", params={"cluster_id": cluster_id}, auth=auth).json() logger.debug(running_cluster) while running_cluster['state'] != "RUNNING": time.sleep(5) logger.info("Waiting for cluster to get in the running state...") running_cluster = requests.get(f"{databricks_url}/clusters/get", params={"cluster_id": cluster_id}, auth=auth).json() logger.debug(running_cluster) if time.time() > timeout: raise TimeoutError logger.info("Cluster is in the running state.") def dbx_library_install(databricks_url, cluster_id, auth, library_data, logger, verify=False, poll_freq = 5, level=0, max_depth=20): ''' Check if library is already installed, uninstall current version, restart cluster, install new version, verify installation. But with recursion :-) ''' if level > max_depth: raise RuntimeError("Recursion depth limit has been reached.") cluster_library_status = requests.get(f"{databricks_url}/libraries/cluster-status", params={"cluster_id": cluster_id}, auth=auth).json() installed_libraries = cluster_library_status.get('library_statuses',[]) library_payload = { "cluster_id": cluster_id, "libraries": [ library_data ] } for library_status in installed_libraries: if library_status['library'] == library_data: if library_status['status'] == 'INSTALLED' and not verify: # uninstall logger.info("Found installed library, uninstalling...") requests.post(f"{databricks_url}/libraries/uninstall", json=library_payload, auth=auth).json() # wait until library changes it's status to uninstalled dbx_library_install(databricks_url, cluster_id, auth, library_data, logger, verify=verify, poll_freq=poll_freq, level=level+1) elif library_status['status'] == 'INSTALLED' and verify: # if verify set to true return elif library_status['status'] == 'UNINSTALL_ON_RESTART': # Restart the cluster logger.info("Found library marked for uninstall, restarting the cluster") response = requests.post(f"{databricks_url}/clusters/restart", json={"cluster_id": cluster_id}, auth=auth).json() logger.debug(response) dbx_cluster_status_check(databricks_url, cluster_id, databricks_auth, logger) # Check whether library has been removed dbx_library_install(databricks_url, cluster_id, auth, library_data, logger, verify=verify, poll_freq=poll_freq, level=level+1) elif library_status['status'] == 'SKIPPED' or library_status['status'] == 'FAILED': logger.error("Uninstall of the libraru has failed") raise RuntimeError("The library deinstallation failed") else: # Wait for status to change logger.info("Waiting for %s library status to stabilize.." % library_status['status']) time.sleep(poll_freq) dbx_library_install(databricks_url, cluster_id, auth, library_data, logger, verify=verify, poll_freq=poll_freq, level=level+1) if len(installed_libraries)==0 and not verify: logger.info("Installing the new version of the library") response = requests.post(f"{databricks_url}/libraries/install", json=library_payload, auth=auth).json() logger.info(response) # Verify that installation was successful dbx_library_install(databricks_url, cluster_id, auth, library_data, logger, verify=True, poll_freq=poll_freq, level=level+1) logger.info("Library have been sucessfully installed") if __name__ == "__main__": args = parse_arguments() FORMAT = '%(asctime)s - %(name)s - %(levelname)s - %(message)s' logging.basicConfig(format=FORMAT, level=args.loglevel) logger = logging.getLogger() databricks_url = f"https://{args.dbxWorkspace}/api/{args.dbxApiVersion}" databricks_auth=BearerAuth(args.dbxAuthToken) clusters = requests.get(f"{databricks_url}/clusters/list", auth=databricks_auth).json() cluster = requests.get(f"{databricks_url}/clusters/get", params={"cluster_id": args.dbxClusterId }, auth=databricks_auth).json() #pprint(clusters) for cluster in clusters['clusters']: print('==============') print(f"Cluster name {cluster['cluster_name']}") print(f"Cluster source {cluster['cluster_source']}") #if cluster.get('cluster_log_conf', None): # print("Cluster log conf") # print(cluster['cluster_log_conf']) #print ("Cluster aws_attributes") #pprint(cluster['aws_attributes']) if cluster.get('custom_tags', None): print ("Cluster custom_tags") pprint(cluster['custom_tags']) else: print("No tags specified") environment_full = { 'dev':'Development', 'qa': 'QA' 'stage': 'Staging', 'prod': 'Production' } platform = 'Delphi' service = 'Databricks' project_short = { 'Apps ETL': 'APPE' 'Simple ETL': 'SETL' 'Exploration Area': 'EXP' 'Undefined': 'UND' } 'plat_env_project_service': 'DLP_DEV_SETL_DBX', 'platform': 'Delphi', 'platform_with_env': 'Delphi Development', 'project': 'Apps ETL', 'service': 'Databricks' def tag_cluster(project, cluster_id, logger): current_tags = requests.get(f"{databricks_url}/clusters/get", params={"cluster_id": args.dbxClusterId }, auth=databricks_auth).json().get(custom_tags, None) # get env from cluster role platform = 'Delphi' service = 'Databricks' plat_env_project_service = f"DLP_{env.upper()}_{project_short.get(project, 'UND')}