engineeringClusterSchedule.py

# Databricks notebook source
import requests
import json

# COMMAND ----------

dbutils.widgets.text("workspaceUrl", "https://adb-123456789101112.13.azuredatabricks.net" ,"Workspace URL")
dbutils.widgets.text("clusterActivity", "start" ,"Cluster Activity")
dbutils.widgets.text("clusterId", "4444-666666-abcdefgh" ,"Cluster ID")

# COMMAND ----------

pat = dbutils.secrets.get("KeyVault", "DatabricksPAT")
dbxWorkspaceURL = dbutils.widgets.get("workspaceUrl")
clusterActivity = dbutils.widgets.get("clusterActivity")
clusterId = dbutils.widgets.get("clusterId")

requestHeader = {'Authorization': 'Bearer ' + pat}

# COMMAND ----------

def getClusterInfo(clusterId):
  endpoint = '/api/2.0/clusters/get'
  response = requests.get(dbxWorkspaceURL + endpoint,headers=requestHeader, params=clusterId)
  return json.loads(response.content)

# COMMAND ----------

def startCluster(clusterId):
  endpoint = '/api/2.0/clusters/start'
  clusterInfo = getClusterInfo(clusterId)
  if clusterInfo['state'] == 'RUNNING':
    return f"{clusterId} is already running"
  else:
    response = requests.post(dbxWorkspaceURL + endpoint,headers=requestHeader, json=clusterId)
    return json.loads(response.content)

# COMMAND ----------

def listClusters():
  endpoint = '/api/2.0/clusters/list'
  response = requests.get(dbxWorkspaceURL + endpoint,headers=requestHeader)
  return json.loads(response.content)

# COMMAND ----------

def terminateCluster(clusterId):
  endpoint = '/api/2.0/clusters/delete'
  clusterInfo = getClusterInfo(clusterId)
  if clusterInfo['state'] == 'TERMINATED':
    return f"{clusterId} is already stopped"
  else:
    response = requests.post(dbxWorkspaceURL + endpoint,headers=requestHeader, json=clusterId)
    return json.loads(response.content)

# COMMAND ----------

def restartCluster(clusterId):
  endpoint = '/api/2.0/clusters/restart'
  clusterInfo = getClusterInfo(clusterId)
  if clusterInfo['state'] == 'RESTARTING':
    return f"{clusterId} is already restarting"
  else:
    response = requests.post(dbxWorkspaceURL + endpoint,headers=requestHeader, json=clusterId)
    return json.loads(response.content)

# COMMAND ----------

if clusterId == '':
  clusters = listClusters()
  for cluster in clusters["clusters"]:
    if "stream" in cluster["cluster_name"]:
      clusterId = {"cluster_id": cluster["cluster_id"]}
      if clusterActivity == "start":
        response = startCluster(clusterId)
        print(response)
      if clusterActivity == "stop":
        response = terminateCluster(clusterId)
        print(response)
      if clusterActivity == "restart":
        response = restartCluster(clusterId)
        print(response)
else:
  clusterId = {"cluster_id": clusterId}
  if clusterActivity == "start":
    response = startCluster(clusterId)
    print(response)
  if clusterActivity == "stop":
    response = terminateCluster(clusterId)
    print(response)
  if clusterActivity == "restart":
    response = restartCluster(clusterId)
    print(response)
  if clusterActivity == "info":
    response = getClusterInfo(clusterId)
    print(response)